Skip to main content

oxi_agent/tools/
memory_recall.rs

1//! `memory_recall` tool — search the memory backend for relevant items.
2
3use super::{AgentTool, AgentToolResult, MemoryItem, ToolContext, ToolError};
4use crate::tools::typed::TypedTool;
5use async_trait::async_trait;
6use schemars::JsonSchema;
7use serde::Deserialize;
8use serde_json::{Value, json};
9
10/// Default number of results returned by [`MemoryRecallTool`] when `limit` is omitted.
11#[allow(dead_code)]
12const DEFAULT_LIMIT: usize = 5;
13/// Maximum number of results [`MemoryRecallTool`] will return.
14const MAX_LIMIT: usize = 20;
15
16/// Tool that searches the configured `MemoryBackend` for memories matching
17/// a query and returns the matches in a compact, model-friendly format.
18#[derive(Deserialize, JsonSchema)]
19pub struct MemoryRecallArgs {
20    query: String,
21    #[serde(default = "default_mem_limit")]
22    limit: u64,
23}
24
25fn default_mem_limit() -> u64 {
26    5
27}
28
29///
30/// Requires `ctx.memory` to be set; otherwise returns an error.
31pub struct MemoryRecallTool;
32
33#[async_trait]
34impl AgentTool for MemoryRecallTool {
35    fn name(&self) -> &str {
36        "memory_recall"
37    }
38
39    fn label(&self) -> &str {
40        "Memory Recall"
41    }
42
43    fn description(&self) -> &str {
44        "Search long-term memory for information relevant to a query. \
45         Returns the most relevant stored memories (facts, preferences, \
46         context, summaries)."
47    }
48
49    fn essential(&self) -> bool {
50        false
51    }
52
53    fn parameters_schema(&self) -> Value {
54        json!({
55            "type": "object",
56            "properties": {
57                "query": {
58                    "type": "string",
59                    "description": "What to search for in memory."
60                },
61                "limit": {
62                    "type": "integer",
63                    "minimum": 1,
64                    "maximum": 20,
65                    "default": 5,
66                    "description": "Maximum number of results to return."
67                }
68            },
69            "required": ["query"]
70        })
71    }
72
73    async fn execute(
74        &self,
75        _tool_call_id: &str,
76        params: Value,
77        _signal: Option<tokio::sync::oneshot::Receiver<()>>,
78        ctx: &ToolContext,
79    ) -> Result<AgentToolResult, ToolError> {
80        let args: MemoryRecallArgs =
81            serde_json::from_value(params).map_err(|e| format!("invalid params: {e}"))?;
82        self.execute_typed(_tool_call_id, args, _signal, ctx).await
83    }
84}
85
86#[async_trait]
87impl TypedTool for MemoryRecallTool {
88    type Args = MemoryRecallArgs;
89
90    async fn execute_typed(
91        &self,
92        _tool_call_id: &str,
93        args: Self::Args,
94        _signal: Option<tokio::sync::oneshot::Receiver<()>>,
95        ctx: &ToolContext,
96    ) -> Result<AgentToolResult, ToolError> {
97        let backend = ctx.memory.as_ref().ok_or("Memory not configured")?;
98        let limit = (args.limit as usize).clamp(1, MAX_LIMIT);
99        let results = backend.search(&args.query, limit).await?;
100        Ok(AgentToolResult::success(format_results(&results)))
101    }
102}
103
104/// Format memory search results into a compact, model-friendly string.
105fn format_results(items: &[MemoryItem]) -> String {
106    if items.is_empty() {
107        return "No matching memories found.".to_string();
108    }
109    let mut out = format!(
110        "Found {} memor{}:\n\n",
111        items.len(),
112        if items.len() == 1 { "y" } else { "ies" }
113    );
114    for (i, item) in items.iter().enumerate() {
115        out.push_str(&format!("{}. [{}] {}\n", i + 1, item.kind, item.content));
116    }
117    out
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::tools::MemoryBackend;
124    use parking_lot::Mutex;
125    use std::future::Future;
126    use std::pin::Pin;
127    use std::sync::Arc;
128
129    /// Returns canned search results and records the requested `k`.
130    #[derive(Debug)]
131    struct MockMemory {
132        items: Vec<MemoryItem>,
133        last_k: Mutex<Option<usize>>,
134    }
135
136    impl MemoryBackend for MockMemory {
137        fn put<'a>(
138            &'a self,
139            _content: &'a str,
140            _kind: &'a str,
141            _subject: &'a str,
142        ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
143            Box::pin(async move { Ok("mem-1".to_string()) })
144        }
145
146        fn search<'a>(
147            &'a self,
148            _query: &'a str,
149            k: usize,
150        ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
151            *self.last_k.lock() = Some(k);
152            let items: Vec<MemoryItem> = self.items.iter().take(k).cloned().collect();
153            Box::pin(async move { Ok(items) })
154        }
155
156        fn list<'a>(
157            &'a self,
158            _subject: &'a str,
159        ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
160            Box::pin(async move { Ok(vec![]) })
161        }
162
163        fn delete<'a>(
164            &'a self,
165            _id: &'a str,
166        ) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
167            Box::pin(async move { Ok(()) })
168        }
169    }
170
171    fn make_item(id: &str, kind: &str, content: &str) -> MemoryItem {
172        MemoryItem {
173            id: id.into(),
174            kind: kind.into(),
175            content: content.into(),
176            subject: "s".into(),
177        }
178    }
179
180    #[tokio::test]
181    async fn recall_returns_formatted_results() {
182        let mock = Arc::new(MockMemory {
183            items: vec![
184                make_item("1", "fact", "Rust is fast"),
185                make_item("2", "preference", "Likes dark mode"),
186            ],
187            last_k: Mutex::new(None),
188        });
189        let ctx = ToolContext::default().with_memory(mock.clone());
190        let result = MemoryRecallTool
191            .execute("c1", json!({"query": "rust", "limit": 5}), None, &ctx)
192            .await
193            .unwrap();
194        assert!(result.success);
195        assert!(result.output.contains("[fact] Rust is fast"));
196        assert!(result.output.contains("[preference] Likes dark mode"));
197        assert_eq!(*mock.last_k.lock(), Some(5));
198    }
199
200    #[tokio::test]
201    async fn recall_reports_empty_results() {
202        let mock = Arc::new(MockMemory {
203            items: vec![],
204            last_k: Mutex::new(None),
205        });
206        let ctx = ToolContext::default().with_memory(mock);
207        let result = MemoryRecallTool
208            .execute("c1", json!({"query": "nothing"}), None, &ctx)
209            .await
210            .unwrap();
211        assert!(result.success);
212        assert_eq!(result.output, "No matching memories found.");
213    }
214
215    #[tokio::test]
216    async fn recall_uses_default_limit() {
217        let mock = Arc::new(MockMemory {
218            items: vec![],
219            last_k: Mutex::new(None),
220        });
221        let ctx = ToolContext::default().with_memory(mock.clone());
222        MemoryRecallTool
223            .execute("c1", json!({"query": "x"}), None, &ctx)
224            .await
225            .unwrap();
226        assert_eq!(*mock.last_k.lock(), Some(DEFAULT_LIMIT));
227    }
228
229    #[tokio::test]
230    async fn recall_clamps_oversized_limit() {
231        let mock = Arc::new(MockMemory {
232            items: vec![],
233            last_k: Mutex::new(None),
234        });
235        let ctx = ToolContext::default().with_memory(mock.clone());
236        MemoryRecallTool
237            .execute("c1", json!({"query": "x", "limit": 100}), None, &ctx)
238            .await
239            .unwrap();
240        assert_eq!(*mock.last_k.lock(), Some(MAX_LIMIT));
241    }
242
243    #[tokio::test]
244    async fn recall_clamps_zero_limit() {
245        let mock = Arc::new(MockMemory {
246            items: vec![],
247            last_k: Mutex::new(None),
248        });
249        let ctx = ToolContext::default().with_memory(mock.clone());
250        MemoryRecallTool
251            .execute("c1", json!({"query": "x", "limit": 0}), None, &ctx)
252            .await
253            .unwrap();
254        assert_eq!(*mock.last_k.lock(), Some(1));
255    }
256
257    #[tokio::test]
258    async fn recall_errors_when_memory_not_configured() {
259        let ctx = ToolContext::default();
260        let err = MemoryRecallTool
261            .execute("c1", json!({"query": "x"}), None, &ctx)
262            .await
263            .unwrap_err();
264        assert_eq!(err, "Memory not configured");
265    }
266
267    #[tokio::test]
268    async fn recall_rejects_missing_query() {
269        let mock = Arc::new(MockMemory {
270            items: vec![],
271            last_k: Mutex::new(None),
272        });
273        let ctx = ToolContext::default().with_memory(mock);
274        let err = MemoryRecallTool
275            .execute("c1", json!({"limit": 3}), None, &ctx)
276            .await
277            .unwrap_err();
278        assert!(err.contains("query"));
279    }
280}