gemini_adk_rs/memory/
mod.rs

1//! Memory service — session-scoped memory for agents.
2//!
3//! Mirrors ADK-JS's `BaseMemoryService`. Provides a trait for storing and
4//! searching memory entries (key-value) with an in-memory default.
5
6mod in_memory;
7mod vertex_ai_memory_bank;
8#[cfg(feature = "vertex-ai-rag")]
9mod vertex_ai_rag;
10
11pub use in_memory::InMemoryMemoryService;
12pub use vertex_ai_memory_bank::{VertexAiMemoryBankConfig, VertexAiMemoryBankService};
13#[cfg(feature = "vertex-ai-rag")]
14pub use vertex_ai_rag::{VertexAiRagMemoryConfig, VertexAiRagMemoryService};
15
16use async_trait::async_trait;
17use serde::{Deserialize, Serialize};
18
19/// A memory entry — a named piece of information stored by an agent.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct MemoryEntry {
22    /// Unique key for this memory.
23    pub key: String,
24    /// The stored value.
25    pub value: serde_json::Value,
26    /// When this entry was created (Unix timestamp seconds).
27    pub created_at: u64,
28    /// When this entry was last updated (Unix timestamp seconds).
29    pub updated_at: u64,
30}
31
32impl MemoryEntry {
33    /// Create a new memory entry.
34    pub fn new(key: impl Into<String>, value: serde_json::Value) -> Self {
35        let now = now_secs();
36        Self {
37            key: key.into(),
38            value,
39            created_at: now,
40            updated_at: now,
41        }
42    }
43}
44
45/// Errors from memory service operations.
46#[derive(Debug, thiserror::Error)]
47pub enum MemoryError {
48    /// The requested memory key was not found.
49    #[error("Memory key not found: {0}")]
50    NotFound(String),
51    /// A storage backend error.
52    #[error("Storage error: {0}")]
53    Storage(String),
54    /// The backend does not support this operation.
55    #[error("Unsupported operation: {0}")]
56    Unsupported(String),
57}
58
59/// Trait for session-scoped memory persistence.
60///
61/// Memory is scoped to a session ID. Implementations must be `Send + Sync`.
62#[async_trait]
63pub trait MemoryService: Send + Sync {
64    /// Store a memory entry for a session.
65    async fn store(&self, session_id: &str, entry: MemoryEntry) -> Result<(), MemoryError>;
66
67    /// Retrieve a memory entry by key.
68    async fn get(&self, session_id: &str, key: &str) -> Result<Option<MemoryEntry>, MemoryError>;
69
70    /// List all memory entries for a session.
71    async fn list(&self, session_id: &str) -> Result<Vec<MemoryEntry>, MemoryError>;
72
73    /// Search memory entries by a query string (simple substring match in default impl).
74    async fn search(&self, session_id: &str, query: &str) -> Result<Vec<MemoryEntry>, MemoryError>;
75
76    /// Delete a memory entry.
77    async fn delete(&self, session_id: &str, key: &str) -> Result<(), MemoryError>;
78
79    /// Clear all memory for a session.
80    async fn clear(&self, session_id: &str) -> Result<(), MemoryError>;
81}
82
83fn now_secs() -> u64 {
84    std::time::SystemTime::now()
85        .duration_since(std::time::UNIX_EPOCH)
86        .unwrap_or_default()
87        .as_secs()
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn memory_entry_new() {
96        let entry = MemoryEntry::new("topic", serde_json::json!("Rust"));
97        assert_eq!(entry.key, "topic");
98        assert_eq!(entry.value, serde_json::json!("Rust"));
99        assert!(entry.created_at > 0);
100    }
101
102    #[test]
103    fn memory_service_is_object_safe() {
104        fn _assert(_: &dyn MemoryService) {}
105    }
106
107    #[tokio::test]
108    async fn store_and_get() {
109        let svc = InMemoryMemoryService::new();
110        let entry = MemoryEntry::new("topic", serde_json::json!("AI"));
111        svc.store("s1", entry).await.unwrap();
112
113        let fetched = svc.get("s1", "topic").await.unwrap();
114        assert!(fetched.is_some());
115        assert_eq!(fetched.unwrap().value, serde_json::json!("AI"));
116    }
117
118    #[tokio::test]
119    async fn get_nonexistent_returns_none() {
120        let svc = InMemoryMemoryService::new();
121        let fetched = svc.get("s1", "missing").await.unwrap();
122        assert!(fetched.is_none());
123    }
124
125    #[tokio::test]
126    async fn list_entries() {
127        let svc = InMemoryMemoryService::new();
128        svc.store("s1", MemoryEntry::new("a", serde_json::json!(1)))
129            .await
130            .unwrap();
131        svc.store("s1", MemoryEntry::new("b", serde_json::json!(2)))
132            .await
133            .unwrap();
134        svc.store("s2", MemoryEntry::new("c", serde_json::json!(3)))
135            .await
136            .unwrap();
137
138        let entries = svc.list("s1").await.unwrap();
139        assert_eq!(entries.len(), 2);
140    }
141
142    #[tokio::test]
143    async fn search_entries() {
144        let svc = InMemoryMemoryService::new();
145        svc.store(
146            "s1",
147            MemoryEntry::new("rust_topic", serde_json::json!("Rust programming")),
148        )
149        .await
150        .unwrap();
151        svc.store(
152            "s1",
153            MemoryEntry::new("python_topic", serde_json::json!("Python scripting")),
154        )
155        .await
156        .unwrap();
157
158        let results = svc.search("s1", "rust").await.unwrap();
159        assert_eq!(results.len(), 1);
160        assert_eq!(results[0].key, "rust_topic");
161    }
162
163    #[tokio::test]
164    async fn delete_entry() {
165        let svc = InMemoryMemoryService::new();
166        svc.store("s1", MemoryEntry::new("k", serde_json::json!(1)))
167            .await
168            .unwrap();
169        svc.delete("s1", "k").await.unwrap();
170        let fetched = svc.get("s1", "k").await.unwrap();
171        assert!(fetched.is_none());
172    }
173
174    #[tokio::test]
175    async fn clear_session() {
176        let svc = InMemoryMemoryService::new();
177        svc.store("s1", MemoryEntry::new("a", serde_json::json!(1)))
178            .await
179            .unwrap();
180        svc.store("s1", MemoryEntry::new("b", serde_json::json!(2)))
181            .await
182            .unwrap();
183        svc.clear("s1").await.unwrap();
184        let entries = svc.list("s1").await.unwrap();
185        assert!(entries.is_empty());
186    }
187}
gemini_adk_rs/memory/mod.rs

gemini_adk_rs/memory/
mod.rs
