Skip to main content

tycode_core/modules/memory/
tool.rs

1//! Memory append tool for storing learnings.
2
3use std::sync::Arc;
4
5use crate::chat::events::{ToolExecutionResult, ToolRequest as ToolRequestEvent, ToolRequestType};
6use crate::tools::r#trait::{
7    ContinuationPreference, ToolCallHandle, ToolCategory, ToolExecutor, ToolOutput, ToolRequest,
8};
9use crate::tools::ToolName;
10
11use super::log::MemoryLog;
12
13pub struct AppendMemoryTool {
14    memory_log: Arc<MemoryLog>,
15}
16
17impl AppendMemoryTool {
18    pub fn new(memory_log: Arc<MemoryLog>) -> Self {
19        Self { memory_log }
20    }
21
22    pub fn tool_name() -> ToolName {
23        ToolName::new("append_memory")
24    }
25}
26
27#[async_trait::async_trait(?Send)]
28impl ToolExecutor for AppendMemoryTool {
29    fn name(&self) -> String {
30        "append_memory".to_string()
31    }
32
33    fn description(&self) -> String {
34        "Appends text to the memory log. Stored memories appear in the model's context in future conversations, helping avoid repeated corrections and follow user preferences. Store when corrected repeatedly or when the user expresses frustration.".to_string()
35    }
36
37    fn input_schema(&self) -> serde_json::Value {
38        serde_json::json!({
39            "type": "object",
40            "properties": {
41                "content": {
42                    "type": "string",
43                    "description": "A concise description of what was learned"
44                },
45                "source": {
46                    "type": "string",
47                    "description": "Optional project name this memory applies to. Omit for global memories."
48                }
49            },
50            "required": ["content"]
51        })
52    }
53
54    fn category(&self) -> ToolCategory {
55        ToolCategory::Execution
56    }
57
58    async fn process(&self, request: &ToolRequest) -> anyhow::Result<Box<dyn ToolCallHandle>> {
59        let content = request.arguments["content"]
60            .as_str()
61            .ok_or_else(|| anyhow::anyhow!("content is required"))?
62            .to_string();
63
64        let source = request
65            .arguments
66            .get("source")
67            .and_then(|v| v.as_str())
68            .map(|s| s.to_string());
69
70        Ok(Box::new(AppendMemoryHandle {
71            content,
72            source,
73            tool_use_id: request.tool_use_id.clone(),
74            memory_log: self.memory_log.clone(),
75        }))
76    }
77}
78
79struct AppendMemoryHandle {
80    content: String,
81    source: Option<String>,
82    tool_use_id: String,
83    memory_log: Arc<MemoryLog>,
84}
85
86#[async_trait::async_trait(?Send)]
87impl ToolCallHandle for AppendMemoryHandle {
88    fn tool_request(&self) -> ToolRequestEvent {
89        ToolRequestEvent {
90            tool_call_id: self.tool_use_id.clone(),
91            tool_name: "append_memory".to_string(),
92            tool_type: ToolRequestType::Other {
93                args: serde_json::json!({
94                    "content": self.content,
95                    "source": self.source
96                }),
97            },
98        }
99    }
100
101    async fn execute(self: Box<Self>) -> ToolOutput {
102        match self
103            .memory_log
104            .append(self.content.clone(), self.source.clone())
105        {
106            Ok(seq) => ToolOutput::Result {
107                content: serde_json::json!({
108                    "seq": seq,
109                    "content": self.content,
110                    "source": self.source,
111                    "success": true
112                })
113                .to_string(),
114                is_error: false,
115                continuation: ContinuationPreference::Continue,
116                ui_result: ToolExecutionResult::Other {
117                    result: serde_json::json!({
118                        "seq": seq,
119                        "success": true
120                    }),
121                },
122            },
123            Err(e) => ToolOutput::Result {
124                content: format!("Failed to append memory: {e:?}"),
125                is_error: true,
126                continuation: ContinuationPreference::Continue,
127                ui_result: ToolExecutionResult::Error {
128                    short_message: "Memory append failed".to_string(),
129                    detailed_message: format!("{e:?}"),
130                },
131            },
132        }
133    }
134}