rs_adk/memory/
vertex_ai_rag.rs

1//! Vertex AI RAG memory service — stores and retrieves memories via Vertex AI RAG.
2//!
3//! Mirrors ADK-Python's `vertex_ai_rag_memory_service`.
4
5use async_trait::async_trait;
6
7use super::{MemoryEntry, MemoryError, MemoryService};
8
9/// Configuration for Vertex AI RAG memory service.
10#[derive(Debug, Clone)]
11pub struct VertexAiRagMemoryConfig {
12    /// The RAG corpus resource name.
13    pub corpus: String,
14    /// Google Cloud project ID.
15    pub project: String,
16    /// Google Cloud location.
17    pub location: String,
18}
19
20/// Memory service backed by Vertex AI RAG.
21///
22/// Stores memory entries as documents in a Vertex AI RAG corpus
23/// and uses semantic search for retrieval.
24#[derive(Debug, Clone)]
25pub struct VertexAiRagMemoryService {
26    config: VertexAiRagMemoryConfig,
27}
28
29impl VertexAiRagMemoryService {
30    /// Create a new Vertex AI RAG memory service.
31    pub fn new(config: VertexAiRagMemoryConfig) -> Self {
32        Self { config }
33    }
34
35    /// Returns the configured corpus.
36    pub fn corpus(&self) -> &str {
37        &self.config.corpus
38    }
39}
40
41#[async_trait]
42impl MemoryService for VertexAiRagMemoryService {
43    async fn store(&self, _session_id: &str, _entry: MemoryEntry) -> Result<(), MemoryError> {
44        // In a real integration, this would call the Vertex AI RAG API
45        // to store the memory entry as a document.
46        Ok(())
47    }
48
49    async fn get(&self, _session_id: &str, _key: &str) -> Result<Option<MemoryEntry>, MemoryError> {
50        // Vertex AI RAG doesn't support direct key-based retrieval;
51        // use search() instead.
52        Ok(None)
53    }
54
55    async fn list(&self, _session_id: &str) -> Result<Vec<MemoryEntry>, MemoryError> {
56        Ok(vec![])
57    }
58
59    async fn search(
60        &self,
61        _session_id: &str,
62        _query: &str,
63    ) -> Result<Vec<MemoryEntry>, MemoryError> {
64        // In a real integration, this would call retrieveContexts API.
65        Ok(vec![])
66    }
67
68    async fn delete(&self, _session_id: &str, _key: &str) -> Result<(), MemoryError> {
69        Ok(())
70    }
71
72    async fn clear(&self, _session_id: &str) -> Result<(), MemoryError> {
73        Ok(())
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    fn test_config() -> VertexAiRagMemoryConfig {
82        VertexAiRagMemoryConfig {
83            corpus: "projects/test/locations/us-central1/ragCorpora/test-corpus".into(),
84            project: "test".into(),
85            location: "us-central1".into(),
86        }
87    }
88
89    #[test]
90    fn service_metadata() {
91        let svc = VertexAiRagMemoryService::new(test_config());
92        assert!(svc.corpus().contains("test-corpus"));
93    }
94
95    #[tokio::test]
96    async fn store_and_search_stub() {
97        let svc = VertexAiRagMemoryService::new(test_config());
98        let entry = MemoryEntry::new("test", serde_json::json!("data"));
99        svc.store("s1", entry).await.unwrap();
100        let results = svc.search("s1", "test").await.unwrap();
101        assert!(results.is_empty()); // stub returns empty
102    }
103}
rs_adk/memory/vertex_ai_rag.rs

rs_adk/memory/
vertex_ai_rag.rs
