rucora_tools/memory/
mod.rs

1//! 记忆工具模块。
2//!
3//! 提供长期记忆存储和检索功能。
4
5use async_trait::async_trait;
6use rucora_core::{
7    error::{MemoryError, ToolError},
8    memory::{Memory, MemoryItem, MemoryQuery},
9    tool::{Tool, ToolCategory},
10};
11use serde_json::{Value, json};
12use std::collections::HashMap;
13use std::sync::{Arc, Mutex};
14
15/// 简单的内存记忆实现
16struct SimpleMemory {
17    records: Mutex<HashMap<String, MemoryItem>>,
18}
19
20impl SimpleMemory {
21    fn new() -> Self {
22        Self {
23            records: Mutex::new(HashMap::new()),
24        }
25    }
26}
27
28#[async_trait]
29impl Memory for SimpleMemory {
30    async fn add(&self, item: MemoryItem) -> Result<(), MemoryError> {
31        self.records.lock().unwrap().insert(item.id.clone(), item);
32        Ok(())
33    }
34
35    async fn query(&self, _query: MemoryQuery) -> Result<Vec<MemoryItem>, MemoryError> {
36        Ok(self.records.lock().unwrap().values().cloned().collect())
37    }
38}
39
40/// 记忆存储工具：存储信息到长期记忆。
41///
42/// 让 Agent 将事实、偏好或笔记存储到长期记忆中。
43/// 支持不同类别：core（永久）、daily（会话）、conversation（对话）。
44///
45/// 输入格式：
46/// ```json
47/// {
48///   "key": "user_lang",
49///   "content": "用户偏好 Rust 语言",
50///   "category": "core"
51/// }
52/// ```
53pub struct MemoryStoreTool {
54    memory: Arc<dyn Memory>,
55}
56
57impl MemoryStoreTool {
58    /// 创建一个新的 MemoryStoreTool 实例。
59    pub fn new() -> Self {
60        Self::from_memory(Arc::new(SimpleMemory::new()))
61    }
62
63    pub fn from_memory(memory: Arc<dyn Memory>) -> Self {
64        Self { memory }
65    }
66}
67
68impl Default for MemoryStoreTool {
69    fn default() -> Self {
70        Self::new()
71    }
72}
73
74#[async_trait]
75impl Tool for MemoryStoreTool {
76    /// 返回工具名称。
77    fn name(&self) -> &str {
78        "memory_store"
79    }
80
81    /// 返回工具描述。
82    fn description(&self) -> Option<&str> {
83        Some(
84            "存储事实、偏好或笔记到长期记忆。使用 category 'core' 表示永久记忆，'daily' 表示会话笔记，'conversation' 表示对话上下文",
85        )
86    }
87
88    /// 返回工具分类。
89    fn categories(&self) -> &'static [ToolCategory] {
90        &[ToolCategory::Memory]
91    }
92
93    /// 返回输入参数的 JSON Schema。
94    fn input_schema(&self) -> Value {
95        json!({
96            "type": "object",
97            "properties": {
98                "key": {
99                    "type": "string",
100                    "description": "记忆的唯一键（如 'user_lang', 'project_stack'）"
101                },
102                "content": {
103                    "type": "string",
104                    "description": "要记忆的信息"
105                },
106                "category": {
107                    "type": "string",
108                    "description": "记忆类别: 'core' (永久), 'daily' (会话), 'conversation' (对话), 或自定义类别。默认为 'core'"
109                }
110            },
111            "required": ["key", "content"]
112        })
113    }
114
115    /// 执行记忆存储。
116    async fn call(&self, input: Value) -> Result<Value, ToolError> {
117        let key = input
118            .get("key")
119            .and_then(|v| v.as_str())
120            .ok_or_else(|| ToolError::Message("缺少必需的 'key' 字段".to_string()))?;
121
122        let content = input
123            .get("content")
124            .and_then(|v| v.as_str())
125            .ok_or_else(|| ToolError::Message("缺少必需的 'content' 字段".to_string()))?;
126
127        let category = input
128            .get("category")
129            .and_then(|v| v.as_str())
130            .unwrap_or("core");
131
132        // 存储
133        let full_key = format!("{category}:{key}");
134        self.memory
135            .add(MemoryItem {
136                id: full_key,
137                content: content.to_string(),
138                metadata: None,
139            })
140            .await
141            .map_err(|e| ToolError::Message(e.to_string()))?;
142
143        Ok(json!({
144            "success": true,
145            "key": key,
146            "category": category,
147            "message": format!("已存储记忆: {}", key)
148        }))
149    }
150}
151
152/// 记忆回忆工具：从长期记忆中检索信息。
153///
154/// 根据键从长期记忆中检索存储的信息。
155///
156/// 输入格式：
157/// ```json
158/// {
159///   "key": "user_lang",
160///   "category": "core"
161/// }
162/// ```
163pub struct MemoryRecallTool {
164    memory: Arc<dyn Memory>,
165}
166
167impl MemoryRecallTool {
168    /// 创建一个新的 MemoryRecallTool 实例。
169    pub fn new() -> Self {
170        Self::from_memory(Arc::new(SimpleMemory::new()))
171    }
172
173    /// 从现有 MemoryStoreTool 创建实例以共享存储。
174    pub fn from_store(store: &MemoryStoreTool) -> Self {
175        Self {
176            memory: store.memory.clone(),
177        }
178    }
179
180    pub fn from_memory(memory: Arc<dyn Memory>) -> Self {
181        Self { memory }
182    }
183}
184
185impl Default for MemoryRecallTool {
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191#[async_trait]
192impl Tool for MemoryRecallTool {
193    /// 返回工具名称。
194    fn name(&self) -> &str {
195        "memory_recall"
196    }
197
198    /// 返回工具描述。
199    fn description(&self) -> Option<&str> {
200        Some("从长期记忆中检索存储的信息")
201    }
202
203    /// 返回工具分类。
204    fn categories(&self) -> &'static [ToolCategory] {
205        &[ToolCategory::Memory]
206    }
207
208    /// 返回输入参数的 JSON Schema。
209    fn input_schema(&self) -> Value {
210        json!({
211            "type": "object",
212            "properties": {
213                "key": {
214                    "type": "string",
215                    "description": "要检索的记忆键"
216                },
217                "category": {
218                    "type": "string",
219                    "description": "记忆类别，默认为 'core'"
220                }
221            },
222            "required": ["key"]
223        })
224    }
225
226    /// 执行记忆检索。
227    async fn call(&self, input: Value) -> Result<Value, ToolError> {
228        let key = input
229            .get("key")
230            .and_then(|v| v.as_str())
231            .ok_or_else(|| ToolError::Message("缺少必需的 'key' 字段".to_string()))?;
232
233        let category = input
234            .get("category")
235            .and_then(|v| v.as_str())
236            .unwrap_or("core");
237
238        // 检索
239        let full_key = format!("{category}:{key}");
240        let mut results = self
241            .memory
242            .query(MemoryQuery {
243                text: full_key,
244                limit: 1,
245            })
246            .await
247            .map_err(|e| ToolError::Message(e.to_string()))?;
248
249        if let Some(item) = results.pop() {
250            Ok(json!({
251                "found": true,
252                "key": key,
253                "category": category,
254                "content": item.content
255            }))
256        } else {
257            Ok(json!({
258                "found": false,
259                "key": key,
260                "category": category,
261                "content": null
262            }))
263        }
264    }
265}
rucora_tools/memory/mod.rs

rucora_tools/memory/
mod.rs
