gemini_adk_rs/memory/
in_memory.rs1use async_trait::async_trait;
4use dashmap::DashMap;
5
6use super::{MemoryEntry, MemoryError, MemoryService};
7
8pub struct InMemoryMemoryService {
13 store: DashMap<String, DashMap<String, MemoryEntry>>,
15}
16
17impl InMemoryMemoryService {
18 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}