agent_sdk/tools/
memory_tools.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde_json::json;
5
6use crate::error::{AgentId, SdkResult};
7use crate::traits::tool::{Tool, ToolDefinition};
8use crate::agent::memory::MemoryStore;
9
10pub struct ReadMemoryTool {
11    pub memory_store: Arc<MemoryStore>,
12}
13
14#[async_trait]
15impl Tool for ReadMemoryTool {
16    fn definition(&self) -> ToolDefinition {
17        ToolDefinition {
18            name: "read_memory".to_string(),
19            description: "Read a value from the shared memory store.".to_string(),
20            parameters: json!({
21                "type": "object",
22                "properties": {
23                    "key": { "type": "string", "description": "The memory key to read" }
24                },
25                "required": ["key"]
26            }),
27        }
28    }
29
30    async fn execute(&self, arguments: serde_json::Value) -> SdkResult<serde_json::Value> {
31        let key = arguments["key"].as_str().unwrap_or("");
32        if key.is_empty() {
33            return Ok(json!({"error": "Missing 'key' argument"}));
34        }
35
36        match self.memory_store.read(key)? {
37            Some(entry) => Ok(json!({
38                "key": entry.key,
39                "value": entry.value,
40                "written_by": entry.written_by.to_string(),
41                "written_at": entry.written_at.to_rfc3339()
42            })),
43            None => Ok(json!({ "key": key, "value": null, "found": false })),
44        }
45    }
46}
47
48pub struct WriteMemoryTool {
49    pub memory_store: Arc<MemoryStore>,
50    pub agent_id: AgentId,
51}
52
53#[async_trait]
54impl Tool for WriteMemoryTool {
55    fn definition(&self) -> ToolDefinition {
56        ToolDefinition {
57            name: "write_memory".to_string(),
58            description: "Write a value to the shared memory store.".to_string(),
59            parameters: json!({
60                "type": "object",
61                "properties": {
62                    "key": { "type": "string", "description": "The memory key" },
63                    "value": { "description": "The value to store" }
64                },
65                "required": ["key", "value"]
66            }),
67        }
68    }
69
70    async fn execute(&self, arguments: serde_json::Value) -> SdkResult<serde_json::Value> {
71        let key = arguments["key"].as_str().unwrap_or("");
72        if key.is_empty() {
73            return Ok(json!({"error": "Missing 'key' argument"}));
74        }
75
76        let value = arguments.get("value").cloned().unwrap_or(serde_json::Value::Null);
77        self.memory_store.write(key, value, self.agent_id)?;
78
79        Ok(json!({ "key": key, "stored": true }))
80    }
81}
82
83pub struct ListMemoryTool {
84    pub memory_store: Arc<MemoryStore>,
85}
86
87#[async_trait]
88impl Tool for ListMemoryTool {
89    fn definition(&self) -> ToolDefinition {
90        ToolDefinition {
91            name: "list_memory".to_string(),
92            description: "List all keys in the shared memory store.".to_string(),
93            parameters: json!({ "type": "object", "properties": {} }),
94        }
95    }
96
97    async fn execute(&self, _arguments: serde_json::Value) -> SdkResult<serde_json::Value> {
98        let keys = self.memory_store.list_keys()?;
99        Ok(json!({ "keys": keys, "count": keys.len() }))
100    }
101}
agent_sdk/tools/memory_tools.rs

agent_sdk/tools/
memory_tools.rs
