robit_agent/tool/
memory.rs

1//! Memory tools for long-term memory.
2//!
3//! Provides tools for the agent to remember, retrieve, and forget information.
4
5use async_trait::async_trait;
6use serde::Deserialize;
7use serde_json::Value;
8
9use crate::storage::{
10    deactivate_memory, delete_memory_permanently, insert_memory, list_memories, recall_memories,
11    update_memory, Memory, MemoryFilter, MemoryType,
12};
13use crate::tool::{Tool, ToolContext, ToolResult};
14use crate::error::Result;
15
16// ============================================================================
17// Memorize tool - store a memory
18// ============================================================================
19
20#[derive(Debug, Deserialize)]
21struct MemorizeArgs {
22    title: String,
23    content: String,
24    #[serde(default = "default_memory_type")]
25    memory_type: String,
26    #[serde(default)]
27    tags: Vec<String>,
28    #[serde(default)]
29    update_if_exists: bool,
30}
31
32fn default_memory_type() -> String {
33    "note".to_string()
34}
35
36pub struct MemorizeTool;
37
38impl MemorizeTool {
39    pub fn new() -> Self {
40        Self
41    }
42}
43
44impl Default for MemorizeTool {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50#[async_trait]
51impl Tool for MemorizeTool {
52    fn name(&self) -> &str {
53        "memorize"
54    }
55
56    fn description(&self) -> &str {
57        "Store important information in long-term memory. Use for user preferences, key facts, project notes, etc."
58    }
59
60    fn parameters_schema(&self) -> Value {
61        serde_json::json!({
62            "type": "object",
63            "properties": {
64                "title": {
65                    "type": "string",
66                    "description": "Short, descriptive title for the memory (for easy retrieval)"
67                },
68                "content": {
69                    "type": "string",
70                    "description": "Full content of the memory"
71                },
72                "memory_type": {
73                    "type": "string",
74                    "description": "Type of memory: fact, preference, note, task, or custom",
75                    "default": "note"
76                },
77                "tags": {
78                    "type": "array",
79                    "items": { "type": "string" },
80                    "description": "Tags for categorization and filtering",
81                    "default": []
82                },
83                "update_if_exists": {
84                    "type": "boolean",
85                    "description": "If true, update existing memory with the same title",
86                    "default": false
87                }
88            },
89            "required": ["title", "content"]
90        })
91    }
92
93    fn requires_confirmation(&self) -> bool {
94        false
95    }
96
97    async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<ToolResult> {
98        let parsed: MemorizeArgs = match serde_json::from_value(args) {
99            Ok(a) => a,
100            Err(e) => return Ok(ToolResult::error(format!("Argument parsing failed: {}", e))),
101        };
102
103        if parsed.title.trim().is_empty() {
104            return Ok(ToolResult::error("Title cannot be empty".to_string()));
105        }
106        if parsed.content.trim().is_empty() {
107            return Ok(ToolResult::error("Content cannot be empty".to_string()));
108        }
109
110        let db_path = match crate::storage::resolve_db_path(&ctx.working_dir, false) {
111            Ok(p) => p,
112            Err(e) => return Ok(ToolResult::error(format!("Failed to resolve DB path: {}", e))),
113        };
114
115        let conn = match rusqlite::Connection::open(&db_path) {
116            Ok(c) => c,
117            Err(e) => return Ok(ToolResult::error(format!("Failed to open DB: {}", e))),
118        };
119
120        let memory_type = MemoryType::from_str(&parsed.memory_type);
121        let mut memory = Memory::new(
122            parsed.title.clone(),
123            parsed.content.clone(),
124            memory_type,
125            parsed.tags.clone(),
126        );
127
128        if ctx.session_id.to_string() != "" {
129            memory = memory.with_session_id(ctx.session_id.to_string());
130        }
131
132        if parsed.update_if_exists {
133            // Look for existing memory with the same title
134            let filter = MemoryFilter {
135                session_id: Some(ctx.session_id.to_string()),
136                only_active: true,
137                ..Default::default()
138            };
139            if let Ok(existing) = list_memories(&conn, &filter, Some(100)) {
140                if let Some(mut existing) = existing.into_iter().find(|m| m.title == parsed.title) {
141                    existing.content = parsed.content;
142                    existing.memory_type = MemoryType::from_str(&parsed.memory_type);
143                    existing.tags = parsed.tags;
144                    if let Err(e) = update_memory(&conn, &existing) {
145                        return Ok(ToolResult::error(format!("Failed to update memory: {}", e)));
146                    }
147                    return Ok(ToolResult::success(format!(
148                        "Updated memory: {} (ID: {})",
149                        existing.title, existing.id
150                    )));
151                }
152            }
153        }
154
155        // Insert as new memory
156        if let Err(e) = insert_memory(&conn, &memory) {
157            return Ok(ToolResult::error(format!("Failed to store memory: {}", e)));
158        }
159
160        Ok(ToolResult::success(format!(
161            "Stored memory: {} (ID: {})",
162            memory.title, memory.id
163        )))
164    }
165}
166
167// ============================================================================
168// Recall tool - retrieve memories
169// ============================================================================
170
171#[derive(Debug, Deserialize)]
172struct RecallArgs {
173    query: Option<String>,
174    memory_type: Option<String>,
175    tags: Option<Vec<String>>,
176    limit: Option<usize>,
177    since: Option<String>,
178    session_id: Option<String>,
179    chat_id: Option<String>,
180}
181
182pub struct RecallTool;
183
184impl RecallTool {
185    pub fn new() -> Self {
186        Self
187    }
188}
189
190impl Default for RecallTool {
191    fn default() -> Self {
192        Self::new()
193    }
194}
195
196#[async_trait]
197impl Tool for RecallTool {
198    fn name(&self) -> &str {
199        "recall"
200    }
201
202    fn description(&self) -> &str {
203        "Retrieve relevant memories from long-term memory. Use when you need context or past information."
204    }
205
206    fn parameters_schema(&self) -> Value {
207        serde_json::json!({
208            "type": "object",
209            "properties": {
210                "query": {
211                    "type": "string",
212                    "description": "Keywords to search in title, content, and tags"
213                },
214                "memory_type": {
215                    "type": "string",
216                    "description": "Filter by memory type: fact, preference, note, task"
217                },
218                "tags": {
219                    "type": "array",
220                    "items": { "type": "string" },
221                    "description": "Filter by tags (any match)"
222                },
223                "limit": {
224                    "type": "integer",
225                    "description": "Maximum number of memories to return",
226                    "default": 10
227                },
228                "since": {
229                    "type": "string",
230                    "description": "Only return memories created after this ISO 8601 timestamp"
231                },
232                "session_id": {
233                    "type": "string",
234                    "description": "Filter by session ID"
235                },
236                "chat_id": {
237                    "type": "string",
238                    "description": "Filter by chat ID (Bot platforms)"
239                }
240            }
241        })
242    }
243
244    fn requires_confirmation(&self) -> bool {
245        false
246    }
247
248    async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<ToolResult> {
249        let parsed: RecallArgs = match serde_json::from_value(args) {
250            Ok(a) => a,
251            Err(e) => return Ok(ToolResult::error(format!("Argument parsing failed: {}", e))),
252        };
253
254        let db_path = match crate::storage::resolve_db_path(&ctx.working_dir, false) {
255            Ok(p) => p,
256            Err(e) => return Ok(ToolResult::error(format!("Failed to resolve DB path: {}", e))),
257        };
258
259        let conn = match rusqlite::Connection::open(&db_path) {
260            Ok(c) => c,
261            Err(e) => return Ok(ToolResult::error(format!("Failed to open DB: {}", e))),
262        };
263
264        let filter = MemoryFilter {
265            memory_type: parsed.memory_type.as_ref().map(|s| MemoryType::from_str(s)),
266            tags: parsed.tags,
267            session_id: parsed.session_id.or(Some(ctx.session_id.to_string())),
268            chat_id: parsed.chat_id,
269            since: parsed.since,
270            only_active: true,
271        };
272
273        let limit = parsed.limit.unwrap_or(10);
274
275        let memories = if let Some(query) = &parsed.query {
276            if query.trim().is_empty() {
277                list_memories(&conn, &filter, Some(limit))
278            } else {
279                recall_memories(&conn, query, &filter, limit)
280            }
281        } else {
282            list_memories(&conn, &filter, Some(limit))
283        };
284
285        match memories {
286            Ok(memories) if memories.is_empty() => Ok(ToolResult::success(
287                "No memories found matching the criteria.".to_string(),
288            )),
289            Ok(memories) => {
290                let mut result = format!("Found {} memories:\n\n", memories.len());
291                for (i, memory) in memories.iter().enumerate() {
292                    result.push_str(&format!(
293                        "{}. [{}] {}\n   {}\n   Tags: {}\n   Created: {}\n\n",
294                        i + 1,
295                        memory.memory_type.as_str(),
296                        memory.title,
297                        memory.content,
298                        if memory.tags.is_empty() {
299                            "(none)".to_string()
300                        } else {
301                            memory.tags.join(", ")
302                        },
303                        memory.created_at
304                    ));
305                }
306                Ok(ToolResult::success(result))
307            }
308            Err(e) => Ok(ToolResult::error(format!("Failed to recall memories: {}", e))),
309        }
310    }
311}
312
313// ============================================================================
314// Forget tool - remove memories
315// ============================================================================
316
317#[derive(Debug, Deserialize)]
318struct ForgetArgs {
319    memory_id: Option<String>,
320    title: Option<String>,
321    #[serde(default)]
322    permanent: bool,
323}
324
325pub struct ForgetTool;
326
327impl ForgetTool {
328    pub fn new() -> Self {
329        Self
330    }
331}
332
333impl Default for ForgetTool {
334    fn default() -> Self {
335        Self::new()
336    }
337}
338
339#[async_trait]
340impl Tool for ForgetTool {
341    fn name(&self) -> &str {
342        "forget"
343    }
344
345    fn description(&self) -> &str {
346        "Remove or deactivate memories you no longer need."
347    }
348
349    fn parameters_schema(&self) -> Value {
350        serde_json::json!({
351            "type": "object",
352            "properties": {
353                "memory_id": {
354                    "type": "string",
355                    "description": "Specific memory ID to remove"
356                },
357                "title": {
358                    "type": "string",
359                    "description": "Remove memories matching this title"
360                },
361                "permanent": {
362                    "type": "boolean",
363                    "description": "If true, permanently delete. If false, just deactivate",
364                    "default": false
365                }
366            }
367        })
368    }
369
370    fn requires_confirmation(&self) -> bool {
371        true
372    }
373
374    async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<ToolResult> {
375        let parsed: ForgetArgs = match serde_json::from_value(args) {
376            Ok(a) => a,
377            Err(e) => return Ok(ToolResult::error(format!("Argument parsing failed: {}", e))),
378        };
379
380        if parsed.memory_id.is_none() && parsed.title.is_none() {
381            return Ok(ToolResult::error(
382                "Either memory_id or title must be provided".to_string(),
383            ));
384        }
385
386        let db_path = match crate::storage::resolve_db_path(&ctx.working_dir, false) {
387            Ok(p) => p,
388            Err(e) => return Ok(ToolResult::error(format!("Failed to resolve DB path: {}", e))),
389        };
390
391        let conn = match rusqlite::Connection::open(&db_path) {
392            Ok(c) => c,
393            Err(e) => return Ok(ToolResult::error(format!("Failed to open DB: {}", e))),
394        };
395
396        if let Some(memory_id) = parsed.memory_id {
397            let result = if parsed.permanent {
398                delete_memory_permanently(&conn, &memory_id)
399            } else {
400                deactivate_memory(&conn, &memory_id)
401            };
402
403            match result {
404                Ok(_) => Ok(ToolResult::success(format!(
405                    "Memory {} has been {}.",
406                    memory_id,
407                    if parsed.permanent {
408                        "permanently deleted"
409                    } else {
410                        "deactivated"
411                    }
412                ))),
413                Err(e) => Ok(ToolResult::error(format!("Failed to remove memory: {}", e))),
414            }
415        } else if let Some(title) = parsed.title {
416            let filter = MemoryFilter {
417                session_id: Some(ctx.session_id.to_string()),
418                only_active: true,
419                ..Default::default()
420            };
421            match list_memories(&conn, &filter, Some(100)) {
422                Ok(memories) => {
423                    let matching: Vec<_> = memories.into_iter().filter(|m| m.title == title).collect();
424                    if matching.is_empty() {
425                        return Ok(ToolResult::error(format!("No memory found with title: {}", title)));
426                    }
427
428                    let mut removed = 0;
429                    for memory in &matching {
430                        let result = if parsed.permanent {
431                            delete_memory_permanently(&conn, &memory.id)
432                        } else {
433                            deactivate_memory(&conn, &memory.id)
434                        };
435                        if result.is_ok() {
436                            removed += 1;
437                        }
438                    }
439
440                    Ok(ToolResult::success(format!(
441                        "Removed {} memory{} with title: {}",
442                        removed,
443                        if removed == 1 { "" } else { "s" },
444                        title
445                    )))
446                }
447                Err(e) => Ok(ToolResult::error(format!("Failed to find memories: {}", e))),
448            }
449        } else {
450            Ok(ToolResult::error(
451                "Internal error: neither memory_id nor title".to_string(),
452            ))
453        }
454    }
455}
456
457// ============================================================================
458// ListMemories tool - list all active memories
459// ============================================================================
460
461#[derive(Debug, Deserialize)]
462struct ListMemoriesArgs {
463    limit: Option<usize>,
464    memory_type: Option<String>,
465}
466
467pub struct ListMemoriesTool;
468
469impl ListMemoriesTool {
470    pub fn new() -> Self {
471        Self
472    }
473}
474
475impl Default for ListMemoriesTool {
476    fn default() -> Self {
477        Self::new()
478    }
479}
480
481#[async_trait]
482impl Tool for ListMemoriesTool {
483    fn name(&self) -> &str {
484        "list_memories"
485    }
486
487    fn description(&self) -> &str {
488        "List all active memories for a quick overview."
489    }
490
491    fn parameters_schema(&self) -> Value {
492        serde_json::json!({
493            "type": "object",
494            "properties": {
495                "limit": {
496                    "type": "integer",
497                    "description": "Maximum number of memories to list",
498                    "default": 20
499                },
500                "memory_type": {
501                    "type": "string",
502                    "description": "Filter by memory type: fact, preference, note, task"
503                }
504            }
505        })
506    }
507
508    fn requires_confirmation(&self) -> bool {
509        false
510    }
511
512    async fn execute(&self, args: Value, ctx: &ToolContext) -> Result<ToolResult> {
513        let parsed: ListMemoriesArgs = match serde_json::from_value(args) {
514            Ok(a) => a,
515            Err(e) => return Ok(ToolResult::error(format!("Argument parsing failed: {}", e))),
516        };
517
518        let db_path = match crate::storage::resolve_db_path(&ctx.working_dir, false) {
519            Ok(p) => p,
520            Err(e) => return Ok(ToolResult::error(format!("Failed to resolve DB path: {}", e))),
521        };
522
523        let conn = match rusqlite::Connection::open(&db_path) {
524            Ok(c) => c,
525            Err(e) => return Ok(ToolResult::error(format!("Failed to open DB: {}", e))),
526        };
527
528        let filter = MemoryFilter {
529            memory_type: parsed.memory_type.as_ref().map(|s| MemoryType::from_str(s)),
530            session_id: Some(ctx.session_id.to_string()),
531            only_active: true,
532            ..Default::default()
533        };
534
535        match list_memories(&conn, &filter, parsed.limit) {
536            Ok(memories) if memories.is_empty() => Ok(ToolResult::success(
537                "No active memories yet. Use `memorize` to store something!".to_string(),
538            )),
539            Ok(memories) => {
540                let mut result = format!("Your memories ({} total):\n\n", memories.len());
541                for (i, memory) in memories.iter().enumerate() {
542                    result.push_str(&format!(
543                        "{}. [{}] {}\n   ID: {}\n   Tags: {}\n\n",
544                        i + 1,
545                        memory.memory_type.as_str(),
546                        memory.title,
547                        memory.id,
548                        if memory.tags.is_empty() {
549                            "(none)".to_string()
550                        } else {
551                            memory.tags.join(", ")
552                        }
553                    ));
554                }
555                Ok(ToolResult::success(result))
556            }
557            Err(e) => Ok(ToolResult::error(format!("Failed to list memories: {}", e))),
558        }
559    }
560}
robit_agent/tool/memory.rs

robit_agent/tool/
memory.rs
