nexo_core/agent/
memory_tool.rs1use super::context::AgentContext;
2use super::tool_registry::ToolHandler;
3use async_trait::async_trait;
4use nexo_llm::ToolDef;
5use nexo_memory::LongTermMemory;
6use serde_json::{json, Value};
7use std::sync::Arc;
8use uuid::Uuid;
9pub struct MemoryTool {
10 memory: Arc<LongTermMemory>,
11 default_recall_mode: String,
12}
13impl MemoryTool {
14 pub fn new(memory: Arc<LongTermMemory>) -> Self {
15 Self::new_with_default_mode(memory, "keyword")
16 }
17 pub fn new_with_default_mode(memory: Arc<LongTermMemory>, mode: impl Into<String>) -> Self {
18 Self {
19 memory,
20 default_recall_mode: normalize_recall_mode(mode.into()),
21 }
22 }
23 pub fn tool_def() -> ToolDef {
24 ToolDef {
25 name: "memory".to_string(),
26 description: "Store and retrieve memories. Actions: remember (save a fact), recall (search memories by keyword), forget (delete by id).".to_string(),
27 parameters: json!({
28 "type": "object",
29 "properties": {
30 "action": {
31 "type": "string",
32 "enum": ["remember", "recall", "forget"]
33 },
34 "content": { "type": "string", "description": "Fact to store (for remember)" },
35 "query": { "type": "string", "description": "Search query (for recall)" },
36 "id": { "type": "string", "description": "Memory UUID to delete (for forget)" },
37 "tags": {
38 "type": "array",
39 "items": { "type": "string" },
40 "description": "Optional tags to categorize the memory"
41 },
42 "limit": { "type": "integer", "description": "Max results to return (for recall, default 5)" },
43 "mode": {
44 "type": "string",
45 "enum": ["keyword", "vector", "hybrid"],
46 "description": "Recall mode. `keyword` (default) = FTS; `vector` = semantic nearest-neighbor; `hybrid` = RRF fusion of both. `vector` and `hybrid` require the memory.vector config."
47 }
48 },
49 "required": ["action"]
50 }),
51 }
52 }
53}
54#[async_trait]
55impl ToolHandler for MemoryTool {
56 async fn call(&self, ctx: &AgentContext, args: Value) -> anyhow::Result<Value> {
57 let action = args["action"].as_str().unwrap_or("");
58 match action {
59 "remember" => {
60 let content = args["content"]
61 .as_str()
62 .ok_or_else(|| anyhow::anyhow!("remember requires 'content'"))?;
63 let tags: Vec<&str> = args["tags"]
64 .as_array()
65 .map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
66 .unwrap_or_default();
67 let id = self.memory.remember(&ctx.agent_id, content, &tags).await?;
68 Ok(json!({ "ok": true, "id": id.to_string() }))
69 }
70 "recall" => {
71 let query = args["query"]
72 .as_str()
73 .ok_or_else(|| anyhow::anyhow!("recall requires 'query'"))?;
74 let limit = args["limit"].as_u64().unwrap_or(5) as usize;
75 let mode_owned = args["mode"]
76 .as_str()
77 .map(str::to_string)
78 .unwrap_or_else(|| self.default_recall_mode.clone());
79 let mode = mode_owned.as_str();
80 let entries = match mode {
81 "vector" => {
82 self.memory
83 .recall_vector(&ctx.agent_id, query, limit)
84 .await?
85 }
86 "hybrid" => {
87 self.memory
88 .recall_hybrid(&ctx.agent_id, query, limit)
89 .await?
90 }
91 "keyword" | "" => self.memory.recall(&ctx.agent_id, query, limit).await?,
92 other => anyhow::bail!("unknown recall mode: {other}"),
93 };
94 for (idx, e) in entries.iter().enumerate() {
97 let score = 1.0 / (idx as f32 + 1.0);
98 if let Err(err) = self
99 .memory
100 .record_recall_event(&ctx.agent_id, e.id, query, score)
101 .await
102 {
103 tracing::warn!(
104 agent_id = %ctx.agent_id,
105 memory_id = %e.id,
106 error = %err,
107 "failed to record recall event"
108 );
109 }
110 }
111 let results: Vec<Value> = entries
112 .iter()
113 .map(|e| {
114 json!({
115 "id": e.id.to_string(),
116 "content": e.content,
117 "tags": e.tags,
118 })
119 })
120 .collect();
121 Ok(json!({ "ok": true, "results": results }))
122 }
123 "forget" => {
124 let id_str = args["id"]
125 .as_str()
126 .ok_or_else(|| anyhow::anyhow!("forget requires 'id'"))?;
127 let id = Uuid::parse_str(id_str)
128 .map_err(|_| anyhow::anyhow!("invalid UUID: {id_str}"))?;
129 let deleted = self.memory.forget(id).await?;
130 Ok(json!({ "ok": deleted }))
131 }
132 other => anyhow::bail!("unknown memory action: {other}"),
133 }
134 }
135}
136
137fn normalize_recall_mode(mode: String) -> String {
138 match mode.trim() {
139 "keyword" | "vector" | "hybrid" => mode.trim().to_string(),
140 _ => "keyword".to_string(),
141 }
142}
143
144#[cfg(test)]
145mod tests {
146 use super::normalize_recall_mode;
147
148 #[test]
149 fn normalize_recall_mode_accepts_supported_values() {
150 assert_eq!(normalize_recall_mode("keyword".to_string()), "keyword");
151 assert_eq!(normalize_recall_mode("vector".to_string()), "vector");
152 assert_eq!(normalize_recall_mode("hybrid".to_string()), "hybrid");
153 }
154
155 #[test]
156 fn normalize_recall_mode_falls_back_to_keyword() {
157 assert_eq!(normalize_recall_mode("".to_string()), "keyword");
158 assert_eq!(normalize_recall_mode("auto".to_string()), "keyword");
159 assert_eq!(normalize_recall_mode(" VECTOR ".to_string()), "keyword");
160 }
161}