use super::traits::{Tool, ToolResult}; use async_trait::async_trait; use serde_json::json; use std::sync::Arc; use crate::memory::{MemoryCategory, MemoryManager}; // ── MemoryStoreTool ────────────────────────────────────────────────── pub struct MemoryStoreTool { memory: Arc, } impl MemoryStoreTool { pub fn new(memory: Arc) -> Self { Self { memory } } } #[async_trait] impl Tool for MemoryStoreTool { fn name(&self) -> &str { "memory_store" } fn description(&self) -> &str { "Store a fact, preference, or insight into long-term memory. \ Use this when the user shares important information you should remember. \ Provide a descriptive key (e.g., 'user_prefers_python', 'project_auth_approach') \ and the full content to remember." } fn read_only(&self) -> bool { false } fn parameters_schema(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "key": { "type": "string", "description": "Semantic identifier for this memory (e.g., 'user_language_pref'). Unique key." }, "content": { "type": "string", "description": "The full content of the memory entry." }, "category": { "type": "string", "enum": ["knowledge", "timeline"], "description": "Memory category. Use 'knowledge' for facts/preferences/insights, 'timeline' for conversation summaries." }, "importance": { "type": "number", "description": "Importance score 0.0-1.0. Higher = more important. Use 0.8+ for critical facts, 0.5 for general info." } }, "required": ["key", "content"] }) } async fn execute(&self, args: serde_json::Value) -> anyhow::Result { let key = args .get("key") .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required parameter: key"))?; let content = args .get("content") .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required parameter: content"))?; let category = args .get("category") .and_then(|v| v.as_str()) .and_then(MemoryCategory::from_str) .unwrap_or(MemoryCategory::Knowledge); let importance = args.get("importance").and_then(|v| v.as_f64()); self.memory .store(key, content, category, None, importance) .await?; Ok(ToolResult { success: true, output: format!("Memory stored: {}", key), error: None, }) } } // ── MemoryRecallTool ───────────────────────────────────────────────── pub struct MemoryRecallTool { memory: Arc, } impl MemoryRecallTool { pub fn new(memory: Arc) -> Self { Self { memory } } } #[async_trait] impl Tool for MemoryRecallTool { fn name(&self) -> &str { "memory_recall" } fn description(&self) -> &str { "Search and retrieve entries from long-term memory. \ Use this to recall previously stored facts, preferences, or conversation history. \ Supports keyword search and optional time-range filtering." } fn read_only(&self) -> bool { true } fn parameters_schema(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "query": { "type": "string", "description": "Search query — keywords to match against memory keys and content." }, "category": { "type": "string", "enum": ["knowledge", "timeline"], "description": "Filter by memory category. Omit to search all categories." }, "since": { "type": "integer", "description": "Start of time range (Unix milliseconds)." }, "until": { "type": "integer", "description": "End of time range (Unix milliseconds)." }, "limit": { "type": "integer", "description": "Max results to return (default 10)." } }, "required": ["query"] }) } async fn execute(&self, args: serde_json::Value) -> anyhow::Result { let query = args .get("query") .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required parameter: query"))?; let category = args .get("category") .and_then(|v| v.as_str()) .and_then(MemoryCategory::from_str); let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(10) as usize; let entries = if args.get("since").is_some() || args.get("until").is_some() { let since = args.get("since").and_then(|v| v.as_i64()).unwrap_or(0); let until = args .get("until") .and_then(|v| v.as_i64()) .unwrap_or(chrono::Utc::now().timestamp_millis()); self.memory .recall_by_time(since, until, limit, category) .await? } else { self.memory.recall(query, limit, category).await? }; if entries.is_empty() { return Ok(ToolResult { success: true, output: "No matching memories found.".to_string(), error: None, }); } let formatted = entries .iter() .map(|e| { format!( "- {} [{}] [importance: {:.1}]: {}", e.key, e.category.as_str(), e.importance, e.content ) }) .collect::>() .join("\n"); Ok(ToolResult { success: true, output: format!("Found {} memories:\n{}", entries.len(), formatted), error: None, }) } } // ── MemoryForgetTool ───────────────────────────────────────────────── pub struct MemoryForgetTool { memory: Arc, } impl MemoryForgetTool { pub fn new(memory: Arc) -> Self { Self { memory } } } #[async_trait] impl Tool for MemoryForgetTool { fn name(&self) -> &str { "memory_forget" } fn description(&self) -> &str { "Delete a memory entry by its key. Use this when information is outdated, \ incorrect, or the user asks to forget something." } fn read_only(&self) -> bool { false } fn parameters_schema(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "key": { "type": "string", "description": "The key of the memory entry to delete." } }, "required": ["key"] }) } async fn execute(&self, args: serde_json::Value) -> anyhow::Result { let key = args .get("key") .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required parameter: key"))?; self.memory.forget(key).await?; Ok(ToolResult { success: true, output: format!("Memory deleted: {}", key), error: None, }) } }