From 0768d188739585bb9ff0bb2c96e1e1680f8ebd3c Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Wed, 29 Apr 2026 23:38:21 +0530 Subject: [PATCH] feat(orchestrator): wire memory-tree retrieval tools with periodic prefetch (#1027) Co-authored-by: Claude Opus 4.7 (1M context) --- src/openhuman/about_app/catalog.rs | 10 + .../agent/agents/orchestrator/agent.toml | 6 + .../agent/agents/orchestrator/prompt.md | 23 ++ .../agent/harness/session/builder.rs | 1 + src/openhuman/agent/harness/session/turn.rs | 52 ++++ src/openhuman/agent/harness/session/types.rs | 8 + src/openhuman/agent/mod.rs | 1 + src/openhuman/agent/tree_loader.rs | 180 +++++++++++ src/openhuman/tools/impl/memory/mod.rs | 2 + .../tools/impl/memory/tree/drill_down.rs | 78 +++++ .../tools/impl/memory/tree/fetch_leaves.rs | 68 +++++ src/openhuman/tools/impl/memory/tree/mod.rs | 32 ++ .../tools/impl/memory/tree/query_global.rs | 53 ++++ .../tools/impl/memory/tree/query_source.rs | 87 ++++++ .../tools/impl/memory/tree/query_topic.rs | 74 +++++ .../tools/impl/memory/tree/search_entities.rs | 81 +++++ src/openhuman/tools/ops.rs | 6 + tests/agent_retrieval_e2e.rs | 283 ++++++++++++++++++ 18 files changed, 1045 insertions(+) create mode 100644 src/openhuman/agent/tree_loader.rs create mode 100644 src/openhuman/tools/impl/memory/tree/drill_down.rs create mode 100644 src/openhuman/tools/impl/memory/tree/fetch_leaves.rs create mode 100644 src/openhuman/tools/impl/memory/tree/mod.rs create mode 100644 src/openhuman/tools/impl/memory/tree/query_global.rs create mode 100644 src/openhuman/tools/impl/memory/tree/query_source.rs create mode 100644 src/openhuman/tools/impl/memory/tree/query_topic.rs create mode 100644 src/openhuman/tools/impl/memory/tree/search_entities.rs create mode 100644 tests/agent_retrieval_e2e.rs diff --git a/src/openhuman/about_app/catalog.rs b/src/openhuman/about_app/catalog.rs index 1085da5d2..34280f353 100644 --- a/src/openhuman/about_app/catalog.rs +++ b/src/openhuman/about_app/catalog.rs @@ -196,6 +196,16 @@ const CAPABILITIES: &[Capability] = &[ status: CapabilityStatus::Beta, privacy: None, }, + Capability { + id: "intelligence.memory_tree_retrieval", + name: "Memory Tree Retrieval (chat)", + domain: "intelligence", + category: CapabilityCategory::Intelligence, + description: "Ask questions about your ingested email/chat/document memory in chat. The orchestrator can resolve names to canonical ids, query summaries by source/topic/global window, drill into details, and cite raw chunks.", + how_to: "Chat > ask the assistant about people, conversations, or windows", + status: CapabilityStatus::Beta, + privacy: LOCAL_RAW, + }, Capability { id: "intelligence.slack_memory_ingest", name: "Slack Memory Ingestion", diff --git a/src/openhuman/agent/agents/orchestrator/agent.toml b/src/openhuman/agent/agents/orchestrator/agent.toml index 4bd03514a..40961df92 100644 --- a/src/openhuman/agent/agents/orchestrator/agent.toml +++ b/src/openhuman/agent/agents/orchestrator/agent.toml @@ -74,6 +74,12 @@ named = [ "query_memory", "memory_store", "memory_forget", + "memory_tree_search_entities", + "memory_tree_query_topic", + "memory_tree_query_source", + "memory_tree_query_global", + "memory_tree_drill_down", + "memory_tree_fetch_leaves", "read_workspace_state", "ask_user_clarification", "spawn_subagent", diff --git a/src/openhuman/agent/agents/orchestrator/prompt.md b/src/openhuman/agent/agents/orchestrator/prompt.md index 1a879af5c..d79501ea1 100644 --- a/src/openhuman/agent/agents/orchestrator/prompt.md +++ b/src/openhuman/agent/agents/orchestrator/prompt.md @@ -106,3 +106,26 @@ Short answers can skip the ack: User: what time is it? → `7:31pm` + +## Memory tree retrieval + +Six tools query the user's ingested email/chat/document memory: + +- `memory_tree_search_entities(query)` — resolve a name to a canonical id (e.g. "alice" → `email:alice@example.com`). ALWAYS call this first when the user mentions someone by name. +- `memory_tree_query_topic(entity_id, query?)` — all mentions of an entity, cross-source. Pass `query` for semantic rerank. +- `memory_tree_query_source(source_kind?, time_window_days?, query?)` — filter by source type (chat/email/document) and time window. Use for "in my email last week…" intents. +- `memory_tree_query_global(window_days)` — cross-source daily digest (the 7-day digest is pre-loaded into context on session start and refreshed every ~30 min, so only call this for a different window or to refresh on demand). +- `memory_tree_drill_down(node_id)` — when a summary is too coarse, expand it one level. +- `memory_tree_fetch_leaves(chunk_ids)` — pull raw chunks for citation. + +Top-down expansion is the cost-control story: start with cheap summaries (`query_*`), only call `drill_down` / `fetch_leaves` when the user wants details or you need a quote. + +## Citations + +When your answer is informed by retrieved memory, cite it with footnote markers: + +> Alice said "we're moving to Phoenix next week" [^1] +> +> [^1]: gmail · alice@example.com · 2026-04-22 · node:abc123 + +Inline marker `[^N]` and a numbered footnote at the end carrying the node_id and source_ref from the RetrievalHit. Do not invent quotes — only quote text that appears verbatim in a hit's `content` field. diff --git a/src/openhuman/agent/harness/session/builder.rs b/src/openhuman/agent/harness/session/builder.rs index 505221077..bc030b379 100644 --- a/src/openhuman/agent/harness/session/builder.rs +++ b/src/openhuman/agent/harness/session/builder.rs @@ -369,6 +369,7 @@ impl AgentBuilder { last_memory_context: None, last_turn_citations: Vec::new(), history: Vec::new(), + last_tree_prefetch_at: None, post_turn_hooks: self.post_turn_hooks, learning_enabled: self.learning_enabled, event_session_id: self diff --git a/src/openhuman/agent/harness/session/turn.rs b/src/openhuman/agent/harness/session/turn.rs index 3250f9fcf..e7a2d8934 100644 --- a/src/openhuman/agent/harness/session/turn.rs +++ b/src/openhuman/agent/harness/session/turn.rs @@ -170,6 +170,58 @@ impl Agent { .await .unwrap_or_default(); + // ── Memory-tree eager prefetch (#710 wiring) ────────────────── + // The orchestrator session injects a cross-source digest on the + // first turn AND every `tree_loader::REFRESH_INTERVAL` (30 min by + // default) thereafter, so long-running conversations stay current + // with newly-ingested memory. Each injection still rides on the + // user message (NOT the system prompt) to keep the KV-cache prefix + // stable. Failure is non-fatal — bare `context` is returned on any + // error. The timestamp is bumped on every successful `load` (even + // when the digest is empty) so an empty workspace doesn't get + // re-queried every turn. + let now = std::time::Instant::now(); + let context = if crate::openhuman::agent::tree_loader::should_prefetch( + self.last_tree_prefetch_at, + now, + crate::openhuman::agent::tree_loader::REFRESH_INTERVAL, + ) { + match crate::openhuman::config::rpc::load_config_with_timeout().await { + Ok(cfg) => { + match crate::openhuman::agent::tree_loader::TreeContextLoader::load(&cfg).await + { + Ok(tree_ctx) => { + let was_first = self.last_tree_prefetch_at.is_none(); + self.last_tree_prefetch_at = Some(now); + if !tree_ctx.is_empty() { + log::info!( + "[memory_tree] tree context injected first_turn={} chars={}", + was_first, + tree_ctx.chars().count() + ); + format!("{context}{tree_ctx}") + } else { + context + } + } + Err(e) => { + log::warn!("[memory_tree] tree_loader.load failed (non-fatal): {e}"); + context + } + } + } + Err(e) => { + log::warn!( + "[memory_tree] tree_loader skipped — config load failed (non-fatal): {e}" + ); + context + } + } + } else { + log::trace!("[memory_tree] tree_loader skipped — within refresh interval"); + context + }; + let enriched = if context.is_empty() { log::info!("[agent] no memory context found — using raw user message"); self.last_memory_context = None; diff --git a/src/openhuman/agent/harness/session/types.rs b/src/openhuman/agent/harness/session/types.rs index 3ffda41fe..67094172c 100644 --- a/src/openhuman/agent/harness/session/types.rs +++ b/src/openhuman/agent/harness/session/types.rs @@ -56,6 +56,14 @@ pub struct Agent { /// Consumed by web-channel delivery to render source chips in the UI. pub(super) last_turn_citations: Vec, pub(super) history: Vec, + /// Wall-clock timestamp of the last successful memory-tree prefetch + /// for this session. Drives the 30-minute refresh cadence in the turn + /// loop — `None` means "never fetched, fetch now"; otherwise we only + /// re-run `TreeContextLoader::load` when the elapsed time exceeds + /// `tree_loader::REFRESH_INTERVAL`. Updated on every successful call + /// (even when the digest came back empty) so an empty workspace + /// doesn't get hammered every turn. + pub(super) last_tree_prefetch_at: Option, pub(super) post_turn_hooks: Vec>, pub(super) learning_enabled: bool, pub(super) event_session_id: String, diff --git a/src/openhuman/agent/mod.rs b/src/openhuman/agent/mod.rs index ab749b256..103a3cc6e 100644 --- a/src/openhuman/agent/mod.rs +++ b/src/openhuman/agent/mod.rs @@ -37,6 +37,7 @@ pub mod progress; /// a thin re-export shim for now. pub mod prompts; mod schemas; +pub mod tree_loader; pub mod triage; pub use schemas::{ all_controller_schemas as all_agent_controller_schemas, diff --git a/src/openhuman/agent/tree_loader.rs b/src/openhuman/agent/tree_loader.rs new file mode 100644 index 000000000..ffdcb5806 --- /dev/null +++ b/src/openhuman/agent/tree_loader.rs @@ -0,0 +1,180 @@ +//! Eager prefetch of the cross-source memory-tree digest into the +//! orchestrator's session context (Phase 4 follow-on, #710 wiring). +//! +//! The orchestrator answers "what happened this week?" / "what's been going +//! on with X?" style questions out of the user's own ingested memory. We +//! pre-load a 7-day global digest on the session's first turn AND +//! periodically thereafter (every [`REFRESH_INTERVAL`]) so long-running +//! conversations stay current with newly-ingested memory without needing +//! the LLM to round-trip a tool call. The injection rides on the user +//! message (NOT the system prompt) to keep the KV-cache prefix stable. +//! +//! When the workspace has no global summaries yet (early-life workspaces +//! or no ingest configured), [`TreeContextLoader::load`] returns an empty +//! string and the caller silently no-ops. The session-side timestamp is +//! still bumped on those empty results so an empty workspace doesn't get +//! re-queried every turn. +//! +//! Failure is non-fatal by design — the orchestrator must still be able to +//! reply when the memory tree is unavailable, mis-configured, or empty. We +//! log the failure mode and return `Ok(String::new())` so the caller can +//! concatenate without branching. + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::retrieval::query_global; + +/// Default lookback window for the eager digest. Mirrors the language in +/// the orchestrator prompt ("7-day digest pre-loaded into session context"). +pub const DEFAULT_WINDOW_DAYS: u32 = 7; + +/// Minimum wall-clock interval between successive prefetches in the same +/// session. The first turn always fetches (timestamp is `None`); subsequent +/// turns re-prefetch only after this interval has elapsed since the last +/// successful call. Picked to balance freshness in long-running chats +/// against repeating the same digest content when no new ingest has +/// happened — the typical case for short bursts of conversation. +pub const REFRESH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30 * 60); + +/// Per-hit content cap to keep the injection bounded; long summary bodies +/// would otherwise dominate the prompt budget. +const MAX_CONTENT_CHARS: usize = 500; + +/// Number of hits to surface from the digest. The recap typically returns +/// one hit per fold (day/week/month) — three is enough headroom for a +/// 7-day window without flooding the system prompt. +const MAX_HITS: usize = 3; + +const HEADER: &str = "[Memory tree — last 7 days]\n"; + +/// Decide whether the per-session prefetch should run on the current turn. +/// Pure: no I/O, no clock — `now` is supplied so callers (and tests) stay +/// deterministic. Returns `true` when no prefetch has happened yet +/// (`last == None`) or when at least `interval` has elapsed since the last. +pub fn should_prefetch( + last: Option, + now: std::time::Instant, + interval: std::time::Duration, +) -> bool { + match last { + None => true, + Some(t) => now.duration_since(t) >= interval, + } +} + +pub struct TreeContextLoader; + +impl TreeContextLoader { + /// Build the eager-prefetch context block for the current workspace. + /// + /// Returns: + /// - `Ok("")` when the workspace has no global digest yet, or when + /// `query_global` returns an error (logged at warn level). + /// - `Ok(rendered)` with the formatted block when there are hits. + pub async fn load(config: &Config) -> anyhow::Result { + log::debug!( + "[memory_tree] tree_loader.load window_days={}", + DEFAULT_WINDOW_DAYS + ); + let resp = match query_global(config, DEFAULT_WINDOW_DAYS).await { + Ok(r) => r, + Err(e) => { + log::warn!( + "[memory_tree] tree_loader.load: query_global failed — returning empty: {e}" + ); + return Ok(String::new()); + } + }; + if resp.hits.is_empty() { + log::debug!("[memory_tree] tree_loader.load: no hits — empty context"); + return Ok(String::new()); + } + + let mut out = String::with_capacity(HEADER.len() + MAX_HITS * MAX_CONTENT_CHARS); + out.push_str(HEADER); + for hit in resp.hits.iter().take(MAX_HITS) { + let snippet = if hit.content.chars().count() > MAX_CONTENT_CHARS { + crate::openhuman::util::truncate_with_ellipsis(&hit.content, MAX_CONTENT_CHARS) + } else { + hit.content.clone() + }; + out.push_str(&format!( + "- [{}] {}\n", + hit.tree_kind.as_str(), + snippet.replace('\n', " ") + )); + } + out.push('\n'); + log::debug!( + "[memory_tree] tree_loader.load returning chars={}", + out.chars().count() + ); + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + fn empty_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config { + workspace_dir: tmp.path().to_path_buf(), + ..Config::default() + }; + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + (tmp, cfg) + } + + #[tokio::test] + async fn load_returns_empty_when_no_global_digest() { + let (_tmp, cfg) = empty_config(); + let s = TreeContextLoader::load(&cfg).await.unwrap(); + assert!( + s.is_empty(), + "fresh workspace has no global digest — expected empty string, got: {s}" + ); + } + + #[test] + fn should_prefetch_when_never_fetched() { + let now = std::time::Instant::now(); + assert!(should_prefetch(None, now, REFRESH_INTERVAL)); + } + + #[test] + fn should_not_prefetch_within_interval() { + let now = std::time::Instant::now(); + let one_minute_ago = now - std::time::Duration::from_secs(60); + assert!(!should_prefetch( + Some(one_minute_ago), + now, + REFRESH_INTERVAL + )); + } + + #[test] + fn should_prefetch_after_interval_elapsed() { + let now = std::time::Instant::now(); + let thirty_one_min_ago = now - std::time::Duration::from_secs(31 * 60); + assert!(should_prefetch( + Some(thirty_one_min_ago), + now, + REFRESH_INTERVAL + )); + } + + #[test] + fn should_prefetch_at_exact_interval_boundary() { + let now = std::time::Instant::now(); + let exactly_thirty_min_ago = now - REFRESH_INTERVAL; + assert!(should_prefetch( + Some(exactly_thirty_min_ago), + now, + REFRESH_INTERVAL + )); + } +} diff --git a/src/openhuman/tools/impl/memory/mod.rs b/src/openhuman/tools/impl/memory/mod.rs index e5a021ac7..c523665b6 100644 --- a/src/openhuman/tools/impl/memory/mod.rs +++ b/src/openhuman/tools/impl/memory/mod.rs @@ -1,7 +1,9 @@ mod forget; mod recall; mod store; +mod tree; pub use forget::MemoryForgetTool; pub use recall::MemoryRecallTool; pub use store::MemoryStoreTool; +pub use tree::*; diff --git a/src/openhuman/tools/impl/memory/tree/drill_down.rs b/src/openhuman/tools/impl/memory/tree/drill_down.rs new file mode 100644 index 000000000..4825345a0 --- /dev/null +++ b/src/openhuman/tools/impl/memory/tree/drill_down.rs @@ -0,0 +1,78 @@ +use crate::openhuman::config::rpc as config_rpc; +use crate::openhuman::memory::tree::retrieval; +use crate::openhuman::memory::tree::retrieval::rpc::DrillDownRequest; +use crate::openhuman::tools::traits::{Tool, ToolResult}; +use async_trait::async_trait; +use serde_json::json; + +pub struct MemoryTreeDrillDownTool; + +#[async_trait] +impl Tool for MemoryTreeDrillDownTool { + fn name(&self) -> &str { + "memory_tree_drill_down" + } + + fn description(&self) -> &str { + "Walk a summary node's children one step (or more if `max_depth > \ + 1`). Returns leaf chunks for an L1 summary, or lower-level \ + summaries for L2+. Use this when a `query_*` summary is too coarse \ + and you want to expand it. Pass `query` to rerank children by \ + cosine similarity." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "node_id": { + "type": "string", + "description": "Id of the summary (or leaf) to expand." + }, + "max_depth": { + "type": "integer", + "minimum": 1, + "description": "How many levels down to walk (default 1)." + }, + "query": { + "type": "string", + "description": "Optional natural-language query — when set, children are reranked by cosine similarity." + }, + "limit": { + "type": "integer", + "minimum": 0, + "description": "Optional cap on returned hits, applied after rerank." + } + }, + "required": ["node_id"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + log::debug!("[tool][memory_tree] drill_down invoked"); + let req: DrillDownRequest = serde_json::from_value(args) + .map_err(|e| anyhow::anyhow!("invalid arguments for memory_tree_drill_down: {e}"))?; + if matches!(req.max_depth, Some(0)) { + return Err(anyhow::anyhow!( + "memory_tree_drill_down: max_depth must be >= 1" + )); + } + let cfg = config_rpc::load_config_with_timeout() + .await + .map_err(|e| anyhow::anyhow!("memory_tree_drill_down: load config failed: {e}"))?; + let hits = retrieval::drill_down( + &cfg, + &req.node_id, + req.max_depth.unwrap_or(1), + req.query.as_deref(), + req.limit, + ) + .await?; + log::debug!( + "[tool][memory_tree] drill_down returning hits={}", + hits.len() + ); + let json = serde_json::to_string(&hits)?; + Ok(ToolResult::success(json)) + } +} diff --git a/src/openhuman/tools/impl/memory/tree/fetch_leaves.rs b/src/openhuman/tools/impl/memory/tree/fetch_leaves.rs new file mode 100644 index 000000000..8d426206b --- /dev/null +++ b/src/openhuman/tools/impl/memory/tree/fetch_leaves.rs @@ -0,0 +1,68 @@ +use crate::openhuman::config::rpc as config_rpc; +use crate::openhuman::memory::tree::retrieval; +use crate::openhuman::memory::tree::retrieval::rpc::FetchLeavesRequest; +use crate::openhuman::tools::traits::{Tool, ToolResult}; +use async_trait::async_trait; +use serde_json::json; + +/// Hard cap on `chunk_ids` enforced at the tool boundary so the tool's +/// behaviour matches the schema description. The retrieval RPC also +/// truncates internally; we mirror that here so excess ids are dropped +/// rather than silently passed through. +const MAX_CHUNK_IDS_PER_CALL: usize = 20; + +pub struct MemoryTreeFetchLeavesTool; + +#[async_trait] +impl Tool for MemoryTreeFetchLeavesTool { + fn name(&self) -> &str { + "memory_tree_fetch_leaves" + } + + fn description(&self) -> &str { + "Batch-fetch raw chunk rows by id (max 20 per call). Use this when \ + you need verbatim content for a citation — the `content` and \ + `source_ref` fields on each hit are the authoritative quote source." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "chunk_ids": { + "type": "array", + "items": {"type": "string"}, + "description": "Chunk ids to hydrate. Capped at 20 per call." + } + }, + "required": ["chunk_ids"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let req: FetchLeavesRequest = serde_json::from_value(args) + .map_err(|e| anyhow::anyhow!("invalid arguments for memory_tree_fetch_leaves: {e}"))?; + log::debug!( + "[rpc][memory_tree] fetch_leaves invoked requested_ids={}", + req.chunk_ids.len() + ); + let cfg = config_rpc::load_config_with_timeout() + .await + .map_err(|e| anyhow::anyhow!("memory_tree_fetch_leaves: load config failed: {e}"))?; + let take = req.chunk_ids.len().min(MAX_CHUNK_IDS_PER_CALL); + if req.chunk_ids.len() > MAX_CHUNK_IDS_PER_CALL { + log::debug!( + "[rpc][memory_tree] fetch_leaves truncating requested_ids={} truncated_to={}", + req.chunk_ids.len(), + MAX_CHUNK_IDS_PER_CALL + ); + } + let hits = retrieval::fetch_leaves(&cfg, &req.chunk_ids[..take]).await?; + log::debug!( + "[rpc][memory_tree] fetch_leaves completed hits={}", + hits.len() + ); + let json = serde_json::to_string(&hits)?; + Ok(ToolResult::success(json)) + } +} diff --git a/src/openhuman/tools/impl/memory/tree/mod.rs b/src/openhuman/tools/impl/memory/tree/mod.rs new file mode 100644 index 000000000..f08bc835d --- /dev/null +++ b/src/openhuman/tools/impl/memory/tree/mod.rs @@ -0,0 +1,32 @@ +//! LLM-callable wrappers for the Phase 4 memory-tree retrieval primitives +//! (issue #710). Each tool is a thin shim over one typed function in +//! [`crate::openhuman::memory::tree::retrieval`]; the `*_rpc` variants are +//! intentionally avoided because they wrap responses in an `RpcOutcome` +//! envelope that is noisier than what the LLM needs. +//! +//! All six tools share the same shape: +//! 1. Deserialize args into the matching `*Request` struct from +//! [`crate::openhuman::memory::tree::retrieval::rpc`]. +//! 2. Load the active workspace `Config` via +//! [`crate::openhuman::config::rpc::load_config_with_timeout`]. +//! 3. Call the typed retrieval function. +//! 4. Serialise the response to JSON and return it as `ToolResult::success`. +//! +//! The tools are stateless unit structs — there is no per-instance state to +//! carry, so they slot directly into `Vec>` without needing +//! constructors. Logs use the `[tool]` / `[memory_tree]` prefixes per the +//! repo's debug-logging conventions. + +mod drill_down; +mod fetch_leaves; +mod query_global; +mod query_source; +mod query_topic; +mod search_entities; + +pub use drill_down::MemoryTreeDrillDownTool; +pub use fetch_leaves::MemoryTreeFetchLeavesTool; +pub use query_global::MemoryTreeQueryGlobalTool; +pub use query_source::MemoryTreeQuerySourceTool; +pub use query_topic::MemoryTreeQueryTopicTool; +pub use search_entities::MemoryTreeSearchEntitiesTool; diff --git a/src/openhuman/tools/impl/memory/tree/query_global.rs b/src/openhuman/tools/impl/memory/tree/query_global.rs new file mode 100644 index 000000000..cc19efa4c --- /dev/null +++ b/src/openhuman/tools/impl/memory/tree/query_global.rs @@ -0,0 +1,53 @@ +use crate::openhuman::config::rpc as config_rpc; +use crate::openhuman::memory::tree::retrieval; +use crate::openhuman::memory::tree::retrieval::rpc::QueryGlobalRequest; +use crate::openhuman::tools::traits::{Tool, ToolResult}; +use async_trait::async_trait; +use serde_json::json; + +pub struct MemoryTreeQueryGlobalTool; + +#[async_trait] +impl Tool for MemoryTreeQueryGlobalTool { + fn name(&self) -> &str { + "memory_tree_query_global" + } + + fn description(&self) -> &str { + "Return the cross-source global digest for the last `window_days`. \ + The 7-day digest is also pre-loaded into the session context at \ + start, so only call this for a different window (e.g. 30 days, \ + 1 day) or to refresh after new ingest." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "window_days": { + "type": "integer", + "minimum": 1, + "description": "Lookback window in days (e.g. 7 for weekly recap)." + } + }, + "required": ["window_days"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + log::debug!("[tool][memory_tree] query_global invoked"); + let req: QueryGlobalRequest = serde_json::from_value(args) + .map_err(|e| anyhow::anyhow!("invalid arguments for memory_tree_query_global: {e}"))?; + let cfg = config_rpc::load_config_with_timeout() + .await + .map_err(|e| anyhow::anyhow!("memory_tree_query_global: load config failed: {e}"))?; + let resp = retrieval::query_global(&cfg, req.window_days).await?; + log::debug!( + "[tool][memory_tree] query_global returning hits={} total={}", + resp.hits.len(), + resp.total + ); + let json = serde_json::to_string(&resp)?; + Ok(ToolResult::success(json)) + } +} diff --git a/src/openhuman/tools/impl/memory/tree/query_source.rs b/src/openhuman/tools/impl/memory/tree/query_source.rs new file mode 100644 index 000000000..c525295f0 --- /dev/null +++ b/src/openhuman/tools/impl/memory/tree/query_source.rs @@ -0,0 +1,87 @@ +use crate::openhuman::config::rpc as config_rpc; +use crate::openhuman::memory::tree::retrieval; +use crate::openhuman::memory::tree::retrieval::rpc::QuerySourceRequest; +use crate::openhuman::memory::tree::types::SourceKind; +use crate::openhuman::tools::traits::{Tool, ToolResult}; +use async_trait::async_trait; +use serde_json::json; + +pub struct MemoryTreeQuerySourceTool; + +#[async_trait] +impl Tool for MemoryTreeQuerySourceTool { + fn name(&self) -> &str { + "memory_tree_query_source" + } + + fn description(&self) -> &str { + "Return summaries from per-source memory trees, optionally filtered \ + by `source_id` (exact), `source_kind` (chat/email/document) and/or \ + `time_window_days`. Use this for intents like \"in my email last \ + week...\" or \"summarise our slack #eng activity\". Newest-first \ + by default; pass `query` for semantic rerank." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "source_id": { + "type": "string", + "description": "Exact source id (e.g. `slack:#eng`, `gmail:abc`)." + }, + "source_kind": { + "type": "string", + "enum": ["chat", "email", "document"], + "description": "Source kind filter when no exact id is known." + }, + "time_window_days": { + "type": "integer", + "minimum": 0, + "description": "Only return summaries whose time range overlaps the last N days." + }, + "query": { + "type": "string", + "description": "Optional natural-language query for cosine-similarity rerank." + }, + "limit": { + "type": "integer", + "minimum": 0, + "description": "Max hits to return (default 10)." + } + } + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + log::debug!("[tool][memory_tree] query_source invoked"); + let req: QuerySourceRequest = serde_json::from_value(args) + .map_err(|e| anyhow::anyhow!("invalid arguments for memory_tree_query_source: {e}"))?; + let cfg = config_rpc::load_config_with_timeout() + .await + .map_err(|e| anyhow::anyhow!("memory_tree_query_source: load config failed: {e}"))?; + let source_kind = match req.source_kind.as_deref() { + Some(s) => Some( + SourceKind::parse(s) + .map_err(|e| anyhow::anyhow!("memory_tree_query_source: {e}"))?, + ), + None => None, + }; + let resp = retrieval::query_source( + &cfg, + req.source_id.as_deref(), + source_kind, + req.time_window_days, + req.query.as_deref(), + req.limit.unwrap_or(10), + ) + .await?; + log::debug!( + "[tool][memory_tree] query_source returning hits={} total={}", + resp.hits.len(), + resp.total + ); + let json = serde_json::to_string(&resp)?; + Ok(ToolResult::success(json)) + } +} diff --git a/src/openhuman/tools/impl/memory/tree/query_topic.rs b/src/openhuman/tools/impl/memory/tree/query_topic.rs new file mode 100644 index 000000000..c0e0c8e65 --- /dev/null +++ b/src/openhuman/tools/impl/memory/tree/query_topic.rs @@ -0,0 +1,74 @@ +use crate::openhuman::config::rpc as config_rpc; +use crate::openhuman::memory::tree::retrieval; +use crate::openhuman::memory::tree::retrieval::rpc::QueryTopicRequest; +use crate::openhuman::tools::traits::{Tool, ToolResult}; +use async_trait::async_trait; +use serde_json::json; + +pub struct MemoryTreeQueryTopicTool; + +#[async_trait] +impl Tool for MemoryTreeQueryTopicTool { + fn name(&self) -> &str { + "memory_tree_query_topic" + } + + fn description(&self) -> &str { + "Return summaries / chunks linked to a canonical entity id (e.g. \ + `email:alice@example.com`, `topic:phoenix`) across every memory \ + tree. Sorted by score then recency, or by cosine similarity if \ + `query` is provided. Use this after `memory_tree_search_entities` \ + resolves a name to a canonical id." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "entity_id": { + "type": "string", + "description": "Canonical entity id (e.g. `email:alice@example.com`, `topic:phoenix`)." + }, + "time_window_days": { + "type": "integer", + "minimum": 0, + "description": "Only return hits whose time range overlaps the last N days." + }, + "query": { + "type": "string", + "description": "Optional natural-language query for cosine-similarity rerank." + }, + "limit": { + "type": "integer", + "minimum": 0, + "description": "Max hits to return (default 10)." + } + }, + "required": ["entity_id"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + log::debug!("[tool][memory_tree] query_topic invoked"); + let req: QueryTopicRequest = serde_json::from_value(args) + .map_err(|e| anyhow::anyhow!("invalid arguments for memory_tree_query_topic: {e}"))?; + let cfg = config_rpc::load_config_with_timeout() + .await + .map_err(|e| anyhow::anyhow!("memory_tree_query_topic: load config failed: {e}"))?; + let resp = retrieval::query_topic( + &cfg, + &req.entity_id, + req.time_window_days, + req.query.as_deref(), + req.limit.unwrap_or(10), + ) + .await?; + log::debug!( + "[tool][memory_tree] query_topic returning hits={} total={}", + resp.hits.len(), + resp.total + ); + let json = serde_json::to_string(&resp)?; + Ok(ToolResult::success(json)) + } +} diff --git a/src/openhuman/tools/impl/memory/tree/search_entities.rs b/src/openhuman/tools/impl/memory/tree/search_entities.rs new file mode 100644 index 000000000..c6cba317d --- /dev/null +++ b/src/openhuman/tools/impl/memory/tree/search_entities.rs @@ -0,0 +1,81 @@ +use crate::openhuman::config::rpc as config_rpc; +use crate::openhuman::memory::tree::retrieval; +use crate::openhuman::memory::tree::retrieval::rpc::SearchEntitiesRequest; +use crate::openhuman::memory::tree::score::extract::EntityKind; +use crate::openhuman::tools::traits::{Tool, ToolResult}; +use async_trait::async_trait; +use serde_json::json; + +pub struct MemoryTreeSearchEntitiesTool; + +#[async_trait] +impl Tool for MemoryTreeSearchEntitiesTool { + fn name(&self) -> &str { + "memory_tree_search_entities" + } + + fn description(&self) -> &str { + "Free-text LIKE search over the entity index — resolve a name or \ + handle to a canonical id (e.g. \"alice\" -> \ + `email:alice@example.com`). ALWAYS call this first when the user \ + mentions someone by name before calling `memory_tree_query_topic`." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Substring to match (case-insensitive)." + }, + "kinds": { + "type": "array", + "items": { + "type": "string", + "enum": [ + "email", "url", "handle", "hashtag", "person", + "organization", "location", "event", "product", + "misc", "topic" + ] + }, + "description": "Optional kind filter — restrict to these entity kinds only." + }, + "limit": { + "type": "integer", + "minimum": 0, + "description": "Max matches (default 5, clamped to 100)." + } + }, + "required": ["query"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + log::debug!("[tool][memory_tree] search_entities invoked"); + let req: SearchEntitiesRequest = serde_json::from_value(args).map_err(|e| { + anyhow::anyhow!("invalid arguments for memory_tree_search_entities: {e}") + })?; + let cfg = config_rpc::load_config_with_timeout() + .await + .map_err(|e| anyhow::anyhow!("memory_tree_search_entities: load config failed: {e}"))?; + let kinds = match req.kinds { + None => None, + Some(list) => { + let parsed: Result, String> = + list.iter().map(|s| EntityKind::parse(s)).collect(); + Some(parsed.map_err(|e| { + anyhow::anyhow!("memory_tree_search_entities: invalid kind: {e}") + })?) + } + }; + let limit = req.limit.unwrap_or(5).min(100); + let matches = retrieval::search_entities(&cfg, &req.query, kinds, limit).await?; + log::debug!( + "[tool][memory_tree] search_entities returning matches={}", + matches.len() + ); + let json = serde_json::to_string(&matches)?; + Ok(ToolResult::success(json)) + } +} diff --git a/src/openhuman/tools/ops.rs b/src/openhuman/tools/ops.rs index d4c331222..7f16b2b07 100644 --- a/src/openhuman/tools/ops.rs +++ b/src/openhuman/tools/ops.rs @@ -119,6 +119,12 @@ pub fn all_tools_with_runtime( Box::new(MemoryStoreTool::new(memory.clone(), security.clone())), Box::new(MemoryRecallTool::new(memory.clone())), Box::new(MemoryForgetTool::new(memory.clone(), security.clone())), + Box::new(MemoryTreeSearchEntitiesTool), + Box::new(MemoryTreeQueryTopicTool), + Box::new(MemoryTreeQuerySourceTool), + Box::new(MemoryTreeQueryGlobalTool), + Box::new(MemoryTreeDrillDownTool), + Box::new(MemoryTreeFetchLeavesTool), Box::new(ScheduleTool::new(security.clone(), root_config.clone())), Box::new(ProxyConfigTool::new(config.clone(), security.clone())), Box::new(GitOperationsTool::new( diff --git a/tests/agent_retrieval_e2e.rs b/tests/agent_retrieval_e2e.rs new file mode 100644 index 000000000..edebef5ad --- /dev/null +++ b/tests/agent_retrieval_e2e.rs @@ -0,0 +1,283 @@ +//! End-to-end coverage for the orchestrator memory-tree retrieval tool +//! wrappers (issue #710 wiring). +//! +//! Goal: prove the `MemoryTree*Tool` instances actually drive the typed +//! retrieval functions against a real ingested workspace and emit JSON the +//! orchestrator LLM can parse + cite from. +//! +//! Why a tool-direct test (and not a full `agent_chat` round-trip): +//! `agent_chat` requires a reachable provider (no provider connection +//! available in unit-test context). The bus-level `mock_agent_run_turn` +//! stub replaces the agent loop wholesale, so it can't observe a tool +//! dispatch happening *inside* the loop. Calling each tool's `execute()` +//! with the same JSON shape the LLM would emit exercises the full +//! deserialise → typed retrieval → serialise pipeline that the orchestrator +//! relies on, and asserts the data round-trips correctly. +//! +//! The orchestrator agent.toml entry registering these tool names is +//! covered by [`orchestrator_lists_memory_tree_tools`] — that catches a +//! regression where the tool wrapper exists but the orchestrator can't see +//! it. + +use chrono::{TimeZone, Utc}; +use openhuman_core::openhuman::config::Config; +use openhuman_core::openhuman::memory::tree::canonicalize::email::{EmailMessage, EmailThread}; +use openhuman_core::openhuman::memory::tree::ingest::ingest_email; +use openhuman_core::openhuman::memory::tree::jobs::drain_until_idle; +use openhuman_core::openhuman::tools::{ + MemoryTreeFetchLeavesTool, MemoryTreeQueryTopicTool, MemoryTreeSearchEntitiesTool, Tool, +}; +use serde_json::{json, Value}; +use tempfile::TempDir; + +/// Build a Config rooted at `tmp/workspace`. The nested `workspace` dir +/// matches what `resolve_config_dir_for_workspace` would derive when +/// `OPENHUMAN_WORKSPACE` points at `tmp` — so the same workspace_dir is +/// used both by the explicit ingest path and by `load_config_with_timeout` +/// inside the tool wrappers. +fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let workspace_dir = tmp.path().join("workspace"); + std::fs::create_dir_all(&workspace_dir).expect("create workspace dir"); + let mut cfg = Config { + workspace_dir: workspace_dir.clone(), + ..Config::default() + }; + // Inert embedder — keeps the test deterministic and avoids any real + // Ollama call. Mirrors `retrieval/integration_test.rs`. + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + (tmp, cfg) +} + +fn alice_phoenix_thread() -> EmailThread { + EmailThread { + provider: "gmail".into(), + thread_subject: "Phoenix migration plan".into(), + messages: vec![ + EmailMessage { + from: "alice@example.com".into(), + to: vec!["bob@example.com".into()], + cc: vec![], + subject: "Phoenix migration plan".into(), + sent_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), + body: "Hey Bob, the phoenix migration runbook is ready for review. \ + I'm coordinating with the infra team and we land Friday." + .into(), + source_ref: Some("".into()), + }, + EmailMessage { + from: "bob@example.com".into(), + to: vec!["alice@example.com".into()], + cc: vec![], + subject: "Re: Phoenix migration plan".into(), + sent_at: Utc.timestamp_millis_opt(1_700_000_060_000).unwrap(), + body: "Confirmed — I'll review the phoenix runbook tonight.".into(), + source_ref: Some("".into()), + }, + ], + } +} + +/// The orchestrator definition must list every memory-tree tool name so +/// the bus filter actually exposes them to the LLM. A wired-up wrapper +/// that's invisible to the orchestrator is dead code. +#[test] +fn orchestrator_lists_memory_tree_tools() { + let toml = include_str!("../src/openhuman/agent/agents/orchestrator/agent.toml"); + for name in [ + "memory_tree_search_entities", + "memory_tree_query_topic", + "memory_tree_query_source", + "memory_tree_query_global", + "memory_tree_drill_down", + "memory_tree_fetch_leaves", + ] { + assert!( + toml.contains(name), + "orchestrator agent.toml must list '{name}' so the LLM can call it" + ); + } +} + +#[tokio::test] +async fn orchestrator_query_topic_tool_returns_alice_phoenix_hits() { + let (tmp, cfg) = test_config(); + + // ── Ingest the email thread + drain async extract jobs so the entity + // index is fully populated before retrieval. + ingest_email( + &cfg, + "gmail:thread-phoenix-1", + "alice", + vec![], + alice_phoenix_thread(), + ) + .await + .expect("ingest_email should succeed"); + drain_until_idle(&cfg) + .await + .expect("job queue should drain cleanly"); + + // ── Set workspace dir so config_rpc::load_config_with_timeout() + // inside the tool resolves to the same workspace we just ingested + // into. The tool wrappers always go through that loader (mirrors + // the production RPC handlers in retrieval/schemas.rs). + struct EnvGuard { + key: &'static str, + prev: Option, + } + impl Drop for EnvGuard { + fn drop(&mut self) { + // SAFETY: see `EnvGuard::set` below — this integration test + // binary owns the env var for its lifetime. + unsafe { + match self.prev.take() { + Some(v) => std::env::set_var(self.key, v), + None => std::env::remove_var(self.key), + } + } + } + } + impl EnvGuard { + fn set(key: &'static str, val: &std::ffi::OsStr) -> Self { + let prev = std::env::var_os(key); + // SAFETY: `cargo test` defaults to running each integration + // test bin in its own process; nothing else in this bin + // mutates `OPENHUMAN_WORKSPACE`. The guard restores the + // previous value on drop. + unsafe { std::env::set_var(key, val) }; + Self { key, prev } + } + } + // Pointing OPENHUMAN_WORKSPACE at `tmp` (not `tmp/workspace`) makes + // `resolve_config_dir_for_workspace` derive `tmp/workspace` as the + // resolved workspace_dir — matching what we already passed into + // `ingest_email` via `cfg.workspace_dir`. + let _ws_guard = EnvGuard::set("OPENHUMAN_WORKSPACE", tmp.path().as_os_str()); + + // ── 1. search_entities resolves "alice" → email:alice@example.com. + // Mirrors the orchestrator prompt's "ALWAYS call this first when + // the user mentions someone by name" flow. + let search = MemoryTreeSearchEntitiesTool; + let search_args = json!({"query": "alice"}); + let search_res = search + .execute(search_args) + .await + .expect("search_entities should not error"); + assert!( + !search_res.is_error, + "search_entities returned an error result: {}", + search_res.output() + ); + let search_json: Value = + serde_json::from_str(&search_res.output()).expect("search output must be valid JSON"); + let matches = search_json + .as_array() + .expect("search_entities returns an array of EntityMatch"); + let alice = matches + .iter() + .find(|m| m.get("canonical_id").and_then(|v| v.as_str()) == Some("email:alice@example.com")) + .unwrap_or_else(|| panic!("search_entities did not return alice; got: {search_json:?}")); + assert!( + alice + .get("mention_count") + .and_then(|v| v.as_u64()) + .unwrap_or(0) + >= 1, + "alice should have at least one mention" + ); + + // ── 2. query_topic on alice's canonical id returns at least one hit + // referencing both her email and the phoenix migration content. + let topic_tool = MemoryTreeQueryTopicTool; + let topic_args = json!({"entity_id": "email:alice@example.com"}); + let topic_res = topic_tool + .execute(topic_args) + .await + .expect("query_topic should not error"); + assert!( + !topic_res.is_error, + "query_topic returned an error result: {}", + topic_res.output() + ); + let topic_json: Value = + serde_json::from_str(&topic_res.output()).expect("topic output must be valid JSON"); + let hits = topic_json + .get("hits") + .and_then(|v| v.as_array()) + .expect("query_topic must include `hits` array"); + assert!( + !hits.is_empty(), + "query_topic returned zero hits — expected at least one for alice" + ); + // Returning ANY hit at all from `query_topic("email:alice@example.com")` + // proves the entity index resolved the canonical id and hydrated nodes + // back. The leaf-level `entities` field on a chunk hit isn't populated + // synchronously by ingest — entity extraction lives in a separate async + // job stage that may not have populated leaf rows. Instead we assert on + // the hydrated content + source_ref so we still catch a regression where + // the chunk lookup returns garbage. + let any_phoenix = hits.iter().any(|h| { + h.get("content") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_lowercase() + .contains("phoenix") + }); + assert!( + any_phoenix, + "expected at least one query_topic hit with phoenix content; got: {topic_json:#}" + ); + let any_source_ref = hits + .iter() + .any(|h| h.get("source_ref").and_then(|v| v.as_str()).is_some()); + assert!( + any_source_ref, + "expected at least one hit to carry a `source_ref` for citation; got: {topic_json:#}" + ); + + // ── 3. fetch_leaves hydrates a leaf chunk — proves the citation path + // (LLM picks an id from a query_* hit, calls fetch_leaves to get + // the verbatim content + source_ref). + let leaf_id = hits + .iter() + .find_map(|h| { + if h.get("node_kind").and_then(|v| v.as_str()) == Some("leaf") { + h.get("node_id") + .and_then(|v| v.as_str()) + .map(str::to_string) + } else { + None + } + }) + .expect("alice's topic hits should include at least one leaf"); + let fetch_tool = MemoryTreeFetchLeavesTool; + let fetch_args = json!({"chunk_ids": [leaf_id.clone()]}); + let fetch_res = fetch_tool + .execute(fetch_args) + .await + .expect("fetch_leaves should not error"); + assert!( + !fetch_res.is_error, + "fetch_leaves returned an error result: {}", + fetch_res.output() + ); + let fetched: Value = + serde_json::from_str(&fetch_res.output()).expect("fetch output must be valid JSON"); + let fetched_arr = fetched.as_array().expect("fetch_leaves returns array"); + assert_eq!( + fetched_arr.len(), + 1, + "fetch_leaves should hydrate exactly the requested chunk" + ); + let content = fetched_arr[0] + .get("content") + .and_then(|v| v.as_str()) + .expect("fetched leaf must carry content"); + assert!( + !content.is_empty(), + "fetched leaf content must not be empty" + ); +}