From 8e8ea4c05e52fb9f7ae59f348980a0a1b3d9a74c Mon Sep 17 00:00:00 2001 From: Steven Enamakel <31011319+senamakel@users.noreply.github.com> Date: Thu, 25 Jun 2026 09:24:13 -0700 Subject: [PATCH] feat(runtime): add runtime Python server (#4106) --- src/openhuman/harness_init/registry.rs | 54 ++- src/openhuman/memory_tree/nlp/client.rs | 227 --------- src/openhuman/memory_tree/nlp/mod.rs | 61 ++- src/openhuman/memory_tree/nlp/provision.rs | 241 ---------- src/openhuman/memory_tree/nlp/service.py | 108 ----- src/openhuman/mod.rs | 1 + src/openhuman/runtime_python_server/mod.rs | 18 + .../runtime_python_server/protocol.rs | 66 +++ .../runtime_python_server/registry.rs | 46 ++ src/openhuman/runtime_python_server/server.py | 111 +++++ src/openhuman/runtime_python_server/server.rs | 453 ++++++++++++++++++ src/openhuman/runtime_python_server/spacy.rs | 306 ++++++++++++ src/openhuman/runtime_python_server/types.rs | 28 ++ 13 files changed, 1107 insertions(+), 613 deletions(-) delete mode 100644 src/openhuman/memory_tree/nlp/client.rs delete mode 100644 src/openhuman/memory_tree/nlp/provision.rs delete mode 100644 src/openhuman/memory_tree/nlp/service.py create mode 100644 src/openhuman/runtime_python_server/mod.rs create mode 100644 src/openhuman/runtime_python_server/protocol.rs create mode 100644 src/openhuman/runtime_python_server/registry.rs create mode 100644 src/openhuman/runtime_python_server/server.py create mode 100644 src/openhuman/runtime_python_server/server.rs create mode 100644 src/openhuman/runtime_python_server/spacy.rs create mode 100644 src/openhuman/runtime_python_server/types.rs diff --git a/src/openhuman/harness_init/registry.rs b/src/openhuman/harness_init/registry.rs index a551b091c..6c85ac211 100644 --- a/src/openhuman/harness_init/registry.rs +++ b/src/openhuman/harness_init/registry.rs @@ -7,7 +7,8 @@ //! Current steps (all non-required — failure degrades to a fallback): //! 1. `python_runtime` — managed CPython (prerequisite for spaCy). //! 2. `spacy` — spaCy venv + `en_core_web_sm` model. -//! 3. `node_runtime` — managed Node.js (skills / MCP). +//! 3. `runtime_python_server` — long-running Python backend host. +//! 4. `node_runtime` — managed Node.js (skills / MCP). //! //! Voice models (Whisper, Piper) and Ollama stay lazy/opt-in and are //! intentionally NOT registered here; they can be added later as steps. @@ -38,7 +39,12 @@ pub struct HarnessInitStep { /// The ordered list of mandatory eager steps. pub fn all_steps() -> Vec { - vec![python_runtime_step(), spacy_step(), node_runtime_step()] + vec![ + python_runtime_step(), + spacy_step(), + runtime_python_server_step(), + node_runtime_step(), + ] } // ── python_runtime ────────────────────────────────────────────────────────── @@ -86,6 +92,36 @@ async fn python_run(config: &Config) -> Result<(), String> { // ── spacy ───────────────────────────────────────────────────────────────── +fn runtime_python_server_step() -> HarnessInitStep { + HarnessInitStep { + id: "runtime_python_server", + label: "Runtime Python server", + required: false, + is_done: |config| Box::pin(runtime_python_server_is_done(config)), + run: |config| Box::pin(runtime_python_server_run(config)), + } +} + +async fn runtime_python_server_is_done(config: &Config) -> bool { + if crate::openhuman::runtime_python_server::enabled_backends(config).is_empty() { + return true; + } + let status = crate::openhuman::runtime_python_server::status().await; + status.running +} + +async fn runtime_python_server_run(config: &Config) -> Result<(), String> { + if crate::openhuman::runtime_python_server::enabled_backends(config).is_empty() { + return Ok(()); + } + crate::openhuman::runtime_python_server::ensure_started(config) + .await + .map(|_| { + log::info!("[harness_init] runtime Python server ready"); + }) + .map_err(|e| format!("{e:#}")) +} + fn spacy_step() -> HarnessInitStep { HarnessInitStep { id: "spacy", @@ -100,14 +136,14 @@ async fn spacy_is_done(config: &Config) -> bool { if !config.runtime_python.enabled || !config.memory_tree.spacy_enabled { return true; } - crate::openhuman::memory_tree::nlp::spacy_provisioned(config) + crate::openhuman::runtime_python_server::spacy_provisioned(config) } async fn spacy_run(config: &Config) -> Result<(), String> { if !config.runtime_python.enabled || !config.memory_tree.spacy_enabled { return Ok(()); } - crate::openhuman::memory_tree::nlp::ensure_spacy(config) + crate::openhuman::runtime_python_server::ensure_spacy(config) .await .map(|_| { log::info!("[harness_init] spaCy provisioned"); @@ -167,7 +203,15 @@ mod tests { fn all_steps_have_stable_ids_and_are_non_required() { let steps = all_steps(); let ids: Vec<_> = steps.iter().map(|s| s.id).collect(); - assert_eq!(ids, vec!["python_runtime", "spacy", "node_runtime"]); + assert_eq!( + ids, + vec![ + "python_runtime", + "spacy", + "runtime_python_server", + "node_runtime" + ] + ); assert!(steps.iter().all(|s| !s.required)); assert!(steps.iter().all(|s| !s.label.is_empty())); } diff --git a/src/openhuman/memory_tree/nlp/client.rs b/src/openhuman/memory_tree/nlp/client.rs deleted file mode 100644 index 15e0a8253..000000000 --- a/src/openhuman/memory_tree/nlp/client.rs +++ /dev/null @@ -1,227 +0,0 @@ -//! Long-lived stdio client for the spaCy NER service. -//! -//! Spawns `service.py` under the provisioned venv interpreter, reads the one -//! `{"ready": true}` handshake line, then issues one request per query over a -//! mutex-guarded stdin/stdout pair. The model loads once for the life of the -//! process; queries are cheap line round-trips. -//! -//! A process-global [`OnceCell`] memoises the (possibly failed) initialisation -//! so the expensive provisioning + model load happens at most once. If init -//! fails (no Python, spaCy install failed, model load error) the cell stores -//! `None` and every caller falls back to the in-Rust extractor for the rest of -//! the process lifetime. - -use std::process::Stdio; -use std::sync::Arc; -use std::time::Duration; - -use anyhow::{bail, Context, Result}; -use serde::Deserialize; -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines}; -use tokio::process::{Child, ChildStdin, ChildStdout, Command}; -use tokio::sync::{Mutex, OnceCell}; - -use crate::openhuman::config::Config; -use crate::openhuman::memory_tree::nlp::provision::{ensure_spacy, SpacyRuntime}; - -/// Per-request timeout for a single extraction round-trip. The model is -/// already loaded; extraction of a short query is sub-100ms, so this is just a -/// guard against a wedged child. -const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); - -/// One named entity span as reported by spaCy. -#[derive(Debug, Clone, Deserialize)] -pub struct SpacyEntity { - pub text: String, - pub label: String, - #[serde(default)] - pub start: u32, - #[serde(default)] - pub end: u32, -} - -/// Parsed extraction response for one query. -#[derive(Debug, Clone, Deserialize)] -pub struct SpacyResponse { - #[serde(default)] - pub entities: Vec, - #[serde(default)] - pub nouns: Vec, - #[serde(default)] - pub id: Option, - #[serde(default)] - pub error: Option, -} - -#[derive(Deserialize)] -struct ReadyLine { - #[serde(default)] - ready: bool, - #[serde(default)] - error: Option, -} - -/// Mutable I/O state for the child, guarded by a mutex so requests serialise. -struct Inner { - // `_child` is retained so the process stays alive (and is killed on drop). - _child: Child, - stdin: ChildStdin, - stdout: Lines>, - next_id: u64, -} - -/// Handle to a running spaCy NER service. -pub struct SpacyNer { - inner: Mutex, -} - -impl SpacyNer { - /// Spawn the service and complete the readiness handshake. - async fn spawn(runtime: &SpacyRuntime) -> Result { - log::debug!( - "[memory_tree::nlp] spawning spaCy service python={} script={}", - runtime.python_bin.display(), - runtime.service_script.display() - ); - let mut child = Command::new(&runtime.python_bin) - .arg("-u") - .arg(&runtime.service_script) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::null()) - .kill_on_drop(true) - .spawn() - .with_context(|| "spawning spaCy service process")?; - - let stdin = child.stdin.take().context("spaCy child stdin missing")?; - let stdout = child.stdout.take().context("spaCy child stdout missing")?; - let mut lines = BufReader::new(stdout).lines(); - - // Handshake: wait for the ready line. - let ready_line = - match tokio::time::timeout(Duration::from_secs(30), lines.next_line()).await { - Ok(Ok(Some(line))) => line, - Ok(Ok(None)) => bail!("spaCy service exited before readiness handshake"), - Ok(Err(e)) => return Err(e).context("reading spaCy readiness line"), - Err(_) => bail!("spaCy service readiness handshake timed out"), - }; - let ready: ReadyLine = serde_json::from_str(&ready_line) - .with_context(|| format!("parsing spaCy ready line: {ready_line}"))?; - if !ready.ready { - bail!( - "spaCy service failed to load model: {}", - ready.error.unwrap_or_else(|| "unknown".into()) - ); - } - - log::info!("[memory_tree::nlp] spaCy service ready"); - Ok(Self { - inner: Mutex::new(Inner { - _child: child, - stdin, - stdout: lines, - next_id: 0, - }), - }) - } - - /// Extract named entities + salient nouns from `text`. - pub async fn extract(&self, text: &str) -> Result { - let mut guard = self.inner.lock().await; - let id = guard.next_id; - guard.next_id += 1; - let id_str = id.to_string(); - - let req = serde_json::json!({ "id": id_str, "text": text }); - let mut line = serde_json::to_string(&req)?; - line.push('\n'); - guard - .stdin - .write_all(line.as_bytes()) - .await - .context("writing spaCy request")?; - guard - .stdin - .flush() - .await - .context("flushing spaCy request")?; - - // Read until the response matching our id (skip stray lines). - loop { - let next = tokio::time::timeout(REQUEST_TIMEOUT, guard.stdout.next_line()).await; - let line = match next { - Ok(Ok(Some(l))) => l, - Ok(Ok(None)) => bail!("spaCy service closed mid-request"), - Ok(Err(e)) => return Err(e).context("reading spaCy response"), - Err(_) => bail!("spaCy request timed out"), - }; - let resp: SpacyResponse = match serde_json::from_str(&line) { - Ok(r) => r, - Err(e) => { - log::warn!("[memory_tree::nlp] unparseable spaCy line skipped: {e}"); - continue; - } - }; - if resp.id.as_deref() == Some(id_str.as_str()) { - if let Some(err) = &resp.error { - bail!("spaCy extraction error: {err}"); - } - return Ok(resp); - } - // Different id — keep reading. - } - } -} - -/// Process-global memoised NER handle. `None` means initialisation was -/// attempted and failed; callers fall back to the Rust extractor. -static SPACY: OnceCell>> = OnceCell::const_new(); - -/// Get the shared spaCy NER handle, initialising (provision + spawn) on first -/// call. Returns `None` when spaCy is unavailable; never errors so callers can -/// branch cleanly to the fallback path. -pub async fn shared_ner(config: &Config) -> Option> { - SPACY - .get_or_init(|| async { - match init_ner(config).await { - Ok(ner) => Some(Arc::new(ner)), - Err(e) => { - log::warn!("[memory_tree::nlp] spaCy unavailable, using Rust fallback: {e:#}"); - None - } - } - }) - .await - .clone() -} - -async fn init_ner(config: &Config) -> Result { - let runtime = ensure_spacy(config).await?; - SpacyNer::spawn(&runtime).await -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn ready_line_parses() { - let r: ReadyLine = - serde_json::from_str(r#"{"ready":true,"model":"en_core_web_sm"}"#).unwrap(); - assert!(r.ready); - let bad: ReadyLine = serde_json::from_str(r#"{"ready":false,"error":"no spacy"}"#).unwrap(); - assert!(!bad.ready); - assert_eq!(bad.error.as_deref(), Some("no spacy")); - } - - #[test] - fn response_parses_entities_and_nouns() { - let resp: SpacyResponse = serde_json::from_str( - r#"{"id":"0","entities":[{"text":"Alice","label":"PERSON","start":0,"end":5}],"nouns":["migration","runbook"]}"#, - ) - .unwrap(); - assert_eq!(resp.entities.len(), 1); - assert_eq!(resp.entities[0].label, "PERSON"); - assert_eq!(resp.nouns, vec!["migration", "runbook"]); - } -} diff --git a/src/openhuman/memory_tree/nlp/mod.rs b/src/openhuman/memory_tree/nlp/mod.rs index 6789e1fdc..43b5c969a 100644 --- a/src/openhuman/memory_tree/nlp/mod.rs +++ b/src/openhuman/memory_tree/nlp/mod.rs @@ -2,7 +2,8 @@ //! //! [`extract_query_entities`] turns a natural-language query into a set of //! canonical entity ids that key into `mem_tree_entity_index` and the -//! co-occurrence graph. It prefers a spaCy NER sidecar (named entities + +//! co-occurrence graph. It prefers the runtime Python server's spaCy backend +//! (named entities + //! salient nouns) and falls back to the in-Rust regex extractor whenever //! spaCy is disabled or unavailable — so retrieval always works offline, just //! with lower person/org recall. @@ -12,11 +13,9 @@ //! same `:` namespace as the indexed chunk entities. No id //! mismatch, no bespoke join. -mod client; -mod provision; - -pub use client::{shared_ner, SpacyNer}; -pub use provision::{ensure_spacy, spacy_provisioned, SpacyRuntime, SPACY_MODEL}; +pub use crate::openhuman::runtime_python_server::{ + ensure_spacy, spacy_provisioned, SpacyResponse, SPACY_MODEL, +}; use crate::openhuman::config::Config; use crate::openhuman::memory_tree::score::extract::{ @@ -51,22 +50,20 @@ pub async fn extract_query_entities(config: &Config, query: &str) -> Vec { - let extracted = spacy_to_extracted(&resp); - let canon = canonicalise(&extracted); - log::debug!( - "[memory_tree::nlp] spaCy query extraction: entities={} nouns={} canonical={}", - resp.entities.len(), - resp.nouns.len(), - canon.len() - ); - return canon; - } - Err(e) => { - log::warn!("[memory_tree::nlp] spaCy extraction failed, falling back: {e:#}"); - } + match crate::openhuman::runtime_python_server::extract_spacy(config, trimmed).await { + Ok(resp) => { + let extracted = spacy_to_extracted(&resp); + let canon = canonicalise(&extracted); + log::debug!( + "[memory_tree::nlp] spaCy query extraction: entities={} nouns={} canonical={}", + resp.entities.len(), + resp.nouns.len(), + canon.len() + ); + return canon; + } + Err(e) => { + log::warn!("[memory_tree::nlp] spaCy extraction failed, falling back: {e:#}"); } } } else { @@ -79,7 +76,7 @@ pub async fn extract_query_entities(config: &Config, query: &str) -> Vec` canonical ids by [`canonicalise`]. -fn spacy_to_extracted(resp: &client::SpacyResponse) -> ExtractedEntities { +fn spacy_to_extracted(resp: &SpacyResponse) -> ExtractedEntities { let entities = resp .entities .iter() @@ -170,16 +167,16 @@ mod tests { #[test] fn spacy_response_maps_nouns_to_topics() { - let resp = client::SpacyResponse { - entities: vec![client::SpacyEntity { - text: "Alice".into(), - label: "PERSON".into(), - start: 0, - end: 5, - }], + let resp = SpacyResponse { + entities: vec![ + crate::openhuman::runtime_python_server::spacy::SpacyEntity { + text: "Alice".into(), + label: "PERSON".into(), + start: 0, + end: 5, + }, + ], nouns: vec!["migration".into()], - id: Some("0".into()), - error: None, }; let extracted = spacy_to_extracted(&resp); let canon = canonicalise(&extracted); diff --git a/src/openhuman/memory_tree/nlp/provision.rs b/src/openhuman/memory_tree/nlp/provision.rs deleted file mode 100644 index 11bb4535a..000000000 --- a/src/openhuman/memory_tree/nlp/provision.rs +++ /dev/null @@ -1,241 +0,0 @@ -//! One-time provisioning of spaCy into a dedicated virtualenv under the -//! managed Python runtime. -//! -//! The deterministic retriever needs spaCy + a small English model to extract -//! entities from a query. The managed CPython distribution -//! (`runtime_python`) ships bare, so the first call here: -//! 1. resolves a Python ≥ 3.12 via [`PythonBootstrap`], -//! 2. creates an isolated venv (so we never mutate system/site packages), -//! 3. `pip install`s spaCy and downloads `en_core_web_sm`, -//! 4. writes a marker file so subsequent launches skip straight to spawning. -//! -//! All of this is network + filesystem heavy but happens at most once per host -//! (guarded by the marker). Any failure propagates as an error so the caller -//! (`nlp::extract_query_entities`) can fall back to the in-Rust extractor. - -use std::path::{Path, PathBuf}; -use std::time::Duration; - -use anyhow::{bail, Context, Result}; -use tokio::process::Command; - -use crate::openhuman::config::Config; -use crate::openhuman::runtime_python::PythonBootstrap; - -/// Embedded stdio service script, written to disk at provision time so the -/// Python interpreter has a real path to execute. -const SERVICE_PY: &str = include_str!("service.py"); - -/// Model spaCy downloads / loads. Small English pipeline — fast to load, ~12MB. -pub const SPACY_MODEL: &str = "en_core_web_sm"; - -/// Timeouts for the one-time install steps. spaCy + model is a few hundred MB -/// of wheels; give pip room on a cold cache / slow link. -const VENV_TIMEOUT: Duration = Duration::from_secs(120); -const PIP_TIMEOUT: Duration = Duration::from_secs(600); - -/// A provisioned spaCy runtime: the venv interpreter plus the service script. -#[derive(Debug, Clone)] -pub struct SpacyRuntime { - /// Python executable inside the dedicated venv (has spaCy + model). - pub python_bin: PathBuf, - /// Path to the written `service.py` stdio server. - pub service_script: PathBuf, -} - -/// Ensure spaCy + the model are installed and return a ready-to-spawn runtime. -/// -/// Idempotent across calls: once the marker file exists we skip venv creation -/// and pip entirely. Errors here are non-fatal to retrieval — the caller falls -/// back to the regex extractor. -pub async fn ensure_spacy(config: &Config) -> Result { - if !config.runtime_python.enabled { - bail!("runtime_python disabled — cannot provision spaCy"); - } - - let root = nlp_cache_root(config); - tokio::fs::create_dir_all(&root) - .await - .with_context(|| format!("creating nlp cache dir {}", root.display()))?; - - let venv_dir = root.join("spacy-venv"); - let marker = venv_dir.join(".openhuman-spacy-ready"); - let service_script = root.join("service.py"); - - // Always (re)write the service script so an upgraded binary ships the - // latest protocol. Cheap (~3KB) and keeps the on-disk copy authoritative. - tokio::fs::write(&service_script, SERVICE_PY) - .await - .with_context(|| format!("writing spaCy service script {}", service_script.display()))?; - - let venv_python = venv_python_path(&venv_dir); - - if marker.exists() && venv_python.exists() { - log::debug!( - "[memory_tree::nlp] spaCy already provisioned at {}", - venv_dir.display() - ); - return Ok(SpacyRuntime { - python_bin: venv_python, - service_script, - }); - } - - log::info!( - "[memory_tree::nlp] provisioning spaCy (one-time): venv={} model={}", - venv_dir.display(), - SPACY_MODEL - ); - - // 1. Resolve a base Python interpreter (managed download or system). - let bootstrap = PythonBootstrap::new(config.runtime_python.clone()); - let base = bootstrap - .resolve() - .await - .context("resolving base python for spaCy venv")?; - log::debug!( - "[memory_tree::nlp] base python resolved version={} bin={}", - base.version, - base.python_bin.display() - ); - - // 2. Create the venv (idempotent — venv is safe to re-run). - run_step( - &base.python_bin, - &["-m", "venv", &venv_dir.to_string_lossy()], - VENV_TIMEOUT, - "create venv", - ) - .await?; - - if !venv_python.exists() { - bail!( - "venv created but interpreter missing at {}", - venv_python.display() - ); - } - - // 3. Upgrade pip + install spaCy. - run_step( - &venv_python, - &["-m", "pip", "install", "--upgrade", "pip", "spacy"], - PIP_TIMEOUT, - "pip install spacy", - ) - .await?; - - // 4. Download the model into the venv. - run_step( - &venv_python, - &["-m", "spacy", "download", SPACY_MODEL], - PIP_TIMEOUT, - "spacy download model", - ) - .await?; - - // 5. Marker — provisioning complete. - tokio::fs::write(&marker, base.version.as_bytes()) - .await - .with_context(|| format!("writing spaCy ready marker {}", marker.display()))?; - - log::info!("[memory_tree::nlp] spaCy provisioning complete"); - Ok(SpacyRuntime { - python_bin: venv_python, - service_script, - }) -} - -/// Run one provisioning subprocess, capturing output and surfacing a useful -/// error on non-zero exit or timeout. -async fn run_step(python_bin: &Path, args: &[&str], timeout: Duration, label: &str) -> Result<()> { - log::debug!( - "[memory_tree::nlp] step `{label}`: {} {:?}", - python_bin.display(), - args - ); - let mut cmd = Command::new(python_bin); - cmd.args(args); - cmd.kill_on_drop(true); - - let output = match tokio::time::timeout(timeout, cmd.output()).await { - Ok(Ok(o)) => o, - Ok(Err(e)) => return Err(e).with_context(|| format!("spawning step `{label}`")), - Err(_) => bail!("step `{label}` timed out after {:?}", timeout), - }; - - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - // Tail only — pip output is verbose and may include paths but no - // secrets; cap to keep logs and the error bounded. - let tail: String = stderr - .chars() - .rev() - .take(800) - .collect::() - .chars() - .rev() - .collect(); - bail!("step `{label}` failed (status {}): {tail}", output.status); - } - Ok(()) -} - -/// Cheap, network-free probe: is spaCy already provisioned on this host? -/// -/// Mirrors the early-return guard in [`ensure_spacy`] (marker file + venv -/// interpreter present) without creating directories, writing the service -/// script, or resolving Python. Used by the harness-init orchestrator to mark -/// the spaCy step `Done` instantly on a warm host. -pub fn spacy_provisioned(config: &Config) -> bool { - let venv_dir = nlp_cache_root(config).join("spacy-venv"); - let marker = venv_dir.join(".openhuman-spacy-ready"); - marker.exists() && venv_python_path(&venv_dir).exists() -} - -/// Resolve the venv's python executable across platforms. -fn venv_python_path(venv_dir: &Path) -> PathBuf { - if cfg!(windows) { - venv_dir.join("Scripts").join("python.exe") - } else { - venv_dir.join("bin").join("python") - } -} - -/// Cache root for NLP artefacts. Honours `runtime_python.cache_dir` when set -/// (keeps all Python state together), else the user cache dir, else a -/// workspace-relative fallback. -fn nlp_cache_root(config: &Config) -> PathBuf { - let configured = config.runtime_python.cache_dir.trim(); - if !configured.is_empty() { - return PathBuf::from(configured).join("memory-nlp"); - } - if let Some(user_cache) = dirs::cache_dir() { - return user_cache.join("openhuman").join("memory-nlp"); - } - config.workspace_dir.join("memory_tree").join("nlp") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn venv_python_path_is_platform_specific() { - let p = venv_python_path(Path::new("/tmp/venv")); - if cfg!(windows) { - assert!(p.ends_with("Scripts/python.exe") || p.ends_with("Scripts\\python.exe")); - } else { - assert_eq!(p, PathBuf::from("/tmp/venv/bin/python")); - } - } - - #[test] - fn cache_root_honours_configured_dir() { - let mut cfg = Config::default(); - cfg.runtime_python.cache_dir = "/custom/py".to_string(); - assert_eq!( - nlp_cache_root(&cfg), - PathBuf::from("/custom/py").join("memory-nlp") - ); - } -} diff --git a/src/openhuman/memory_tree/nlp/service.py b/src/openhuman/memory_tree/nlp/service.py deleted file mode 100644 index 28ddc933f..000000000 --- a/src/openhuman/memory_tree/nlp/service.py +++ /dev/null @@ -1,108 +0,0 @@ -#!/usr/bin/env python3 -"""spaCy NER stdio service for the OpenHuman memory retriever. - -Long-lived line-oriented JSON protocol over stdin/stdout. The model is loaded -exactly once at startup, then the process answers extraction requests until its -stdin is closed (the Rust parent drops the child on shutdown via kill_on_drop). - -Protocol --------- -On startup, after the model loads, emit exactly one line: - {"ready": true, "model": "en_core_web_sm"} -or, on fatal load failure: - {"ready": false, "error": ""} (then exit non-zero) - -For each request line `{"id": , "text": }`, reply with one line: - {"id": , - "entities": [{"text": str, "label": str, "start": int, "end": int}, ...], - "nouns": [str, ...]} - -Entities are spaCy named entities (doc.ents). `nouns` are deduplicated lower- -case lemmas of common/proper nouns — E2GraphRAG keys retrieval on both named -entities and salient nouns, so a query like "migration runbook" still yields -graph anchors even with no PERSON/ORG spans. - -All output is a single compact JSON object per line, flushed immediately -(the parent runs us with `python -u`, but we flush defensively anyway). -""" - -import json -import sys - -MODEL_NAME = "en_core_web_sm" -# Disable the parser/lemmatizer pipes we do not need for speed; keep the -# tagger (POS, needed for noun selection) and ner. `lemma` falls back to the -# surface form when the lemmatizer is absent, which is fine for our keys. -_DISABLE = ["parser"] - - -def _emit(obj): - sys.stdout.write(json.dumps(obj, ensure_ascii=False)) - sys.stdout.write("\n") - sys.stdout.flush() - - -def _load(): - import spacy - - try: - return spacy.load(MODEL_NAME, disable=_DISABLE) - except Exception: - # Fall back to loading with the full pipeline if the disable list is - # incompatible with the installed model build. - return spacy.load(MODEL_NAME) - - -def _extract(nlp, text): - doc = nlp(text) - entities = [ - { - "text": ent.text, - "label": ent.label_, - "start": int(ent.start_char), - "end": int(ent.end_char), - } - for ent in doc.ents - ] - seen = set() - nouns = [] - for tok in doc: - if tok.pos_ in ("NOUN", "PROPN") and not tok.is_stop and tok.is_alpha: - key = (tok.lemma_ or tok.text).lower().strip() - if len(key) >= 2 and key not in seen: - seen.add(key) - nouns.append(key) - return entities, nouns - - -def main(): - try: - nlp = _load() - except Exception as exc: # pragma: no cover - exercised only without spaCy - _emit({"ready": False, "error": f"{type(exc).__name__}: {exc}"}) - return 1 - - _emit({"ready": True, "model": MODEL_NAME}) - - for line in sys.stdin: - line = line.strip() - if not line: - continue - try: - req = json.loads(line) - except Exception as exc: - _emit({"id": None, "error": f"bad request json: {exc}"}) - continue - req_id = req.get("id") - text = req.get("text") or "" - try: - entities, nouns = _extract(nlp, text) - _emit({"id": req_id, "entities": entities, "nouns": nouns}) - except Exception as exc: # pragma: no cover - defensive - _emit({"id": req_id, "error": f"{type(exc).__name__}: {exc}"}) - - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/src/openhuman/mod.rs b/src/openhuman/mod.rs index dc44c8dbc..a5fafce3a 100644 --- a/src/openhuman/mod.rs +++ b/src/openhuman/mod.rs @@ -96,6 +96,7 @@ pub mod referral; pub mod routing; pub mod runtime_node; pub mod runtime_python; +pub mod runtime_python_server; pub mod sandbox; pub mod scheduler_gate; pub mod screen_intelligence; diff --git a/src/openhuman/runtime_python_server/mod.rs b/src/openhuman/runtime_python_server/mod.rs new file mode 100644 index 000000000..fe5c4d2d0 --- /dev/null +++ b/src/openhuman/runtime_python_server/mod.rs @@ -0,0 +1,18 @@ +//! Long-running Python backend host. +//! +//! `runtime_python` owns interpreter resolution. This module owns the +//! process-level server that keeps Python-backed model modules warm and serves +//! Rust callers over a private JSONL stdio protocol. + +pub mod protocol; +pub mod registry; +pub mod server; +pub mod spacy; +pub mod types; + +pub use registry::{enabled_backends, RuntimePythonBackend}; +pub use server::{ensure_started, status, RuntimePythonServer}; +pub use spacy::{ + ensure_spacy, extract as extract_spacy, spacy_provisioned, SpacyResponse, SPACY_MODEL, +}; +pub use types::{BackendStatus, RuntimePythonServerStatus}; diff --git a/src/openhuman/runtime_python_server/protocol.rs b/src/openhuman/runtime_python_server/protocol.rs new file mode 100644 index 000000000..3a76f6ec8 --- /dev/null +++ b/src/openhuman/runtime_python_server/protocol.rs @@ -0,0 +1,66 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +pub const PROTOCOL_VERSION: u32 = 1; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ReadyLine { + #[serde(default)] + pub ready: bool, + #[serde(default)] + pub protocol: Option, + #[serde(default)] + pub backends: Vec, + #[serde(default)] + pub error: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PythonServerRequest { + pub id: String, + pub method: String, + #[serde(default)] + pub params: Value, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PythonServerError { + pub code: String, + pub message: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PythonServerResponse { + pub id: Option, + #[serde(default)] + pub ok: bool, + #[serde(default)] + pub result: Option, + #[serde(default)] + pub error: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ready_line_parses() { + let ready: ReadyLine = + serde_json::from_str(r#"{"ready":true,"protocol":1,"backends":["spacy"]}"#).unwrap(); + assert!(ready.ready); + assert_eq!(ready.protocol, Some(PROTOCOL_VERSION)); + assert_eq!(ready.backends, vec!["spacy"]); + } + + #[test] + fn response_parses_error_envelope() { + let response: PythonServerResponse = serde_json::from_str( + r#"{"id":"7","ok":false,"error":{"code":"bad_request","message":"missing text"}}"#, + ) + .unwrap(); + assert!(!response.ok); + assert_eq!(response.id.as_deref(), Some("7")); + assert_eq!(response.error.unwrap().code, "bad_request"); + } +} diff --git a/src/openhuman/runtime_python_server/registry.rs b/src/openhuman/runtime_python_server/registry.rs new file mode 100644 index 000000000..c5846f276 --- /dev/null +++ b/src/openhuman/runtime_python_server/registry.rs @@ -0,0 +1,46 @@ +use crate::openhuman::config::Config; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RuntimePythonBackend { + Spacy, +} + +impl RuntimePythonBackend { + pub fn id(self) -> &'static str { + match self { + Self::Spacy => "spacy", + } + } +} + +pub fn enabled_backends(config: &Config) -> Vec { + if !config.runtime_python.enabled { + return Vec::new(); + } + + let mut backends = Vec::new(); + if config.memory_tree.spacy_enabled { + backends.push(RuntimePythonBackend::Spacy); + } + backends +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn registry_respects_runtime_and_spacy_flags() { + let mut config = Config::default(); + config.runtime_python.enabled = false; + config.memory_tree.spacy_enabled = true; + assert!(enabled_backends(&config).is_empty()); + + config.runtime_python.enabled = true; + config.memory_tree.spacy_enabled = false; + assert!(enabled_backends(&config).is_empty()); + + config.memory_tree.spacy_enabled = true; + assert_eq!(enabled_backends(&config), vec![RuntimePythonBackend::Spacy]); + } +} diff --git a/src/openhuman/runtime_python_server/server.py b/src/openhuman/runtime_python_server/server.py new file mode 100644 index 000000000..79bd1ef41 --- /dev/null +++ b/src/openhuman/runtime_python_server/server.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""OpenHuman runtime Python server. + +Private JSONL stdio protocol. Rust owns the process and sends one compact JSON +request per line. This server keeps expensive Python backends warm for the +life of the Rust core process. +""" + +import json +import sys + +PROTOCOL = 1 +SPACY_MODEL = "en_core_web_sm" + +_spacy_nlp = None + + +def _emit(obj): + sys.stdout.write(json.dumps(obj, ensure_ascii=False, separators=(",", ":"))) + sys.stdout.write("\n") + sys.stdout.flush() + + +def _error(req_id, code, message): + return {"id": req_id, "ok": False, "error": {"code": code, "message": str(message)}} + + +def _configure_stdio(): + if hasattr(sys.stdin, "reconfigure"): + sys.stdin.reconfigure(encoding="utf-8") + if hasattr(sys.stdout, "reconfigure"): + sys.stdout.reconfigure(encoding="utf-8") + + +def _load_spacy(): + global _spacy_nlp + if _spacy_nlp is not None: + return _spacy_nlp + + import spacy + + try: + _spacy_nlp = spacy.load(SPACY_MODEL, disable=["parser"]) + except Exception: + _spacy_nlp = spacy.load(SPACY_MODEL) + return _spacy_nlp + + +def _spacy_extract(params): + text = (params or {}).get("text") or "" + nlp = _load_spacy() + doc = nlp(text) + entities = [ + { + "text": ent.text, + "label": ent.label_, + "start": int(ent.start_char), + "end": int(ent.end_char), + } + for ent in doc.ents + ] + seen = set() + nouns = [] + for tok in doc: + if tok.pos_ in ("NOUN", "PROPN") and not tok.is_stop and tok.is_alpha: + key = (tok.lemma_ or tok.text).lower().strip() + if len(key) >= 2 and key not in seen: + seen.add(key) + nouns.append(key) + return {"entities": entities, "nouns": nouns} + + +def _handle(req): + req_id = req.get("id") + method = req.get("method") + params = req.get("params") or {} + if method == "spacy.extract": + return {"id": req_id, "ok": True, "result": _spacy_extract(params)} + return _error(req_id, "unknown_method", f"unknown runtime_python_server method: {method}") + + +def main(): + _configure_stdio() + try: + _load_spacy() + except Exception as exc: + _emit({"ready": False, "error": f"{type(exc).__name__}: {exc}"}) + return 1 + + _emit({"ready": True, "protocol": PROTOCOL, "backends": ["spacy"]}) + for line in sys.stdin: + line = line.strip() + if not line: + continue + try: + req = json.loads(line) + except Exception as exc: + _emit(_error(None, "bad_json", exc)) + continue + if not isinstance(req, dict): + _emit(_error(None, "bad_request", "request must be a JSON object")) + continue + try: + _emit(_handle(req)) + except Exception as exc: + _emit(_error(req.get("id"), type(exc).__name__, exc)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/openhuman/runtime_python_server/server.rs b/src/openhuman/runtime_python_server/server.rs new file mode 100644 index 000000000..b8508a46d --- /dev/null +++ b/src/openhuman/runtime_python_server/server.rs @@ -0,0 +1,453 @@ +use std::path::PathBuf; +use std::sync::{Arc, OnceLock}; +use std::time::{Duration, Instant}; + +use anyhow::{bail, Context, Result}; +use serde::de::DeserializeOwned; +use serde_json::{json, Value}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines}; +use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout}; +use tokio::sync::Mutex; + +use super::protocol::{PythonServerRequest, PythonServerResponse, ReadyLine, PROTOCOL_VERSION}; +use super::registry::{enabled_backends, RuntimePythonBackend}; +use super::types::{BackendStatus, RuntimePythonServerStatus}; +use crate::openhuman::config::Config; +use crate::openhuman::runtime_python::process::PythonLaunchSpec; + +const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30); +const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); +const START_FAILURE_BACKOFF: Duration = Duration::from_secs(300); + +static SERVER: OnceLock> = OnceLock::new(); + +fn server_slot() -> &'static Mutex { + SERVER.get_or_init(|| Mutex::new(ServerCache::Empty)) +} + +#[derive(Clone)] +enum ServerCache { + Empty, + Ready(Arc), + Failed { + message: String, + retry_after: Instant, + }, +} + +#[derive(Debug, Clone)] +struct ServerLaunch { + python_bin: PathBuf, + script_path: PathBuf, + backends: Vec, +} + +struct ServerInner { + _child: Child, + stdin: ChildStdin, + stdout: Lines>, + next_id: u64, + ready_backends: Vec, +} + +fn drain_server_stderr(stderr: ChildStderr) { + tokio::spawn(async move { + let mut reader = BufReader::new(stderr); + let mut buf = Vec::with_capacity(1024); + let mut line_count = 0u64; + let mut byte_count = 0u64; + + loop { + buf.clear(); + match reader.read_until(b'\n', &mut buf).await { + Ok(0) => { + log::debug!( + "[runtime_python_server] stderr drain closed lines={} bytes={}", + line_count, + byte_count + ); + break; + } + Ok(n) => { + line_count += 1; + byte_count += n as u64; + log::trace!( + "[runtime_python_server] drained stderr line bytes={} total_lines={} total_bytes={}", + n, + line_count, + byte_count + ); + } + Err(error) => { + log::debug!( + "[runtime_python_server] stderr drain failed after lines={} bytes={}: {error}", + line_count, + byte_count + ); + break; + } + } + } + }); +} + +pub struct RuntimePythonServer { + launch: ServerLaunch, + inner: Mutex>, +} + +impl RuntimePythonServer { + async fn new(config: &Config) -> Result { + let launch = prepare_launch(config).await?; + Ok(Self { + launch, + inner: Mutex::new(None), + }) + } + + pub async fn start(&self) -> Result<()> { + let mut guard = self.inner.lock().await; + if guard.is_some() { + return Ok(()); + } + let inner = spawn_inner(&self.launch).await?; + *guard = Some(inner); + Ok(()) + } + + pub async fn request(&self, method: &str, params: Value) -> Result + where + T: DeserializeOwned, + { + match self.request_once(method, params.clone()).await { + Ok(value) => Ok(value), + Err(err) => { + log::warn!( + "[runtime_python_server] request failed; restarting server before retry: {err:#}" + ); + self.reset().await; + self.request_once(method, params).await + } + } + } + + async fn request_once(&self, method: &str, params: Value) -> Result + where + T: DeserializeOwned, + { + let mut guard = self.inner.lock().await; + if guard.is_none() { + *guard = Some(spawn_inner(&self.launch).await?); + } + let inner = guard.as_mut().context("runtime python server missing")?; + let id = inner.next_id.to_string(); + inner.next_id += 1; + + let request = PythonServerRequest { + id: id.clone(), + method: method.to_string(), + params, + }; + let mut line = serde_json::to_string(&request)?; + line.push('\n'); + log::debug!( + "[runtime_python_server] sending request id={} method={}", + id, + method + ); + inner + .stdin + .write_all(line.as_bytes()) + .await + .context("writing runtime python server request")?; + inner + .stdin + .flush() + .await + .context("flushing runtime python server request")?; + + loop { + let next = tokio::time::timeout(REQUEST_TIMEOUT, inner.stdout.next_line()).await; + let line = match next { + Ok(Ok(Some(line))) => line, + Ok(Ok(None)) => bail!("runtime python server closed stdout"), + Ok(Err(error)) => { + return Err(error).context("reading runtime python server response") + } + Err(_) => bail!("runtime python server request timed out"), + }; + let response: PythonServerResponse = match serde_json::from_str(&line) { + Ok(response) => response, + Err(error) => { + log::warn!( + "[runtime_python_server] unparseable response skipped: {error}; line_len={}", + line.len() + ); + continue; + } + }; + if response.id.as_deref() != Some(id.as_str()) { + log::debug!( + "[runtime_python_server] skipped response for different id={:?}", + response.id + ); + continue; + } + if !response.ok { + let message = response + .error + .map(|error| format!("{}: {}", error.code, error.message)) + .unwrap_or_else(|| "unknown python server error".to_string()); + bail!("runtime python server `{method}` failed: {message}"); + } + let result = response.result.unwrap_or(Value::Null); + return serde_json::from_value(result) + .with_context(|| format!("decoding runtime python server `{method}` result")); + } + } + + async fn reset(&self) { + let mut guard = self.inner.lock().await; + *guard = None; + } + + fn status_from_inner(&self, inner: Option<&ServerInner>) -> RuntimePythonServerStatus { + let running = inner.is_some(); + let ready_backends = inner + .map(|inner| inner.ready_backends.as_slice()) + .unwrap_or(&[]); + RuntimePythonServerStatus { + enabled: true, + running, + backends: self + .launch + .backends + .iter() + .map(|backend| BackendStatus { + id: backend.id().to_string(), + enabled: true, + ready: ready_backends.iter().any(|id| id == backend.id()), + message: None, + }) + .collect(), + message: None, + } + } + + pub async fn status(&self) -> RuntimePythonServerStatus { + let guard = self.inner.lock().await; + self.status_from_inner(guard.as_ref()) + } +} + +pub async fn ensure_started(config: &Config) -> Result> { + let mut guard = server_slot().lock().await; + match &*guard { + ServerCache::Ready(existing) => { + let existing = existing.clone(); + if let Err(error) = existing.start().await { + let message = format!("{error:#}"); + log::warn!( + "[runtime_python_server] cached server failed to start; backing off: {message}" + ); + *guard = ServerCache::Failed { + message: message.clone(), + retry_after: Instant::now() + START_FAILURE_BACKOFF, + }; + bail!("runtime python server unavailable: {message}"); + } + return Ok(existing); + } + ServerCache::Failed { + message, + retry_after, + } if Instant::now() < *retry_after => { + bail!("runtime python server unavailable after previous startup failure: {message}"); + } + ServerCache::Failed { .. } | ServerCache::Empty => {} + } + + match start_new_server(config).await { + Ok(server) => { + *guard = ServerCache::Ready(server.clone()); + Ok(server) + } + Err(error) => { + let message = format!("{error:#}"); + log::warn!( + "[runtime_python_server] startup failed; caching fallback state for {:?}: {message}", + START_FAILURE_BACKOFF + ); + *guard = ServerCache::Failed { + message: message.clone(), + retry_after: Instant::now() + START_FAILURE_BACKOFF, + }; + bail!("runtime python server unavailable: {message}"); + } + } +} + +async fn start_new_server(config: &Config) -> Result> { + let server = Arc::new(RuntimePythonServer::new(config).await?); + server.start().await?; + Ok(server) +} + +pub async fn status() -> RuntimePythonServerStatus { + let cached = { + let guard = server_slot().lock().await; + guard.clone() + }; + match cached { + ServerCache::Ready(server) => server.status().await, + ServerCache::Failed { message, .. } => RuntimePythonServerStatus { + enabled: true, + running: false, + backends: Vec::new(), + message: Some(format!("runtime python server unavailable: {message}")), + }, + ServerCache::Empty => { + RuntimePythonServerStatus::disabled("runtime python server has not started") + } + } +} + +async fn prepare_launch(config: &Config) -> Result { + let backends = enabled_backends(config); + if backends.is_empty() { + bail!("no runtime python server backends enabled"); + } + + let spacy_runtime = if backends.contains(&RuntimePythonBackend::Spacy) { + Some(super::spacy::ensure_spacy(config).await?) + } else { + None + }; + + let python_bin = if let Some(spacy_runtime) = spacy_runtime { + spacy_runtime.python_bin + } else { + crate::openhuman::runtime_python::PythonBootstrap::new(config.runtime_python.clone()) + .resolve() + .await? + .python_bin + }; + let script_path = write_server_script(config).await?; + + Ok(ServerLaunch { + python_bin, + script_path, + backends, + }) +} + +async fn write_server_script(config: &Config) -> Result { + let root = super::spacy::python_server_cache_root(config); + tokio::fs::create_dir_all(&root) + .await + .with_context(|| format!("creating runtime python server cache {}", root.display()))?; + let script_path = root.join("runtime_python_server.py"); + tokio::fs::write(&script_path, include_str!("server.py")) + .await + .with_context(|| { + format!( + "writing runtime python server script {}", + script_path.display() + ) + })?; + Ok(script_path) +} + +async fn spawn_inner(launch: &ServerLaunch) -> Result { + log::info!( + "[runtime_python_server] starting server python={} script={} backends={:?}", + launch.python_bin.display(), + launch.script_path.display(), + launch.backends + ); + let resolved = crate::openhuman::runtime_python::ResolvedPython { + bin_dir: launch + .python_bin + .parent() + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(".")), + python_bin: launch.python_bin.clone(), + version: "runtime-backend".to_string(), + source: crate::openhuman::runtime_python::PythonSource::Managed, + }; + let spec = PythonLaunchSpec::new(launch.script_path.clone()); + let mut child = + crate::openhuman::runtime_python::process::spawn_stdio_process(&resolved, &spec) + .context("spawning runtime python server")?; + let stdin = child + .stdin + .take() + .context("runtime python server stdin missing")?; + let stdout = child + .stdout + .take() + .context("runtime python server stdout missing")?; + if let Some(stderr) = child.stderr.take() { + drain_server_stderr(stderr); + } else { + log::debug!("[runtime_python_server] stderr pipe missing; continuing without drain"); + } + let mut lines = BufReader::new(stdout).lines(); + + let ready_line = match tokio::time::timeout(HANDSHAKE_TIMEOUT, lines.next_line()).await { + Ok(Ok(Some(line))) => line, + Ok(Ok(None)) => bail!("runtime python server exited before readiness handshake"), + Ok(Err(error)) => return Err(error).context("reading runtime python server handshake"), + Err(_) => bail!("runtime python server readiness handshake timed out"), + }; + let ready: ReadyLine = serde_json::from_str(&ready_line) + .with_context(|| format!("parsing runtime python server ready line: {ready_line}"))?; + if !ready.ready { + bail!( + "runtime python server failed to start: {}", + ready.error.unwrap_or_else(|| "unknown".to_string()) + ); + } + if ready.protocol != Some(PROTOCOL_VERSION) { + bail!( + "runtime python server protocol mismatch: expected {}, got {:?}", + PROTOCOL_VERSION, + ready.protocol + ); + } + log::info!( + "[runtime_python_server] server ready backends={:?}", + ready.backends + ); + + Ok(ServerInner { + _child: child, + stdin, + stdout: lines, + next_id: 0, + ready_backends: ready.backends, + }) +} + +pub async fn request_spacy_extract( + config: &Config, + text: &str, +) -> Result { + let server = ensure_started(config).await?; + server + .request("spacy.extract", json!({ "text": text })) + .await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn prepare_launch_rejects_disabled_backends() { + let mut config = Config::default(); + config.runtime_python.enabled = false; + let err = prepare_launch(&config).await.unwrap_err().to_string(); + assert!(err.contains("no runtime python server backends enabled")); + } +} diff --git a/src/openhuman/runtime_python_server/spacy.rs b/src/openhuman/runtime_python_server/spacy.rs new file mode 100644 index 000000000..1d5448f25 --- /dev/null +++ b/src/openhuman/runtime_python_server/spacy.rs @@ -0,0 +1,306 @@ +use std::path::{Path, PathBuf}; +use std::sync::OnceLock; +use std::time::Duration; + +use anyhow::{bail, Context, Result}; +use serde::{Deserialize, Serialize}; +use tokio::process::Command; +use tokio::sync::Mutex; + +use crate::openhuman::config::Config; +use crate::openhuman::runtime_python::PythonBootstrap; + +pub const SPACY_MODEL: &str = "en_core_web_sm"; + +const VENV_TIMEOUT: Duration = Duration::from_secs(120); +const PIP_TIMEOUT: Duration = Duration::from_secs(600); + +static SPACY_PROVISION_LOCK: OnceLock> = OnceLock::new(); + +fn spacy_provision_lock() -> &'static Mutex<()> { + SPACY_PROVISION_LOCK.get_or_init(|| Mutex::new(())) +} + +#[derive(Debug, Clone)] +pub struct SpacyRuntime { + pub python_bin: PathBuf, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SpacyEntity { + pub text: String, + pub label: String, + #[serde(default)] + pub start: u32, + #[serde(default)] + pub end: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SpacyResponse { + #[serde(default)] + pub entities: Vec, + #[serde(default)] + pub nouns: Vec, +} + +pub async fn extract(config: &Config, text: &str) -> Result { + super::server::request_spacy_extract(config, text).await +} + +pub async fn ensure_spacy(config: &Config) -> Result { + let _guard = spacy_provision_lock().lock().await; + if !config.runtime_python.enabled { + bail!("runtime_python disabled — cannot provision spaCy"); + } + + let root = python_server_cache_root(config); + tokio::fs::create_dir_all(&root).await.with_context(|| { + format!( + "creating runtime python server cache dir {}", + root.display() + ) + })?; + + let venv_dir = runtime_spacy_venv_dir(config); + let venv_python = venv_python_path(&venv_dir); + + if spacy_venv_ready(&venv_dir) { + log::debug!( + "[runtime_python_server::spacy] spaCy already provisioned at {}", + venv_dir.display() + ); + return Ok(SpacyRuntime { + python_bin: venv_python, + }); + } + + if let Some(existing_venv) = migrate_or_reuse_legacy_spacy_venv(config, &venv_dir).await? { + return Ok(SpacyRuntime { + python_bin: venv_python_path(&existing_venv), + }); + } + + log::info!( + "[runtime_python_server::spacy] provisioning spaCy venv={} model={}", + venv_dir.display(), + SPACY_MODEL + ); + + let bootstrap = PythonBootstrap::new(config.runtime_python.clone()); + let base = bootstrap + .resolve() + .await + .context("resolving base python for runtime python server spaCy venv")?; + log::debug!( + "[runtime_python_server::spacy] base python resolved version={} bin={}", + base.version, + base.python_bin.display() + ); + + run_step( + &base.python_bin, + &["-m", "venv", &venv_dir.to_string_lossy()], + VENV_TIMEOUT, + "create venv", + ) + .await?; + + if !venv_python.exists() { + bail!( + "venv created but interpreter missing at {}", + venv_python.display() + ); + } + + run_step( + &venv_python, + &["-m", "pip", "install", "--upgrade", "pip", "spacy"], + PIP_TIMEOUT, + "pip install spacy", + ) + .await?; + + run_step( + &venv_python, + &["-m", "spacy", "download", SPACY_MODEL], + PIP_TIMEOUT, + "spacy download model", + ) + .await?; + + let marker = venv_dir.join(".openhuman-spacy-ready"); + tokio::fs::write(&marker, base.version.as_bytes()) + .await + .with_context(|| format!("writing spaCy ready marker {}", marker.display()))?; + + log::info!("[runtime_python_server::spacy] spaCy provisioning complete"); + Ok(SpacyRuntime { + python_bin: venv_python, + }) +} + +async fn run_step(python_bin: &Path, args: &[&str], timeout: Duration, label: &str) -> Result<()> { + log::debug!( + "[runtime_python_server::spacy] step `{label}`: {} {:?}", + python_bin.display(), + args + ); + let mut cmd = Command::new(python_bin); + cmd.args(args); + cmd.kill_on_drop(true); + + let output = match tokio::time::timeout(timeout, cmd.output()).await { + Ok(Ok(output)) => output, + Ok(Err(error)) => return Err(error).with_context(|| format!("spawning step `{label}`")), + Err(_) => bail!("step `{label}` timed out after {:?}", timeout), + }; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + let tail: String = stderr + .chars() + .rev() + .take(800) + .collect::() + .chars() + .rev() + .collect(); + bail!("step `{label}` failed (status {}): {tail}", output.status); + } + Ok(()) +} + +pub fn spacy_provisioned(config: &Config) -> bool { + spacy_venv_ready(&runtime_spacy_venv_dir(config)) + || legacy_spacy_venv_dirs(config) + .into_iter() + .any(|venv_dir| spacy_venv_ready(&venv_dir)) +} + +pub(crate) fn python_server_cache_root(config: &Config) -> PathBuf { + let configured = config.runtime_python.cache_dir.trim(); + if !configured.is_empty() { + return PathBuf::from(configured).join("runtime-python-server"); + } + if let Some(user_cache) = dirs::cache_dir() { + return user_cache.join("openhuman").join("runtime-python-server"); + } + config.workspace_dir.join("runtime_python_server") +} + +fn runtime_spacy_venv_dir(config: &Config) -> PathBuf { + python_server_cache_root(config).join("spacy-venv") +} + +async fn migrate_or_reuse_legacy_spacy_venv( + config: &Config, + target_venv: &Path, +) -> Result> { + for legacy_venv in legacy_spacy_venv_dirs(config) { + if legacy_venv == target_venv || !spacy_venv_ready(&legacy_venv) { + continue; + } + + if !target_venv.exists() { + if let Some(parent) = target_venv.parent() { + tokio::fs::create_dir_all(parent) + .await + .with_context(|| format!("creating spaCy venv parent {}", parent.display()))?; + } + match tokio::fs::rename(&legacy_venv, target_venv).await { + Ok(()) => { + log::info!( + "[runtime_python_server::spacy] migrated legacy spaCy venv {} -> {}", + legacy_venv.display(), + target_venv.display() + ); + return Ok(Some(target_venv.to_path_buf())); + } + Err(error) => { + log::warn!( + "[runtime_python_server::spacy] could not migrate legacy spaCy venv {} -> {}; reusing legacy path: {error}", + legacy_venv.display(), + target_venv.display() + ); + return Ok(Some(legacy_venv)); + } + } + } + + log::info!( + "[runtime_python_server::spacy] reusing legacy spaCy venv {} because target {} is not ready", + legacy_venv.display(), + target_venv.display() + ); + return Ok(Some(legacy_venv)); + } + + Ok(None) +} + +fn legacy_spacy_venv_dirs(config: &Config) -> Vec { + let mut roots = Vec::new(); + let configured = config.runtime_python.cache_dir.trim(); + if !configured.is_empty() { + roots.push(PathBuf::from(configured).join("memory-nlp")); + } else if let Some(user_cache) = dirs::cache_dir() { + roots.push(user_cache.join("openhuman").join("memory-nlp")); + } + roots.push(config.workspace_dir.join("memory_tree").join("nlp")); + roots + .into_iter() + .map(|root| root.join("spacy-venv")) + .collect() +} + +fn spacy_venv_ready(venv_dir: &Path) -> bool { + venv_dir.join(".openhuman-spacy-ready").exists() && venv_python_path(venv_dir).exists() +} + +fn venv_python_path(venv_dir: &Path) -> PathBuf { + if cfg!(windows) { + venv_dir.join("Scripts").join("python.exe") + } else { + venv_dir.join("bin").join("python") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cache_root_honours_runtime_python_cache_dir() { + let mut config = Config::default(); + config.runtime_python.cache_dir = "/tmp/openhuman-python".to_string(); + assert_eq!( + python_server_cache_root(&config), + PathBuf::from("/tmp/openhuman-python").join("runtime-python-server") + ); + } + + #[test] + fn legacy_configured_cache_is_considered_provisioned() { + let temp = tempfile::tempdir().unwrap(); + let mut config = Config::default(); + config.runtime_python.cache_dir = temp.path().to_string_lossy().to_string(); + let legacy_venv = temp.path().join("memory-nlp").join("spacy-venv"); + std::fs::create_dir_all(legacy_venv.join(if cfg!(windows) { "Scripts" } else { "bin" })) + .unwrap(); + std::fs::write(legacy_venv.join(".openhuman-spacy-ready"), "test").unwrap(); + std::fs::write(venv_python_path(&legacy_venv), "").unwrap(); + + assert!(spacy_provisioned(&config)); + } + + #[test] + fn spacy_response_parses() { + let response: SpacyResponse = serde_json::from_str( + r#"{"entities":[{"text":"Alice","label":"PERSON","start":0,"end":5}],"nouns":["migration"]}"#, + ) + .unwrap(); + assert_eq!(response.entities[0].label, "PERSON"); + assert_eq!(response.nouns, vec!["migration"]); + } +} diff --git a/src/openhuman/runtime_python_server/types.rs b/src/openhuman/runtime_python_server/types.rs new file mode 100644 index 000000000..ad7459e79 --- /dev/null +++ b/src/openhuman/runtime_python_server/types.rs @@ -0,0 +1,28 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct BackendStatus { + pub id: String, + pub enabled: bool, + pub ready: bool, + pub message: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RuntimePythonServerStatus { + pub enabled: bool, + pub running: bool, + pub backends: Vec, + pub message: Option, +} + +impl RuntimePythonServerStatus { + pub fn disabled(message: impl Into) -> Self { + Self { + enabled: false, + running: false, + backends: Vec::new(), + message: Some(message.into()), + } + } +}