From 8ee89acdaf96a02d1978aee8b5eab78584dc0596 Mon Sep 17 00:00:00 2001 From: Steven Enamakel <31011319+senamakel@users.noreply.github.com> Date: Tue, 23 Jun 2026 00:53:10 -0700 Subject: [PATCH] feat(memory): deterministic E2GraphRAG retrieval (replace agentic walk) (#3947) --- scripts/test-rust-e2e.sh | 2 +- src/openhuman/agent_memory/README.md | 2 +- src/openhuman/agent_memory/agent/agent.toml | 2 - src/openhuman/agent_memory/agent/prompt.md | 16 +- src/openhuman/agent_memory/ops.rs | 100 +- src/openhuman/agent_memory/tools.rs | 10 +- .../config/schema/load/env_overlay.rs | 5 + src/openhuman/config/schema/storage_memory.rs | 15 + src/openhuman/memory/query/fast_walk.rs | 82 ++ src/openhuman/memory/query/mod.rs | 25 +- .../memory/query/smart_walk/dispatch.rs | 687 ------------- src/openhuman/memory/query/smart_walk/mod.rs | 29 - .../memory/query/smart_walk/prompts.rs | 292 ------ .../memory/query/smart_walk/runner.rs | 182 ---- .../query/smart_walk/smart_walk_tests.rs | 693 ------------- src/openhuman/memory/query/smart_walk/tool.rs | 188 ---- .../memory/query/smart_walk/types.rs | 66 -- src/openhuman/memory/query/walk.rs | 966 ------------------ src/openhuman/memory/schema/definitions.rs | 72 +- src/openhuman/memory/schema/handlers.rs | 129 +-- src/openhuman/memory_search/tools/mod.rs | 15 +- src/openhuman/memory_store/chunks/store.rs | 21 + src/openhuman/memory_tree/graph/bfs.rs | 215 ++++ src/openhuman/memory_tree/graph/mod.rs | 19 + src/openhuman/memory_tree/graph/store.rs | 230 +++++ src/openhuman/memory_tree/mod.rs | 2 + src/openhuman/memory_tree/nlp/client.rs | 227 ++++ src/openhuman/memory_tree/nlp/mod.rs | 189 ++++ src/openhuman/memory_tree/nlp/provision.rs | 229 +++++ src/openhuman/memory_tree/nlp/service.py | 108 ++ src/openhuman/memory_tree/retrieval/fast.rs | 446 ++++++++ src/openhuman/memory_tree/retrieval/mod.rs | 2 + src/openhuman/memory_tree/score/mod.rs | 35 + src/openhuman/tools/ops.rs | 2 - tests/memory_fast_retrieve_e2e.rs | 131 +++ tests/memory_threads_raw_coverage_e2e.rs | 2 - tests/memory_tree_sync_raw_coverage_e2e.rs | 117 +-- tests/memory_tree_walk_e2e.rs | 535 ---------- 38 files changed, 2089 insertions(+), 3999 deletions(-) create mode 100644 src/openhuman/memory/query/fast_walk.rs delete mode 100644 src/openhuman/memory/query/smart_walk/dispatch.rs delete mode 100644 src/openhuman/memory/query/smart_walk/mod.rs delete mode 100644 src/openhuman/memory/query/smart_walk/prompts.rs delete mode 100644 src/openhuman/memory/query/smart_walk/runner.rs delete mode 100644 src/openhuman/memory/query/smart_walk/smart_walk_tests.rs delete mode 100644 src/openhuman/memory/query/smart_walk/tool.rs delete mode 100644 src/openhuman/memory/query/smart_walk/types.rs delete mode 100644 src/openhuman/memory/query/walk.rs create mode 100644 src/openhuman/memory_tree/graph/bfs.rs create mode 100644 src/openhuman/memory_tree/graph/mod.rs create mode 100644 src/openhuman/memory_tree/graph/store.rs create mode 100644 src/openhuman/memory_tree/nlp/client.rs create mode 100644 src/openhuman/memory_tree/nlp/mod.rs create mode 100644 src/openhuman/memory_tree/nlp/provision.rs create mode 100644 src/openhuman/memory_tree/nlp/service.py create mode 100644 src/openhuman/memory_tree/retrieval/fast.rs create mode 100644 tests/memory_fast_retrieve_e2e.rs delete mode 100644 tests/memory_tree_walk_e2e.rs diff --git a/scripts/test-rust-e2e.sh b/scripts/test-rust-e2e.sh index 0b6590ccc..4007cf772 100755 --- a/scripts/test-rust-e2e.sh +++ b/scripts/test-rust-e2e.sh @@ -52,7 +52,7 @@ ALL_E2E_SUITES=( memory_roundtrip_e2e memory_sources_e2e memory_tree_summarizer_e2e - memory_tree_walk_e2e + memory_fast_retrieve_e2e ollama_embeddings_fallback_e2e screen_intelligence_vision_e2e skill_registry_e2e diff --git a/src/openhuman/agent_memory/README.md b/src/openhuman/agent_memory/README.md index c47f56a98..eb246d659 100644 --- a/src/openhuman/agent_memory/README.md +++ b/src/openhuman/agent_memory/README.md @@ -65,4 +65,4 @@ The built-in agent is registered at `src/openhuman/agent_memory/agent/`: - `prompt.rs` — dynamic prompt builder - `prompt.md` — system prompt archetype -The agent has access to the full memory retrieval tool surface: `memory_tree`, `memory_recall`, `memory_smart_walk`, `memory_tree_walk`, plus entity and source querying tools. +The agent has access to the full memory retrieval tool surface: `memory_tree` (with deterministic E2GraphRAG `walk`/`smart_walk` modes plus `search_entities`/`query_source`/`cover_window`/`drill_down`/`fetch_leaves`), `memory_recall`, and `query_memory`. diff --git a/src/openhuman/agent_memory/agent/agent.toml b/src/openhuman/agent_memory/agent/agent.toml index cdc89a475..4729652d8 100644 --- a/src/openhuman/agent_memory/agent/agent.toml +++ b/src/openhuman/agent_memory/agent/agent.toml @@ -20,8 +20,6 @@ hint = "chat" named = [ "memory_recall", "memory_tree", - "memory_tree_walk", - "memory_smart_walk", "query_memory", "memory_doctor", "ask_user_clarification", diff --git a/src/openhuman/agent_memory/agent/prompt.md b/src/openhuman/agent_memory/agent/prompt.md index 7de2fa39f..8b3c70ac2 100644 --- a/src/openhuman/agent_memory/agent/prompt.md +++ b/src/openhuman/agent_memory/agent/prompt.md @@ -6,24 +6,22 @@ You are a memory retrieval specialist. Your job is to find and return relevant i Use the right tool for the job: -1. **`memory_smart_walk`** — your primary tool. Combines vector search, keyword matching, entity lookup, and tree browsing. Use for open-ended queries ("what do I know about X?", "find conversations about Y"). -2. **`memory_tree`** — unified dispatcher with modes: +1. **`memory_tree`** — your primary tool. Unified dispatcher with modes: + - `walk` / `smart_walk` — deterministic E2GraphRAG retrieval. Extracts query entities, routes between entity-graph (local) and dense-summary (global) search with no LLM, and returns ranked evidence hits. Use for open-ended queries ("what do I know about X?", "find conversations about Y"). - `search_entities` — find canonical entity IDs first (call before filtering by entity) - `query_source` — filter by source kind (chat, email, document) + time window - `drill_down` — expand a summary node one level deeper - `fetch_leaves` — pull raw chunks for citation -3. **`memory_tree_walk`** — basic tree navigation. Use when you need to explore the hierarchical summary structure step by step. -4. **`memory_recall`** — legacy key-value memory search. Good for exact preference/fact lookups. -5. **`query_memory`** — simple text search across stored memories. -6. **`memory_doctor`** — diagnose tree health issues. +2. **`memory_recall`** — legacy key-value memory search. Good for exact preference/fact lookups. +3. **`query_memory`** — simple text search across stored memories. +4. **`memory_doctor`** — diagnose tree health issues. ## Performance contract -- Start broad, then narrow. Use `search_entities` or `memory_smart_walk` first, then drill down. -- Avoid redundant walks. If `memory_smart_walk` already found the answer, don't re-walk with `memory_tree_walk`. +- Start broad, then narrow. Use `memory_tree` mode `walk` (or `search_entities`) first, then `drill_down` / `fetch_leaves` for detail. +- `walk`/`smart_walk` are deterministic and cheap — a single call returns ranked evidence; you do the synthesis. No multi-turn walking. - Cite sources. Every fact in your answer should trace back to a specific chunk or summary node. - Report what you didn't find. If the memory tree has gaps, say so explicitly rather than guessing. -- Prefer fewer turns. A 3-turn retrieval is better than an 8-turn one if both reach the same answer. ## Output format diff --git a/src/openhuman/agent_memory/ops.rs b/src/openhuman/agent_memory/ops.rs index 939cc5194..5c82db9f2 100644 --- a/src/openhuman/agent_memory/ops.rs +++ b/src/openhuman/agent_memory/ops.rs @@ -1,70 +1,56 @@ -//! Memory agent operations — benchmarking harness for memory tree walking -//! and retrieval performance measurement. +//! Memory agent operations — benchmarking harness for memory retrieval +//! performance measurement. +//! +//! Now measures the deterministic [`fast_retrieve`] retriever (E2GraphRAG). +//! There is no LLM in the loop, so the trace is a single retrieval "step" +//! rather than a multi-turn walk; `total_turns` stays 0 and the benchmark +//! focuses on wall-clock latency + hit count. use crate::openhuman::agent_memory::types::{BenchmarkSummary, RetrievalStep, WalkBenchmark}; use crate::openhuman::config::Config; -use crate::openhuman::inference::provider::traits::Provider; -use crate::openhuman::memory::query::smart_walk::{ - run_smart_walk, SmartWalkOptions, SmartWalkStopReason, -}; +use crate::openhuman::memory_tree::retrieval::{fast_retrieve, FastRetrieveOptions}; use std::path::PathBuf; use std::time::Instant; -/// Run a single benchmarked smart walk against the memory tree. +/// Run a single benchmarked deterministic retrieval against the memory tree. pub async fn bench_walk( config: &Config, - provider: &dyn Provider, query: &str, namespace: &str, content_root: Option, - max_turns: usize, - model: Option, + limit: usize, ) -> anyhow::Result { let effective_root = content_root .clone() .unwrap_or_else(|| config.memory_tree_content_root()); log::info!( - "[agent_memory::bench] query_len={} namespace={} content_root={} max_turns={}", + "[agent_memory::bench] query_len={} namespace={} content_root={} limit={}", query.len(), namespace, effective_root.display(), - max_turns + limit ); - let opts = SmartWalkOptions { - max_turns, - namespace: namespace.to_string(), - model, - content_root: content_root.clone(), + let opts = FastRetrieveOptions { + limit, + ..FastRetrieveOptions::default() }; let start = Instant::now(); - let outcome = run_smart_walk(config, provider, query, opts).await?; + let resp = fast_retrieve(config, query, opts).await?; let total_elapsed = start.elapsed(); - let steps: Vec = outcome - .trace - .iter() - .map(|step| RetrievalStep { - turn: step.turn, - action: step.action.clone(), - args_summary: step.args_summary.clone(), - result_preview: step.result_preview.clone(), - elapsed: std::time::Duration::ZERO, // per-step timing not available from smart_walk yet - chunks_returned: 0, - bytes_scanned: 0, - }) - .collect(); - - let total_chunks = outcome.evidence.len(); - - let stop_reason = match &outcome.stopped_reason { - SmartWalkStopReason::Answered => "answered".to_string(), - SmartWalkStopReason::MaxTurnsReached => "max_turns_reached".to_string(), - SmartWalkStopReason::LlmGaveUp => "llm_gave_up".to_string(), - SmartWalkStopReason::Error(e) => format!("error: {e}"), - }; + let total_bytes_scanned: u64 = resp.hits.iter().map(|h| h.content.len() as u64).sum(); + let steps: Vec = vec![RetrievalStep { + turn: 1, + action: "fast_retrieve".to_string(), + args_summary: format!("limit={limit}"), + result_preview: format!("{} hits (total {})", resp.hits.len(), resp.total), + elapsed: total_elapsed, + chunks_returned: resp.hits.len(), + bytes_scanned: total_bytes_scanned, + }]; let benchmark = WalkBenchmark { query: query.to_string(), @@ -72,24 +58,18 @@ pub async fn bench_walk( content_root: effective_root.display().to_string(), total_elapsed, steps, - total_turns: outcome.turns_used, - total_chunks_retrieved: total_chunks, - total_bytes_scanned: outcome - .evidence - .iter() - .map(|e| e.snippet.len() as u64) - .sum(), - answer: outcome.answer, - stop_reason, + total_turns: 0, // deterministic — no LLM turns + total_chunks_retrieved: resp.hits.len(), + total_bytes_scanned, + answer: String::new(), // synthesis is the high-level agent's job + stop_reason: "deterministic".to_string(), }; log::info!( - "[agent_memory::bench] completed query_len={} elapsed={:?} turns={} chunks={} stop={}", + "[agent_memory::bench] completed query_len={} elapsed={:?} chunks={}", query.len(), total_elapsed, - benchmark.total_turns, benchmark.total_chunks_retrieved, - benchmark.stop_reason ); Ok(benchmark) @@ -98,27 +78,15 @@ pub async fn bench_walk( /// Run a batch of queries and produce a summary. pub async fn bench_batch( config: &Config, - provider: &dyn Provider, queries: &[&str], namespace: &str, content_root: Option, - max_turns: usize, - model: Option, + limit: usize, ) -> anyhow::Result<(Vec, BenchmarkSummary)> { let mut results = Vec::with_capacity(queries.len()); for query in queries { - match bench_walk( - config, - provider, - query, - namespace, - content_root.clone(), - max_turns, - model.clone(), - ) - .await - { + match bench_walk(config, query, namespace, content_root.clone(), limit).await { Ok(bench) => results.push(bench), Err(e) => { log::warn!( diff --git a/src/openhuman/agent_memory/tools.rs b/src/openhuman/agent_memory/tools.rs index d8bbba13d..088357a01 100644 --- a/src/openhuman/agent_memory/tools.rs +++ b/src/openhuman/agent_memory/tools.rs @@ -1,11 +1,11 @@ //! Tool: `call_memory_agent` — invoke the memory retrieval agent to walk //! the memory tree and return context for a query. //! -//! Unlike the lower-level `memory_tree`, `memory_smart_walk`, and -//! `memory_tree_walk` tools which are individual retrieval primitives, -//! this tool spawns the full `agent_memory` sub-agent which autonomously -//! decides which retrieval strategies to combine, navigates the tree -//! across multiple turns, and returns a synthesised, cited answer. +//! Unlike the lower-level `memory_tree` tool (whose `walk`/`smart_walk` modes +//! now run the deterministic E2GraphRAG retriever and the other modes are +//! individual retrieval primitives), this tool spawns the full `agent_memory` +//! sub-agent which decides which retrieval strategies to combine and returns a +//! synthesised, cited answer. //! //! Supports both sync (blocking) and async (fire-and-forget) modes. diff --git a/src/openhuman/config/schema/load/env_overlay.rs b/src/openhuman/config/schema/load/env_overlay.rs index c30bc4f1c..8e2c84b40 100644 --- a/src/openhuman/config/schema/load/env_overlay.rs +++ b/src/openhuman/config/schema/load/env_overlay.rs @@ -684,6 +684,11 @@ impl Config { self.memory_tree.cloud_summarization_opt_in = val; } } + if let Some(raw) = env.get("OPENHUMAN_MEMORY_TREE_SPACY_ENABLED") { + if let Some(val) = parse_env_bool("OPENHUMAN_MEMORY_TREE_SPACY_ENABLED", &raw) { + self.memory_tree.spacy_enabled = val; + } + } } fn apply_update_env(&mut self, env: &E) { diff --git a/src/openhuman/config/schema/storage_memory.rs b/src/openhuman/config/schema/storage_memory.rs index 541bb760e..7f9974e11 100644 --- a/src/openhuman/config/schema/storage_memory.rs +++ b/src/openhuman/config/schema/storage_memory.rs @@ -357,6 +357,20 @@ pub struct MemoryTreeConfig { /// memory content will be sent to an external service. #[serde(default)] pub cloud_summarization_opt_in: bool, + + /// Enable the spaCy NER sidecar used by the deterministic (E2GraphRAG) + /// retriever to extract entities from a query. When `true` (default), the + /// managed Python runtime provisions spaCy on first use and serves entity + /// extraction over stdio. When `false` — or whenever Python/spaCy is + /// unavailable — query-entity extraction falls back to the in-Rust + /// regex+LLM extractor (`score::extract`). Env override: + /// `OPENHUMAN_MEMORY_TREE_SPACY_ENABLED`. + #[serde(default = "default_memory_tree_spacy_enabled")] + pub spacy_enabled: bool, +} + +fn default_memory_tree_spacy_enabled() -> bool { + true } /// Returns `None` so that existing installs that never opted into Phase 4 @@ -459,6 +473,7 @@ impl Default for MemoryTreeConfig { cloud_llm_model: default_cloud_llm_model(), smart_walk_model: None, cloud_summarization_opt_in: false, + spacy_enabled: default_memory_tree_spacy_enabled(), } } } diff --git a/src/openhuman/memory/query/fast_walk.rs b/src/openhuman/memory/query/fast_walk.rs new file mode 100644 index 000000000..549c1e39c --- /dev/null +++ b/src/openhuman/memory/query/fast_walk.rs @@ -0,0 +1,82 @@ +//! Deterministic replacement for the former agentic `walk` / `smart_walk` +//! tool modes. +//! +//! Both modes now resolve to [`fast_retrieve`] — the E2GraphRAG, LLM-free +//! retriever. It returns a structured [`QueryResponse`] of ranked evidence +//! (no synthesized prose); a higher-level context agent composes the answer. + +use crate::openhuman::config::rpc as config_rpc; +use crate::openhuman::memory_tree::retrieval::{fast_retrieve, FastRetrieveOptions}; +use crate::openhuman::tools::traits::ToolResult; + +/// Parse the shared `memory_tree` args and run deterministic retrieval. +/// Accepts `query` (required), `limit`, `time_window_days`, and `max_hops`. +pub async fn run_fast_walk(args: serde_json::Value) -> anyhow::Result { + let query = args + .get("query") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if query.trim().is_empty() { + return Err(anyhow::anyhow!("memory_tree walk: `query` is required")); + } + + let limit = args + .get("limit") + .and_then(|v| v.as_u64()) + .map(|n| n as usize) + .unwrap_or(10); + let time_window_days = args + .get("time_window_days") + .and_then(|v| v.as_u64()) + .map(|n| n as u32); + let max_hops = args + .get("max_hops") + .and_then(|v| v.as_u64()) + .map(|n| n as u32) + .unwrap_or(2); + + log::debug!( + "[tool][memory_tree] walk (deterministic) query_len={} limit={} max_hops={} window={:?}", + query.len(), + limit, + max_hops, + time_window_days + ); + + let cfg = config_rpc::load_config_with_timeout() + .await + .map_err(|e| anyhow::anyhow!("memory_tree walk: load config failed: {e}"))?; + + let opts = FastRetrieveOptions { + limit, + max_hops, + time_window_days, + }; + let resp = fast_retrieve(&cfg, &query, opts).await?; + log::debug!( + "[tool][memory_tree] walk returning hits={} total={}", + resp.hits.len(), + resp.total + ); + let json = serde_json::to_string(&resp)?; + Ok(ToolResult::success(json)) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[tokio::test] + async fn missing_query_errors() { + let err = run_fast_walk(json!({})).await.unwrap_err(); + assert!(err.to_string().contains("`query` is required")); + } + + #[tokio::test] + async fn blank_query_errors() { + let err = run_fast_walk(json!({"query": " "})).await.unwrap_err(); + assert!(err.to_string().contains("`query` is required")); + } +} diff --git a/src/openhuman/memory/query/mod.rs b/src/openhuman/memory/query/mod.rs index 3529684c6..1efef4af0 100644 --- a/src/openhuman/memory/query/mod.rs +++ b/src/openhuman/memory/query/mod.rs @@ -9,12 +9,11 @@ mod backend; mod cover_window; mod drill_down; +mod fast_walk; mod fetch_leaves; mod ingest_document; mod query_source; mod search_entities; -pub mod smart_walk; -pub mod walk; // Re-export individual tool types for callers that need them directly // (e.g. tool registration in ops.rs). @@ -24,12 +23,6 @@ pub use fetch_leaves::MemoryTreeFetchLeavesTool; pub use ingest_document::MemoryTreeIngestDocumentTool; pub use query_source::MemoryTreeQuerySourceTool; pub use search_entities::MemoryTreeSearchEntitiesTool; -pub use smart_walk::{ - run_smart_walk, SmartMemoryWalkTool, SmartWalkOptions, SmartWalkOutcome, SmartWalkStep, - SmartWalkStopReason, -}; -pub use walk::MemoryTreeWalkTool as MemoryQueryWalkTool; -pub use walk::{run_walk, MemoryTreeWalkTool, WalkOptions, WalkOutcome, WalkStep, WalkStopReason}; pub use MemoryTreeTool as MemoryQueryTool; use crate::openhuman::tools::traits::{Tool, ToolResult}; @@ -55,9 +48,9 @@ impl Tool for MemoryTreeTool { `drill_down` (expand a coarse summary one level), \ `cover_window` (minimum node set covering a time window [since_ms, until_ms] — use for last-24h / time-bounded recaps), \ `fetch_leaves` (pull raw chunks for citation), `ingest_document` (write a document into the tree for future retrieval), \ - `walk` (agentic multi-turn walk — LLM navigates summaries and returns a synthesized answer for a natural-language query), \ - `smart_walk` (multi-strategy retrieval — combines vector search, keyword search, entity lookup, \ - and tree browsing across raw files, wiki summaries, documents, and episodic memories)." + `walk` / `smart_walk` (deterministic E2GraphRAG retrieval — extracts query entities, routes between \ + entity-graph (local) and dense-summary (global) search with no LLM, and returns ranked evidence \ + hits for a natural-language query)." } fn parameters_schema(&self) -> serde_json::Value { @@ -97,7 +90,12 @@ impl Tool for MemoryTreeTool { }, "time_window_days": { "type": "integer", - "description": "query_source: look-back window in days." + "description": "query_source / walk / smart_walk: look-back window in days (applied to the dense/global branch for walk)." + }, + // walk / smart_walk params + "max_hops": { + "type": "integer", + "description": "walk / smart_walk: entity-graph relatedness hop threshold for E2GraphRAG routing (default 2, capped at 4)." }, // drill_down params "node_id": { @@ -158,8 +156,7 @@ impl Tool for MemoryTreeTool { "cover_window" => MemoryTreeCoverWindowTool.execute(args).await, "fetch_leaves" => MemoryTreeFetchLeavesTool.execute(args).await, "ingest_document" => MemoryTreeIngestDocumentTool.execute(args).await, - "walk" => MemoryTreeWalkTool.execute(args).await, - "smart_walk" => SmartMemoryWalkTool.execute(args).await, + "walk" | "smart_walk" => fast_walk::run_fast_walk(args).await, other => { log::debug!("[tool][memory_tree] unknown_mode mode={other}"); Err(anyhow::anyhow!( diff --git a/src/openhuman/memory/query/smart_walk/dispatch.rs b/src/openhuman/memory/query/smart_walk/dispatch.rs deleted file mode 100644 index 777ff8451..000000000 --- a/src/openhuman/memory/query/smart_walk/dispatch.rs +++ /dev/null @@ -1,687 +0,0 @@ -//! Tool call dispatch for the smart_walk inner loop. -//! -//! Each `dispatch_*` function handles one named inner tool and returns -//! `(args_summary, result_text, is_final_answer, answer_text)`. - -use crate::openhuman::config::Config; -use crate::openhuman::memory::query::smart_walk::prompts::InnerCall; -use crate::openhuman::memory::query::smart_walk::types::{ - Evidence, MAX_EVIDENCE_ITEMS, MAX_FILE_READ_BYTES, MAX_KEYWORD_RESULTS, -}; -use crate::openhuman::memory_store::chunks::types::SourceKind; -use crate::openhuman::memory_tree::retrieval; -use crate::openhuman::memory_tree::score::extract::EntityKind; -use crate::openhuman::memory_tree::tree_runtime::store::{read_children, read_node}; -use std::path::{Path, PathBuf}; - -// ── Top-level dispatcher ───────────────────────────────────────────────────── - -pub(crate) async fn dispatch_call( - config: &Config, - namespace: &str, - content_root: &Path, - call: &InnerCall, - evidence: &mut Vec, -) -> (String, String, bool, String) { - match call.name.as_str() { - "keyword_search" => { - let cr = content_root.to_path_buf(); - let c = call.clone(); - tokio::task::spawn_blocking(move || dispatch_keyword_search(&cr, &c)) - .await - .unwrap_or_else(|e| (String::new(), format!("error: {e}"), false, String::new())) - } - "entity_search" => dispatch_entity_search(config, call).await, - "list_sources" => { - let cr = content_root.to_path_buf(); - let c = call.clone(); - tokio::task::spawn_blocking(move || dispatch_list_sources(&cr, &c)) - .await - .unwrap_or_else(|e| (String::new(), format!("error: {e}"), false, String::new())) - } - "read_content" => { - let cr = content_root.to_path_buf(); - let c = call.clone(); - tokio::task::spawn_blocking(move || dispatch_read_content(&cr, &c)) - .await - .unwrap_or_else(|e| (String::new(), format!("error: {e}"), false, String::new())) - } - "browse_tree" => dispatch_browse_tree(config, namespace, call).await, - "collect_evidence" => dispatch_collect_evidence(call, evidence), - "answer" => dispatch_answer(call), - "vector_search" => dispatch_vector_search(config, call).await, - other => { - log::warn!("[smart_walk] unknown action: {other}"); - ( - format!("action={other}"), - format!( - "unknown action '{other}'. Valid: keyword_search, entity_search, \ - list_sources, read_content, browse_tree, vector_search, \ - collect_evidence, answer" - ), - false, - String::new(), - ) - } - } -} - -// ── keyword_search ─────────────────────────────────────────────────────────── - -pub(crate) fn dispatch_keyword_search( - content_root: &Path, - call: &InnerCall, -) -> (String, String, bool, String) { - let pattern = call - .args - .get("pattern") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - let content_type = call - .args - .get("content_type") - .and_then(|v| v.as_str()) - .unwrap_or("all"); - - if pattern.is_empty() { - return ( - "pattern=".into(), - "error: keyword_search requires a non-empty pattern".into(), - false, - String::new(), - ); - } - - log::debug!( - "[smart_walk] keyword_search pattern={} content_type={}", - pattern, - content_type - ); - - let args_summary = format!("pattern=\"{}\" type={}", pattern, content_type); - - let search_dirs: Vec = match content_type { - "raw" => vec![content_root.join("raw")], - "wiki" => vec![content_root.join("wiki")], - "document" => vec![content_root.join("document")], - "episodic" => vec![content_root.join("episodic")], - _ => vec![ - content_root.join("raw"), - content_root.join("wiki"), - content_root.join("document"), - content_root.join("episodic"), - ], - }; - - let pattern_lower = pattern.to_lowercase(); - let mut results: Vec = Vec::new(); - - for dir in &search_dirs { - if !dir.exists() { - continue; - } - search_dir_recursive(dir, &pattern_lower, &mut results, content_root); - if results.len() >= MAX_KEYWORD_RESULTS { - break; - } - } - - results.truncate(MAX_KEYWORD_RESULTS); - - if results.is_empty() { - ( - args_summary, - format!("no matches for pattern \"{}\"", pattern), - false, - String::new(), - ) - } else { - let count = results.len(); - ( - args_summary, - format!("{count} matches:\n{}", results.join("\n")), - false, - String::new(), - ) - } -} - -pub(crate) fn search_dir_recursive( - dir: &Path, - pattern: &str, - results: &mut Vec, - content_root: &Path, -) { - let entries = match std::fs::read_dir(dir) { - Ok(e) => e, - Err(_) => return, - }; - - for entry in entries.flatten() { - if results.len() >= MAX_KEYWORD_RESULTS { - return; - } - - let path = entry.path(); - if path.is_dir() { - search_dir_recursive(&path, pattern, results, content_root); - } else if path.extension().map_or(false, |e| e == "md") { - if let Ok(content) = std::fs::read_to_string(&path) { - if content.to_lowercase().contains(pattern) { - let rel = path - .strip_prefix(content_root) - .unwrap_or(&path) - .to_string_lossy() - .to_string(); - - let line_match = content - .lines() - .find(|l| l.to_lowercase().contains(pattern)) - .unwrap_or("") - .trim(); - let preview: String = line_match.chars().take(120).collect(); - results.push(format!(" [{rel}] {preview}")); - } - } - } - } -} - -// ── entity_search ──────────────────────────────────────────────────────────── - -async fn dispatch_entity_search( - config: &Config, - call: &InnerCall, -) -> (String, String, bool, String) { - let query = call - .args - .get("query") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - let kinds: Option> = - call.args - .get("kinds") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str()) - .filter_map(|s| EntityKind::parse(s).ok()) - .collect() - }); - - if query.is_empty() { - return ( - "query=".into(), - "error: entity_search requires a non-empty query".into(), - false, - String::new(), - ); - } - - log::debug!( - "[smart_walk] entity_search query={} kinds={:?}", - query, - kinds - .as_ref() - .map(|ks| ks.iter().map(|k| k.as_str()).collect::>()) - ); - let args_summary = format!( - "query=\"{}\" kinds={:?}", - query, - kinds - .as_ref() - .map(|ks| ks.iter().map(|k| k.as_str()).collect::>()) - ); - - match retrieval::search_entities(config, &query, kinds, 10).await { - Ok(matches) => { - if matches.is_empty() { - ( - args_summary, - format!("no entities matching \"{}\"", query), - false, - String::new(), - ) - } else { - let formatted: Vec = matches - .iter() - .map(|m| { - format!( - " [{}] kind={} surface=\"{}\" mentions={} last_seen={}", - m.canonical_id, - m.kind.as_str(), - m.surface, - m.mention_count, - m.last_seen_ms - ) - }) - .collect(); - ( - args_summary, - format!( - "{} entities found:\n{}", - formatted.len(), - formatted.join("\n") - ), - false, - String::new(), - ) - } - } - Err(e) => ( - args_summary, - format!("entity search error: {e}"), - false, - String::new(), - ), - } -} - -// ── list_sources ───────────────────────────────────────────────────────────── - -pub(crate) fn dispatch_list_sources( - content_root: &Path, - call: &InnerCall, -) -> (String, String, bool, String) { - let content_type = call - .args - .get("content_type") - .and_then(|v| v.as_str()) - .unwrap_or("all"); - - log::debug!("[smart_walk] list_sources type={}", content_type); - let args_summary = format!("type={}", content_type); - - let mut listing = Vec::new(); - - let types_to_scan: Vec<&str> = match content_type { - "all" => vec!["raw", "wiki", "document", "episodic"], - t => vec![t], - }; - - for ctype in types_to_scan { - let dir = content_root.join(ctype); - if !dir.exists() { - listing.push(format!(" {ctype}/: (empty)")); - continue; - } - - match std::fs::read_dir(&dir) { - Ok(entries) => { - let mut subdirs: Vec = entries - .flatten() - .filter(|e| e.path().is_dir()) - .filter_map(|e| e.file_name().into_string().ok()) - .collect(); - subdirs.sort(); - - if subdirs.is_empty() { - listing.push(format!(" {ctype}/: (no subdirectories)")); - } else { - let count = subdirs.len(); - let preview: Vec<&str> = subdirs.iter().map(|s| s.as_str()).take(10).collect(); - listing.push(format!( - " {ctype}/ ({count} sources): {}{}", - preview.join(", "), - if count > 10 { ", ..." } else { "" } - )); - } - } - Err(e) => listing.push(format!(" {ctype}/: error: {e}")), - } - } - - ( - args_summary, - format!("Content sources:\n{}", listing.join("\n")), - false, - String::new(), - ) -} - -// ── read_content ───────────────────────────────────────────────────────────── - -pub(crate) fn dispatch_read_content( - content_root: &Path, - call: &InnerCall, -) -> (String, String, bool, String) { - let path_str = call - .args - .get("path") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - if path_str.is_empty() { - return ( - "path=".into(), - "error: read_content requires a non-empty path".into(), - false, - String::new(), - ); - } - - let requested = Path::new(&path_str); - if requested.is_absolute() || path_str.contains("..") { - return ( - format!("path={path_str}"), - "error: path must stay within the content root".into(), - false, - String::new(), - ); - } - - log::debug!("[smart_walk] read_content path={}", path_str); - - let full_path = content_root.join(requested); - if !full_path.exists() { - return ( - format!("path={path_str}"), - format!("file not found: {path_str}"), - false, - String::new(), - ); - } - - let canonical_root = match content_root.canonicalize() { - Ok(p) => p, - Err(e) => { - return ( - format!("path={path_str}"), - format!("error resolving content root: {e}"), - false, - String::new(), - ); - } - }; - let canonical_path = match full_path.canonicalize() { - Ok(p) => p, - Err(e) => { - return ( - format!("path={path_str}"), - format!("error resolving path: {e}"), - false, - String::new(), - ); - } - }; - if !canonical_path.starts_with(&canonical_root) { - return ( - format!("path={path_str}"), - "error: path escapes content root".into(), - false, - String::new(), - ); - } - - match std::fs::read_to_string(&canonical_path) { - Ok(content) => { - let truncated: String = content.chars().take(MAX_FILE_READ_BYTES).collect(); - let was_truncated = content.len() > MAX_FILE_READ_BYTES; - let suffix = if was_truncated { - format!("\n\n[...truncated, {} total chars]", content.len()) - } else { - String::new() - }; - ( - format!("path={path_str}"), - format!("{truncated}{suffix}"), - false, - String::new(), - ) - } - Err(e) => ( - format!("path={path_str}"), - format!("error reading: {e}"), - false, - String::new(), - ), - } -} - -// ── browse_tree ────────────────────────────────────────────────────────────── - -async fn dispatch_browse_tree( - config: &Config, - namespace: &str, - call: &InnerCall, -) -> (String, String, bool, String) { - let node_id = call - .args - .get("node_id") - .and_then(|v| v.as_str()) - .unwrap_or("root") - .to_string(); - - log::debug!("[smart_walk] browse_tree node_id={}", node_id); - - let config_owned = config.clone(); - let ns_owned = namespace.to_string(); - let id_owned = node_id.clone(); - - let result = tokio::task::spawn_blocking(move || { - let node = match read_node(&config_owned, &ns_owned, &id_owned) { - Ok(Some(n)) => n, - Ok(None) => return format!("unknown node: {id_owned}"), - Err(e) => return format!("error reading node {id_owned}: {e}"), - }; - - let children = match read_children(&config_owned, &ns_owned, &id_owned) { - Ok(c) => c, - Err(_) => vec![], - }; - - let mut out = format!( - "Node: {} (level={:?})\nSummary: {}\n", - node.node_id, node.level, node.summary - ); - - if children.is_empty() { - out.push_str("Children: (none — leaf node)\n"); - } else { - out.push_str(&format!("Children ({}):\n", children.len())); - for c in &children { - let preview: String = c.summary.chars().take(100).collect(); - out.push_str(&format!( - " - id={} level={:?}: {}\n", - c.node_id, c.level, preview - )); - } - } - out - }) - .await - .unwrap_or_else(|_| format!("error building context for node {node_id}")); - - (format!("node_id={node_id}"), result, false, String::new()) -} - -// ── vector_search ──────────────────────────────────────────────────────────── - -async fn dispatch_vector_search( - config: &Config, - call: &InnerCall, -) -> (String, String, bool, String) { - use crate::openhuman::memory::query::smart_walk::types::truncate_chars; - - let query = call - .args - .get("query") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - let source_kind = call - .args - .get("source_kind") - .and_then(|v| v.as_str()) - .and_then(|s| match s { - "chat" => Some(SourceKind::Chat), - "email" => Some(SourceKind::Email), - "document" => Some(SourceKind::Document), - _ => None, - }); - - let time_window_days = call - .args - .get("time_window_days") - .and_then(|v| v.as_u64()) - .map(|n| n as u32); - - if query.is_empty() { - return ( - "query=".into(), - "error: vector_search requires a non-empty query".into(), - false, - String::new(), - ); - } - - log::debug!( - "[smart_walk] vector_search query={} source_kind={:?} window_days={:?}", - query, - source_kind, - time_window_days - ); - let args_summary = format!( - "query=\"{}\" kind={:?} window={:?}", - truncate_chars(&query, 40), - source_kind, - time_window_days - ); - - match retrieval::query_source( - config, - None, - source_kind, - time_window_days, - Some(&query), - 10, - ) - .await - { - Ok(resp) => { - if resp.hits.is_empty() { - ( - args_summary, - format!("no vector matches for \"{}\"", query), - false, - String::new(), - ) - } else { - let formatted: Vec = resp - .hits - .iter() - .map(|h| { - let preview: String = h.content.chars().take(120).collect(); - format!(" [{}] (score={:.2}) {}", h.node_id, h.score, preview) - }) - .collect(); - ( - args_summary, - format!( - "{} semantic matches:\n{}", - formatted.len(), - formatted.join("\n") - ), - false, - String::new(), - ) - } - } - Err(e) => ( - args_summary, - format!("vector search error: {e}"), - false, - String::new(), - ), - } -} - -// ── collect_evidence ───────────────────────────────────────────────────────── - -pub(crate) fn dispatch_collect_evidence( - call: &InnerCall, - evidence: &mut Vec, -) -> (String, String, bool, String) { - let items = call - .args - .get("items") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - - if items.is_empty() { - return ( - "items=[]".into(), - "error: collect_evidence requires non-empty items array".into(), - false, - String::new(), - ); - } - - let mut added = 0; - for item in &items { - if evidence.len() >= MAX_EVIDENCE_ITEMS { - break; - } - let source_path = item - .get("source") - .and_then(|v| v.as_str()) - .unwrap_or("unknown") - .to_string(); - let snippet = item - .get("snippet") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let relevance = item - .get("relevance") - .and_then(|v| v.as_str()) - .unwrap_or("relevant") - .to_string(); - - if !snippet.is_empty() { - evidence.push(Evidence { - source_path, - snippet, - relevance, - }); - added += 1; - } - } - - log::debug!( - "[smart_walk] collect_evidence added={} total={}", - added, - evidence.len() - ); - - ( - format!("{added} items"), - format!( - "collected {added} evidence items (total: {})", - evidence.len() - ), - false, - String::new(), - ) -} - -// ── answer ─────────────────────────────────────────────────────────────────── - -pub(crate) fn dispatch_answer(call: &InnerCall) -> (String, String, bool, String) { - let text = call - .args - .get("text") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - log::debug!("[smart_walk] answer text_len={}", text.len()); - ("(final answer)".into(), text.clone(), true, text) -} diff --git a/src/openhuman/memory/query/smart_walk/mod.rs b/src/openhuman/memory/query/smart_walk/mod.rs deleted file mode 100644 index 7c896fa64..000000000 --- a/src/openhuman/memory/query/smart_walk/mod.rs +++ /dev/null @@ -1,29 +0,0 @@ -//! E2GraphRAG-inspired smart memory retrieval. -//! -//! Unlike the basic `walk` module which only navigates the time-based summary -//! tree, smart_walk combines multiple retrieval strategies: -//! -//! 1. **Vector search** — semantic similarity across all stored content -//! 2. **Keyword search** — pattern matching across raw content files on disk -//! 3. **Entity search** — find entities and follow relationships -//! 4. **Tree browse** — navigate wiki summary hierarchies -//! 5. **Content read** — read specific files (raw/wiki/document/episodic) -//! 6. **Source listing** — discover available sources and content types -//! -//! The walker LLM (defaulting to DeepSeek Flash) plans which strategies to -//! use, collects evidence snippets, then synthesizes a cited answer. - -mod dispatch; -mod prompts; -mod runner; -mod tool; -pub mod types; - -#[cfg(test)] -mod smart_walk_tests; - -// ── Public re-exports ──────────────────────────────────────────────────────── - -pub use runner::run_smart_walk; -pub use tool::SmartMemoryWalkTool; -pub use types::{Evidence, SmartWalkOptions, SmartWalkOutcome, SmartWalkStep, SmartWalkStopReason}; diff --git a/src/openhuman/memory/query/smart_walk/prompts.rs b/src/openhuman/memory/query/smart_walk/prompts.rs deleted file mode 100644 index 6bdc35937..000000000 --- a/src/openhuman/memory/query/smart_walk/prompts.rs +++ /dev/null @@ -1,292 +0,0 @@ -//! Prompt construction, content inventory, model resolution, tool-call parsing, -//! and fallback synthesis for smart_walk. - -use crate::openhuman::config::Config; -use crate::openhuman::memory::query::smart_walk::types::{truncate_chars, Evidence, SmartWalkStep}; -use std::path::Path; - -// ── Inner call type (used by parser and dispatch) ─────────────────────────── - -#[derive(Clone)] -pub(crate) struct InnerCall { - pub(crate) name: String, - pub(crate) args: serde_json::Value, -} - -// ── System prompt ──────────────────────────────────────────────────────────── - -pub(crate) fn build_system_prompt() -> String { - r#"You are a smart memory retrieval agent. Your task is to answer queries by -searching through a user's personal memory — which includes raw files (emails, -chats, commits, documents), wiki summaries, episodic conversation memories, -and document archives. - -## Strategy - -Use a multi-strategy approach inspired by graph-based retrieval: - -1. **Start broad**: Use `list_sources` to understand what content is available, - then `keyword_search` or `vector_search` to find relevant starting points. - -2. **Follow connections**: When you find a relevant entity or topic, use - `entity_search` to find related entities and follow the connections. - -3. **Drill into details**: Use `read_content` to read specific files for - full context. Use `browse_tree` to navigate wiki summary hierarchies. - -4. **Collect evidence**: As you find relevant information, use `collect_evidence` - to save snippets. This builds your citation buffer for the final answer. - -5. **Synthesize**: When you have enough evidence, use `answer` to provide a - comprehensive response with citations. - -## Rules - -- Be efficient: don't re-search for things you already found. -- Prefer vector_search for semantic/conceptual queries. -- Prefer keyword_search for specific names, IDs, or exact phrases. -- Use entity_search when the query mentions people, projects, or organizations. -- Always collect_evidence before answering, so your answer has citations. -- Use tags with JSON content for actions. Format: - {"name":"tool_name","arguments":{"param":"value"}} -- You can call multiple tools in one turn by including multiple blocks. - -## Example turn - -I'll search for recent emails about the project. - -{"name":"list_sources","arguments":{"content_type":"all"}} -{"name":"keyword_search","arguments":{"pattern":"project","content_type":"raw"}} -"# - .into() -} - -pub(crate) fn build_inner_tools_text() -> String { - r#"## Available tools - -**keyword_search** `{"pattern": "", "content_type": "all|raw|wiki|document|episodic"}` -Search for a text pattern (case-insensitive) across memory files. Returns matching file paths and line previews. - -**vector_search** `{"query": "", "source_kind": "chat|email|document", "time_window_days": 30}` -Semantic similarity search over indexed summaries. All params except query are optional. - -**entity_search** `{"query": "", "kinds": ["person", "email", "url", "handle"]}` -Find entities (people, emails, URLs, handles) in the entity index. kinds is optional. - -**list_sources** `{"content_type": "all|raw|wiki|document|episodic"}` -List available content sources and their subdirectories. - -**read_content** `{"path": ""}` -Read a specific content file. Path is relative to the content root (e.g. "raw/github-com-example/commits/123.md"). - -**browse_tree** `{"node_id": "root"}` -Navigate the wiki summary tree. Returns node summary and children. Use "root" to start. - -**collect_evidence** `{"items": [{"source": "", "snippet": "", "relevance": ""}]}` -Save evidence snippets for citation in your final answer. Call this as you find relevant information. - -**answer** `{"text": ""}` -Return your final answer. Reference collected evidence by source path."# - .into() -} - -// ── Content inventory ──────────────────────────────────────────────────────── - -pub(crate) fn build_content_inventory(content_root: &Path) -> String { - let mut parts = Vec::new(); - - for (label, subdir) in &[ - ("Raw content", "raw"), - ("Wiki summaries", "wiki"), - ("Documents", "document"), - ("Episodic memories", "episodic"), - ] { - let dir = content_root.join(subdir); - if dir.exists() { - let count = count_files_recursive(&dir); - if count > 0 { - parts.push(format!("- **{label}** ({subdir}/): {count} files")); - } - } - } - - if parts.is_empty() { - "No content files found.".into() - } else { - parts.join("\n") - } -} - -pub(crate) fn count_files_recursive(dir: &Path) -> usize { - let mut count = 0; - if let Ok(entries) = std::fs::read_dir(dir) { - for entry in entries.flatten() { - let path = entry.path(); - if path.is_dir() { - count += count_files_recursive(&path); - } else if path.extension().map_or(false, |e| e == "md") { - count += 1; - } - } - } - count -} - -// ── Model resolution ───────────────────────────────────────────────────────── - -const DEFAULT_SMART_WALK_MODEL: &str = "hint:summarization"; - -pub(crate) fn resolve_walk_model(config: &Config) -> String { - // 1. Explicit smart_walk_model config takes priority - if let Some(ref swm) = config.memory_tree.smart_walk_model { - if !swm.is_empty() { - return swm.clone(); - } - } - // 2. Default to summarization-v1 (routed through the OpenHuman backend) - DEFAULT_SMART_WALK_MODEL.to_string() -} - -// ── Tool call parser ───────────────────────────────────────────────────────── - -pub(crate) fn parse_tool_calls(response: &str) -> (String, Vec) { - let mut calls: Vec = Vec::new(); - let mut text_parts: Vec<&str> = Vec::new(); - let mut remaining: &str = response; - - const OPEN: &str = ""; - const CLOSE: &str = ""; - - loop { - match remaining.find(OPEN) { - None => { - if !remaining.trim().is_empty() && calls.is_empty() { - text_parts.push(remaining); - } - break; - } - Some(start) => { - let before = &remaining[..start]; - if !before.trim().is_empty() { - text_parts.push(before); - } - let after_open = &remaining[start + OPEN.len()..]; - match after_open.find(CLOSE) { - None => break, - Some(close_idx) => { - let inner = after_open[..close_idx].trim(); - if let Some(call) = parse_single_tool_call(inner) { - calls.push(call); - } - remaining = &after_open[close_idx + CLOSE.len()..]; - } - } - } - } - } - - let text_before = text_parts.concat(); - (text_before, calls) -} - -fn parse_single_tool_call(inner: &str) -> Option { - // Primary: JSON format {"name":"...","arguments":{...}} - if let Ok(val) = serde_json::from_str::(inner) { - if let Some(name) = val.get("name").and_then(|v| v.as_str()) { - let args = val - .get("arguments") - .cloned() - .unwrap_or(serde_json::Value::Object(Default::default())); - log::debug!( - "[smart_walk::parse_single_tool_call] json path: tool={} args_keys={}", - name, - args.as_object().map(|m| m.len()).unwrap_or(0) - ); - return Some(InnerCall { - name: name.to_string(), - args, - }); - } - } - // Fallback: XML-style nameJSON - if let (Some(name), args) = ( - extract_xml_tag(inner, "tool_name"), - extract_xml_tag(inner, "parameters"), - ) { - log::debug!( - "[smart_walk::parse_single_tool_call] xml fallback path: tool={} has_params={}", - name.trim(), - args.is_some() - ); - let parsed_args = args - .and_then(|a| serde_json::from_str::(a.trim()).ok()) - .unwrap_or_else(|| { - let mut map = serde_json::Map::new(); - for line in inner.lines() { - let trimmed = line.trim(); - if trimmed.starts_with('<') - && !trimmed.starts_with("') { - let tag = &trimmed[1..tag_end]; - if let Some(close) = trimmed.find(&format!("")) { - let value = &trimmed[tag_end + 1..close]; - map.insert( - tag.to_string(), - serde_json::Value::String(value.to_string()), - ); - } - } - } - } - serde_json::Value::Object(map) - }); - return Some(InnerCall { - name: name.trim().to_string(), - args: parsed_args, - }); - } - None -} - -fn extract_xml_tag<'a>(text: &'a str, tag: &str) -> Option<&'a str> { - let open = format!("<{tag}>"); - let close = format!(""); - let start = text.find(&open)? + open.len(); - let end = text[start..].find(&close)? + start; - Some(&text[start..end]) -} - -// ── Fallback synthesis ─────────────────────────────────────────────────────── - -pub(crate) fn synthesize_fallback(trace: &[SmartWalkStep], evidence: &[Evidence]) -> String { - let mut out = String::new(); - - if !evidence.is_empty() { - out.push_str("Based on the evidence collected:\n\n"); - for (i, ev) in evidence.iter().enumerate() { - out.push_str(&format!( - "{}. [{}] {}: {}\n", - i + 1, - ev.source_path, - ev.relevance, - truncate_chars(&ev.snippet, 150) - )); - } - } else if !trace.is_empty() { - out.push_str("Could not converge on an answer. Steps taken:\n\n"); - for s in trace { - out.push_str(&format!( - "- Turn {}: {} → {}\n", - s.turn, - s.action, - truncate_chars(&s.result_preview, 100) - )); - } - } else { - out.push_str("Could not converge on an answer — no steps taken."); - } - out -} diff --git a/src/openhuman/memory/query/smart_walk/runner.rs b/src/openhuman/memory/query/smart_walk/runner.rs deleted file mode 100644 index 33b3bfb24..000000000 --- a/src/openhuman/memory/query/smart_walk/runner.rs +++ /dev/null @@ -1,182 +0,0 @@ -//! Main agentic loop for smart_walk. - -use crate::openhuman::config::Config; -use crate::openhuman::inference::provider::traits::{ChatMessage, Provider}; -use crate::openhuman::memory::query::smart_walk::dispatch::dispatch_call; -use crate::openhuman::memory::query::smart_walk::prompts::{ - build_content_inventory, build_inner_tools_text, build_system_prompt, parse_tool_calls, - resolve_walk_model, synthesize_fallback, -}; -use crate::openhuman::memory::query::smart_walk::types::{ - truncate_chars, Evidence, SmartWalkOptions, SmartWalkOutcome, SmartWalkStep, - SmartWalkStopReason, HARD_MAX_TURNS, SMART_WALK_TEMP, -}; - -pub async fn run_smart_walk( - config: &Config, - provider: &dyn Provider, - query: &str, - opts: SmartWalkOptions, -) -> anyhow::Result { - let max_turns = opts.max_turns.min(HARD_MAX_TURNS); - let model = opts - .model - .clone() - .unwrap_or_else(|| resolve_walk_model(config)); - - let content_root = opts - .content_root - .clone() - .unwrap_or_else(|| config.memory_tree_content_root()); - - log::debug!( - "[smart_walk] starting query_len={} namespace={} max_turns={} model={} content_root={}", - query.len(), - opts.namespace, - max_turns, - model, - content_root.display() - ); - - let system = build_system_prompt(); - let inner_tools = build_inner_tools_text(); - - let cr = content_root.clone(); - let inventory = tokio::task::spawn_blocking(move || build_content_inventory(&cr)) - .await - .unwrap_or_else(|_| "error building content inventory".into()); - - let mut history: Vec = vec![ - ChatMessage::system(format!("{system}\n\n{inner_tools}")), - ChatMessage::user(format!( - "Query: {query}\n\n## Available content\n{inventory}" - )), - ]; - - let mut trace: Vec = Vec::new(); - let mut evidence: Vec = Vec::new(); - - for turn in 1..=max_turns { - log::debug!("[smart_walk] turn={turn} evidence_count={}", evidence.len()); - - let response = match provider - .chat_with_history(&history, &model, SMART_WALK_TEMP) - .await - { - Ok(r) => r, - Err(e) => { - log::warn!("[smart_walk] provider error on turn={turn}: {e:#}"); - let err_msg = format!("Provider error on turn {turn}: {e}"); - return Ok(SmartWalkOutcome { - answer: format!( - "Walk failed: {err_msg}\n\nPartial from {} turn(s).", - trace.len() - ), - evidence, - trace, - turns_used: turn, - stopped_reason: SmartWalkStopReason::Error(err_msg), - }); - } - }; - - log::debug!("[smart_walk] turn={turn} response_len={}", response.len()); - - let (text_before, calls) = parse_tool_calls(&response); - - if calls.is_empty() { - let trimmed = response.trim().to_string(); - if trimmed.is_empty() { - log::debug!("[smart_walk] turn={turn} LLM gave up (empty response)"); - return Ok(SmartWalkOutcome { - answer: synthesize_fallback(&trace, &evidence), - evidence, - trace, - turns_used: turn, - stopped_reason: SmartWalkStopReason::LlmGaveUp, - }); - } - log::debug!("[smart_walk] turn={turn} no tool calls — treating as answer"); - return Ok(SmartWalkOutcome { - answer: trimmed, - evidence, - trace, - turns_used: turn, - stopped_reason: SmartWalkStopReason::Answered, - }); - } - - history.push(ChatMessage::assistant(response.clone())); - - // Process ALL tool calls in this turn (not just the first). - let mut combined_results = Vec::new(); - for call in &calls { - log::debug!( - "[smart_walk] turn={turn} action={} args={}", - call.name, - call.args - ); - - let (args_summary, tool_result, is_answer, answer_text) = - dispatch_call(config, &opts.namespace, &content_root, call, &mut evidence).await; - - let result_preview: String = tool_result.chars().take(200).collect(); - trace.push(SmartWalkStep { - turn, - action: call.name.clone(), - args_summary, - result_preview: result_preview.clone(), - }); - - if is_answer { - log::debug!("[smart_walk] turn={turn} answer action — stopping"); - return Ok(SmartWalkOutcome { - answer: answer_text, - evidence, - trace, - turns_used: turn, - stopped_reason: SmartWalkStopReason::Answered, - }); - } - - combined_results.push(format!( - "{}", - call.name, tool_result - )); - } - - let evidence_summary = if evidence.is_empty() { - String::new() - } else { - format!( - "\n\nEvidence collected so far ({} items):\n{}", - evidence.len(), - evidence - .iter() - .enumerate() - .map(|(i, e)| format!(" {}. [{}] {}", i + 1, e.source_path, e.relevance)) - .collect::>() - .join("\n") - ) - }; - - let result_msg = format!("{}{}", combined_results.join("\n"), evidence_summary); - history.push(ChatMessage::user(result_msg)); - - if !text_before.trim().is_empty() { - log::debug!( - "[smart_walk] turn={turn} text before tool calls: {}", - truncate_chars(&text_before, 80) - ); - } - } - - log::debug!("[smart_walk] max_turns={max_turns} reached"); - Ok(SmartWalkOutcome { - answer: synthesize_fallback(&trace, &evidence), - evidence, - trace, - turns_used: max_turns, - stopped_reason: SmartWalkStopReason::MaxTurnsReached, - }) -} diff --git a/src/openhuman/memory/query/smart_walk/smart_walk_tests.rs b/src/openhuman/memory/query/smart_walk/smart_walk_tests.rs deleted file mode 100644 index b851303c8..000000000 --- a/src/openhuman/memory/query/smart_walk/smart_walk_tests.rs +++ /dev/null @@ -1,693 +0,0 @@ -//! Tests for the smart_walk module. - -#[cfg(test)] -mod tests { - use crate::openhuman::config::Config; - use crate::openhuman::inference::provider::traits::{ChatMessage, Provider}; - use crate::openhuman::memory::query::smart_walk::dispatch::{ - dispatch_keyword_search, dispatch_list_sources, dispatch_read_content, search_dir_recursive, - }; - use crate::openhuman::memory::query::smart_walk::prompts::{ - build_content_inventory, parse_tool_calls, InnerCall, - }; - use crate::openhuman::memory::query::smart_walk::runner::run_smart_walk; - use crate::openhuman::memory::query::smart_walk::types::{ - SmartWalkOptions, SmartWalkStopReason, - }; - use async_trait::async_trait; - use std::path::Path; - use std::sync::Mutex; - use tempfile::TempDir; - - struct StubProvider { - responses: Mutex>, - } - - impl StubProvider { - fn new(responses: Vec<&str>) -> Self { - Self { - responses: Mutex::new(responses.into_iter().map(|s| s.to_string()).collect()), - } - } - } - - #[async_trait] - impl Provider for StubProvider { - async fn chat_with_system( - &self, - _system: Option<&str>, - _message: &str, - _model: &str, - _temp: f64, - ) -> anyhow::Result { - let mut responses = self.responses.lock().unwrap(); - if responses.is_empty() { - return Err(anyhow::anyhow!("StubProvider: no more responses")); - } - Ok(responses.remove(0)) - } - - async fn chat_with_history( - &self, - _messages: &[ChatMessage], - _model: &str, - _temp: f64, - ) -> anyhow::Result { - let mut responses = self.responses.lock().unwrap(); - if responses.is_empty() { - return Err(anyhow::anyhow!("StubProvider: no more responses")); - } - Ok(responses.remove(0)) - } - } - - fn test_config(tmp: &TempDir) -> Config { - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().join("workspace"); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - cfg - } - - fn seed_content(content_root: &Path) { - let raw_dir = content_root.join("raw").join("test-source").join("commits"); - std::fs::create_dir_all(&raw_dir).unwrap(); - std::fs::write( - raw_dir.join("123_abc.md"), - "---\nsource_kind: document\n---\n# Test Commit\nFixed the login bug in auth module.\n", - ) - .unwrap(); - - let doc_dir = content_root.join("document").join("test-doc"); - std::fs::create_dir_all(&doc_dir).unwrap(); - std::fs::write( - doc_dir.join("readme.md"), - "---\nsource_kind: document\n---\n# README\nProject documentation for the auth system.\n", - ) - .unwrap(); - - let wiki_dir = content_root - .join("wiki") - .join("summaries") - .join("source-test"); - std::fs::create_dir_all(wiki_dir.join("L1")).unwrap(); - std::fs::write( - wiki_dir.join("L1").join("summary-001.md"), - "---\nkind: summary\nlevel: 1\n---\nSummary of auth changes in May 2026.\n", - ) - .unwrap(); - } - - #[tokio::test] - async fn smart_walk_keyword_search_and_answer() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_content(&content_root); - - let provider = StubProvider::new(vec![ - // Turn 1: keyword search for "login" - r#"{"name":"keyword_search","arguments":{"pattern":"login","content_type":"all"}}"#, - // Turn 2: read the matching file - r#"{"name":"read_content","arguments":{"path":"raw/test-source/commits/123_abc.md"}}"#, - // Turn 3: collect evidence and answer - r#"{"name":"collect_evidence","arguments":{"items":[{"source":"raw/test-source/commits/123_abc.md","snippet":"Fixed the login bug in auth module.","relevance":"directly mentions login fix"}]}} -{"name":"answer","arguments":{"text":"The login bug was fixed in the auth module, as documented in commit 123_abc."}}"#, - ]); - - let opts = SmartWalkOptions { - max_turns: 10, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk(&cfg, &provider, "What happened with the login bug?", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::Answered); - assert!(outcome.answer.contains("login")); - assert_eq!(outcome.evidence.len(), 1); - assert!(outcome.evidence[0].snippet.contains("login bug")); - } - - #[tokio::test] - async fn smart_walk_list_sources() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_content(&content_root); - - let provider = StubProvider::new(vec![ - // Turn 1: list sources - r#"{"name":"list_sources","arguments":{"content_type":"all"}}"#, - // Turn 2: answer - r#"{"name":"answer","arguments":{"text":"Found raw, document, and wiki content."}}"#, - ]); - - let opts = SmartWalkOptions { - max_turns: 5, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk(&cfg, &provider, "What sources are available?", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::Answered); - assert!(outcome.answer.contains("raw")); - } - - #[tokio::test] - async fn smart_walk_max_turns() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_content(&content_root); - - let provider = StubProvider::new(vec![ - r#"{"name":"list_sources","arguments":{"content_type":"all"}}"#, - r#"{"name":"list_sources","arguments":{"content_type":"raw"}}"#, - r#"{"name":"list_sources","arguments":{"content_type":"wiki"}}"#, - ]); - - let opts = SmartWalkOptions { - max_turns: 3, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk(&cfg, &provider, "loop test", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::MaxTurnsReached); - assert_eq!(outcome.turns_used, 3); - } - - #[test] - fn parse_multiple_tool_calls() { - let response = r#"Let me search. -{"name":"keyword_search","arguments":{"pattern":"test"}} -{"name":"entity_search","arguments":{"query":"Alice"}}"#; - - let (text, calls) = parse_tool_calls(response); - assert_eq!(calls.len(), 2); - assert_eq!(calls[0].name, "keyword_search"); - assert_eq!(calls[1].name, "entity_search"); - assert!(text.contains("Let me search")); - } - - #[test] - fn content_inventory_counts_files() { - let tmp = TempDir::new().unwrap(); - let content_root = tmp.path().join("content"); - seed_content(&content_root); - - let inventory = build_content_inventory(&content_root); - assert!(inventory.contains("Raw content")); - assert!(inventory.contains("Documents")); - assert!(inventory.contains("Wiki summaries")); - } - - // ── Staging integration tests (run with --ignored) ──────────────── - - fn staging_content_root() -> Option { - let path = std::path::PathBuf::from( - "/Users/enamakel/.openhuman-staging/users/69d9cb73e61f755583c3671f/workspace/memory_tree/content", - ); - if path.exists() { - Some(path) - } else { - None - } - } - - #[test] - #[ignore] - fn staging_keyword_search_finds_steven() { - let content_root = staging_content_root().expect("staging content not available"); - let mut results = Vec::new(); - search_dir_recursive( - &content_root.join("raw"), - "steven", - &mut results, - &content_root, - ); - println!("keyword 'steven': {} results", results.len()); - for r in results.iter().take(5) { - println!(" {}", r); - } - assert!( - !results.is_empty(), - "should find 'steven' in staging raw content" - ); - } - - #[test] - #[ignore] - fn staging_content_inventory() { - let content_root = staging_content_root().expect("staging content not available"); - let inventory = build_content_inventory(&content_root); - println!("Inventory:\n{}", inventory); - assert!(inventory.contains("Raw content")); - assert!(inventory.contains("Documents")); - } - - #[test] - #[ignore] - fn staging_list_sources_shows_github() { - let content_root = staging_content_root().expect("staging content not available"); - let call = InnerCall { - name: "list_sources".into(), - args: serde_json::json!({"content_type": "all"}), - }; - let (_, result, _, _) = dispatch_list_sources(&content_root, &call); - println!("list_sources:\n{}", result); - assert!(result.contains("raw/"), "should list raw sources"); - } - - #[test] - #[ignore] - fn staging_read_wiki_summary() { - let content_root = staging_content_root().expect("staging content not available"); - let wiki_dir = content_root.join("wiki").join("summaries"); - if !wiki_dir.exists() { - println!("no wiki summaries found — skipping"); - return; - } - // Find first summary file - let first = walkdir_first_md(&wiki_dir); - if let Some(path) = first { - let rel = path - .strip_prefix(&content_root) - .unwrap() - .to_string_lossy() - .to_string(); - println!("Reading wiki: {}", rel); - let call = InnerCall { - name: "read_content".into(), - args: serde_json::json!({"path": rel}), - }; - let (_, result, _, _) = dispatch_read_content(&content_root, &call); - println!("Content preview: {}", &result[..result.len().min(300)]); - assert!( - !result.starts_with("error"), - "should read wiki file without error" - ); - } - } - - #[test] - #[ignore] - fn staging_read_episodic_memory() { - let content_root = staging_content_root().expect("staging content not available"); - let ep_dir = content_root.join("episodic"); - if !ep_dir.exists() { - println!("no episodic memories — skipping"); - return; - } - let first = walkdir_first_md(&ep_dir); - if let Some(path) = first { - let rel = path - .strip_prefix(&content_root) - .unwrap() - .to_string_lossy() - .to_string(); - println!("Reading episodic: {}", rel); - let call = InnerCall { - name: "read_content".into(), - args: serde_json::json!({"path": rel}), - }; - let (_, result, _, _) = dispatch_read_content(&content_root, &call); - println!("Content preview: {}", &result[..result.len().min(300)]); - assert!( - !result.starts_with("error"), - "should read episodic file without error" - ); - } - } - - #[test] - #[ignore] - fn staging_full_smart_walk_keyword_pipeline() { - let content_root = staging_content_root().expect("staging content not available"); - - // Simulate the pipeline: list_sources → keyword_search → read_content - let call = InnerCall { - name: "list_sources".into(), - args: serde_json::json!({"content_type": "raw"}), - }; - let (_, sources, _, _) = dispatch_list_sources(&content_root, &call); - println!("Step 1 - Sources:\n{}", sources); - - let call = InnerCall { - name: "keyword_search".into(), - args: serde_json::json!({"pattern": "memory", "content_type": "all"}), - }; - let (_, search_result, _, _) = dispatch_keyword_search(&content_root, &call); - println!("Step 2 - Search 'memory':\n{}", search_result); - - if search_result.contains('[') { - // Extract first file path from results - if let Some(path_start) = search_result.find('[') { - if let Some(path_end) = search_result[path_start + 1..].find(']') { - let file_path = &search_result[path_start + 1..path_start + 1 + path_end]; - println!("Step 3 - Reading: {}", file_path); - let call = InnerCall { - name: "read_content".into(), - args: serde_json::json!({"path": file_path}), - }; - let (_, content, _, _) = dispatch_read_content(&content_root, &call); - println!( - "Step 3 - Content ({} chars): {}", - content.len(), - &content[..content.len().min(200)] - ); - assert!( - !content.starts_with("error"), - "pipeline should complete without errors" - ); - } - } - } - } - - fn walkdir_first_md(dir: &std::path::Path) -> Option { - fn recurse(dir: &std::path::Path) -> Option { - for entry in std::fs::read_dir(dir).ok()?.flatten() { - let path = entry.path(); - if path.is_dir() { - if let Some(found) = recurse(&path) { - return Some(found); - } - } else if path.extension().map_or(false, |e| e == "md") { - return Some(path); - } - } - None - } - recurse(dir) - } - - fn seed_synced_memory(content_root: &Path) { - // Raw email content - let email_dir = content_root.join("raw").join("email").join("inbox"); - std::fs::create_dir_all(&email_dir).unwrap(); - std::fs::write( - email_dir.join("001_meeting.md"), - "---\nsource_kind: email\nauthor: alice@example.com\ndate: 2026-06-01\n---\n\ - # Team standup notes\n\n\ - Action items:\n\ - - Deploy the auth service refactor by Friday\n\ - - Review PR #342 for the billing module\n\ - - Schedule security audit with external team\n", - ) - .unwrap(); - std::fs::write( - email_dir.join("002_project.md"), - "---\nsource_kind: email\nauthor: bob@example.com\ndate: 2026-06-02\n---\n\ - # Project Phoenix status update\n\n\ - The migration is 80% complete. Remaining:\n\ - - Database schema changes (blocked on DBA review)\n\ - - API versioning for backward compatibility\n\ - - Load testing the new endpoints\n", - ) - .unwrap(); - std::fs::write( - email_dir.join("003_personal.md"), - "---\nsource_kind: email\nauthor: carol@example.com\ndate: 2026-06-03\n---\n\ - # Lunch plans\n\n\ - Hey, want to grab sushi on Thursday? The new place on 5th street \ - got great reviews.\n", - ) - .unwrap(); - - // Episodic memories - let ep_dir = content_root.join("episodic").join("daily"); - std::fs::create_dir_all(&ep_dir).unwrap(); - std::fs::write( - ep_dir.join("2026-06-01.md"), - "---\nkind: episodic\ndate: 2026-06-01\n---\n\ - Worked on the auth service refactor. Had a productive standup.\n\ - Identified three blockers for Project Phoenix.\n", - ) - .unwrap(); - - // Wiki summaries - let wiki_dir = content_root - .join("wiki") - .join("summaries") - .join("email-inbox"); - std::fs::create_dir_all(wiki_dir.join("L1")).unwrap(); - std::fs::write( - wiki_dir.join("L1").join("summary-week-22.md"), - "---\nkind: summary\nlevel: 1\n---\n\ - Week 22 summary: Team focused on Project Phoenix migration \ - and auth service refactor. Key contacts: alice@example.com (standup), \ - bob@example.com (project status), carol@example.com (social).\n", - ) - .unwrap(); - - // Document content - let doc_dir = content_root.join("document").join("notes"); - std::fs::create_dir_all(&doc_dir).unwrap(); - std::fs::write( - doc_dir.join("project-phoenix.md"), - "---\nsource_kind: document\n---\n\ - # Project Phoenix\n\n\ - ## Overview\n\ - Migration from legacy monolith to microservices.\n\n\ - ## Status\n\ - Phase 2 of 3 — data migration and API versioning.\n\n\ - ## Key risks\n\ - - Data integrity during cutover\n\ - - Backward compatibility for mobile clients\n", - ) - .unwrap(); - } - - #[tokio::test] - async fn walk_synced_email_with_keyword_and_evidence() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_synced_memory(&content_root); - - let provider = StubProvider::new(vec![ - // Turn 1: keyword search - r#"{"name":"keyword_search","arguments":{"pattern":"auth service","content_type":"all"}}"#, - // Turn 2: collect evidence - concat!( - r#"{"name":"collect_evidence","arguments":{"items":["#, - r#"{"source":"raw/email/inbox/001_meeting.md","snippet":"Deploy the auth service refactor by Friday","relevance":"action item"}"#, - r#"]}}"#, - ), - // Turn 3: answer - r#"{"name":"answer","arguments":{"text":"The auth service refactor needs to be deployed by Friday."}}"#, - ]); - - let opts = SmartWalkOptions { - max_turns: 5, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk( - &cfg, - &provider, - "What's happening with the auth service?", - opts, - ) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::Answered); - assert!(outcome.answer.contains("auth service")); - assert!(!outcome.evidence.is_empty()); - } - - #[tokio::test] - async fn walk_with_xml_format_tool_calls() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_synced_memory(&content_root); - - let provider = StubProvider::new(vec![ - // Turn 1: XML-formatted tool call - concat!( - "", - "keyword_search", - "{\"pattern\": \"project phoenix\", \"content_type\": \"all\"}", - "", - ), - // Turn 2: JSON-formatted answer - r#"{"name":"answer","arguments":{"text":"Project Phoenix is in phase 2 of 3."}}"#, - ]); - - let opts = SmartWalkOptions { - max_turns: 5, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk( - &cfg, - &provider, - "What's the status of Project Phoenix?", - opts, - ) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::Answered); - assert!(outcome.answer.contains("phase 2")); - } - - #[tokio::test] - async fn walk_reads_across_content_types() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_synced_memory(&content_root); - - let provider = StubProvider::new(vec![ - // Turn 1: list sources - r#"{"name":"list_sources","arguments":{"content_type":"all"}}"#, - // Turn 2: read document - r#"{"name":"read_content","arguments":{"path":"document/notes/project-phoenix.md"}}"#, - // Turn 3: read episodic - r#"{"name":"read_content","arguments":{"path":"episodic/daily/2026-06-01.md"}}"#, - // Turn 4: collect + answer - concat!( - r#"{"name":"collect_evidence","arguments":{"items":["#, - r#"{"source":"document/notes/project-phoenix.md","snippet":"Phase 2 of 3","relevance":"status"},"#, - r#"{"source":"episodic/daily/2026-06-01.md","snippet":"Identified three blockers","relevance":"context"}"#, - r#"]}}"#, - r#"{"name":"answer","arguments":{"text":"Project Phoenix: Phase 2/3 with 3 blockers identified."}}"#, - ), - ]); - - let opts = SmartWalkOptions { - max_turns: 5, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk(&cfg, &provider, "Summarize Project Phoenix status", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::Answered); - assert_eq!(outcome.evidence.len(), 2); - } - - #[tokio::test] - async fn walk_llm_gives_up_uses_fallback() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_synced_memory(&content_root); - - let provider = StubProvider::new(vec![ - // Turn 1: search finds nothing - r#"{"name":"keyword_search","arguments":{"pattern":"quantum computing","content_type":"all"}}"#, - // Turn 2: LLM gives up with empty response - "", - ]); - - let opts = SmartWalkOptions { - max_turns: 5, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk(&cfg, &provider, "Tell me about quantum computing", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::LlmGaveUp); - assert!(outcome.evidence.is_empty()); - assert!(outcome.answer.contains("Could not converge")); - } - - #[tokio::test] - async fn walk_direct_answer_without_tools() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_synced_memory(&content_root); - - let provider = StubProvider::new(vec![ - // LLM directly answers without using any tools - "I don't have enough context to answer that question from your memory.", - ]); - - let opts = SmartWalkOptions { - max_turns: 5, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk(&cfg, &provider, "What's the meaning of life?", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::Answered); - assert!(outcome.answer.contains("don't have enough context")); - assert_eq!(outcome.turns_used, 1); - assert!(outcome.evidence.is_empty()); - } - - #[tokio::test] - async fn walk_collect_evidence_deduplicates_within_limit() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let content_root = cfg.workspace_dir.join("memory_tree").join("content"); - seed_synced_memory(&content_root); - - let provider = StubProvider::new(vec![ - // Turn 1: collect a batch of evidence - concat!( - r#"{"name":"collect_evidence","arguments":{"items":["#, - r#"{"source":"raw/email/inbox/001_meeting.md","snippet":"Deploy auth","relevance":"task"},"#, - r#"{"source":"raw/email/inbox/002_project.md","snippet":"Migration 80%","relevance":"status"}"#, - r#"]}}"#, - ), - // Turn 2: collect more evidence (including a duplicate of the first source) - concat!( - r#"{"name":"collect_evidence","arguments":{"items":["#, - r#"{"source":"document/notes/project-phoenix.md","snippet":"Phase 2 of 3","relevance":"doc"},"#, - r#"{"source":"raw/email/inbox/001_meeting.md","snippet":"Deploy auth (duplicate)","relevance":"task"}"#, - r#"]}}"#, - ), - // Turn 3: answer - r#"{"name":"answer","arguments":{"text":"Summary with evidence items including duplicate source."}}"#, - ]); - - let opts = SmartWalkOptions { - max_turns: 10, - namespace: "default".into(), - model: Some("test-model".into()), - content_root: Some(content_root), - }; - - let outcome = run_smart_walk(&cfg, &provider, "Summarize everything", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, SmartWalkStopReason::Answered); - // 2 items from turn 1 + 2 items from turn 2 (one of which duplicates a turn-1 source); - // collect_evidence does not deduplicate, so all 4 items are present. - assert_eq!(outcome.evidence.len(), 4); - } -} diff --git a/src/openhuman/memory/query/smart_walk/tool.rs b/src/openhuman/memory/query/smart_walk/tool.rs deleted file mode 100644 index f2ecfd7a9..000000000 --- a/src/openhuman/memory/query/smart_walk/tool.rs +++ /dev/null @@ -1,188 +0,0 @@ -//! `SmartMemoryWalkTool` — the agent-facing tool wrapper, plus the -//! `ChatProviderAdapter` that bridges the memory chat provider to the -//! inference `Provider` trait. - -use crate::openhuman::config::rpc as config_rpc; -use crate::openhuman::inference::provider::traits::{ChatMessage, Provider}; -use crate::openhuman::memory::chat::{build_chat_provider, ChatPrompt}; -use crate::openhuman::memory::query::smart_walk::runner::run_smart_walk; -use crate::openhuman::memory::query::smart_walk::types::{ - truncate_chars, SmartWalkOptions, HARD_MAX_TURNS, -}; -use crate::openhuman::tools::traits::{PermissionLevel, Tool, ToolCategory, ToolResult}; -use async_trait::async_trait; -use serde_json::json; - -// ── Tool ───────────────────────────────────────────────────────────────────── - -pub struct SmartMemoryWalkTool; - -#[async_trait] -impl Tool for SmartMemoryWalkTool { - fn name(&self) -> &str { - "memory_smart_walk" - } - - fn description(&self) -> &str { - "Smart memory retrieval — combines vector search, keyword search, \ - entity lookup, and tree browsing to answer queries about the user's \ - memory. More capable than the basic walk: searches across raw files, \ - wiki summaries, documents, and episodic memories." - } - - fn parameters_schema(&self) -> serde_json::Value { - json!({ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Natural-language question to answer by searching memory." - }, - "namespace": { - "type": "string", - "description": "Memory namespace. Default: \"default\"." - }, - "max_turns": { - "type": "integer", - "description": "Max LLM turns. Default 12, hard cap 25." - }, - "model": { - "type": "string", - "description": "Provider:model override (e.g. 'deepseek:deepseek-chat')." - } - }, - "required": ["query"] - }) - } - - fn category(&self) -> ToolCategory { - ToolCategory::System - } - - fn permission_level(&self) -> PermissionLevel { - PermissionLevel::ReadOnly - } - - fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { - true - } - - async fn execute(&self, args: serde_json::Value) -> anyhow::Result { - let query = args - .get("query") - .and_then(|v| v.as_str()) - .ok_or_else(|| anyhow::anyhow!("memory_smart_walk: `query` is required"))? - .to_string(); - - let namespace = args - .get("namespace") - .and_then(|v| v.as_str()) - .unwrap_or("default") - .to_string(); - - let max_turns = args - .get("max_turns") - .and_then(|v| v.as_u64()) - .map(|n| (n as usize).min(HARD_MAX_TURNS)) - .unwrap_or(12); - - let model = args - .get("model") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - let cfg = config_rpc::load_config_with_timeout() - .await - .map_err(|e| anyhow::anyhow!("memory_smart_walk: load config failed: {e}"))?; - - let opts = SmartWalkOptions { - max_turns, - namespace, - model, - content_root: None, - }; - - let chat_provider = build_chat_provider(&cfg) - .map_err(|e| anyhow::anyhow!("memory_smart_walk: build chat provider failed: {e}"))?; - let adapter = ChatProviderAdapter { - inner: chat_provider, - }; - - let outcome = run_smart_walk(&cfg, &adapter, &query, opts).await?; - - let mut out = format!("{}\n", outcome.answer); - - if !outcome.evidence.is_empty() { - out.push_str("\n## Evidence\n"); - for (i, ev) in outcome.evidence.iter().enumerate() { - out.push_str(&format!( - "{}. **{}** — {}\n > {}\n", - i + 1, - ev.source_path, - ev.relevance, - truncate_chars(&ev.snippet, 200) - )); - } - } - - out.push_str("\n## Trace\n"); - for step in &outcome.trace { - out.push_str(&format!( - "- **Turn {}** `{}` {}: {}\n", - step.turn, step.action, step.args_summary, step.result_preview - )); - } - out.push_str(&format!( - "\n*Stop reason: {:?}, turns used: {}*\n", - outcome.stopped_reason, outcome.turns_used - )); - - Ok(ToolResult::success(out)) - } -} - -// ── ChatProviderAdapter ─────────────────────────────────────────────────────── - -pub(crate) struct ChatProviderAdapter { - pub(crate) inner: std::sync::Arc, -} - -#[async_trait] -impl Provider for ChatProviderAdapter { - async fn chat_with_system( - &self, - system: Option<&str>, - message: &str, - _model: &str, - temperature: f64, - ) -> anyhow::Result { - let prompt = ChatPrompt { - system: system.unwrap_or("").to_string(), - user: message.to_string(), - temperature, - kind: "memory_smart_walk", - max_tokens: None, - }; - self.inner.chat_for_text(&prompt).await - } - - async fn chat_with_history( - &self, - messages: &[ChatMessage], - model: &str, - temperature: f64, - ) -> anyhow::Result { - let system = messages - .iter() - .find(|m| m.role == "system") - .map(|m| m.content.as_str()); - let user: String = messages - .iter() - .filter(|m| m.role != "system") - .map(|m| m.content.as_str()) - .collect::>() - .join("\n"); - self.chat_with_system(system, &user, model, temperature) - .await - } -} diff --git a/src/openhuman/memory/query/smart_walk/types.rs b/src/openhuman/memory/query/smart_walk/types.rs deleted file mode 100644 index 94e55109b..000000000 --- a/src/openhuman/memory/query/smart_walk/types.rs +++ /dev/null @@ -1,66 +0,0 @@ -//! Public output types and shared constants for smart_walk. - -pub(crate) const SMART_WALK_TEMP: f64 = 0.2; -pub(crate) const HARD_MAX_TURNS: usize = 25; -pub(crate) const MAX_EVIDENCE_ITEMS: usize = 30; -pub(crate) const MAX_KEYWORD_RESULTS: usize = 15; -pub(crate) const MAX_FILE_READ_BYTES: usize = 8000; - -pub(crate) fn truncate_chars(value: &str, max_chars: usize) -> String { - value.chars().take(max_chars).collect() -} - -// ── Public output types ───────────────────────────────────────────────────── - -#[derive(Debug, Clone)] -pub struct SmartWalkOptions { - pub max_turns: usize, - pub namespace: String, - /// Provider string override (e.g. "deepseek:deepseek-chat"). - pub model: Option, - /// Content root override. Defaults to config.memory_tree_content_root(). - pub content_root: Option, -} - -impl Default for SmartWalkOptions { - fn default() -> Self { - Self { - max_turns: 12, - namespace: "default".into(), - model: None, - content_root: None, - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum SmartWalkStopReason { - Answered, - MaxTurnsReached, - LlmGaveUp, - Error(String), -} - -#[derive(Debug, Clone)] -pub struct SmartWalkStep { - pub turn: usize, - pub action: String, - pub args_summary: String, - pub result_preview: String, -} - -#[derive(Debug, Clone)] -pub struct Evidence { - pub source_path: String, - pub snippet: String, - pub relevance: String, -} - -#[derive(Debug, Clone)] -pub struct SmartWalkOutcome { - pub answer: String, - pub evidence: Vec, - pub trace: Vec, - pub turns_used: usize, - pub stopped_reason: SmartWalkStopReason, -} diff --git a/src/openhuman/memory/query/walk.rs b/src/openhuman/memory/query/walk.rs deleted file mode 100644 index 7ce1a8274..000000000 --- a/src/openhuman/memory/query/walk.rs +++ /dev/null @@ -1,966 +0,0 @@ -//! Agentic memory-tree walk tool. -//! -//! Given a free-text query, a lightweight LLM navigates the summary tree in -//! a turn-based inner loop — calling `descend`, `peek`, `fetch_leaves`, or -//! `answer` each turn — and returns a synthesised answer with a trace. -//! -//! The inner loop uses `Provider::chat_with_history` (prompt-guided tool -//! calling via XML tags) because the `Provider::chat()` default does not -//! surface native `tool_calls` for prompt-guided backends. The response text -//! is parsed for `…` blocks, matching the harness -//! convention established in `agent/harness/parse.rs`. -//! -//! For the `Tool::execute` path, a thin `ChatProviderAdapter` wraps the -//! memory-tree's `ChatProvider` (available from `build_chat_provider`) to -//! satisfy the `Provider` trait — avoiding a dependency on the full routing -//! stack which requires a configured remote backend. - -use crate::openhuman::config::rpc as config_rpc; -use crate::openhuman::config::Config; -use crate::openhuman::inference::provider::traits::{ChatMessage, Provider}; -use crate::openhuman::memory::chat::{build_chat_provider, ChatPrompt}; -use crate::openhuman::memory_tree::retrieval; -use crate::openhuman::memory_tree::retrieval::fetch::fetch_leaves as do_fetch_leaves; -use crate::openhuman::memory_tree::tree_runtime::store::{read_children, read_node}; -use crate::openhuman::tools::traits::{PermissionLevel, Tool, ToolCategory, ToolResult}; -use async_trait::async_trait; -use serde_json::json; - -// ── Temperature (matches SUMMARIZATION_TEMP convention) ──────────────────── -const WALK_TEMP: f64 = 0.3; -/// Hard cap on LLM turns, even if the caller requests more. -const HARD_MAX_TURNS: usize = 20; - -// ── Public output types ───────────────────────────────────────────────────── - -#[derive(Debug, Clone)] -pub struct WalkOptions { - /// Maximum number of LLM turns before giving up. Default: 6. - pub max_turns: usize, - /// Node id to start from. `None` → namespace root. - pub start_node_id: Option, - /// Memory namespace. Default: `"default"`. - pub namespace: String, - /// Model override. `None` → `config.local_ai.chat_model_id`. - pub model: Option, -} - -impl Default for WalkOptions { - fn default() -> Self { - Self { - max_turns: 6, - start_node_id: None, - namespace: "default".into(), - model: None, - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum WalkStopReason { - /// LLM called `answer { text }`. - Answered, - /// Loop exhausted `max_turns` without an answer action. - MaxTurnsReached, - /// LLM returned no tool call and no meaningful text — treated as giving up. - LlmGaveUp, - /// A hard error prevented the walk from completing. - Error(String), -} - -#[derive(Debug, Clone)] -pub struct WalkStep { - pub turn: usize, - pub action: String, - pub args_summary: String, - pub result_preview: String, -} - -#[derive(Debug, Clone)] -pub struct WalkOutcome { - pub answer: String, - pub trace: Vec, - pub turns_used: usize, - pub stopped_reason: WalkStopReason, -} - -// ── Public API ────────────────────────────────────────────────────────────── - -pub struct MemoryTreeWalkTool; - -#[async_trait] -impl Tool for MemoryTreeWalkTool { - fn name(&self) -> &str { - "memory_tree_walk" - } - - fn description(&self) -> &str { - "Agentically walk the memory tree to answer a query — a lightweight \ - LLM navigates summaries, drills into relevant branches, and returns \ - a synthesized answer with citations." - } - - fn parameters_schema(&self) -> serde_json::Value { - json!({ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Natural-language question to answer by walking the memory tree." - }, - "namespace": { - "type": "string", - "description": "Memory namespace. Default: \"default\"." - }, - "start_node_id": { - "type": "string", - "description": "Optional starting node id. Default: namespace root." - }, - "max_turns": { - "type": "integer", - "description": "Max LLM turns. Default 6, hard cap 20." - } - }, - "required": ["query"] - }) - } - - fn category(&self) -> ToolCategory { - ToolCategory::System - } - - fn permission_level(&self) -> PermissionLevel { - PermissionLevel::ReadOnly - } - - fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { - true - } - - async fn execute(&self, args: serde_json::Value) -> anyhow::Result { - let query = args - .get("query") - .and_then(|v| v.as_str()) - .ok_or_else(|| anyhow::anyhow!("memory_tree_walk: `query` is required"))? - .to_string(); - - let namespace = args - .get("namespace") - .and_then(|v| v.as_str()) - .unwrap_or("default") - .to_string(); - - let start_node_id = args - .get("start_node_id") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - let max_turns = args - .get("max_turns") - .and_then(|v| v.as_u64()) - .map(|n| (n as usize).min(HARD_MAX_TURNS)) - .unwrap_or(6); - - let cfg = config_rpc::load_config_with_timeout() - .await - .map_err(|e| anyhow::anyhow!("memory_tree_walk: load config failed: {e}"))?; - - let model = args - .get("model") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - let opts = WalkOptions { - max_turns, - start_node_id, - namespace, - model, - }; - - // Build a chat provider from config (same path used by the summariser) - // and wrap it in the thin `ChatProviderAdapter` that satisfies `Provider`. - let chat_provider = build_chat_provider(&cfg) - .map_err(|e| anyhow::anyhow!("memory_tree_walk: build chat provider failed: {e}"))?; - let adapter = ChatProviderAdapter { - inner: chat_provider, - }; - - let outcome = run_walk(&cfg, &adapter, &query, opts).await?; - - // Format output as markdown with trace. - let mut out = format!("{}\n\n## Trace\n", outcome.answer); - for step in &outcome.trace { - out.push_str(&format!( - "- **Turn {}** `{}` {}: {}\n", - step.turn, step.action, step.args_summary, step.result_preview - )); - } - out.push_str(&format!( - "\n*Stop reason: {:?}, turns used: {}*\n", - outcome.stopped_reason, outcome.turns_used - )); - - Ok(ToolResult::success(out)) - } -} - -/// Drive the walk without going through the Tool trait. -/// Useful for tests and callers that already hold a `Config` and `Provider`. -pub async fn run_walk( - config: &Config, - provider: &dyn Provider, - query: &str, - opts: WalkOptions, -) -> anyhow::Result { - let max_turns = opts.max_turns.min(HARD_MAX_TURNS); - let model = opts - .model - .clone() - .unwrap_or_else(|| config.local_ai.chat_model_id.clone()); - - log::debug!( - "[memory_tree_walk] starting walk query_len={} namespace={} max_turns={} model={}", - query.len(), - opts.namespace, - max_turns, - model - ); - - // Determine the starting node. - let start_id = opts - .start_node_id - .clone() - .unwrap_or_else(|| "root".to_string()); - - // Load the starting node summary + children to build the first context message. - let initial_context = build_node_context(config, &opts.namespace, &start_id).await; - log::debug!( - "[memory_tree_walk] initial_context node_id={} context_len={}", - start_id, - initial_context.len() - ); - - let system = build_system_prompt(); - let inner_tools_text = build_inner_tools_text(); - - // Chat history: system → tool instructions injected → user query → context. - let mut history: Vec = vec![ - ChatMessage::system(format!("{system}\n\n{inner_tools_text}")), - ChatMessage::user(format!( - "Query: {query}\n\nCurrent position in memory tree:\n{initial_context}" - )), - ]; - - let mut trace: Vec = Vec::new(); - let mut current_node_id = start_id.clone(); - - for turn in 1..=max_turns { - log::debug!("[memory_tree_walk] turn={turn} current_node={current_node_id}"); - - let response = match provider - .chat_with_history(&history, &model, WALK_TEMP) - .await - { - Ok(r) => r, - Err(e) => { - log::warn!("[memory_tree_walk] provider error on turn={turn}: {e:#}"); - let err_msg = format!("Provider error on turn {turn}: {e}"); - return Ok(WalkOutcome { - answer: format!( - "Walk failed: {err_msg}\n\nPartial trace from {} turn(s).", - trace.len() - ), - trace, - turns_used: turn, - stopped_reason: WalkStopReason::Error(err_msg), - }); - } - }; - - log::debug!( - "[memory_tree_walk] turn={turn} response_len={}", - response.len() - ); - - // Parse tool calls from the response text. - let (text_before, calls) = parse_walk_tool_calls(&response); - - if calls.is_empty() { - // No tool call — treat as final answer if there's meaningful text. - let trimmed = response.trim().to_string(); - if trimmed.is_empty() { - log::debug!("[memory_tree_walk] turn={turn} LLM gave up (empty response)"); - return Ok(WalkOutcome { - answer: synthesize_fallback_answer(&trace), - trace, - turns_used: turn, - stopped_reason: WalkStopReason::LlmGaveUp, - }); - } - log::debug!("[memory_tree_walk] turn={turn} no tool calls — treating as final answer"); - return Ok(WalkOutcome { - answer: trimmed, - trace, - turns_used: turn, - stopped_reason: WalkStopReason::Answered, - }); - } - - // Process the first tool call (walk is serial). - let call = &calls[0]; - log::debug!( - "[memory_tree_walk] turn={turn} action={} args={}", - call.name, - call.args - ); - - // Append assistant turn to history. - history.push(ChatMessage::assistant(response.clone())); - - // Dispatch inner walk primitive. - let (step_args_summary, tool_result, is_answer, answer_text) = - dispatch_inner_call(config, &opts.namespace, call, &mut current_node_id).await; - - let result_preview: String = tool_result.chars().take(200).collect(); - trace.push(WalkStep { - turn, - action: call.name.clone(), - args_summary: step_args_summary, - result_preview: result_preview.clone(), - }); - - if is_answer { - log::debug!("[memory_tree_walk] turn={turn} answer action — stopping"); - return Ok(WalkOutcome { - answer: answer_text, - trace, - turns_used: turn, - stopped_reason: WalkStopReason::Answered, - }); - } - - // Append tool result as user message (prompt-guided protocol). - let tool_msg = format!( - "{}\n\nCurrent position: {current_node_id}\n", - tool_result - ); - history.push(ChatMessage::user(tool_msg)); - - // Drop the preamble text if there was any (don't lose context). - if !text_before.trim().is_empty() { - log::debug!( - "[memory_tree_walk] turn={turn} text before tool call: {}", - &text_before[..text_before.len().min(80)] - ); - } - } - - // Max turns reached. - log::debug!("[memory_tree_walk] max_turns={max_turns} reached — synthesising fallback"); - Ok(WalkOutcome { - answer: synthesize_fallback_answer(&trace), - trace, - turns_used: max_turns, - stopped_reason: WalkStopReason::MaxTurnsReached, - }) -} - -// ── ChatProviderAdapter ───────────────────────────────────────────────────── -// -// Bridges the memory-tree's lightweight `ChatProvider` into the top-level -// `Provider` trait so `run_walk` can accept both production adapters and -// unit-test stubs that implement `Provider` directly. - -struct ChatProviderAdapter { - inner: std::sync::Arc, -} - -#[async_trait] -impl Provider for ChatProviderAdapter { - async fn chat_with_system( - &self, - system: Option<&str>, - message: &str, - _model: &str, - temperature: f64, - ) -> anyhow::Result { - let prompt = ChatPrompt { - system: system.unwrap_or("").to_string(), - user: message.to_string(), - temperature, - kind: "memory_tree_walk", - max_tokens: None, - }; - self.inner.chat_for_text(&prompt).await - } - - async fn chat_with_history( - &self, - messages: &[ChatMessage], - model: &str, - temperature: f64, - ) -> anyhow::Result { - let system = messages - .iter() - .find(|m| m.role == "system") - .map(|m| m.content.as_str()); - // Combine all non-system messages into the user turn. - let user: String = messages - .iter() - .filter(|m| m.role != "system") - .map(|m| m.content.as_str()) - .collect::>() - .join("\n"); - self.chat_with_system(system, &user, model, temperature) - .await - } -} - -// ── Inner helpers ─────────────────────────────────────────────────────────── - -/// A parsed tool call from the inner walk loop. -struct InnerCall { - name: String, - args: serde_json::Value, -} - -/// Parse `…` blocks from a response string. -/// Returns `(text_before_first_call, calls)`. -fn parse_walk_tool_calls(response: &str) -> (String, Vec) { - let mut calls: Vec = Vec::new(); - let mut text_parts: Vec<&str> = Vec::new(); - let mut remaining: &str = response; - - const OPEN: &str = ""; - const CLOSE: &str = ""; - - loop { - match remaining.find(OPEN) { - None => { - // No more tags; collect trailing text. - if !remaining.trim().is_empty() && calls.is_empty() { - text_parts.push(remaining); - } - break; - } - Some(start) => { - let before = &remaining[..start]; - if !before.trim().is_empty() { - text_parts.push(before); - } - let after_open = &remaining[start + OPEN.len()..]; - match after_open.find(CLOSE) { - None => break, // malformed — stop - Some(close_idx) => { - let inner = &after_open[..close_idx]; - if let Ok(val) = serde_json::from_str::(inner.trim()) { - if let Some(name) = val.get("name").and_then(|v| v.as_str()) { - let args = val - .get("arguments") - .cloned() - .unwrap_or(serde_json::Value::Object(Default::default())); - calls.push(InnerCall { - name: name.to_string(), - args, - }); - } - } - remaining = &after_open[close_idx + CLOSE.len()..]; - } - } - } - } - } - - let text_before = text_parts.concat(); - (text_before, calls) -} - -/// Dispatch an inner walk primitive and return -/// `(args_summary, result_text, is_final_answer, answer_text)`. -async fn dispatch_inner_call( - config: &Config, - namespace: &str, - call: &InnerCall, - current_node_id: &mut String, -) -> (String, String, bool, String) { - match call.name.as_str() { - "descend" => { - let node_id = call - .args - .get("node_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - log::debug!( - "[memory_tree_walk] descend node_id={node_id} from={current_node_id} namespace={namespace}" - ); - - if node_id.is_empty() { - return ( - "node_id=".into(), - "error: descend requires a non-empty node_id".into(), - false, - String::new(), - ); - } - - // Move to the target node. - let ctx = build_node_context(config, namespace, &node_id).await; - if ctx.starts_with("unknown node") { - ( - format!("node_id={node_id}"), - format!("unknown node: {node_id}"), - false, - String::new(), - ) - } else { - *current_node_id = node_id.clone(); - (format!("node_id={node_id}"), ctx, false, String::new()) - } - } - - "peek" => { - let node_ids: Vec = call - .args - .get("node_ids") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str()) - .map(|s| s.to_string()) - .collect() - }) - .unwrap_or_default(); - - log::debug!( - "[memory_tree_walk] peek node_ids={} namespace={namespace}", - node_ids.len() - ); - - let args_summary = format!("node_ids=[{}]", node_ids.join(", ")); - - let config_owned = config.clone(); - let ns_owned = namespace.to_string(); - let ids_owned = node_ids.clone(); - - let result = tokio::task::spawn_blocking(move || -> Vec { - ids_owned - .iter() - .map(|id| match read_node(&config_owned, &ns_owned, id) { - Ok(Some(node)) => { - format!( - "id={} level={:?} summary={}", - id, - node.level, - &node.summary[..node.summary.len().min(120)] - ) - } - Ok(None) => format!("id={id} unknown node"), - Err(e) => format!("id={id} error: {e}"), - }) - .collect() - }) - .await - .unwrap_or_default(); - - (args_summary, result.join("\n"), false, String::new()) - } - - "fetch_leaves" => { - let node_id = call - .args - .get("node_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - - log::debug!("[memory_tree_walk] fetch_leaves node_id={node_id} namespace={namespace}"); - - if node_id.is_empty() { - return ( - "node_id=".into(), - "error: fetch_leaves requires a non-empty node_id".into(), - false, - String::new(), - ); - } - - // fetch_leaves in retrieval takes a list of chunk ids. Reuse - // drill_down to get the leaf hits under this node, then return content. - let hits = match retrieval::drill_down(config, &node_id, 1, None, Some(10)).await { - Ok(h) => h, - Err(e) => { - return ( - format!("node_id={node_id}"), - format!("error fetching leaves: {e}"), - false, - String::new(), - ); - } - }; - - let text = if hits.is_empty() { - // Try to fetch the node itself as a leaf (chunk). - let chunk_ids = vec![node_id.clone()]; - match do_fetch_leaves(config, &chunk_ids).await { - Ok(leaf_hits) if !leaf_hits.is_empty() => leaf_hits - .iter() - .map(|h| format!("[{}] {}", h.node_id, h.content)) - .collect::>() - .join("\n---\n"), - _ => format!("no leaves found under node_id={node_id}"), - } - } else { - hits.iter() - .map(|h| format!("[{}] {}", h.node_id, h.content)) - .collect::>() - .join("\n---\n") - }; - - (format!("node_id={node_id}"), text, false, String::new()) - } - - "answer" => { - let text = call - .args - .get("text") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - log::debug!("[memory_tree_walk] answer action text_len={}", text.len()); - ("(final answer)".into(), text.clone(), true, text) - } - - other => { - log::warn!("[memory_tree_walk] unknown inner action: {other}"); - ( - format!("action={other}"), - format!( - "unknown walk action '{other}'. Valid: descend, peek, fetch_leaves, answer" - ), - false, - String::new(), - ) - } - } -} - -/// Build a human-readable context string for a node: its summary + children list. -async fn build_node_context(config: &Config, namespace: &str, node_id: &str) -> String { - let config_owned = config.clone(); - let ns_owned = namespace.to_string(); - let id_owned = node_id.to_string(); - - tokio::task::spawn_blocking(move || { - let node = match read_node(&config_owned, &ns_owned, &id_owned) { - Ok(Some(n)) => n, - Ok(None) => return format!("unknown node: {id_owned}"), - Err(e) => return format!("error reading node {id_owned}: {e}"), - }; - - let children = match read_children(&config_owned, &ns_owned, &id_owned) { - Ok(c) => c, - Err(_) => vec![], - }; - - let mut out = format!( - "Node: {} (level={:?})\nSummary: {}\n", - node.node_id, node.level, node.summary - ); - - if children.is_empty() { - out.push_str("Children: (none — this is a leaf)\n"); - } else { - out.push_str(&format!("Children ({}):\n", children.len())); - for c in &children { - out.push_str(&format!( - " - id={} level={:?} summary_preview={}\n", - c.node_id, - c.level, - &c.summary[..c.summary.len().min(80)] - )); - } - } - - out - }) - .await - .unwrap_or_else(|_| format!("error building context for node {node_id}")) -} - -fn build_system_prompt() -> String { - "You are a memory-tree navigator. Your job is to answer user queries \ - by walking a hierarchical summary tree.\n\ - Use the provided tools to navigate: `descend` to move into a child node, \ - `peek` to preview multiple children without descending, \ - `fetch_leaves` to retrieve raw content from a node, \ - and `answer` when you have enough information to respond.\n\ - Be efficient — prefer `peek` to survey options before `descend`.\n\ - Always end with `answer { \"text\": \"...\" }` when ready.\n\ - Use XML tool_call tags:\n\ - {\"name\": \"descend\", \"arguments\": {\"node_id\": \"some/id\"}}" - .into() -} - -fn build_inner_tools_text() -> String { - "## Inner walk tools\n\n\ - **descend** `{\"node_id\": \"\"}` — move to a child node and see its summary and children.\n\ - **peek** `{\"node_ids\": [\"\", \"\"]}` — preview summaries for a list of nodes without descending.\n\ - **fetch_leaves** `{\"node_id\": \"\"}` — retrieve raw chunk text under a node for citation.\n\ - **answer** `{\"text\": \"\"}` — stop and return your synthesised answer." - .into() -} - -fn synthesize_fallback_answer(trace: &[WalkStep]) -> String { - if trace.is_empty() { - return "Could not converge on an answer — no steps taken.".into(); - } - let preview: Vec = trace - .iter() - .map(|s| { - format!( - "Turn {}: {} → {}", - s.turn, - s.action, - &s.result_preview[..s.result_preview.len().min(100)] - ) - }) - .collect(); - format!( - "Could not converge on an answer within the turn limit. Here is what I saw:\n\n{}", - preview.join("\n") - ) -} - -// ── Tests ─────────────────────────────────────────────────────────────────── - -#[cfg(test)] -mod tests { - use super::*; - use crate::openhuman::config::Config; - use crate::openhuman::inference::provider::traits::ChatMessage; - use crate::openhuman::memory_tree::tree_runtime::store::write_node; - use crate::openhuman::memory_tree::tree_runtime::types::{NodeLevel, TreeNode}; - use async_trait::async_trait; - use chrono::Utc; - use std::sync::Mutex; - use tempfile::TempDir; - - // ── Stub provider ────────────────────────────────────────────────── - - /// A scripted stub provider that returns predefined responses in sequence. - /// Each `chat_with_history` call pops the next response from the queue. - struct StubProvider { - responses: Mutex>, - } - - impl StubProvider { - fn new(responses: Vec<&str>) -> Self { - Self { - responses: Mutex::new(responses.into_iter().map(|s| s.to_string()).collect()), - } - } - } - - #[async_trait] - impl Provider for StubProvider { - async fn chat_with_system( - &self, - _system: Option<&str>, - _message: &str, - _model: &str, - _temp: f64, - ) -> anyhow::Result { - let mut responses = self.responses.lock().unwrap(); - if responses.is_empty() { - return Err(anyhow::anyhow!("StubProvider: no more scripted responses")); - } - Ok(responses.remove(0)) - } - - async fn chat_with_history( - &self, - _messages: &[ChatMessage], - _model: &str, - _temp: f64, - ) -> anyhow::Result { - let mut responses = self.responses.lock().unwrap(); - if responses.is_empty() { - return Err(anyhow::anyhow!("StubProvider: no more scripted responses")); - } - Ok(responses.remove(0)) - } - } - - // ── Tree helpers ─────────────────────────────────────────────────── - - fn test_config(tmp: &TempDir) -> Config { - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().join("workspace"); - std::fs::create_dir_all(&cfg.workspace_dir).unwrap(); - cfg - } - - fn make_node(namespace: &str, node_id: &str, summary: &str, child_count: u32) -> TreeNode { - let level = crate::openhuman::memory_tree::tree_runtime::types::level_from_node_id(node_id); - let parent_id = - crate::openhuman::memory_tree::tree_runtime::types::derive_parent_id(node_id); - let ts = Utc::now(); - TreeNode { - node_id: node_id.to_string(), - namespace: namespace.to_string(), - level, - parent_id, - summary: summary.to_string(), - token_count: crate::openhuman::memory_tree::tree_runtime::types::estimate_tokens( - summary, - ), - child_count, - created_at: ts, - updated_at: ts, - metadata: None, - } - } - - /// Seed: root → 2024 (child of root) → 2024/01 (leaf). - fn seed_tree(cfg: &Config, ns: &str) { - write_node( - cfg, - &make_node(ns, "root", "All-time summary: project logs 2024", 1), - ) - .unwrap(); - write_node( - cfg, - &make_node(ns, "2024", "Year 2024: shipped v1, v2, v3", 1), - ) - .unwrap(); - write_node( - cfg, - &make_node(ns, "2024/01", "January 2024: initial project launch", 0), - ) - .unwrap(); - } - - // ── Test 1: walks_and_answers ────────────────────────────────────── - - /// Script: turn1=descend(2024), turn2=fetch_leaves(2024/01), turn3=answer. - #[tokio::test] - async fn walks_and_answers() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let ns = "default"; - seed_tree(&cfg, ns); - - let provider = StubProvider::new(vec![ - // Turn 1: descend into the year node. - r#"{"name":"descend","arguments":{"node_id":"2024"}}"#, - // Turn 2: fetch leaves under the month node. - r#"{"name":"fetch_leaves","arguments":{"node_id":"2024/01"}}"#, - // Turn 3: answer. - r#"{"name":"answer","arguments":{"text":"The project launched in January 2024."}}"#, - ]); - - let opts = WalkOptions { - max_turns: 6, - start_node_id: None, - namespace: ns.to_string(), - model: Some("test-model".into()), - }; - - let outcome = run_walk(&cfg, &provider, "When did the project launch?", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, WalkStopReason::Answered); - assert!( - outcome.answer.contains("January 2024"), - "answer should mention January 2024, got: {}", - outcome.answer - ); - assert_eq!(outcome.trace.len(), 3, "expected 3 steps"); - assert_eq!(outcome.trace[0].action, "descend"); - assert_eq!(outcome.trace[1].action, "fetch_leaves"); - assert_eq!(outcome.trace[2].action, "answer"); - assert_eq!(outcome.turns_used, 3); - } - - // ── Test 2: max_turns_cap ───────────────────────────────────────── - - /// Script: always `descend` in a loop — should stop at max_turns. - #[tokio::test] - async fn max_turns_cap() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let ns = "default"; - seed_tree(&cfg, ns); - - // Provide more `descend` responses than max_turns (3) so the cap fires. - let provider = StubProvider::new(vec![ - r#"{"name":"descend","arguments":{"node_id":"2024"}}"#, - r#"{"name":"descend","arguments":{"node_id":"2024"}}"#, - r#"{"name":"descend","arguments":{"node_id":"2024"}}"#, - r#"{"name":"descend","arguments":{"node_id":"2024"}}"#, - ]); - - let opts = WalkOptions { - max_turns: 3, - start_node_id: None, - namespace: ns.to_string(), - model: Some("test-model".into()), - }; - - let outcome = run_walk(&cfg, &provider, "infinite loop query", opts) - .await - .unwrap(); - - assert_eq!(outcome.stopped_reason, WalkStopReason::MaxTurnsReached); - assert_eq!(outcome.turns_used, 3); - assert!( - outcome.answer.contains("Could not converge"), - "expected fallback answer, got: {}", - outcome.answer - ); - } - - // ── Test 3: unknown_node_recovers ───────────────────────────────── - - /// Script: descend into a non-existent node → result says "unknown node" → loop continues. - #[tokio::test] - async fn unknown_node_recovers() { - let tmp = TempDir::new().unwrap(); - let cfg = test_config(&tmp); - let ns = "default"; - seed_tree(&cfg, ns); - - let provider = StubProvider::new(vec![ - // Turn 1: descend into a node that does not exist. - r#"{"name":"descend","arguments":{"node_id":"does_not_exist"}}"#, - // Turn 2: answer (loop continues after bad descend). - r#"{"name":"answer","arguments":{"text":"I could not find that node."}}"#, - ]); - - let opts = WalkOptions { - max_turns: 6, - start_node_id: None, - namespace: ns.to_string(), - model: Some("test-model".into()), - }; - - let outcome = run_walk(&cfg, &provider, "find nonexistent data", opts) - .await - .unwrap(); - - // The first trace step should indicate "unknown node". - assert_eq!(outcome.trace.len(), 2); - assert!( - outcome.trace[0].result_preview.contains("unknown node"), - "expected 'unknown node' in preview, got: {}", - outcome.trace[0].result_preview - ); - // The walk should continue and eventually answer. - assert_eq!(outcome.stopped_reason, WalkStopReason::Answered); - assert!(outcome.answer.contains("could not find that node")); - } -} diff --git a/src/openhuman/memory/schema/definitions.rs b/src/openhuman/memory/schema/definitions.rs index 4b37c1c1a..d13aa62a1 100644 --- a/src/openhuman/memory/schema/definitions.rs +++ b/src/openhuman/memory/schema/definitions.rs @@ -896,10 +896,11 @@ pub fn schemas(function: &str) -> ControllerSchema { "smart_walk" => ControllerSchema { namespace: NAMESPACE, function: "smart_walk", - description: "Multi-strategy memory retrieval — combines vector \ - search, keyword search, entity lookup, and tree browsing to \ - answer natural-language queries across raw files, wiki \ - summaries, documents, and episodic memories.", + description: "Deterministic E2GraphRAG memory retrieval — extracts \ + query entities (spaCy, with regex fallback), routes between \ + entity-graph (local) and dense-summary (global) search with no \ + LLM, and returns ranked evidence hits for a natural-language \ + query.", inputs: vec![ FieldSchema { name: "query", @@ -908,59 +909,42 @@ pub fn schemas(function: &str) -> ControllerSchema { required: true, }, FieldSchema { - name: "namespace", - ty: TypeSchema::String, - comment: "Memory namespace. Default: \"default\".", - required: false, - }, - FieldSchema { - name: "max_turns", + name: "limit", ty: TypeSchema::U64, - comment: "Max LLM turns. Default 12, hard cap 25.", + comment: "Max evidence hits to return. Default 10.", required: false, }, FieldSchema { - name: "model", - ty: TypeSchema::String, - comment: "Provider:model override (e.g. 'deepseek:deepseek-chat').", + name: "time_window_days", + ty: TypeSchema::U64, + comment: "Restrict the global/dense branch to the last N days.", + required: false, + }, + FieldSchema { + name: "max_hops", + ty: TypeSchema::U64, + comment: "Entity-graph relatedness hop threshold. Default 2.", required: false, }, ], outputs: vec![ FieldSchema { - name: "answer", - ty: TypeSchema::String, - comment: "Synthesized answer with evidence citations.", - required: true, - }, - FieldSchema { - name: "turns_used", - ty: TypeSchema::U64, - comment: "Number of LLM turns consumed.", - required: true, - }, - FieldSchema { - name: "evidence_count", - ty: TypeSchema::U64, - comment: "Number of evidence items collected.", - required: true, - }, - FieldSchema { - name: "stopped_reason", - ty: TypeSchema::String, - comment: "Why the walk stopped (answered/max_turns/llm_gave_up/error).", - required: true, - }, - FieldSchema { - name: "evidence", + name: "hits", ty: TypeSchema::Array(Box::new(TypeSchema::Json)), - comment: "Array of {source_path, snippet, relevance} evidence items.", + comment: "Ranked RetrievalHit evidence (node_id, content, \ + entities, score, time range, ...).", required: true, }, FieldSchema { - name: "trace", - ty: TypeSchema::Array(Box::new(TypeSchema::Json)), - comment: "Array of {turn, action, args_summary, result_preview} trace steps.", + name: "total", + ty: TypeSchema::U64, + comment: "Pre-truncation match count.", + required: true, + }, + FieldSchema { + name: "truncated", + ty: TypeSchema::Bool, + comment: "True when total exceeds the returned hit count.", required: true, }, ], diff --git a/src/openhuman/memory/schema/handlers.rs b/src/openhuman/memory/schema/handlers.rs index 940023810..9190b2fac 100644 --- a/src/openhuman/memory/schema/handlers.rs +++ b/src/openhuman/memory/schema/handlers.rs @@ -252,114 +252,63 @@ pub(super) fn handle_set_enabled(params: Map) -> ControllerFuture pub(super) fn handle_smart_walk(params: Map) -> ControllerFuture { Box::pin(async move { - use crate::openhuman::memory::chat::build_chat_provider; - use crate::openhuman::memory::query::smart_walk::{ - run_smart_walk, SmartWalkOptions, SmartWalkStopReason, - }; + use crate::openhuman::memory_tree::retrieval::{fast_retrieve, FastRetrieveOptions}; + // `max_turns`/`model` are accepted for backwards compatibility but + // ignored — retrieval is now deterministic (E2GraphRAG), so there are + // no LLM turns or model to select. `namespace` is NOT silently + // ignored: fast-retrieve operates over the whole leaf/summary store + // (leaf storage is intentionally not namespace-scoped), so a caller + // that previously relied on `namespace` as a retrieval boundary must + // fail closed rather than receive unscoped hits. #[derive(serde::Deserialize)] struct Req { query: String, - #[serde(default = "default_namespace")] - namespace: String, #[serde(default)] + limit: Option, + #[serde(default)] + time_window_days: Option, + #[serde(default)] + max_hops: Option, + #[serde(default)] + namespace: Option, + #[serde(default)] + #[allow(dead_code)] max_turns: Option, #[serde(default)] + #[allow(dead_code)] model: Option, } - fn default_namespace() -> String { - "default".into() - } let req = parse_value::(Value::Object(params))?; + + // Fail closed for a non-default namespace: deterministic retrieval has + // no namespace boundary, so honouring it silently would leak + // cross-namespace hits. The default namespace is the whole store. + if let Some(ns) = req.namespace.as_deref() { + if !ns.is_empty() && ns != "default" { + return Err(format!( + "smart_walk: namespace `{ns}` is not supported — deterministic \ + retrieval is not namespace-scoped (leaf storage is global). \ + Omit `namespace` or pass \"default\"." + )); + } + } + let config = config_rpc::load_config_with_timeout().await?; - let chat_provider = build_chat_provider(&config) - .map_err(|e| format!("smart_walk: build chat provider failed: {e}"))?; - - struct Adapter { - inner: std::sync::Arc, - } - - #[async_trait::async_trait] - impl crate::openhuman::inference::provider::traits::Provider for Adapter { - async fn chat_with_system( - &self, - system: Option<&str>, - message: &str, - _model: &str, - temperature: f64, - ) -> anyhow::Result { - let prompt = crate::openhuman::memory::chat::ChatPrompt { - system: system.unwrap_or("").to_string(), - user: message.to_string(), - temperature, - kind: "memory_smart_walk_rpc", - max_tokens: None, - }; - self.inner.chat_for_text(&prompt).await - } - - async fn chat_with_history( - &self, - messages: &[crate::openhuman::inference::provider::traits::ChatMessage], - model: &str, - temperature: f64, - ) -> anyhow::Result { - let system = messages - .iter() - .find(|m| m.role == "system") - .map(|m| m.content.as_str()); - let user: String = messages - .iter() - .filter(|m| m.role != "system") - .map(|m| m.content.as_str()) - .collect::>() - .join("\n"); - self.chat_with_system(system, &user, model, temperature) - .await - } - } - - let adapter = Adapter { - inner: chat_provider, + let opts = FastRetrieveOptions { + limit: req.limit.map(|n| n as usize).unwrap_or(10), + max_hops: req.max_hops.unwrap_or(2), + time_window_days: req.time_window_days, }; - let opts = SmartWalkOptions { - max_turns: req.max_turns.map(|n| n as usize).unwrap_or(12), - namespace: req.namespace, - model: req.model, - content_root: None, - }; - - let outcome = run_smart_walk(&config, &adapter, &req.query, opts) + let resp = fast_retrieve(&config, &req.query, opts) .await .map_err(|e| format!("smart_walk error: {e}"))?; - let stopped = match outcome.stopped_reason { - SmartWalkStopReason::Answered => "answered", - SmartWalkStopReason::MaxTurnsReached => "max_turns", - SmartWalkStopReason::LlmGaveUp => "llm_gave_up", - SmartWalkStopReason::Error(_) => "error", - }; - - let result = serde_json::json!({ - "answer": outcome.answer, - "turns_used": outcome.turns_used, - "evidence_count": outcome.evidence.len(), - "stopped_reason": stopped, - "evidence": outcome.evidence.iter().map(|e| serde_json::json!({ - "source_path": e.source_path, - "snippet": e.snippet, - "relevance": e.relevance, - })).collect::>(), - "trace": outcome.trace.iter().map(|s| serde_json::json!({ - "turn": s.turn, - "action": s.action, - "args_summary": s.args_summary, - "result_preview": s.result_preview, - })).collect::>(), - }); + let result = serde_json::to_value(&resp) + .map_err(|e| format!("smart_walk: serialize response failed: {e}"))?; to_json(RpcOutcome::new(result, vec![])) }) } diff --git a/src/openhuman/memory_search/tools/mod.rs b/src/openhuman/memory_search/tools/mod.rs index 5d5757949..cfb100538 100644 --- a/src/openhuman/memory_search/tools/mod.rs +++ b/src/openhuman/memory_search/tools/mod.rs @@ -17,14 +17,11 @@ pub use crate::openhuman::memory_store::tools::{ MemoryStoreKindsTool, MemoryStoreRawChunksTool, MemoryStoreRawSearchTool, }; -// Re-export existing tools from memory::query -pub use crate::openhuman::memory::query::smart_walk::run_smart_walk; -pub use crate::openhuman::memory::query::walk::{ - run_walk, WalkOptions, WalkOutcome, WalkStep, WalkStopReason, -}; +// Re-export existing tools from memory::query. The former agentic `walk` / +// `smart_walk` tools are gone — retrieval is now the deterministic +// `fast_retrieve` exposed via the `memory_tree` tool's `walk`/`smart_walk` +// modes (see `memory_tree::retrieval::fast`). pub use crate::openhuman::memory::query::{ - MemoryQueryWalkTool, MemoryTreeDrillDownTool, MemoryTreeFetchLeavesTool, - MemoryTreeIngestDocumentTool, MemoryTreeQuerySourceTool, MemoryTreeSearchEntitiesTool, - MemoryTreeWalkTool, SmartMemoryWalkTool, SmartWalkOptions, SmartWalkOutcome, SmartWalkStep, - SmartWalkStopReason, + MemoryTreeDrillDownTool, MemoryTreeFetchLeavesTool, MemoryTreeIngestDocumentTool, + MemoryTreeQuerySourceTool, MemoryTreeSearchEntitiesTool, }; diff --git a/src/openhuman/memory_store/chunks/store.rs b/src/openhuman/memory_store/chunks/store.rs index e466d524f..9799773b0 100644 --- a/src/openhuman/memory_store/chunks/store.rs +++ b/src/openhuman/memory_store/chunks/store.rs @@ -190,6 +190,27 @@ CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_node CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_index_timestamp ON mem_tree_entity_index(timestamp_ms); +-- E2GraphRAG: undirected weighted entity co-occurrence graph. One row per +-- unordered entity pair that has been extracted together within the same +-- chunk; `weight` accumulates co-occurrence frequency. Canonical ordering +-- (`entity_a < entity_b`) keeps the pair unique without a second row, and +-- the `entity_b` index makes neighbour lookups symmetric (a row matches a +-- query entity whether it appears as `entity_a` or `entity_b`). Read at +-- query time by `memory_tree::graph` for bounded-hop shortest-path filtering +-- during deterministic (LLM-free) retrieval. +CREATE TABLE IF NOT EXISTS mem_tree_entity_edges ( + entity_a TEXT NOT NULL, + entity_b TEXT NOT NULL, + weight INTEGER NOT NULL DEFAULT 1, + updated_ms INTEGER NOT NULL, + PRIMARY KEY (entity_a, entity_b) +); + +CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_edges_a + ON mem_tree_entity_edges(entity_a); +CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_edges_b + ON mem_tree_entity_edges(entity_b); + -- Phase 3a (#709): summary trees / bucket-seal. -- `mem_tree_trees` tracks one tree per scope (source/topic/global). CREATE TABLE IF NOT EXISTS mem_tree_trees ( diff --git a/src/openhuman/memory_tree/graph/bfs.rs b/src/openhuman/memory_tree/graph/bfs.rs new file mode 100644 index 000000000..9aba760f6 --- /dev/null +++ b/src/openhuman/memory_tree/graph/bfs.rs @@ -0,0 +1,215 @@ +//! Bounded shortest-path (hop distance) over the entity co-occurrence graph. +//! +//! This is the E2GraphRAG "graph filter" step: given the entities extracted +//! from a query, decide which *pairs* are semantically related by checking +//! whether they sit within `h` hops of each other in the co-occurrence graph. +//! Pairs within range route retrieval to the *local* (index-intersection) +//! branch; an empty result routes to the *global* (dense-tree) branch. +//! +//! Everything here is pure code — no LLM. Lookups go through +//! [`super::store::neighbors`], cached per query so a node is expanded at most +//! once even when it lies between several query-entity pairs. + +use std::collections::{HashMap, HashSet, VecDeque}; + +use anyhow::Result; + +use crate::openhuman::config::Config; +use crate::openhuman::memory_tree::graph::store; + +/// One related query-entity pair and the hop distance between them. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PairDistance { + pub a: String, + pub b: String, + pub dist: u32, +} + +/// Per-query neighbour cache so each entity's adjacency is read from SQLite at +/// most once across all pair BFS runs. +struct NeighborCache<'a> { + config: &'a Config, + cache: HashMap>, +} + +impl<'a> NeighborCache<'a> { + fn new(config: &'a Config) -> Self { + Self { + config, + cache: HashMap::new(), + } + } + + fn neighbors(&mut self, node: &str) -> Result<&[String]> { + if !self.cache.contains_key(node) { + let adj = store::neighbors(self.config, node)? + .into_iter() + .map(|(other, _weight)| other) + .collect(); + self.cache.insert(node.to_string(), adj); + } + // Safe: just inserted if missing. + Ok(self.cache.get(node).unwrap()) + } +} + +/// Compute hop distances between every unordered pair of `entity_ids`, keeping +/// only pairs reachable within `max_h` hops. `max_h == 0` short-circuits to an +/// empty result (no pair can be 0 hops apart — distinct entities are ≥1 hop). +/// +/// A breadth-first search runs from each entity, stopping once it has either +/// found all its peers in the query set or exhausted the `max_h` frontier. +/// Distance 1 means the two entities co-occurred directly in some chunk. +pub fn pair_distances( + config: &Config, + entity_ids: &[String], + max_h: u32, +) -> Result> { + if max_h == 0 || entity_ids.len() < 2 { + return Ok(Vec::new()); + } + // Deduplicate while preserving a stable canonical order so output pairs are + // deterministic (a < b). + let mut unique: Vec = entity_ids.to_vec(); + unique.sort(); + unique.dedup(); + if unique.len() < 2 { + return Ok(Vec::new()); + } + let targets: HashSet<&str> = unique.iter().map(|s| s.as_str()).collect(); + + let mut cache = NeighborCache::new(config); + let mut out: Vec = Vec::new(); + // Track pairs already recorded so a BFS from both endpoints doesn't emit + // the same pair twice. Keyed canonically (a < b). + let mut seen: HashSet<(String, String)> = HashSet::new(); + + for src in &unique { + // Peers we still need to reach from `src`; once empty we can stop the + // BFS early instead of draining the whole frontier (latency then + // scales with unresolved pairs, not graph size). + let mut remaining: HashSet<&str> = targets + .iter() + .copied() + .filter(|t| *t != src.as_str()) + .collect(); + if remaining.is_empty() { + continue; + } + // BFS frontier from `src` bounded to `max_h`. + let mut visited: HashSet = HashSet::new(); + visited.insert(src.clone()); + let mut queue: VecDeque<(String, u32)> = VecDeque::new(); + queue.push_back((src.clone(), 0)); + + while let Some((node, dist)) = queue.pop_front() { + if dist >= max_h || remaining.is_empty() { + continue; + } + let neighbors: Vec = cache.neighbors(&node)?.to_vec(); + for nb in neighbors { + if !visited.insert(nb.clone()) { + continue; + } + let nd = dist + 1; + // Record a pair the moment we reach another query target. BFS + // guarantees `nd` is the shortest distance. + if nb.as_str() != src.as_str() && remaining.remove(nb.as_str()) { + let (a, b) = if src.as_str() < nb.as_str() { + (src.clone(), nb.clone()) + } else { + (nb.clone(), src.clone()) + }; + if seen.insert((a.clone(), b.clone())) { + out.push(PairDistance { a, b, dist: nd }); + } + } + queue.push_back((nb, nd)); + } + } + } + + out.sort_by(|x, y| { + x.dist + .cmp(&y.dist) + .then_with(|| x.a.cmp(&y.a)) + .then_with(|| x.b.cmp(&y.b)) + }); + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory_tree::graph::store::{pairs_from_entities, upsert_edges}; + use tempfile::TempDir; + + fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + (tmp, cfg) + } + + /// Seed a path graph alice — bob — carol (two chunks). + fn seed_path(cfg: &Config) { + upsert_edges( + cfg, + &pairs_from_entities(&["person:alice".into(), "person:bob".into()]), + 1_000, + ) + .unwrap(); + upsert_edges( + cfg, + &pairs_from_entities(&["person:bob".into(), "person:carol".into()]), + 1_000, + ) + .unwrap(); + } + + #[test] + fn direct_cooccurrence_is_distance_one() { + let (_tmp, cfg) = test_config(); + seed_path(&cfg); + let pairs = pair_distances(&cfg, &["person:alice".into(), "person:bob".into()], 2).unwrap(); + assert_eq!(pairs.len(), 1); + assert_eq!(pairs[0].dist, 1); + assert_eq!(pairs[0].a, "person:alice"); + assert_eq!(pairs[0].b, "person:bob"); + } + + #[test] + fn two_hop_pair_found_within_h2_not_h1() { + let (_tmp, cfg) = test_config(); + seed_path(&cfg); + let ids = vec!["person:alice".to_string(), "person:carol".to_string()]; + + let h1 = pair_distances(&cfg, &ids, 1).unwrap(); + assert!( + h1.is_empty(), + "alice-carol is 2 hops apart; h=1 finds nothing" + ); + + let h2 = pair_distances(&cfg, &ids, 2).unwrap(); + assert_eq!(h2.len(), 1); + assert_eq!(h2[0].dist, 2); + } + + #[test] + fn disconnected_entities_yield_no_pairs() { + let (_tmp, cfg) = test_config(); + seed_path(&cfg); + // dave is not in the graph at all. + let pairs = + pair_distances(&cfg, &["person:alice".into(), "person:dave".into()], 3).unwrap(); + assert!(pairs.is_empty()); + } + + #[test] + fn single_entity_yields_no_pairs() { + let (_tmp, cfg) = test_config(); + seed_path(&cfg); + let pairs = pair_distances(&cfg, &["person:alice".into()], 2).unwrap(); + assert!(pairs.is_empty()); + } +} diff --git a/src/openhuman/memory_tree/graph/mod.rs b/src/openhuman/memory_tree/graph/mod.rs new file mode 100644 index 000000000..37fc92e89 --- /dev/null +++ b/src/openhuman/memory_tree/graph/mod.rs @@ -0,0 +1,19 @@ +//! Entity co-occurrence graph for E2GraphRAG-style deterministic retrieval. +//! +//! Two pieces: +//! - [`store`] — persistence for `mem_tree_entity_edges`, the undirected +//! weighted co-occurrence graph built incrementally at ingest time. +//! - [`bfs`] — bounded shortest-path (hop distance) over that graph, used as +//! the query-time "graph filter" that routes retrieval between the local +//! (entity-index intersection) and global (dense summary-tree) branches. +//! +//! The graph bridges the entity index and the summary tree without any LLM in +//! the loop: query entities → hop-distance filter → candidate chunk lookup. + +pub mod bfs; +pub mod store; + +pub use bfs::{pair_distances, PairDistance}; +pub use store::{ + clear_edges_for_entities_tx, neighbors, pairs_from_entities, upsert_edges, upsert_edges_tx, +}; diff --git a/src/openhuman/memory_tree/graph/store.rs b/src/openhuman/memory_tree/graph/store.rs new file mode 100644 index 000000000..f14410c36 --- /dev/null +++ b/src/openhuman/memory_tree/graph/store.rs @@ -0,0 +1,230 @@ +//! Persistence for the entity co-occurrence graph (`mem_tree_entity_edges`). +//! +//! The edge table is an undirected weighted graph: each row is one unordered +//! pair of canonical entity ids that co-occurred inside the same chunk, with +//! `weight` accumulating the number of co-occurrences. Canonical ordering +//! (`entity_a < entity_b`) keeps a pair on a single row; neighbour lookups +//! union both columns so a query entity matches regardless of which side it +//! was stored on. +//! +//! The schema is declared in `memory_store/chunks/store.rs::SCHEMA`; this file +//! only owns the CRUD operations. Writes happen inside the same transaction +//! that indexes a chunk's entities (see `score::persist_score_tx`) so the +//! graph never diverges from `mem_tree_entity_index`. + +use anyhow::Result; +use rusqlite::{params, Transaction}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory_store::chunks::store::with_connection; + +/// Order a pair canonically so `(a, b)` and `(b, a)` collapse onto one row. +/// Returns `None` for a self-pair (`a == b`) — an entity never edges itself. +fn order_pair<'a>(a: &'a str, b: &'a str) -> Option<(&'a str, &'a str)> { + use std::cmp::Ordering; + match a.cmp(b) { + Ordering::Less => Some((a, b)), + Ordering::Greater => Some((b, a)), + Ordering::Equal => None, + } +} + +/// Emit every unordered pair from a set of canonical entity ids. Deduplicates +/// the input first (a chunk can mention the same entity in several spans) so a +/// repeated id doesn't inflate a single chunk's contribution to one edge. +pub fn pairs_from_entities(entity_ids: &[String]) -> Vec<(String, String)> { + use std::collections::BTreeSet; + let unique: Vec<&String> = entity_ids + .iter() + .collect::>() + .into_iter() + .collect(); + let mut out = Vec::new(); + for i in 0..unique.len() { + for j in (i + 1)..unique.len() { + if let Some((a, b)) = order_pair(unique[i], unique[j]) { + out.push((a.to_string(), b.to_string())); + } + } + } + out +} + +/// Upsert a batch of co-occurrence pairs inside an existing transaction, +/// incrementing `weight` for pairs that already exist. Idempotent on the +/// `(entity_a, entity_b)` primary key — re-running adds to the weight, which +/// is the intended accumulation semantics (more co-occurrences = stronger +/// edge). `pairs` are expected to already be canonically ordered via +/// [`pairs_from_entities`]. +pub fn upsert_edges_tx( + tx: &Transaction<'_>, + pairs: &[(String, String)], + timestamp_ms: i64, +) -> Result { + // Defensively canonicalise (`a < b`) and dedup here too, so a caller that + // passes reversed or duplicate pairs can never split one edge across two + // rows (`a,b` + `b,a`) or double-count a single chunk's contribution. + // `pairs_from_entities` already does this, but the table invariant must not + // depend on callers getting it right. + use std::collections::BTreeSet; + let canonical: BTreeSet<(String, String)> = pairs + .iter() + .filter_map(|(a, b)| order_pair(a, b).map(|(x, y)| (x.to_string(), y.to_string()))) + .collect(); + if canonical.is_empty() { + return Ok(0); + } + let mut stmt = tx.prepare( + "INSERT INTO mem_tree_entity_edges (entity_a, entity_b, weight, updated_ms) + VALUES (?1, ?2, 1, ?3) + ON CONFLICT(entity_a, entity_b) + DO UPDATE SET weight = weight + 1, updated_ms = ?3", + )?; + for (a, b) in &canonical { + stmt.execute(params![a, b, timestamp_ms])?; + } + Ok(canonical.len()) +} + +/// Convenience wrapper that opens its own connection + transaction. Prefer +/// [`upsert_edges_tx`] when an enclosing transaction already exists (the +/// ingest hook does), so the graph write commits atomically with the entity +/// index write. +pub fn upsert_edges( + config: &Config, + pairs: &[(String, String)], + timestamp_ms: i64, +) -> Result { + if pairs.is_empty() { + return Ok(0); + } + with_connection(config, |conn| { + let tx = conn.unchecked_transaction()?; + let n = upsert_edges_tx(&tx, pairs, timestamp_ms)?; + tx.commit()?; + Ok(n) + }) +} + +/// Return the immediate neighbours of `entity_id` with their edge weights. +/// Unions both columns so the lookup is symmetric. Ordered by weight DESC so +/// the strongest associations come first (callers can cap traversal breadth). +pub fn neighbors(config: &Config, entity_id: &str) -> Result> { + with_connection(config, |conn| { + let mut stmt = conn.prepare( + "SELECT entity_b AS other, weight FROM mem_tree_entity_edges WHERE entity_a = ?1 + UNION ALL + SELECT entity_a AS other, weight FROM mem_tree_entity_edges WHERE entity_b = ?1 + ORDER BY weight DESC", + )?; + let rows = stmt + .query_map(params![entity_id], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)) + })? + .collect::>>()?; + Ok(rows) + }) +} + +/// Remove every edge touching any of `entity_ids`. Used before re-indexing a +/// re-scored chunk so co-occurrences that no longer hold don't leak through as +/// stale edges (mirrors `clear_entity_index_for_node`). Note this is coarse — +/// it drops the entity's edges to *all* peers, not just the ones from this +/// chunk; the subsequent re-index rebuilds the surviving edges from the +/// chunk's current entities. Edges contributed by other chunks are recovered +/// on their next re-score; for the common append-only ingest path this branch +/// is not exercised. +pub fn clear_edges_for_entities_tx(tx: &Transaction<'_>, entity_ids: &[String]) -> Result { + if entity_ids.is_empty() { + return Ok(0); + } + let mut removed = 0; + let mut stmt = + tx.prepare("DELETE FROM mem_tree_entity_edges WHERE entity_a = ?1 OR entity_b = ?1")?; + for id in entity_ids { + removed += stmt.execute(params![id])?; + } + Ok(removed) +} + +/// Count edge rows (for tests / diagnostics). +pub fn count_edges(config: &Config) -> Result { + with_connection(config, |conn| { + let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_entity_edges", [], |r| { + r.get(0) + })?; + Ok(n.max(0) as u64) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + (tmp, cfg) + } + + #[test] + fn pairs_are_canonically_ordered_and_deduped() { + let pairs = pairs_from_entities(&[ + "person:bob".into(), + "person:alice".into(), + "person:alice".into(), // dup — must not double an edge + ]); + assert_eq!(pairs, vec![("person:alice".into(), "person:bob".into())]); + } + + #[test] + fn self_pairs_are_skipped() { + let pairs = pairs_from_entities(&["x:1".into(), "x:1".into()]); + assert!(pairs.is_empty()); + } + + #[test] + fn upsert_increments_weight_and_neighbors_is_symmetric() { + let (_tmp, cfg) = test_config(); + let p = vec![("person:alice".to_string(), "person:bob".to_string())]; + upsert_edges(&cfg, &p, 1_000).unwrap(); + upsert_edges(&cfg, &p, 2_000).unwrap(); + + // Symmetric lookup from either endpoint. + let from_alice = neighbors(&cfg, "person:alice").unwrap(); + assert_eq!(from_alice, vec![("person:bob".to_string(), 2)]); + let from_bob = neighbors(&cfg, "person:bob").unwrap(); + assert_eq!(from_bob, vec![("person:alice".to_string(), 2)]); + assert_eq!(count_edges(&cfg).unwrap(), 1); + } + + #[test] + fn clear_edges_removes_all_touching() { + let (_tmp, cfg) = test_config(); + upsert_edges( + &cfg, + &pairs_from_entities(&[ + "person:alice".into(), + "person:bob".into(), + "topic:phoenix".into(), + ]), + 1_000, + ) + .unwrap(); + assert_eq!(count_edges(&cfg).unwrap(), 3); + + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + let n = clear_edges_for_entities_tx(&tx, &["person:alice".into()]).unwrap(); + tx.commit()?; + // alice-bob and alice-phoenix removed; bob-phoenix survives. + assert_eq!(n, 2); + Ok(()) + }) + .unwrap(); + assert_eq!(count_edges(&cfg).unwrap(), 1); + assert!(neighbors(&cfg, "person:alice").unwrap().is_empty()); + } +} diff --git a/src/openhuman/memory_tree/mod.rs b/src/openhuman/memory_tree/mod.rs index 145de6e6a..5f6e579b3 100644 --- a/src/openhuman/memory_tree/mod.rs +++ b/src/openhuman/memory_tree/mod.rs @@ -5,9 +5,11 @@ //! It is flavor-agnostic; the specific tree instances (global, topic, //! source) and their policies live in [`crate::openhuman::memory`]. +pub mod graph; pub mod health; pub mod ingest; pub mod io; +pub mod nlp; pub mod retrieval; pub mod score; pub mod summarise; diff --git a/src/openhuman/memory_tree/nlp/client.rs b/src/openhuman/memory_tree/nlp/client.rs new file mode 100644 index 000000000..15e0a8253 --- /dev/null +++ b/src/openhuman/memory_tree/nlp/client.rs @@ -0,0 +1,227 @@ +//! 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 new file mode 100644 index 000000000..abab366d5 --- /dev/null +++ b/src/openhuman/memory_tree/nlp/mod.rs @@ -0,0 +1,189 @@ +//! Query-side NLP for the deterministic (E2GraphRAG) retriever. +//! +//! [`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 + +//! 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. +//! +//! Output is intentionally `Vec`: it reuses +//! [`score::resolver::canonicalise`] so query entity ids land in the exact +//! 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, SpacyRuntime, SPACY_MODEL}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory_tree::score::extract::{ + EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic, +}; +use crate::openhuman::memory_tree::score::resolver::{canonicalise, CanonicalEntity}; + +/// Map a spaCy entity label to our [`EntityKind`]. Unknown labels collapse to +/// [`EntityKind::Misc`] so they still participate as graph anchors. +fn map_spacy_label(label: &str) -> EntityKind { + match label { + "PERSON" => EntityKind::Person, + "ORG" | "NORP" => EntityKind::Organization, + "GPE" | "LOC" | "FAC" => EntityKind::Location, + "PRODUCT" => EntityKind::Product, + "EVENT" => EntityKind::Event, + "DATE" | "TIME" => EntityKind::Datetime, + "MONEY" | "QUANTITY" | "PERCENT" | "CARDINAL" | "ORDINAL" => EntityKind::Quantity, + "LANGUAGE" => EntityKind::Technology, + _ => EntityKind::Misc, + } +} + +/// Extract canonical query entities, preferring spaCy and falling back to the +/// in-Rust regex extractor. Never fails: an unavailable sidecar degrades to +/// the fallback rather than erroring, because an empty/partial entity set just +/// routes retrieval toward the global (dense) branch. +pub async fn extract_query_entities(config: &Config, query: &str) -> Vec { + let trimmed = query.trim(); + if trimmed.is_empty() { + return Vec::new(); + } + + if config.memory_tree.spacy_enabled { + if let Some(ner) = shared_ner(config).await { + match ner.extract(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 { + log::debug!("[memory_tree::nlp] spaCy disabled by config — using regex fallback"); + } + + fallback_extract(trimmed).await +} + +/// Build [`ExtractedEntities`] from a spaCy response: named entities become +/// entity spans, salient nouns become topics. Topics are promoted to +/// `topic:` canonical ids by [`canonicalise`]. +fn spacy_to_extracted(resp: &client::SpacyResponse) -> ExtractedEntities { + let entities = resp + .entities + .iter() + .map(|e| ExtractedEntity { + kind: map_spacy_label(&e.label), + text: e.text.clone(), + span_start: e.start, + span_end: e.end, + score: 1.0, + }) + .collect(); + let topics = resp + .nouns + .iter() + .map(|n| ExtractedTopic { + label: n.clone(), + score: 1.0, + }) + .collect(); + ExtractedEntities { + entities, + topics, + llm_importance: None, + llm_importance_reason: None, + } +} + +/// Regex-only fallback. Deterministic, no network, no LLM — catches +/// emails/urls/handles/hashtags in the query. Person/org recall is lost +/// (spaCy's job), which simply biases retrieval toward the global branch. +async fn fallback_extract(query: &str) -> Vec { + use crate::openhuman::memory_tree::score::extract::{CompositeExtractor, EntityExtractor}; + let extractor = CompositeExtractor::regex_only(); + match extractor.extract(query).await { + Ok(extracted) => { + let canon = canonicalise(&extracted); + log::debug!( + "[memory_tree::nlp] regex fallback query extraction: canonical={}", + canon.len() + ); + canon + } + Err(e) => { + log::warn!("[memory_tree::nlp] regex fallback failed: {e:#}"); + Vec::new() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cfg_spacy_off() -> Config { + let mut c = Config::default(); + c.memory_tree.spacy_enabled = false; + c + } + + #[test] + fn label_mapping_covers_common_kinds() { + assert_eq!(map_spacy_label("PERSON"), EntityKind::Person); + assert_eq!(map_spacy_label("ORG"), EntityKind::Organization); + assert_eq!(map_spacy_label("GPE"), EntityKind::Location); + assert_eq!(map_spacy_label("WHATEVER"), EntityKind::Misc); + } + + #[tokio::test] + async fn fallback_used_when_spacy_disabled_extracts_mechanical_entities() { + let cfg = cfg_spacy_off(); + let ents = extract_query_entities(&cfg, "ping alice@example.com about #launch").await; + assert!( + ents.iter() + .any(|e| e.canonical_id == "email:alice@example.com"), + "regex fallback should find the email; got {ents:?}" + ); + assert!( + ents.iter().any(|e| e.kind == EntityKind::Hashtag), + "regex fallback should find the hashtag; got {ents:?}" + ); + } + + #[tokio::test] + async fn empty_query_yields_no_entities() { + let cfg = cfg_spacy_off(); + assert!(extract_query_entities(&cfg, " ").await.is_empty()); + } + + #[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, + }], + nouns: vec!["migration".into()], + id: Some("0".into()), + error: None, + }; + let extracted = spacy_to_extracted(&resp); + let canon = canonicalise(&extracted); + assert!(canon.iter().any(|c| c.canonical_id == "person:alice")); + assert!(canon.iter().any(|c| c.canonical_id == "topic:migration")); + } +} diff --git a/src/openhuman/memory_tree/nlp/provision.rs b/src/openhuman/memory_tree/nlp/provision.rs new file mode 100644 index 000000000..1d7373ea1 --- /dev/null +++ b/src/openhuman/memory_tree/nlp/provision.rs @@ -0,0 +1,229 @@ +//! 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(()) +} + +/// 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 new file mode 100644 index 000000000..28ddc933f --- /dev/null +++ b/src/openhuman/memory_tree/nlp/service.py @@ -0,0 +1,108 @@ +#!/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/memory_tree/retrieval/fast.rs b/src/openhuman/memory_tree/retrieval/fast.rs new file mode 100644 index 000000000..831f35f50 --- /dev/null +++ b/src/openhuman/memory_tree/retrieval/fast.rs @@ -0,0 +1,446 @@ +//! `fast_retrieve` — deterministic, LLM-free memory retrieval (E2GraphRAG). +//! +//! This replaces the agentic `walk` / `smart_walk` loops. Instead of an LLM +//! navigating the summary tree turn-by-turn, retrieval routing is decided +//! purely by spaCy query entities + co-occurrence-graph hop distance: +//! +//! 1. Extract query entities `Eq` (spaCy, with regex fallback). +//! 2. `Eq` empty → **global**: dense rerank over the summary tree. +//! 3. Otherwise compute related entity pairs `Ph` within `h` hops. +//! - `Ph` empty → **global with occurrence ranking**: dense top-2k, then +//! re-rank by how many `Eq` entities each summary mentions. +//! - `Ph` non-empty → **local (index mapping)**: intersect the entity-index +//! node sets of each related pair; tighten `h` while the candidate set is +//! larger than `k`; rank survivors by entity coverage then recency. +//! +//! Output is a structured [`QueryResponse`] of [`RetrievalHit`]s (no +//! synthesized prose) for a higher-level context agent to consume. + +use std::cmp::Reverse; +use std::collections::{HashMap, HashSet}; + +use anyhow::Result; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::source_scope::current_source_scope; +use crate::openhuman::memory_store::content::read as content_read; +use crate::openhuman::memory_store::trees::store as tree_store; +use crate::openhuman::memory_tree::graph; +use crate::openhuman::memory_tree::nlp; +use crate::openhuman::memory_tree::retrieval::fetch::{self, fetch_leaves}; +use crate::openhuman::memory_tree::retrieval::source::query_source; +use crate::openhuman::memory_tree::retrieval::types::{ + hit_from_summary, QueryResponse, RetrievalHit, +}; +use crate::openhuman::memory_tree::score::store::lookup_entity; + +/// Tunables for [`fast_retrieve`]. Defaults match the E2GraphRAG paper's +/// small-`k` regime and a 2-hop relatedness threshold. +#[derive(Clone, Debug)] +pub struct FastRetrieveOptions { + /// `k` — max hits returned. + pub limit: usize, + /// `h` — initial hop threshold for relatedness. Tightened (decremented) + /// during the local branch when too many candidates match. + pub max_hops: u32, + /// Look-back window (days) applied on the global/dense branch. `None` = + /// unbounded. + pub time_window_days: Option, +} + +impl Default for FastRetrieveOptions { + fn default() -> Self { + Self { + limit: 10, + max_hops: 2, + time_window_days: None, + } + } +} + +/// Per-node-entity lookup cap. Popular entities can touch many nodes; this +/// bounds the intersection work while staying well above any realistic `k`. +const LOOKUP_LIMIT: usize = 500; + +/// Default / ceiling for `limit` (`k`). The tool and RPC paths expose this, so +/// a huge value must not be able to request oversized dense/local result sets. +const DEFAULT_LIMIT: usize = 10; +const MAX_RETRIEVE_LIMIT: usize = 100; +/// Default / ceiling for the hop threshold. A large `max_hops` would force many +/// bounded-BFS passes through the tightening loop; cap it. +const DEFAULT_MAX_HOPS: u32 = 2; +const MAX_GRAPH_HOPS: u32 = 4; + +/// Run deterministic retrieval for `query`. Never invokes an LLM. +pub async fn fast_retrieve( + config: &Config, + query: &str, + opts: FastRetrieveOptions, +) -> Result { + let k = match opts.limit { + 0 => DEFAULT_LIMIT, + n => n.min(MAX_RETRIEVE_LIMIT), + }; + let max_hops = match opts.max_hops { + 0 => DEFAULT_MAX_HOPS, + n => n.min(MAX_GRAPH_HOPS), + }; + let trimmed = query.trim(); + if trimmed.is_empty() { + return Ok(QueryResponse::empty()); + } + + // 1. Query entities. + let entities = nlp::extract_query_entities(config, trimmed).await; + let eq: Vec = dedup_ids(entities.iter().map(|e| e.canonical_id.clone())); + log::debug!( + "[retrieval::fast] query_len={} eq={} k={} h={}", + trimmed.len(), + eq.len(), + k, + max_hops + ); + + // 2. No entities → pure global dense retrieval. + if eq.is_empty() { + log::debug!("[retrieval::fast] branch=global_dense (no query entities)"); + return query_source(config, None, None, opts.time_window_days, Some(trimmed), k).await; + } + + // 3. Graph filter: related entity pairs within `h` hops. + let pairs = graph::pair_distances(config, &eq, max_hops)?; + if pairs.is_empty() { + log::debug!("[retrieval::fast] branch=global_occurrence (no related pairs)"); + return global_occurrence(config, trimmed, &eq, k, opts.time_window_days).await; + } + + // 4. Local branch: index-mapping intersection with `h` tightening. + let mut h = max_hops; + let mut cands = local_candidates(config, &eq, &pairs)?; + let mut last_non_empty = cands.clone(); + while cands.len() > k && h > 1 { + h -= 1; + let tightened = graph::pair_distances(config, &eq, h)?; + let next = local_candidates(config, &eq, &tightened)?; + if next.is_empty() { + // Tightening removed everything — keep the looser, non-empty set. + break; + } + last_non_empty = next.clone(); + cands = next; + } + if cands.is_empty() { + cands = last_non_empty; + } + + if cands.is_empty() { + // Related pairs existed but never co-occurred in an indexed node + // (e.g. only 2-hop links). Fall back to global occurrence ranking. + log::debug!("[retrieval::fast] branch=local->global_occurrence (empty intersection)"); + return global_occurrence(config, trimmed, &eq, k, opts.time_window_days).await; + } + + log::debug!( + "[retrieval::fast] branch=local candidates={} final_h={}", + cands.len(), + h + ); + resolve_local(config, cands, k).await +} + +/// Candidate node accumulated during the local branch. +#[derive(Clone, Debug)] +struct Candidate { + node_kind: String, + /// Distinct query entities that landed on this node (coverage signal). + matched: HashSet, + latest_ts: i64, +} + +/// Intersect the entity-index node sets of each related pair and union the +/// results. A node enters the candidate set when it is indexed against *both* +/// members of some related pair. +fn local_candidates( + config: &Config, + eq: &[String], + pairs: &[graph::PairDistance], +) -> Result> { + let _ = eq; // `eq` documents intent; matched entities come from the pairs. + let mut out: HashMap = HashMap::new(); + for pair in pairs { + let hits_a = lookup_entity(config, &pair.a, Some(LOOKUP_LIMIT))?; + if hits_a.is_empty() { + continue; + } + let hits_b = lookup_entity(config, &pair.b, Some(LOOKUP_LIMIT))?; + if hits_b.is_empty() { + continue; + } + let b_nodes: HashMap<&str, i64> = hits_b + .iter() + .map(|h| (h.node_id.as_str(), h.timestamp_ms)) + .collect(); + for ha in &hits_a { + let Some(&b_ts) = b_nodes.get(ha.node_id.as_str()) else { + continue; + }; + let entry = out.entry(ha.node_id.clone()).or_insert_with(|| Candidate { + node_kind: ha.node_kind.clone(), + matched: HashSet::new(), + latest_ts: 0, + }); + entry.matched.insert(pair.a.clone()); + entry.matched.insert(pair.b.clone()); + entry.latest_ts = entry.latest_ts.max(ha.timestamp_ms).max(b_ts); + } + } + Ok(out) +} + +/// Rank candidates by entity coverage (desc) then recency (desc), resolve to +/// hits, apply the profile source-scope gate **before** counting/truncating +/// (so an out-of-scope top hit never displaces an allowed lower-ranked one, +/// and `total` never reveals out-of-scope match counts), then truncate to `k`. +async fn resolve_local( + config: &Config, + cands: HashMap, + k: usize, +) -> Result { + let mut ordered: Vec<(String, Candidate)> = cands.into_iter().collect(); + ordered.sort_by(|a, b| { + b.1.matched + .len() + .cmp(&a.1.matched.len()) + .then_with(|| b.1.latest_ts.cmp(&a.1.latest_ts)) + .then_with(|| a.0.cmp(&b.0)) + }); + + // Coverage score per node id so resolved hits carry the ranking signal. + let coverage: HashMap = ordered + .iter() + .map(|(id, c)| (id.clone(), c.matched.len() as f32)) + .collect(); + + let leaf_ids: Vec = ordered + .iter() + .filter(|(_, c)| c.node_kind == "leaf") + .map(|(id, _)| id.clone()) + .collect(); + let summary_ids: Vec = ordered + .iter() + .filter(|(_, c)| c.node_kind != "leaf") + .map(|(id, _)| id.clone()) + .collect(); + + let mut by_id: HashMap = HashMap::new(); + // `fetch_leaves` caps each batch at MAX_BATCH and would silently drop the + // rest, so chunk to resolve every candidate leaf before the scope gate. + for chunk in leaf_ids.chunks(fetch::MAX_BATCH) { + for hit in fetch_leaves(config, chunk).await? { + by_id.insert(hit.node_id.clone(), hit); + } + } + if !summary_ids.is_empty() { + for hit in resolve_summaries(config, &summary_ids)? { + by_id.insert(hit.node_id.clone(), hit); + } + } + + // Scope-gate the full ranked set first; `total` reflects only in-scope hits. + let scope = current_source_scope(); + let mut hits: Vec = Vec::with_capacity(ordered.len()); + for (id, _) in &ordered { + if let Some(mut hit) = by_id.remove(id) { + if !scope_allows(scope.as_ref(), &hit.tree_scope) { + continue; + } + if let Some(score) = coverage.get(id) { + hit.score = *score; + } + hits.push(hit); + } + } + let total = hits.len(); + hits.truncate(k); + Ok(QueryResponse::new(hits, total)) +} + +/// Resolve summary node ids to hits, hydrating the full body from disk and +/// threading the owning tree's scope for provenance. +fn resolve_summaries(config: &Config, summary_ids: &[String]) -> Result> { + let nodes = tree_store::get_summaries_batch(config, summary_ids)?; + let mut scope_cache: HashMap = HashMap::new(); + let mut out = Vec::with_capacity(nodes.len()); + for (_, mut node) in nodes { + // Hydrate full body (the `content` column is a ≤500-char preview). + if let Ok(body) = content_read::read_summary_body(config, &node.id) { + node.content = body; + } + let scope = match scope_cache.get(&node.tree_id) { + Some(s) => s.clone(), + None => { + let s = tree_store::get_tree(config, &node.tree_id) + .ok() + .flatten() + .map(|t| t.scope) + .unwrap_or_default(); + scope_cache.insert(node.tree_id.clone(), s.clone()); + s + } + }; + out.push(hit_from_summary(&node, &scope)); + } + Ok(out) +} + +/// Global branch with occurrence ranking: dense-retrieve top-`2k`, then re-rank +/// by how many `Eq` entities each summary mentions (occurrence). Stable sort +/// preserves the semantic order on ties. +async fn global_occurrence( + config: &Config, + query: &str, + eq: &[String], + k: usize, + window: Option, +) -> Result { + let resp = query_source(config, None, None, window, Some(query), k.saturating_mul(2)).await?; + let eq_set: HashSet<&str> = eq.iter().map(|s| s.as_str()).collect(); + let mut hits = resp.hits; + let total = resp.total; + hits.sort_by_key(|h| { + let occ = h + .entities + .iter() + .filter(|e| eq_set.contains(e.as_str())) + .count(); + Reverse(occ) + }); + hits.truncate(k); + Ok(QueryResponse::new(hits, total)) +} + +/// Deduplicate ids preserving first-seen order. +fn dedup_ids(ids: impl Iterator) -> Vec { + let mut seen = HashSet::new(); + let mut out = Vec::new(); + for id in ids { + if seen.insert(id.clone()) { + out.push(id); + } + } + out +} + +/// Profile source-scope gate: `None` = unrestricted; otherwise the scope must +/// be on the allowlist. Empty scopes (unknown provenance) are allowed through +/// only when unrestricted. +fn scope_allows(scope: Option<&HashSet>, tree_scope: &str) -> bool { + match scope { + None => true, + Some(set) => set.contains(tree_scope), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::ingest_pipeline::ingest_chat; + use crate::openhuman::memory_sync::canonicalize::chat::{ChatBatch, ChatMessage}; + use chrono::{TimeZone, Utc}; + use tempfile::TempDir; + + fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + // spaCy off in CI — exercise the regex fallback + graph routing + // deterministically without a Python runtime. + cfg.memory_tree.spacy_enabled = false; + (tmp, cfg) + } + + async fn seed_chat(cfg: &Config, source: &str, text: &str) { + let batch = ChatBatch { + platform: "slack".into(), + channel_label: source.into(), + messages: vec![ChatMessage { + author: "alice".into(), + timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), + text: text.into(), + source_ref: Some("slack://x".into()), + }], + }; + ingest_chat(cfg, source, "alice", vec![], batch) + .await + .unwrap(); + } + + #[tokio::test] + async fn empty_query_returns_empty() { + let (_tmp, cfg) = test_config(); + let resp = fast_retrieve(&cfg, " ", FastRetrieveOptions::default()) + .await + .unwrap(); + assert!(resp.hits.is_empty()); + } + + #[tokio::test] + async fn no_entities_routes_global_without_panicking() { + let (_tmp, cfg) = test_config(); + // A query with no mechanical entities (regex fallback finds nothing) + // must take the global branch and return cleanly on an empty store. + let resp = fast_retrieve( + &cfg, + "what happened recently", + FastRetrieveOptions::default(), + ) + .await + .unwrap(); + assert!(resp.hits.is_empty()); + } + + #[tokio::test] + async fn local_branch_finds_cooccurring_entities() { + let (_tmp, cfg) = test_config(); + // Two emails co-occur in one message → an edge + both indexed on the + // same leaf. A query mentioning both routes local and returns the leaf. + seed_chat( + &cfg, + "slack:#eng", + "Sync between alice@example.com and bob@example.com about the runbook.", + ) + .await; + + let resp = fast_retrieve( + &cfg, + "alice@example.com and bob@example.com", + FastRetrieveOptions::default(), + ) + .await + .unwrap(); + assert!( + !resp.hits.is_empty(), + "co-occurring entities should yield a local hit; got {resp:?}" + ); + // Coverage score = 2 distinct query entities matched the leaf. + assert!(resp.hits.iter().any(|h| h.score >= 2.0)); + } + + #[test] + fn scope_allows_respects_allowlist() { + assert!(scope_allows(None, "slack:#eng")); + let mut set = HashSet::new(); + set.insert("slack:#eng".to_string()); + assert!(scope_allows(Some(&set), "slack:#eng")); + assert!(!scope_allows(Some(&set), "gmail:alice@example.com")); + } + + #[test] + fn dedup_ids_preserves_first_seen_order() { + let out = dedup_ids(["a".to_string(), "b".into(), "a".into(), "c".into()].into_iter()); + assert_eq!(out, vec!["a", "b", "c"]); + } +} diff --git a/src/openhuman/memory_tree/retrieval/mod.rs b/src/openhuman/memory_tree/retrieval/mod.rs index ed49dad13..a41543fae 100644 --- a/src/openhuman/memory_tree/retrieval/mod.rs +++ b/src/openhuman/memory_tree/retrieval/mod.rs @@ -19,6 +19,7 @@ pub mod cover; pub mod drill_down; +pub mod fast; pub mod fetch; pub mod rpc; pub mod schemas; @@ -33,6 +34,7 @@ mod integration_tests; pub use cover::cover_window; pub use drill_down::drill_down; +pub use fast::{fast_retrieve, FastRetrieveOptions}; pub use fetch::fetch_leaves; pub use schemas::{ all_controller_schemas as all_retrieval_controller_schemas, diff --git a/src/openhuman/memory_tree/score/mod.rs b/src/openhuman/memory_tree/score/mod.rs index 00525db52..fc64d628d 100644 --- a/src/openhuman/memory_tree/score/mod.rs +++ b/src/openhuman/memory_tree/score/mod.rs @@ -392,12 +392,36 @@ pub fn persist_score( timestamp_ms, tree_id, )?; + // E2GraphRAG: accumulate the chunk's entity co-occurrence edges so + // query-time hop-distance filtering has a graph to walk. + persist_cooccurrence_edges(config, &result.canonical_entities, timestamp_ms); } } Ok(()) } +/// Build and persist the undirected co-occurrence edges implied by a chunk's +/// canonical entities. Best-effort: a graph write failure must never block the +/// (already-committed) score/entity-index write, so errors are logged and +/// swallowed. Co-occurrence = two entities mentioned in the same chunk. +fn persist_cooccurrence_edges( + config: &crate::openhuman::config::Config, + entities: &[crate::openhuman::memory_tree::score::resolver::CanonicalEntity], + timestamp_ms: i64, +) { + use crate::openhuman::memory_tree::graph; + let ids: Vec = entities.iter().map(|e| e.canonical_id.clone()).collect(); + let pairs = graph::pairs_from_entities(&ids); + if pairs.is_empty() { + return; + } + match graph::upsert_edges(config, &pairs, timestamp_ms) { + Ok(n) => log::debug!("[memory_tree::graph] indexed cooccurrence edges count={n}"), + Err(e) => log::warn!("[memory_tree::graph] edge upsert failed (non-fatal): {e:#}"), + } +} + pub(crate) fn persist_score_tx( tx: &Transaction<'_>, result: &ScoreResult, @@ -419,6 +443,17 @@ pub(crate) fn persist_score_tx( timestamp_ms, tree_id, )?; + // E2GraphRAG: accumulate co-occurrence edges in the SAME transaction + // so the graph never diverges from the entity index. + let ids: Vec = result + .canonical_entities + .iter() + .map(|e| e.canonical_id.clone()) + .collect(); + let pairs = crate::openhuman::memory_tree::graph::pairs_from_entities(&ids); + let n = + crate::openhuman::memory_tree::graph::upsert_edges_tx(tx, &pairs, timestamp_ms)?; + log::debug!("[memory_tree::graph] indexed cooccurrence edges (tx) count={n}"); } } diff --git a/src/openhuman/tools/ops.rs b/src/openhuman/tools/ops.rs index 299da44b5..e3492e54e 100644 --- a/src/openhuman/tools/ops.rs +++ b/src/openhuman/tools/ops.rs @@ -249,8 +249,6 @@ pub fn all_tools_with_runtime( // can explain an empty/stalled wiki + the fix. Box::new(MemoryDoctorTool::new(config.clone())), Box::new(MemoryQueryTool), - Box::new(MemoryQueryWalkTool), - Box::new(SmartMemoryWalkTool), // memory_search tools — vector search, chunk context, hybrid search, // and previously unregistered raw store tools. Box::new(MemoryVectorSearchTool), diff --git a/tests/memory_fast_retrieve_e2e.rs b/tests/memory_fast_retrieve_e2e.rs new file mode 100644 index 000000000..e1578229e --- /dev/null +++ b/tests/memory_fast_retrieve_e2e.rs @@ -0,0 +1,131 @@ +//! E2E tests for the deterministic E2GraphRAG retriever (`fast_retrieve`). +//! +//! These replace the old agentic `memory_tree_walk_e2e.rs`. There is no LLM in +//! the retrieval loop, so no mock server is needed — we ingest a small chat +//! corpus, then assert that: +//! - an entity-relationship query routes to the *local* branch and returns +//! the chunk where the two entities co-occur, ranked by entity coverage; +//! - a query with no extractable entities routes to the *global* branch and +//! returns cleanly (no panic) over the same store; +//! - the output is structured `QueryResponse` evidence (hits), not prose. +//! +//! spaCy is disabled here so the run is deterministic and Python-free in CI — +//! query-entity extraction uses the regex fallback (emails/handles/hashtags), +//! which is enough to exercise both routing branches. +//! +//! Run with: +//! cargo test --test memory_fast_retrieve_e2e +//! or via the project wrapper: +//! bash scripts/test-rust-with-mock.sh --test memory_fast_retrieve_e2e + +use chrono::{TimeZone, Utc}; +use tempfile::TempDir; + +use openhuman_core::openhuman::config::Config; +use openhuman_core::openhuman::memory::ingest_pipeline::ingest_chat; +use openhuman_core::openhuman::memory_sync::canonicalize::chat::{ChatBatch, ChatMessage}; +use openhuman_core::openhuman::memory_tree::retrieval::{fast_retrieve, FastRetrieveOptions}; + +fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + // Inert embedder — no Ollama/cloud in CI. + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + // Deterministic, Python-free entity extraction (regex fallback). + cfg.memory_tree.spacy_enabled = false; + (tmp, cfg) +} + +async fn seed_chat(cfg: &Config, source: &str, text: &str) { + let batch = ChatBatch { + platform: "slack".into(), + channel_label: source.into(), + messages: vec![ChatMessage { + author: "alice".into(), + timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), + text: text.into(), + source_ref: Some("slack://x".into()), + }], + }; + ingest_chat(cfg, source, "alice", vec![], batch) + .await + .expect("ingest_chat should succeed"); +} + +#[tokio::test] +async fn local_branch_returns_cooccurring_evidence() { + let (_tmp, cfg) = test_config(); + // alice + bob co-occur in one message → graph edge + both indexed on the + // same leaf chunk. + seed_chat( + &cfg, + "slack:#eng", + "Sync between alice@example.com and bob@example.com on the runbook.", + ) + .await; + // An unrelated message that should NOT surface for the alice+bob query. + seed_chat( + &cfg, + "slack:#random", + "Lunch plans for friday with the team.", + ) + .await; + + let resp = fast_retrieve( + &cfg, + "what did alice@example.com and bob@example.com discuss", + FastRetrieveOptions::default(), + ) + .await + .expect("fast_retrieve should succeed"); + + assert!( + !resp.hits.is_empty(), + "co-occurring entities should yield a local hit; got {resp:?}" + ); + // Coverage score = both query entities matched the same node. + assert!( + resp.hits.iter().any(|h| h.score >= 2.0), + "top local hit should have entity-coverage score >= 2; got {:?}", + resp.hits.iter().map(|h| h.score).collect::>() + ); + // Structured evidence, not prose — every hit has a node id + content. + assert!(resp.hits.iter().all(|h| !h.node_id.is_empty())); +} + +#[tokio::test] +async fn global_branch_handles_entity_free_query() { + let (_tmp, cfg) = test_config(); + seed_chat( + &cfg, + "slack:#eng", + "Sync between alice@example.com and bob@example.com on the runbook.", + ) + .await; + + // No mechanical entities in the query → global/dense branch. With the inert + // embedder this returns recency-ordered summaries (possibly empty), and + // crucially must not panic or error. + let resp = fast_retrieve( + &cfg, + "give me a recap of everything important", + FastRetrieveOptions::default(), + ) + .await + .expect("global branch should succeed"); + // total/truncated are well-formed regardless of hit count. + assert_eq!(resp.truncated, resp.total > resp.hits.len()); +} + +#[tokio::test] +async fn empty_store_returns_no_hits() { + let (_tmp, cfg) = test_config(); + let resp = fast_retrieve(&cfg, "anything at all", FastRetrieveOptions::default()) + .await + .expect("retrieval over empty store should succeed"); + assert!(resp.hits.is_empty()); + assert_eq!(resp.total, 0); +} diff --git a/tests/memory_threads_raw_coverage_e2e.rs b/tests/memory_threads_raw_coverage_e2e.rs index b49b4ccb0..9a51d8fe0 100644 --- a/tests/memory_threads_raw_coverage_e2e.rs +++ b/tests/memory_threads_raw_coverage_e2e.rs @@ -24,7 +24,6 @@ use openhuman_core::openhuman::embeddings::NoopEmbedding; use openhuman_core::openhuman::memory::query::{ MemoryQueryTool, MemoryTreeDrillDownTool, MemoryTreeFetchLeavesTool, MemoryTreeIngestDocumentTool, MemoryTreeQuerySourceTool, MemoryTreeSearchEntitiesTool, - MemoryTreeWalkTool, }; use openhuman_core::openhuman::memory::tools::{ MemoryForgetTool, MemoryRecallTool, MemoryStoreTool, @@ -1060,7 +1059,6 @@ fn memory_schema_registries_and_query_tool_metadata_cover_public_surfaces() { &MemoryTreeDrillDownTool, &MemoryTreeFetchLeavesTool, &MemoryTreeIngestDocumentTool, - &MemoryTreeWalkTool, ] { assert!(!tool.name().is_empty()); assert!(!tool.description().is_empty()); diff --git a/tests/memory_tree_sync_raw_coverage_e2e.rs b/tests/memory_tree_sync_raw_coverage_e2e.rs index 821e75fae..6a91c82c3 100644 --- a/tests/memory_tree_sync_raw_coverage_e2e.rs +++ b/tests/memory_tree_sync_raw_coverage_e2e.rs @@ -16,9 +16,6 @@ use tempfile::TempDir; use openhuman_core::core::event_bus::{DomainEvent, EventHandler}; use openhuman_core::openhuman::config::Config; use openhuman_core::openhuman::inference::provider::traits::{ChatMessage, Provider}; -use openhuman_core::openhuman::memory::query::{ - run_walk, MemoryTreeWalkTool, WalkOptions, WalkStopReason, -}; use openhuman_core::openhuman::memory_store::chunks::store::upsert_chunks; use openhuman_core::openhuman::memory_store::chunks::types::{ approx_token_count, chunk_id, Chunk, Metadata, SourceKind as ChunkSourceKind, SourceRef, @@ -45,10 +42,8 @@ use openhuman_core::openhuman::memory_tree::tree::{ append_leaf_deferred, get_or_create_tree, store as tree_store, LabelStrategy, LeafRef, }; use openhuman_core::openhuman::memory_tree::tree_runtime::{ - derive_parent_id, engine, estimate_tokens, level_from_node_id, rpc as tree_runtime_rpc, - store as runtime_store, TreeNode, + engine, rpc as tree_runtime_rpc, store as runtime_store, }; -use openhuman_core::openhuman::tools::traits::Tool; struct EnvVarGuard { key: &'static str, @@ -102,22 +97,6 @@ fn config_in(tmp: &TempDir) -> Config { cfg } -fn runtime_node(namespace: &str, node_id: &str, summary: &str) -> TreeNode { - let ts = Utc.with_ymd_and_hms(2026, 5, 29, 13, 45, 0).unwrap(); - TreeNode { - node_id: node_id.to_string(), - namespace: namespace.to_string(), - level: level_from_node_id(node_id), - parent_id: derive_parent_id(node_id), - summary: summary.to_string(), - token_count: estimate_tokens(summary), - child_count: 0, - created_at: ts, - updated_at: ts, - metadata: None, - } -} - fn staged_chunk(cfg: &Config, source_id: &str, seq: u32, tokens: u32) -> Chunk { let ts = Utc .timestamp_millis_opt(1_700_000_000_000 + seq as i64) @@ -280,34 +259,6 @@ async fn tree_runtime_engine_rpc_and_walk_cover_success_and_edge_paths() { .expect("rebuild tree"); assert_eq!(rebuilt.total_nodes, 6); assert_eq!(runtime_store::buffer_read(&cfg, ns).unwrap().len(), 1); - - let walk_provider = ScriptedProvider::new([ - r#"Surveying children {"name":"peek","arguments":{"node_ids":["2026","missing"]}}"#, - r#"{"name":"descend","arguments":{"node_id":"2026/05/29"}}"#, - r#"{"name":"fetch_leaves","arguments":{"node_id":"2026/05/29"}}"#, - r#"{"name":"answer","arguments":{"text":"Alice and Bob discussed launch cleanup."}}"#, - ]); - let outcome = run_walk( - &cfg, - &walk_provider, - "Who discussed launch cleanup?", - WalkOptions { - max_turns: 8, - start_node_id: None, - namespace: ns.into(), - model: Some("scripted".into()), - }, - ) - .await - .expect("walk"); - assert_eq!(outcome.stopped_reason, WalkStopReason::Answered); - assert_eq!(outcome.turns_used, 4); - assert!(outcome.answer.contains("Alice and Bob")); - - let tool = MemoryTreeWalkTool; - assert_eq!(tool.name(), "memory_tree_walk"); - let missing_query = tool.execute(json!({ "namespace": ns })).await.unwrap_err(); - assert!(missing_query.to_string().contains("`query` is required")); } #[tokio::test] @@ -377,72 +328,6 @@ async fn bucket_seal_deferred_and_fallback_paths_preserve_buffers_and_labels() { assert_eq!(parent.item_ids, sealed); } -#[tokio::test] -async fn memory_walk_provider_errors_and_unknown_actions_are_reported() { - let tmp = TempDir::new().expect("tempdir"); - let cfg = config_in(&tmp); - let ns = "walk-round14"; - for node in [ - runtime_node(ns, "root", "root has one 2026 child"), - runtime_node(ns, "2026", "year node"), - ] { - runtime_store::write_node(&cfg, &node).expect("write node"); - } - - let provider = ScriptedProvider::new([ - r#"{"name":"dance","arguments":{"node_id":"2026"}}"#, - "", - ]); - let outcome = run_walk( - &cfg, - &provider, - "exercise unknown action", - WalkOptions { - max_turns: 2, - start_node_id: Some("root".into()), - namespace: ns.into(), - model: Some("scripted".into()), - }, - ) - .await - .expect("unknown action walk"); - assert_eq!(outcome.stopped_reason, WalkStopReason::LlmGaveUp); - assert!(outcome.trace[0] - .result_preview - .contains("unknown walk action")); - - struct FailingProvider; - #[async_trait] - impl Provider for FailingProvider { - async fn chat_with_system( - &self, - system_prompt: Option<&str>, - message: &str, - model: &str, - temperature: f64, - ) -> anyhow::Result { - let _ = (system_prompt, message, model, temperature); - anyhow::bail!("scripted provider failure") - } - } - - let failed = run_walk( - &cfg, - &FailingProvider, - "force provider error", - WalkOptions { - max_turns: 1, - start_node_id: Some("missing".into()), - namespace: ns.into(), - model: None, - }, - ) - .await - .expect("provider errors become walk outcome"); - assert!(matches!(failed.stopped_reason, WalkStopReason::Error(_))); - assert!(failed.answer.contains("Walk failed")); -} - #[tokio::test] async fn composio_providers_sync_state_and_bus_surfaces_cover_read_write_edges() { let _lock = env_lock(); diff --git a/tests/memory_tree_walk_e2e.rs b/tests/memory_tree_walk_e2e.rs deleted file mode 100644 index 2f042fd9e..000000000 --- a/tests/memory_tree_walk_e2e.rs +++ /dev/null @@ -1,535 +0,0 @@ -//! E2E tests for the `memory_tree_walk` agentic tool. -//! -//! These tests exercise `run_walk` end-to-end against a real HTTP mock -//! LLM server (wiremock) to prove that: -//! - The `OpenAiCompatibleProvider` → `run_walk` chain works over HTTP. -//! - Tool-call XML parsing, trace assembly, and stop-reason detection -//! all behave correctly with per-turn scripted LLM responses. -//! - Turn-cap enforcement fires correctly. -//! - Unknown node descends are handled gracefully. -//! -//! Tree nodes are seeded directly via `store::write_node` so the tests -//! are isolated from the summariser pipeline and focus purely on walk. -//! -//! Run with: -//! cargo test --test memory_tree_walk_e2e -//! or via the project wrapper: -//! bash scripts/test-rust-with-mock.sh --test memory_tree_walk_e2e - -use std::collections::VecDeque; -use std::sync::{Arc, Mutex, OnceLock}; - -use chrono::Utc; -use serde_json::json; -use tempfile::TempDir; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; - -use openhuman_core::openhuman::config::Config; -use openhuman_core::openhuman::inference::provider::compatible::{ - AuthStyle, OpenAiCompatibleProvider, -}; -use openhuman_core::openhuman::memory_tree::tools::walk::{run_walk, WalkOptions, WalkStopReason}; -use openhuman_core::openhuman::memory_tree::tree_runtime::store::write_node; -use openhuman_core::openhuman::memory_tree::tree_runtime::types::{ - derive_parent_id, estimate_tokens, level_from_node_id, TreeNode, -}; - -// ── Environment serialisation lock ────────────────────────────────────────── -// -// `OPENHUMAN_WORKSPACE` is a process-global env var. Tests that set it -// must serialise with this lock. - -static ENV_LOCK: OnceLock> = OnceLock::new(); - -fn env_lock() -> std::sync::MutexGuard<'static, ()> { - let m = ENV_LOCK.get_or_init(|| Mutex::new(())); - match m.lock() { - Ok(g) => g, - Err(p) => p.into_inner(), - } -} - -// ── Per-turn sequential mock responder ────────────────────────────────────── -// -// `ScriptedResponder` pops one canned OpenAI-compatible JSON response from its -// queue on each `respond` call. This allows a single wiremock `Mock` to hand -// back turn-1, turn-2, turn-3 responses in sequence. -// -// The queue is `Arc>>` so it can be shared between the -// `Mock` registration (which clones the arc) and the test body for inspection. - -#[derive(Clone)] -struct ScriptedResponder { - queue: Arc>>, - call_count: Arc>, -} - -impl ScriptedResponder { - /// Create a new responder pre-loaded with `content_strings`. - /// - /// Each string becomes the `choices[0].message.content` of an - /// OpenAI-compatible non-streaming response. - fn new(content_strings: Vec<&str>) -> Self { - log::debug!( - "[memory_tree_walk_e2e] ScriptedResponder created with {} turns", - content_strings.len() - ); - let queue: VecDeque = content_strings.iter().map(|s| s.to_string()).collect(); - Self { - queue: Arc::new(Mutex::new(queue)), - call_count: Arc::new(Mutex::new(0)), - } - } - - fn call_count(&self) -> usize { - *self.call_count.lock().unwrap() - } -} - -impl Respond for ScriptedResponder { - fn respond(&self, _request: &Request) -> ResponseTemplate { - let mut queue = self.queue.lock().unwrap(); - let mut count = self.call_count.lock().unwrap(); - *count += 1; - let current_turn = *count; - - let content = queue - .pop_front() - .unwrap_or_else(|| "ScriptedResponder: queue exhausted".to_string()); - - log::debug!( - "[memory_tree_walk_e2e] ScriptedResponder turn={current_turn} content_preview={}", - &content[..content.len().min(80)] - ); - - let body = json!({ - "id": format!("chatcmpl-walk-e2e-{current_turn}"), - "object": "chat.completion", - "created": 1_700_000_000_u64, - "model": "e2e-walk-model", - "choices": [{ - "index": 0, - "message": { - "role": "assistant", - "content": content - }, - "finish_reason": "stop" - }], - "usage": { - "prompt_tokens": 20, - "completion_tokens": 10, - "total_tokens": 30 - } - }); - - ResponseTemplate::new(200).set_body_json(body) - } -} - -// ── Tree helpers ───────────────────────────────────────────────────────────── - -/// Build a minimal `Config` pointing at a temp workspace. -fn test_config(tmp: &TempDir) -> Config { - let mut cfg = Config::default(); - cfg.workspace_dir = tmp.path().join("workspace"); - std::fs::create_dir_all(&cfg.workspace_dir).expect("create workspace dir"); - cfg -} - -/// Create a `TreeNode` with the correct level and parent derived from its id. -fn make_node(namespace: &str, node_id: &str, summary: &str, child_count: u32) -> TreeNode { - let level = level_from_node_id(node_id); - let parent_id = derive_parent_id(node_id); - let ts = Utc::now(); - TreeNode { - node_id: node_id.to_string(), - namespace: namespace.to_string(), - level, - parent_id, - summary: summary.to_string(), - token_count: estimate_tokens(summary), - child_count, - created_at: ts, - updated_at: ts, - metadata: None, - } -} - -/// Seed a 3-level tree: -/// root → 2024 (year) → 2024/01 (month leaf) -/// -/// Returns the node ids for easy reference in tests. -fn seed_tree(cfg: &Config, ns: &str) -> (&'static str, &'static str, &'static str) { - let root_id = "root"; - let year_id = "2024"; - let month_id = "2024/01"; - - write_node( - cfg, - &make_node(ns, root_id, "All-time summary: project activity 2024.", 1), - ) - .expect("write root node"); - - write_node( - cfg, - &make_node( - ns, - year_id, - "Year 2024: major deployment cycle, shipped v1 release.", - 1, - ), - ) - .expect("write year node"); - - write_node( - cfg, - &make_node( - ns, - month_id, - "January 2024: user discussed deployment X on the Slack channel.", - 0, - ), - ) - .expect("write month node"); - - log::debug!( - "[memory_tree_walk_e2e] seeded tree namespace={ns} root={root_id} year={year_id} month={month_id}" - ); - - (root_id, year_id, month_id) -} - -// ── Provider helper ─────────────────────────────────────────────────────────── - -/// Build an `OpenAiCompatibleProvider` pointing at `wiremock_uri/v1`. -fn make_provider(wiremock_uri: &str) -> OpenAiCompatibleProvider { - log::debug!( - "[memory_tree_walk_e2e] building provider base_url={}/v1", - wiremock_uri - ); - OpenAiCompatibleProvider::new( - "e2e-walk-test", - &format!("{}/v1", wiremock_uri), - Some("test-key"), - AuthStyle::Bearer, - ) -} - -// ── Test 1: walks_descend_fetch_answer ─────────────────────────────────────── - -/// Happy-path walk: descend → fetch_leaves → answer in 3 turns. -/// -/// Validates: -/// - `WalkStopReason::Answered` -/// - `answer` contains expected text -/// - trace has 3 steps with correct action names -/// - `turns_used == 3` -/// - mock LLM received exactly 3 HTTP calls -#[tokio::test] -async fn walks_descend_fetch_answer() { - let _lock = env_lock(); - - log::debug!("[memory_tree_walk_e2e] test=walks_descend_fetch_answer starting"); - - // ── Start wiremock ── - let server = MockServer::start().await; - log::debug!( - "[memory_tree_walk_e2e] wiremock listening at {}", - server.uri() - ); - - // ── Seed tree ── - let tmp = TempDir::new().expect("tempdir"); - let cfg = test_config(&tmp); - let ns = "e2e-walk-test"; - let (_, year_id, month_id) = seed_tree(&cfg, ns); - - // ── Script 3 turns ── - // Turn 1: descend into the year node. - // Turn 2: fetch_leaves on the month node. - // Turn 3: answer. - let turn1 = format!( - r#"{{"name":"descend","arguments":{{"node_id":"{year_id}"}}}}"# - ); - let turn2 = format!( - r#"{{"name":"fetch_leaves","arguments":{{"node_id":"{month_id}"}}}}"# - ); - let turn3 = - r#"{"name":"answer","arguments":{"text":"The user discussed deployment X"}}"#.to_string(); - - let responder = ScriptedResponder::new(vec![&turn1, &turn2, &turn3]); - let responder_clone = responder.clone(); - - Mock::given(method("POST")) - .and(path("/v1/chat/completions")) - .respond_with(responder) - .mount(&server) - .await; - - // ── Run walk ── - let provider = make_provider(&server.uri()); - let opts = WalkOptions { - max_turns: 6, - start_node_id: None, - namespace: ns.to_string(), - model: Some("e2e-walk-model".into()), - }; - - log::debug!("[memory_tree_walk_e2e] calling run_walk query='deployment query'"); - let outcome = run_walk( - &cfg, - &provider, - "What was discussed about deployment?", - opts, - ) - .await - .expect("run_walk should succeed"); - - log::debug!( - "[memory_tree_walk_e2e] outcome stopped_reason={:?} turns_used={} trace_len={}", - outcome.stopped_reason, - outcome.turns_used, - outcome.trace.len() - ); - - // ── Assertions ── - assert_eq!( - outcome.stopped_reason, - WalkStopReason::Answered, - "walk should stop with Answered, got {:?}", - outcome.stopped_reason - ); - assert!( - outcome.answer.contains("The user discussed deployment X"), - "answer should contain expected text, got: {}", - outcome.answer - ); - assert_eq!( - outcome.turns_used, 3, - "expected 3 turns, got {}", - outcome.turns_used - ); - assert_eq!( - outcome.trace.len(), - 3, - "expected trace of 3 steps, got {}", - outcome.trace.len() - ); - assert_eq!( - outcome.trace[0].action, "descend", - "step 0 should be descend" - ); - assert_eq!( - outcome.trace[1].action, "fetch_leaves", - "step 1 should be fetch_leaves" - ); - assert_eq!(outcome.trace[2].action, "answer", "step 2 should be answer"); - - // Verify LLM was called exactly 3 times over HTTP. - let llm_calls = responder_clone.call_count(); - assert_eq!( - llm_calls, 3, - "LLM mock should have received exactly 3 HTTP requests, got {llm_calls}" - ); - - log::debug!("[memory_tree_walk_e2e] test=walks_descend_fetch_answer PASSED"); -} - -// ── Test 2: respects_max_turns_cap_with_mock ───────────────────────────────── - -/// Turn-cap enforcement: the mock always returns `descend` (never `answer`), -/// so the walk must stop at `max_turns` with `MaxTurnsReached`. -/// -/// Validates: -/// - `WalkStopReason::MaxTurnsReached` -/// - `turns_used == max_turns` (3) -/// - fallback answer contains "converge" or similar marker -/// - mock received exactly 3 HTTP calls (== max_turns) -#[tokio::test] -async fn respects_max_turns_cap_with_mock() { - let _lock = env_lock(); - - log::debug!("[memory_tree_walk_e2e] test=respects_max_turns_cap_with_mock starting"); - - let server = MockServer::start().await; - log::debug!( - "[memory_tree_walk_e2e] wiremock listening at {}", - server.uri() - ); - - let tmp = TempDir::new().expect("tempdir"); - let cfg = test_config(&tmp); - let ns = "e2e-walk-cap-test"; - let (_, year_id, _) = seed_tree(&cfg, ns); - - // Always descend (never answer) — more entries than max_turns so the cap fires. - let forever_descend = format!( - r#"{{"name":"descend","arguments":{{"node_id":"{year_id}"}}}}"# - ); - - let responder = ScriptedResponder::new(vec![ - &forever_descend, - &forever_descend, - &forever_descend, - &forever_descend, - &forever_descend, - ]); - let responder_clone = responder.clone(); - - Mock::given(method("POST")) - .and(path("/v1/chat/completions")) - .respond_with(responder) - .mount(&server) - .await; - - let provider = make_provider(&server.uri()); - let opts = WalkOptions { - max_turns: 3, - start_node_id: None, - namespace: ns.to_string(), - model: Some("e2e-walk-model".into()), - }; - - log::debug!("[memory_tree_walk_e2e] calling run_walk with max_turns=3"); - let outcome = run_walk(&cfg, &provider, "infinite loop query", opts) - .await - .expect("run_walk should succeed (not error)"); - - log::debug!( - "[memory_tree_walk_e2e] outcome stopped_reason={:?} turns_used={} trace_len={}", - outcome.stopped_reason, - outcome.turns_used, - outcome.trace.len() - ); - - assert_eq!( - outcome.stopped_reason, - WalkStopReason::MaxTurnsReached, - "walk should stop with MaxTurnsReached, got {:?}", - outcome.stopped_reason - ); - assert_eq!( - outcome.turns_used, 3, - "turns_used should be max_turns=3, got {}", - outcome.turns_used - ); - assert!( - outcome.answer.to_lowercase().contains("converge") - || outcome.answer.to_lowercase().contains("turn limit") - || outcome.answer.to_lowercase().contains("could not"), - "fallback answer should indicate failure to converge, got: {}", - outcome.answer - ); - - // Mock should have been called exactly max_turns (3) times. - let llm_calls = responder_clone.call_count(); - assert_eq!( - llm_calls, 3, - "LLM mock should have received exactly 3 HTTP requests (max_turns), got {llm_calls}" - ); - - log::debug!("[memory_tree_walk_e2e] test=respects_max_turns_cap_with_mock PASSED"); -} - -// ── Test 3: handles_unknown_node_gracefully ─────────────────────────────────── - -/// Unknown-node recovery: turn 1 descends into a non-existent node; -/// the walk reports "unknown node" in the trace but continues. -/// Turn 2 answers, so the walk completes with `Answered`. -/// -/// Validates: -/// - `WalkStopReason::Answered` -/// - `trace[0].result_preview` contains "unknown node" -/// - `trace.len() == 2` -/// - answer from turn 2 is preserved -#[tokio::test] -async fn handles_unknown_node_gracefully() { - let _lock = env_lock(); - - log::debug!("[memory_tree_walk_e2e] test=handles_unknown_node_gracefully starting"); - - let server = MockServer::start().await; - log::debug!( - "[memory_tree_walk_e2e] wiremock listening at {}", - server.uri() - ); - - let tmp = TempDir::new().expect("tempdir"); - let cfg = test_config(&tmp); - let ns = "e2e-walk-unknown-test"; - seed_tree(&cfg, ns); - - // Turn 1: descend into a node that does not exist. - let turn1 = - r#"{"name":"descend","arguments":{"node_id":"does_not_exist"}}"#; - // Turn 2: answer (walk continues after bad descend). - let turn2 = r#"{"name":"answer","arguments":{"text":"I gave up: node not found"}}"#; - - let responder = ScriptedResponder::new(vec![turn1, turn2]); - let responder_clone = responder.clone(); - - Mock::given(method("POST")) - .and(path("/v1/chat/completions")) - .respond_with(responder) - .mount(&server) - .await; - - let provider = make_provider(&server.uri()); - let opts = WalkOptions { - max_turns: 6, - start_node_id: None, - namespace: ns.to_string(), - model: Some("e2e-walk-model".into()), - }; - - log::debug!("[memory_tree_walk_e2e] calling run_walk with unknown-node script"); - let outcome = run_walk(&cfg, &provider, "find nonexistent data", opts) - .await - .expect("run_walk should succeed"); - - log::debug!( - "[memory_tree_walk_e2e] outcome stopped_reason={:?} turns_used={} trace_len={}", - outcome.stopped_reason, - outcome.turns_used, - outcome.trace.len() - ); - - // Walk should complete despite the bad descend. - assert_eq!( - outcome.stopped_reason, - WalkStopReason::Answered, - "walk should eventually answer, got {:?}", - outcome.stopped_reason - ); - - assert_eq!( - outcome.trace.len(), - 2, - "expected 2 trace steps, got {}", - outcome.trace.len() - ); - - // The first step's result_preview should indicate the unknown node. - let step0_preview = &outcome.trace[0].result_preview; - assert!( - step0_preview.contains("unknown node"), - "step 0 result_preview should contain 'unknown node', got: {step0_preview}" - ); - - // The final answer should be from turn 2. - assert!( - outcome.answer.contains("I gave up"), - "answer should contain turn-2 text, got: {}", - outcome.answer - ); - - // 2 HTTP calls made. - let llm_calls = responder_clone.call_count(); - assert_eq!( - llm_calls, 2, - "LLM mock should have received exactly 2 HTTP requests, got {llm_calls}" - ); - - log::debug!("[memory_tree_walk_e2e] test=handles_unknown_node_gracefully PASSED"); -}