Skip to main content

agent_sdk/agent/
memory.rs

1use std::path::PathBuf;
2
3use chrono::Utc;
4
5use crate::error::{AgentId, SdkError, SdkResult};
6use crate::types::memory::MemoryEntry;
7
8pub struct MemoryStore {
9    base_dir: PathBuf,
10}
11
12impl MemoryStore {
13    pub fn new(base_dir: PathBuf) -> SdkResult<Self> {
14        std::fs::create_dir_all(&base_dir).map_err(SdkError::Io)?;
15        Ok(Self { base_dir })
16    }
17
18    fn key_path(&self, key: &str) -> PathBuf {
19        let safe_key: String = key
20            .chars()
21            .map(|c| {
22                if c.is_alphanumeric() || c == '_' || c == '-' || c == '.' {
23                    c
24                } else {
25                    '_'
26                }
27            })
28            .collect();
29        self.base_dir.join(format!("{}.json", safe_key))
30    }
31
32    pub fn read(&self, key: &str) -> SdkResult<Option<MemoryEntry>> {
33        let path = self.key_path(key);
34        if !path.exists() {
35            return Ok(None);
36        }
37        let content = std::fs::read_to_string(&path)?;
38        let entry: MemoryEntry = serde_json::from_str(&content)?;
39        Ok(Some(entry))
40    }
41
42    pub fn write(&self, key: &str, value: serde_json::Value, agent_id: AgentId) -> SdkResult<()> {
43        let entry = MemoryEntry {
44            key: key.to_string(),
45            value,
46            written_by: agent_id,
47            written_at: Utc::now(),
48        };
49        let content = serde_json::to_string_pretty(&entry)?;
50        let path = self.key_path(key);
51        std::fs::write(&path, content)?;
52        Ok(())
53    }
54
55    pub fn list_keys(&self) -> SdkResult<Vec<String>> {
56        let mut keys = Vec::new();
57        if self.base_dir.exists() {
58            for entry in std::fs::read_dir(&self.base_dir)? {
59                let entry = entry?;
60                let name = entry.file_name().to_string_lossy().to_string();
61                if let Some(key) = name.strip_suffix(".json") {
62                    keys.push(key.to_string());
63                }
64            }
65        }
66        keys.sort();
67        Ok(keys)
68    }
69}