Skip to main content

nexo_core/agent/
memory_tool.rs

1use 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                // Log one recall event per hit for Phase 10.5 signal tracking.
95                // Reciprocal rank as score — position 1 → 1.0, position 2 → 0.5, …
96                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}