agent_memory_store/
lib.rs

1/*!
2agent-memory-store: simple key-value memory store for AI agents.
3
4Persist facts, preferences, and summaries across conversation turns.
5Supports tagging for grouped retrieval and tracks last-access time.
6
7```rust
8use agent_memory_store::MemoryStore;
9use serde_json::json;
10
11let mut store = MemoryStore::new();
12store.store("user_name", json!("Alice"), &["user", "profile"]);
13let entry = store.get("user_name").unwrap();
14assert_eq!(entry.value, json!("Alice"));
15assert_eq!(store.len(), 1);
16```
17*/
18
19use serde_json::Value;
20use std::collections::HashMap;
21use std::time::{SystemTime, UNIX_EPOCH};
22
23fn now_f64() -> f64 {
24    SystemTime::now()
25        .duration_since(UNIX_EPOCH)
26        .map(|d| d.as_secs_f64())
27        .unwrap_or(0.0)
28}
29
30/// A stored memory entry.
31#[derive(Debug, Clone)]
32pub struct MemoryEntry {
33    pub key: String,
34    pub value: Value,
35    pub tags: Vec<String>,
36    pub created_at: f64,
37    pub accessed_at: f64,
38    pub access_count: usize,
39}
40
41impl MemoryEntry {
42    pub fn to_json(&self) -> Value {
43        serde_json::json!({
44            "key": self.key,
45            "value": self.value,
46            "tags": self.tags,
47            "created_at": self.created_at,
48            "accessed_at": self.accessed_at,
49            "access_count": self.access_count,
50        })
51    }
52}
53
54/// Append-and-update key-value memory store.
55#[derive(Debug, Default, Clone)]
56pub struct MemoryStore {
57    entries: HashMap<String, MemoryEntry>,
58}
59
60impl MemoryStore {
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    /// Store `value` under `key`. Overwrites any existing entry.
66    pub fn store(&mut self, key: &str, value: Value, tags: &[&str]) {
67        let now = now_f64();
68        let existing = self.entries.get(key);
69        let created_at = existing.map(|e| e.created_at).unwrap_or(now);
70        let access_count = existing.map(|e| e.access_count).unwrap_or(0);
71        self.entries.insert(
72            key.to_owned(),
73            MemoryEntry {
74                key: key.to_owned(),
75                value,
76                tags: tags.iter().map(|s| s.to_string()).collect(),
77                created_at,
78                accessed_at: now,
79                access_count,
80            },
81        );
82    }
83
84    /// Get an entry by key, updating its access metadata.
85    pub fn get(&mut self, key: &str) -> Option<&MemoryEntry> {
86        let e = self.entries.get_mut(key)?;
87        e.accessed_at = now_f64();
88        e.access_count += 1;
89        Some(e)
90    }
91
92    /// Get an entry by key without updating access metadata.
93    pub fn peek(&self, key: &str) -> Option<&MemoryEntry> {
94        self.entries.get(key)
95    }
96
97    /// Delete an entry. Returns `true` if the key existed.
98    pub fn delete(&mut self, key: &str) -> bool {
99        self.entries.remove(key).is_some()
100    }
101
102    /// All entries, sorted by key.
103    pub fn all(&self) -> Vec<&MemoryEntry> {
104        let mut v: Vec<&MemoryEntry> = self.entries.values().collect();
105        v.sort_by_key(|e| e.key.as_str());
106        v
107    }
108
109    /// Entries matching a specific tag.
110    pub fn by_tag(&self, tag: &str) -> Vec<&MemoryEntry> {
111        self.entries
112            .values()
113            .filter(|e| e.tags.iter().any(|t| t == tag))
114            .collect()
115    }
116
117    /// Keys that currently exist.
118    pub fn keys(&self) -> Vec<&str> {
119        let mut k: Vec<&str> = self.entries.keys().map(|s| s.as_str()).collect();
120        k.sort();
121        k
122    }
123
124    pub fn len(&self) -> usize {
125        self.entries.len()
126    }
127
128    pub fn is_empty(&self) -> bool {
129        self.entries.is_empty()
130    }
131
132    /// Remove all entries.
133    pub fn clear(&mut self) {
134        self.entries.clear();
135    }
136
137    /// Serialize all entries to a JSON array.
138    pub fn to_json(&self) -> Value {
139        Value::Array(self.all().iter().map(|e| e.to_json()).collect())
140    }
141
142    /// True if `key` exists.
143    pub fn contains(&self, key: &str) -> bool {
144        self.entries.contains_key(key)
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use serde_json::json;
152
153    #[test]
154    fn store_and_peek() {
155        let mut s = MemoryStore::new();
156        s.store("k", json!("val"), &["tag1"]);
157        let e = s.peek("k").unwrap();
158        assert_eq!(e.value, json!("val"));
159    }
160
161    #[test]
162    fn get_increments_access() {
163        let mut s = MemoryStore::new();
164        s.store("k", json!(1), &[]);
165        s.get("k").unwrap();
166        s.get("k").unwrap();
167        assert_eq!(s.peek("k").unwrap().access_count, 2);
168    }
169
170    #[test]
171    fn missing_key_returns_none() {
172        let mut s = MemoryStore::new();
173        assert!(s.get("missing").is_none());
174        assert!(s.peek("missing").is_none());
175    }
176
177    #[test]
178    fn overwrite_preserves_created_at() {
179        let mut s = MemoryStore::new();
180        s.store("k", json!(1), &[]);
181        let t1 = s.peek("k").unwrap().created_at;
182        std::thread::sleep(std::time::Duration::from_millis(5));
183        s.store("k", json!(2), &[]);
184        let t2 = s.peek("k").unwrap().created_at;
185        assert!((t1 - t2).abs() < 0.1);
186    }
187
188    #[test]
189    fn delete_removes_entry() {
190        let mut s = MemoryStore::new();
191        s.store("k", json!(1), &[]);
192        assert!(s.delete("k"));
193        assert!(!s.contains("k"));
194    }
195
196    #[test]
197    fn delete_returns_false_when_missing() {
198        let mut s = MemoryStore::new();
199        assert!(!s.delete("nope"));
200    }
201
202    #[test]
203    fn by_tag_filters() {
204        let mut s = MemoryStore::new();
205        s.store("a", json!(1), &["user"]);
206        s.store("b", json!(2), &["user", "profile"]);
207        s.store("c", json!(3), &["system"]);
208        assert_eq!(s.by_tag("user").len(), 2);
209        assert_eq!(s.by_tag("system").len(), 1);
210        assert_eq!(s.by_tag("unknown").len(), 0);
211    }
212
213    #[test]
214    fn all_sorted_by_key() {
215        let mut s = MemoryStore::new();
216        s.store("b", json!(2), &[]);
217        s.store("a", json!(1), &[]);
218        let all = s.all();
219        assert_eq!(all[0].key, "a");
220        assert_eq!(all[1].key, "b");
221    }
222
223    #[test]
224    fn len_increments() {
225        let mut s = MemoryStore::new();
226        assert_eq!(s.len(), 0);
227        s.store("k", json!(1), &[]);
228        assert_eq!(s.len(), 1);
229    }
230
231    #[test]
232    fn contains_returns_correct() {
233        let mut s = MemoryStore::new();
234        assert!(!s.contains("k"));
235        s.store("k", json!(1), &[]);
236        assert!(s.contains("k"));
237    }
238
239    #[test]
240    fn clear_empties_store() {
241        let mut s = MemoryStore::new();
242        s.store("k", json!(1), &[]);
243        s.clear();
244        assert!(s.is_empty());
245    }
246
247    #[test]
248    fn to_json_is_array() {
249        let mut s = MemoryStore::new();
250        s.store("k", json!(1), &[]);
251        let j = s.to_json();
252        assert!(j.is_array());
253        assert_eq!(j.as_array().unwrap().len(), 1);
254    }
255
256    #[test]
257    fn entry_to_json_has_fields() {
258        let mut s = MemoryStore::new();
259        s.store("k", json!("v"), &["t"]);
260        let e = s.peek("k").unwrap();
261        let j = e.to_json();
262        assert_eq!(j["key"], "k");
263        assert_eq!(j["value"], "v");
264        assert_eq!(j["tags"][0], "t");
265    }
266
267    #[test]
268    fn keys_sorted() {
269        let mut s = MemoryStore::new();
270        s.store("c", json!(3), &[]);
271        s.store("a", json!(1), &[]);
272        s.store("b", json!(2), &[]);
273        assert_eq!(s.keys(), vec!["a", "b", "c"]);
274    }
275}
agent_memory_store/lib.rs

agent_memory_store/
lib.rs
