trusty_memory/chat/
tools.rs

1//! Chat tool definitions + the `execute_*` dispatcher set.
2//!
3//! Why: the chat assistant's tool surface (`all_tools`) and the in-process
4//! dispatcher that runs each tool (`execute_tool` + the per-tool `execute_*`
5//! functions) form one cohesive concern split out of the former monolithic
6//! `chat.rs` (issue #607).
7//! What: `ChatBody`, `MAX_TOOL_ROUNDS`, `all_tools`, `execute_tool`, and every
8//! `execute_*` helper, moved verbatim. Visibility unchanged.
9//! Test: `all_tools_returns_expected_set`, `execute_tool_dispatches_known_tools`
10//! in `web::tests`.
11
12use crate::web::{load_user_config, palace_info_from, DreamStatusPayload};
13use crate::AppState;
14use serde::Deserialize;
15use serde_json::{json, Value};
16use trusty_common::memory_core::dream::PersistedDreamStats;
17use trusty_common::memory_core::palace::{PalaceId, RoomType};
18use trusty_common::memory_core::retrieval::{
19    recall_across_palaces_with_default_embedder, recall_with_default_embedder,
20};
21use trusty_common::memory_core::store::kg::Triple;
22use trusty_common::memory_core::PalaceRegistry;
23use trusty_common::{ChatMessage, ToolDef};
24
25// ---------------------------------------------------------------------------
26
27#[derive(Deserialize)]
28pub(crate) struct ChatBody {
29    #[serde(default)]
30    pub(crate) palace_id: Option<String>,
31    pub(crate) message: String,
32    #[serde(default)]
33    pub(crate) history: Vec<ChatMessage>,
34    /// Optional existing chat-session id; when provided we load+append+save.
35    #[serde(default)]
36    pub(crate) session_id: Option<String>,
37}
38
39/// Hard cap on the number of `tool -> assistant` round trips per chat turn.
40///
41/// Why: Without a bound, a malicious or confused model could request tools
42/// indefinitely; 10 is generous enough for any realistic plan-and-act loop
43/// while still terminating quickly when the model gets stuck.
44pub(crate) const MAX_TOOL_ROUNDS: usize = 10;
45
46/// Build the complete set of tool definitions the chat assistant can call.
47///
48/// Why: Centralizing the tool surface keeps the wire schema, the dispatcher in
49/// `execute_tool`, and the system prompt in lock-step — adding a new tool means
50/// editing this one function plus a match arm.
51/// What: Returns the 11 read/write tools spanning palace introspection,
52/// memory recall/create, KG read/write, and daemon status.
53/// Test: `all_tools_returns_expected_set` asserts names and required-arg shape.
54pub(crate) fn all_tools() -> Vec<ToolDef> {
55    vec![
56        ToolDef {
57            name: "list_palaces".into(),
58            description: "List all memory palaces on this machine with their metadata (id, name, description, counts).".into(),
59            parameters: json!({ "type": "object", "properties": {}, "required": [] }),
60        },
61        ToolDef {
62            name: "get_palace".into(),
63            description: "Get details for a specific palace by id.".into(),
64            parameters: json!({
65                "type": "object",
66                "properties": { "palace_id": { "type": "string", "description": "Palace id (kebab-case)" } },
67                "required": ["palace_id"],
68            }),
69        },
70        ToolDef {
71            name: "recall_memories".into(),
72            description: "Semantic search for memories in a palace. Returns the top-k most relevant drawers ranked by similarity to the query.".into(),
73            parameters: json!({
74                "type": "object",
75                "properties": {
76                    "palace_id": { "type": "string" },
77                    "query": { "type": "string", "description": "Free-text query" },
78                    "top_k": { "type": "integer", "minimum": 1, "maximum": 50, "default": 5 }
79                },
80                "required": ["palace_id", "query"],
81            }),
82        },
83        ToolDef {
84            name: "list_drawers".into(),
85            description: "List all drawers (memories) in a palace, most recent first.".into(),
86            parameters: json!({
87                "type": "object",
88                "properties": { "palace_id": { "type": "string" } },
89                "required": ["palace_id"],
90            }),
91        },
92        ToolDef {
93            name: "kg_query".into(),
94            description: "Query the temporal knowledge graph for all currently-active triples whose subject matches.".into(),
95            parameters: json!({
96                "type": "object",
97                "properties": {
98                    "palace_id": { "type": "string" },
99                    "subject": { "type": "string" }
100                },
101                "required": ["palace_id", "subject"],
102            }),
103        },
104        ToolDef {
105            name: "get_config".into(),
106            description: "Get the trusty-memory daemon's configuration (provider, model, data root). API keys are masked.".into(),
107            parameters: json!({ "type": "object", "properties": {}, "required": [] }),
108        },
109        ToolDef {
110            name: "get_status".into(),
111            description: "Get daemon health: version, palace count, totals for drawers/vectors/triples.".into(),
112            parameters: json!({ "type": "object", "properties": {}, "required": [] }),
113        },
114        ToolDef {
115            name: "get_dream_status".into(),
116            description: "Get aggregated dreamer activity across all palaces (merged/pruned/compacted counts, last run timestamp).".into(),
117            parameters: json!({ "type": "object", "properties": {}, "required": [] }),
118        },
119        ToolDef {
120            name: "get_palace_dream_status".into(),
121            description: "Get dreamer activity stats for a specific palace.".into(),
122            parameters: json!({
123                "type": "object",
124                "properties": { "palace_id": { "type": "string" } },
125                "required": ["palace_id"],
126            }),
127        },
128        ToolDef {
129            name: "create_memory".into(),
130            description: "Store a new memory (drawer) in a palace. The content is embedded and inserted into the vector index plus the drawer table.".into(),
131            parameters: json!({
132                "type": "object",
133                "properties": {
134                    "palace_id": { "type": "string" },
135                    "content": { "type": "string", "description": "Verbatim memory text" },
136                    "room": { "type": "string", "description": "Room name (Frontend/Backend/Testing/Planning/Documentation/Research/Configuration/Meetings/General or a custom name); defaults to General." },
137                    "tags": { "type": "array", "items": { "type": "string" } },
138                    "importance": { "type": "number", "minimum": 0.0, "maximum": 1.0, "default": 0.5 }
139                },
140                "required": ["palace_id", "content"],
141            }),
142        },
143        ToolDef {
144            name: "kg_assert".into(),
145            description: "Assert a knowledge-graph triple. Any prior active triple with the same (subject, predicate) is closed out (valid_to set to now) before the new one is inserted.".into(),
146            parameters: json!({
147                "type": "object",
148                "properties": {
149                    "palace_id": { "type": "string" },
150                    "subject": { "type": "string" },
151                    "predicate": { "type": "string" },
152                    "object": { "type": "string" },
153                    "confidence": { "type": "number", "minimum": 0.0, "maximum": 1.0, "default": 1.0 }
154                },
155                "required": ["palace_id", "subject", "predicate", "object"],
156            }),
157        },
158        ToolDef {
159            name: "memory_recall_all".into(),
160            description: "Semantic search across ALL palaces simultaneously. Returns the top-k most relevant drawers ranked by similarity, regardless of which palace they belong to. Each result includes a `palace_id` field identifying its source.".into(),
161            parameters: json!({
162                "type": "object",
163                "properties": {
164                    "q": { "type": "string", "description": "Free-text query" },
165                    "top_k": { "type": "integer", "minimum": 1, "maximum": 50, "default": 10 },
166                    "deep": { "type": "boolean", "default": false }
167                },
168                "required": ["q"],
169            }),
170        },
171    ]
172}
173
174/// Execute a tool call against the live `AppState`.
175///
176/// Why: We want the model's tool invocations to call the same Rust paths the
177/// HTTP handlers use — no extra HTTP round-trip, no JSON re-parsing, and the
178/// results always reflect this daemon's view of the world.
179/// What: Parses `arguments` as JSON, dispatches by tool name, returns a JSON
180/// value that becomes the `role: "tool"` message content. Errors are caught
181/// and returned as `{"error": "..."}` JSON so the model can react.
182/// Test: `execute_tool_dispatches_known_tools` covers the dispatch path and
183/// the unknown-tool error case.
184pub(crate) async fn execute_tool(name: &str, args: &str, state: &AppState) -> Value {
185    let parsed: Value = serde_json::from_str(args).unwrap_or(json!({}));
186    match name {
187        "list_palaces" => execute_list_palaces(state).await,
188        "get_palace" => match parsed.get("palace_id").and_then(|v| v.as_str()) {
189            Some(id) => execute_get_palace(state, id).await,
190            None => json!({ "error": "missing required argument: palace_id" }),
191        },
192        "recall_memories" => {
193            let pid = parsed.get("palace_id").and_then(|v| v.as_str());
194            let q = parsed.get("query").and_then(|v| v.as_str());
195            let top_k = parsed.get("top_k").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
196            match (pid, q) {
197                (Some(p), Some(q)) => execute_recall(state, p, q, top_k).await,
198                _ => json!({ "error": "missing required argument(s): palace_id, query" }),
199            }
200        }
201        "list_drawers" => match parsed.get("palace_id").and_then(|v| v.as_str()) {
202            Some(id) => execute_list_drawers(state, id).await,
203            None => json!({ "error": "missing required argument: palace_id" }),
204        },
205        "kg_query" => {
206            let pid = parsed.get("palace_id").and_then(|v| v.as_str());
207            let subj = parsed.get("subject").and_then(|v| v.as_str());
208            match (pid, subj) {
209                (Some(p), Some(s)) => execute_kg_query(state, p, s).await,
210                _ => json!({ "error": "missing required argument(s): palace_id, subject" }),
211            }
212        }
213        "get_config" => execute_get_config(state),
214        "get_status" => execute_get_status(state).await,
215        "get_dream_status" => execute_get_dream_status(state).await,
216        "get_palace_dream_status" => match parsed.get("palace_id").and_then(|v| v.as_str()) {
217            Some(id) => execute_get_palace_dream_status(state, id).await,
218            None => json!({ "error": "missing required argument: palace_id" }),
219        },
220        "create_memory" => {
221            let pid = parsed.get("palace_id").and_then(|v| v.as_str());
222            let content = parsed.get("content").and_then(|v| v.as_str());
223            let room = parsed.get("room").and_then(|v| v.as_str());
224            let tags: Vec<String> = parsed
225                .get("tags")
226                .and_then(|v| v.as_array())
227                .map(|arr| {
228                    arr.iter()
229                        .filter_map(|t| t.as_str().map(|s| s.to_string()))
230                        .collect()
231                })
232                .unwrap_or_default();
233            let importance = parsed
234                .get("importance")
235                .and_then(|v| v.as_f64())
236                .map(|f| f as f32)
237                .unwrap_or(0.5);
238            match (pid, content) {
239                (Some(p), Some(c)) => {
240                    execute_create_memory(state, p, c, room, tags, importance).await
241                }
242                _ => json!({ "error": "missing required argument(s): palace_id, content" }),
243            }
244        }
245        "kg_assert" => {
246            let pid = parsed.get("palace_id").and_then(|v| v.as_str());
247            let subj = parsed.get("subject").and_then(|v| v.as_str());
248            let pred = parsed.get("predicate").and_then(|v| v.as_str());
249            let obj = parsed.get("object").and_then(|v| v.as_str());
250            let conf = parsed
251                .get("confidence")
252                .and_then(|v| v.as_f64())
253                .map(|f| f as f32)
254                .unwrap_or(1.0);
255            match (pid, subj, pred, obj) {
256                (Some(p), Some(s), Some(pr), Some(o)) => {
257                    execute_kg_assert(state, p, s, pr, o, conf).await
258                }
259                _ => json!({
260                    "error": "missing required argument(s): palace_id, subject, predicate, object"
261                }),
262            }
263        }
264        "memory_recall_all" => {
265            let q = parsed.get("q").and_then(|v| v.as_str());
266            let top_k = parsed.get("top_k").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
267            let deep = parsed
268                .get("deep")
269                .and_then(|v| v.as_bool())
270                .unwrap_or(false);
271            match q {
272                Some(q) => execute_recall_all(state, q, top_k, deep).await,
273                None => json!({ "error": "missing required argument: q" }),
274            }
275        }
276        _ => json!({ "error": format!("unknown tool: {name}") }),
277    }
278}
279
280async fn execute_list_palaces(state: &AppState) -> Value {
281    let palaces = match PalaceRegistry::list_palaces(&state.data_root) {
282        Ok(v) => v,
283        Err(e) => return json!({ "error": format!("list palaces: {e:#}") }),
284    };
285    let out: Vec<Value> = palaces
286        .into_iter()
287        .map(|p| {
288            let handle = state.registry.open_palace(&state.data_root, &p.id).ok();
289            let info = palace_info_from(&p, handle.as_ref());
290            serde_json::to_value(info).unwrap_or(json!({}))
291        })
292        .collect();
293    json!(out)
294}
295
296async fn execute_get_palace(state: &AppState, id: &str) -> Value {
297    let palaces = match PalaceRegistry::list_palaces(&state.data_root) {
298        Ok(v) => v,
299        Err(e) => return json!({ "error": format!("list palaces: {e:#}") }),
300    };
301    match palaces.into_iter().find(|p| p.id.0 == id) {
302        Some(p) => {
303            let handle = state.registry.open_palace(&state.data_root, &p.id).ok();
304            serde_json::to_value(palace_info_from(&p, handle.as_ref())).unwrap_or(json!({}))
305        }
306        None => json!({ "error": format!("palace not found: {id}") }),
307    }
308}
309
310async fn execute_recall(state: &AppState, palace_id: &str, query: &str, top_k: usize) -> Value {
311    let handle = match state
312        .registry
313        .open_palace(&state.data_root, &PalaceId::new(palace_id))
314    {
315        Ok(h) => h,
316        Err(e) => return json!({ "error": format!("open palace {palace_id}: {e:#}") }),
317    };
318    match recall_with_default_embedder(&handle, query, top_k).await {
319        Ok(hits) => json!(hits
320            .into_iter()
321            .map(|r| json!({
322                "drawer_id": r.drawer.id.to_string(),
323                "content": r.drawer.content,
324                "importance": r.drawer.importance,
325                "tags": r.drawer.tags,
326                "score": r.score,
327                "layer": r.layer,
328            }))
329            .collect::<Vec<_>>()),
330        Err(e) => json!({ "error": format!("recall: {e:#}") }),
331    }
332}
333
334/// Execute a cross-palace recall and return JSON results tagged with palace id.
335///
336/// Why: Both the MCP `memory_recall_all` tool and the `GET /api/v1/recall`
337/// HTTP route share the same wiring — list palaces, open handles, fan out via
338/// `recall_across_palaces_with_default_embedder`, and serialize.
339/// What: Lists every palace on disk, opens each (skipping any that fail with
340/// a `tracing::warn!`), and delegates to the core fan-out. On success returns
341/// a JSON array; on listing failure returns `{ "error": "..." }`.
342/// Test: Indirectly via `recall_across_palaces_merges_results` (core merge
343/// logic) and the HTTP/MCP integration paths.
344pub(crate) async fn execute_recall_all(
345    state: &AppState,
346    query: &str,
347    top_k: usize,
348    deep: bool,
349) -> Value {
350    let palaces = match PalaceRegistry::list_palaces(&state.data_root) {
351        Ok(v) => v,
352        Err(e) => return json!({ "error": format!("list palaces: {e:#}") }),
353    };
354    let mut handles = Vec::with_capacity(palaces.len());
355    for p in &palaces {
356        match state.registry.open_palace(&state.data_root, &p.id) {
357            Ok(h) => handles.push(h),
358            Err(e) => {
359                tracing::warn!(palace = %p.id, "execute_recall_all: open failed: {e:#}");
360            }
361        }
362    }
363    if handles.is_empty() {
364        return json!([]);
365    }
366    match recall_across_palaces_with_default_embedder(&handles, query, top_k, deep).await {
367        Ok(results) => json!(results
368            .into_iter()
369            .map(|r| json!({
370                "palace_id": r.palace_id,
371                "drawer_id": r.result.drawer.id.to_string(),
372                "content": r.result.drawer.content,
373                "importance": r.result.drawer.importance,
374                "tags": r.result.drawer.tags,
375                "score": r.result.score,
376                "layer": r.result.layer,
377            }))
378            .collect::<Vec<_>>()),
379        Err(e) => json!({ "error": format!("recall_across_palaces: {e:#}") }),
380    }
381}
382
383async fn execute_list_drawers(state: &AppState, palace_id: &str) -> Value {
384    let handle = match state
385        .registry
386        .open_palace(&state.data_root, &PalaceId::new(palace_id))
387    {
388        Ok(h) => h,
389        Err(e) => return json!({ "error": format!("open palace {palace_id}: {e:#}") }),
390    };
391    let drawers = handle.list_drawers(None, None, 200);
392    serde_json::to_value(drawers).unwrap_or(json!([]))
393}
394
395async fn execute_kg_query(state: &AppState, palace_id: &str, subject: &str) -> Value {
396    let handle = match state
397        .registry
398        .open_palace(&state.data_root, &PalaceId::new(palace_id))
399    {
400        Ok(h) => h,
401        Err(e) => return json!({ "error": format!("open palace {palace_id}: {e:#}") }),
402    };
403    match handle.kg.query_active(subject).await {
404        Ok(triples) => serde_json::to_value(triples).unwrap_or(json!([])),
405        Err(e) => json!({ "error": format!("kg query: {e:#}") }),
406    }
407}
408
409fn execute_get_config(state: &AppState) -> Value {
410    let cfg = load_user_config().unwrap_or_default();
411    json!({
412        "openrouter_configured": !cfg.openrouter_api_key.is_empty(),
413        "openrouter_model": cfg.openrouter_model,
414        "local_model": {
415            "enabled": cfg.local_model.enabled,
416            "base_url": cfg.local_model.base_url,
417            "model": cfg.local_model.model,
418        },
419        "data_root": state.data_root.display().to_string(),
420    })
421}
422
423async fn execute_get_status(state: &AppState) -> Value {
424    let palaces = PalaceRegistry::list_palaces(&state.data_root).unwrap_or_default();
425    let (mut total_drawers, mut total_vectors, mut total_kg_triples) = (0usize, 0usize, 0usize);
426    for p in &palaces {
427        if let Ok(handle) = state.registry.open_palace(&state.data_root, &p.id) {
428            total_drawers = total_drawers.saturating_add(handle.drawers.read().len());
429            total_vectors = total_vectors.saturating_add(handle.vector_store.index_size());
430            total_kg_triples = total_kg_triples.saturating_add(handle.kg.count_active_triples());
431        }
432    }
433    json!({
434        "version": state.version,
435        "palace_count": palaces.len(),
436        "default_palace": state.default_palace,
437        "data_root": state.data_root.display().to_string(),
438        "total_drawers": total_drawers,
439        "total_vectors": total_vectors,
440        "total_kg_triples": total_kg_triples,
441    })
442}
443
444pub(crate) async fn execute_get_dream_status(state: &AppState) -> Value {
445    let palaces = PalaceRegistry::list_palaces(&state.data_root).unwrap_or_default();
446    let mut out = DreamStatusPayload::default();
447    let mut latest: Option<chrono::DateTime<chrono::Utc>> = None;
448    for p in palaces {
449        let data_dir = state.data_root.join(p.id.as_str());
450        let snap = match PersistedDreamStats::load(&data_dir) {
451            Ok(Some(s)) => s,
452            _ => continue,
453        };
454        out.merged = out.merged.saturating_add(snap.stats.merged);
455        out.pruned = out.pruned.saturating_add(snap.stats.pruned);
456        out.compacted = out.compacted.saturating_add(snap.stats.compacted);
457        out.closets_updated = out
458            .closets_updated
459            .saturating_add(snap.stats.closets_updated);
460        out.duration_ms = out.duration_ms.saturating_add(snap.stats.duration_ms);
461        latest = match latest {
462            Some(t) if t >= snap.last_run_at => Some(t),
463            _ => Some(snap.last_run_at),
464        };
465    }
466    out.last_run_at = latest;
467    serde_json::to_value(out).unwrap_or(json!({}))
468}
469
470async fn execute_get_palace_dream_status(state: &AppState, palace_id: &str) -> Value {
471    let data_dir = state.data_root.join(palace_id);
472    if !data_dir.exists() {
473        return json!({ "error": format!("palace not found: {palace_id}") });
474    }
475    match PersistedDreamStats::load(&data_dir) {
476        Ok(Some(s)) => serde_json::to_value(DreamStatusPayload::from(s)).unwrap_or(json!({})),
477        Ok(None) => serde_json::to_value(DreamStatusPayload::default()).unwrap_or(json!({})),
478        Err(e) => json!({ "error": format!("read dream stats: {e:#}") }),
479    }
480}
481
482async fn execute_create_memory(
483    state: &AppState,
484    palace_id: &str,
485    content: &str,
486    room: Option<&str>,
487    tags: Vec<String>,
488    importance: f32,
489) -> Value {
490    let handle = match state
491        .registry
492        .open_palace(&state.data_root, &PalaceId::new(palace_id))
493    {
494        Ok(h) => h,
495        Err(e) => return json!({ "error": format!("open palace {palace_id}: {e:#}") }),
496    };
497    let room = room.map(RoomType::parse).unwrap_or(RoomType::General);
498    match handle
499        .remember(content.to_string(), room, tags, importance)
500        .await
501    {
502        Ok(id) => json!({ "drawer_id": id.to_string(), "status": "stored" }),
503        Err(e) => json!({ "error": format!("remember: {e:#}") }),
504    }
505}
506
507async fn execute_kg_assert(
508    state: &AppState,
509    palace_id: &str,
510    subject: &str,
511    predicate: &str,
512    object: &str,
513    confidence: f32,
514) -> Value {
515    let handle = match state
516        .registry
517        .open_palace(&state.data_root, &PalaceId::new(palace_id))
518    {
519        Ok(h) => h,
520        Err(e) => return json!({ "error": format!("open palace {palace_id}: {e:#}") }),
521    };
522    let triple = Triple {
523        subject: subject.to_string(),
524        predicate: predicate.to_string(),
525        object: object.to_string(),
526        valid_from: chrono::Utc::now(),
527        valid_to: None,
528        confidence,
529        provenance: Some("chat:assistant".to_string()),
530    };
531    match handle.kg.assert(triple).await {
532        Ok(()) => json!({ "status": "asserted" }),
533        Err(e) => json!({ "error": format!("kg assert: {e:#}") }),
534    }
535}
trusty_memory/chat/tools.rs

trusty_memory/chat/
tools.rs
