oxicode_sdk/ports/inmem/
memory.rs

1//! In-memory `MemoryStore` for tests and ephemeral sessions.
2//!
3//! Use this when:
4//! - You need an in-process memory store with no I/O
5//! - You're writing tests and want a fast fake
6//! - You're prototyping and persistence doesn't matter yet
7//!
8//! For durable memory, implement `oxicode_sdk::ports::MemoryStore` against
9//! SQLite/Redis/vector DB.
10
11use parking_lot::Mutex;
12use std::collections::HashMap;
13use std::future::Future;
14use std::pin::Pin;
15
16use crate::SdkError;
17use crate::ports::{MemoryEntry, MemoryStore};
18
19/// Thread-safe in-memory memory store.
20pub struct InMemoryMemoryStore {
21    inner: Mutex<HashMap<String, MemoryEntry>>,
22}
23
24impl std::fmt::Debug for InMemoryMemoryStore {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        f.debug_struct("InMemoryMemoryStore").finish()
27    }
28}
29
30impl Default for InMemoryMemoryStore {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl InMemoryMemoryStore {
37    /// Create a new empty store.
38    pub fn new() -> Self {
39        Self {
40            inner: Mutex::new(HashMap::new()),
41        }
42    }
43}
44
45impl MemoryStore for InMemoryMemoryStore {
46    fn put(
47        &self,
48        entry: MemoryEntry,
49    ) -> Pin<Box<dyn Future<Output = Result<(), SdkError>> + Send + '_>> {
50        self.inner.lock().insert(entry.id.clone(), entry);
51        Box::pin(async { Ok(()) })
52    }
53
54    fn list(
55        &self,
56        subject: &str,
57    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryEntry>, SdkError>> + Send + '_>> {
58        let g = self.inner.lock();
59        let result: Vec<MemoryEntry> = g
60            .values()
61            .filter(|e| e.subject == subject)
62            .cloned()
63            .collect();
64        Box::pin(async { Ok(result) })
65    }
66
67    fn search(
68        &self,
69        query: &[f32],
70        k: usize,
71    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryEntry>, SdkError>> + Send + '_>> {
72        let g = self.inner.lock();
73        let mut scored: Vec<_> = g
74            .values()
75            .filter_map(|e| e.embedding.as_ref().map(|emb| (e, cosine(query, emb))))
76            .collect();
77        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
78        let result: Vec<MemoryEntry> = scored.into_iter().take(k).map(|(e, _)| e.clone()).collect();
79        Box::pin(async { Ok(result) })
80    }
81    fn delete(&self, id: &str) -> Pin<Box<dyn Future<Output = Result<(), SdkError>> + Send + '_>> {
82        self.inner.lock().remove(id);
83        Box::pin(async { Ok(()) })
84    }
85}
86
87fn cosine(a: &[f32], b: &[f32]) -> f32 {
88    if a.len() != b.len() || a.is_empty() {
89        return 0.0;
90    }
91    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
92    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
93    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
94    if na == 0.0 || nb == 0.0 {
95        0.0
96    } else {
97        dot / (na * nb)
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use serde_json::json;
105
106    fn entry(id: &str, subject: &str, emb: Option<Vec<f32>>) -> MemoryEntry {
107        MemoryEntry {
108            id: id.into(),
109            subject: subject.into(),
110            kind: "episodic".into(),
111            embedding: emb,
112            content: json!({}),
113            created_at: chrono::Utc::now(),
114        }
115    }
116
117    #[tokio::test]
118    async fn put_and_list() {
119        let s = InMemoryMemoryStore::new();
120        s.put(entry("a", "agent-1", None)).await.unwrap();
121        s.put(entry("b", "agent-2", None)).await.unwrap();
122        s.put(entry("c", "agent-1", None)).await.unwrap();
123        let list = s.list("agent-1").await.unwrap();
124        assert_eq!(list.len(), 2);
125    }
126
127    #[tokio::test]
128    async fn cosine_search_returns_top_k() {
129        let s = InMemoryMemoryStore::new();
130        s.put(entry("a", "x", Some(vec![1.0, 0.0, 0.0])))
131            .await
132            .unwrap();
133        s.put(entry("b", "x", Some(vec![0.0, 1.0, 0.0])))
134            .await
135            .unwrap();
136        s.put(entry("c", "x", Some(vec![0.9, 0.1, 0.0])))
137            .await
138            .unwrap();
139        let top = s.search(&[1.0, 0.0, 0.0], 2).await.unwrap();
140        assert_eq!(top.len(), 2);
141        assert_eq!(top[0].id, "a");
142        assert_eq!(top[1].id, "c");
143    }
144}
oxicode_sdk/ports/inmem/memory.rs

oxicode_sdk/ports/inmem/
memory.rs
