Skip to main content

gemini_adk_rs/memory/
in_memory.rs

1//! In-memory memory service using DashMap.
2
3use async_trait::async_trait;
4use dashmap::DashMap;
5
6use super::{MemoryEntry, MemoryError, MemoryService};
7
8/// In-memory memory service backed by [`DashMap`] for lock-free concurrent access.
9///
10/// Memory entries are scoped by session ID. Suitable for testing, prototyping,
11/// and single-process deployments. Data is lost on process restart.
12pub struct InMemoryMemoryService {
13    /// session_id → (key → MemoryEntry)
14    store: DashMap<String, DashMap<String, MemoryEntry>>,
15}
16
17impl InMemoryMemoryService {
18    /// Create a new in-memory memory service.
19    pub fn new() -> Self {
20        Self {
21            store: DashMap::new(),
22        }
23    }
24}
25
26impl Default for InMemoryMemoryService {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32#[async_trait]
33impl MemoryService for InMemoryMemoryService {
34    async fn store(&self, session_id: &str, entry: MemoryEntry) -> Result<(), MemoryError> {
35        let session = self.store.entry(session_id.to_string()).or_default();
36        session.insert(entry.key.clone(), entry);
37        Ok(())
38    }
39
40    async fn get(&self, session_id: &str, key: &str) -> Result<Option<MemoryEntry>, MemoryError> {
41        let result = self
42            .store
43            .get(session_id)
44            .and_then(|session| session.get(key).map(|e| e.clone()));
45        Ok(result)
46    }
47
48    async fn list(&self, session_id: &str) -> Result<Vec<MemoryEntry>, MemoryError> {
49        let entries = self
50            .store
51            .get(session_id)
52            .map(|session| session.iter().map(|e| e.value().clone()).collect())
53            .unwrap_or_default();
54        Ok(entries)
55    }
56
57    async fn search(&self, session_id: &str, query: &str) -> Result<Vec<MemoryEntry>, MemoryError> {
58        let query_lower = query.to_lowercase();
59        let entries = self
60            .store
61            .get(session_id)
62            .map(|session| {
63                session
64                    .iter()
65                    .filter(|e| {
66                        e.key().to_lowercase().contains(&query_lower)
67                            || e.value()
68                                .value
69                                .to_string()
70                                .to_lowercase()
71                                .contains(&query_lower)
72                    })
73                    .map(|e| e.value().clone())
74                    .collect()
75            })
76            .unwrap_or_default();
77        Ok(entries)
78    }
79
80    async fn delete(&self, session_id: &str, key: &str) -> Result<(), MemoryError> {
81        if let Some(session) = self.store.get(session_id) {
82            session.remove(key);
83        }
84        Ok(())
85    }
86
87    async fn clear(&self, session_id: &str) -> Result<(), MemoryError> {
88        if let Some(session) = self.store.get(session_id) {
89            session.clear();
90        }
91        Ok(())
92    }
93}