gemini_adk_rs/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//! * `add_session_to_memory` ingests session events by uploading them as a RAG
5//!   file (one JSON line per event) into the configured corpus, encoding the
6//!   `app_name` / `user_id` / `session_id` into the file's display name.
7//! * `search_memory` queries the corpus via the `:retrieveContexts` endpoint,
8//!   then filters results by the encoded `app_name` / `user_id`.
9//!
10//! The generic [`MemoryService`] trait (key-value oriented) is implemented on
11//! top of these primitives: `store` uploads a single entry, `search` queries
12//! the corpus.
13
14use std::sync::Arc;
15
16use async_trait::async_trait;
17use base64::Engine as _;
18use base64::engine::general_purpose::URL_SAFE_NO_PAD;
19use serde::Deserialize;
20use serde_json::{Value, json};
21
22use super::{MemoryEntry, MemoryError, MemoryService};
23
24/// Display-name prefix used to tag uploaded session files, mirroring ADK.
25const SOURCE_DISPLAY_NAME_PREFIX: &str = "adk-memory-v1.";
26
27// ──────────────────────────────────────────────────────────────────────────────
28// Configuration
29// ──────────────────────────────────────────────────────────────────────────────
30
31/// Configuration for the Vertex AI RAG memory service.
32#[derive(Debug, Clone)]
33pub struct VertexAiRagMemoryConfig {
34    /// The RAG corpus resource name.
35    /// Format: `projects/{project}/locations/{location}/ragCorpora/{corpus_id}`
36    pub corpus: String,
37    /// Google Cloud project ID.
38    pub project: String,
39    /// Google Cloud location (e.g. `us-central1`).
40    pub location: String,
41    /// Number of contexts to retrieve (maps to `similarity_top_k`).
42    pub similarity_top_k: Option<u32>,
43    /// Only return contexts with a vector distance smaller than this threshold.
44    pub vector_distance_threshold: Option<f64>,
45}
46
47impl VertexAiRagMemoryConfig {
48    /// Create a new config from a corpus resource name, deriving the
49    /// project / location from it when fully-qualified.
50    pub fn from_corpus(corpus: impl Into<String>) -> Self {
51        let corpus = corpus.into();
52        let (project, location) = parse_project_location(&corpus);
53        Self {
54            corpus,
55            project,
56            location,
57            similarity_top_k: None,
58            vector_distance_threshold: Some(10.0),
59        }
60    }
61
62    /// Build the `:retrieveContexts` endpoint URL.
63    fn retrieve_contexts_url(&self) -> String {
64        format!(
65            "https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}:retrieveContexts",
66            project = self.project,
67            location = self.location,
68        )
69    }
70
71    /// Build the RAG file upload endpoint URL for the configured corpus.
72    ///
73    /// Format:
74    /// `https://{location}-aiplatform.googleapis.com/upload/v1beta1/{corpus}/ragFiles:upload`
75    fn upload_rag_file_url(&self) -> String {
76        format!(
77            "https://{location}-aiplatform.googleapis.com/upload/v1beta1/{corpus}/ragFiles:upload",
78            location = self.location,
79            corpus = self.corpus,
80        )
81    }
82}
83
84/// Best-effort extraction of `{project}` / `{location}` from a corpus resource
85/// name of the form `projects/{project}/locations/{location}/ragCorpora/{id}`.
86fn parse_project_location(corpus: &str) -> (String, String) {
87    let mut project = String::new();
88    let mut location = String::new();
89    let mut parts = corpus.split('/');
90    while let Some(seg) = parts.next() {
91        match seg {
92            "projects" => project = parts.next().unwrap_or_default().to_string(),
93            "locations" => location = parts.next().unwrap_or_default().to_string(),
94            _ => {}
95        }
96    }
97    (project, location)
98}
99
100// ──────────────────────────────────────────────────────────────────────────────
101// source_display_name encode / decode (mirrors ADK)
102// ──────────────────────────────────────────────────────────────────────────────
103
104fn encode_part(value: &str) -> String {
105    URL_SAFE_NO_PAD.encode(value.as_bytes())
106}
107
108fn decode_part(value: &str) -> Option<String> {
109    URL_SAFE_NO_PAD
110        .decode(value.as_bytes())
111        .ok()
112        .and_then(|bytes| String::from_utf8(bytes).ok())
113}
114
115/// Encode `app_name` / `user_id` / `session_id` into a display name, matching
116/// ADK's `_build_source_display_name`.
117pub(crate) fn build_source_display_name(app_name: &str, user_id: &str, session_id: &str) -> String {
118    format!(
119        "{prefix}{a}.{u}.{s}",
120        prefix = SOURCE_DISPLAY_NAME_PREFIX,
121        a = encode_part(app_name),
122        u = encode_part(user_id),
123        s = encode_part(session_id),
124    )
125}
126
127/// Decode a display name back into `(app_name, user_id, session_id)`, matching
128/// ADK's `_parse_source_display_name`. Returns `None` if it doesn't match the
129/// expected three-part encoded form.
130pub(crate) fn parse_source_display_name(name: &str) -> Option<(String, String, String)> {
131    if let Some(rest) = name.strip_prefix(SOURCE_DISPLAY_NAME_PREFIX) {
132        let parts: Vec<&str> = rest.split('.').collect();
133        if parts.len() != 3 {
134            return None;
135        }
136        return Some((
137            decode_part(parts[0])?,
138            decode_part(parts[1])?,
139            decode_part(parts[2])?,
140        ));
141    }
142    // Legacy dot-delimited (plaintext) form.
143    let parts: Vec<&str> = name.split('.').collect();
144    if parts.len() != 3 {
145        return None;
146    }
147    Some((
148        parts[0].to_string(),
149        parts[1].to_string(),
150        parts[2].to_string(),
151    ))
152}
153
154// ──────────────────────────────────────────────────────────────────────────────
155// Token provider (mirrors session/vertex_ai.rs)
156// ──────────────────────────────────────────────────────────────────────────────
157
158enum TokenProvider {
159    None,
160    Static(String),
161    Refresher(Arc<dyn Fn() -> String + Send + Sync>),
162}
163
164impl TokenProvider {
165    fn get(&self) -> Result<String, MemoryError> {
166        match self {
167            TokenProvider::None => Err(MemoryError::Storage(
168                "missing auth token: call .with_token() or .with_token_refresher()".into(),
169            )),
170            TokenProvider::Static(t) => Ok(t.clone()),
171            TokenProvider::Refresher(f) => Ok(f()),
172        }
173    }
174}
175
176// ──────────────────────────────────────────────────────────────────────────────
177// Wire DTOs — `:retrieveContexts` response
178// ──────────────────────────────────────────────────────────────────────────────
179
180#[derive(Debug, Default, Deserialize)]
181#[serde(rename_all = "camelCase")]
182struct RetrieveContextsResponse {
183    #[serde(default)]
184    contexts: RagContexts,
185}
186
187#[derive(Debug, Default, Deserialize)]
188#[serde(rename_all = "camelCase")]
189struct RagContexts {
190    #[serde(default)]
191    contexts: Vec<RagContext>,
192}
193
194#[derive(Debug, Default, Deserialize)]
195#[serde(rename_all = "camelCase")]
196struct RagContext {
197    #[serde(default)]
198    text: String,
199    #[serde(default)]
200    source_display_name: String,
201}
202
203/// Build the request body for a `:retrieveContexts` call against this corpus.
204fn build_retrieve_body(
205    query: &str,
206    corpus: &str,
207    similarity_top_k: Option<u32>,
208    vector_distance_threshold: Option<f64>,
209) -> Value {
210    let mut vertex_rag_store = json!({
211        "ragResources": [ { "ragCorpus": corpus } ],
212    });
213    if let Some(threshold) = vector_distance_threshold {
214        vertex_rag_store["vectorDistanceThreshold"] = json!(threshold);
215    }
216    let mut query_obj = json!({ "text": query });
217    if let Some(top_k) = similarity_top_k {
218        query_obj["ragRetrievalConfig"] = json!({ "topK": top_k });
219    }
220    json!({
221        "vertexRagStore": vertex_rag_store,
222        "query": query_obj,
223    })
224}
225
226// ──────────────────────────────────────────────────────────────────────────────
227// Service struct
228// ──────────────────────────────────────────────────────────────────────────────
229
230/// Memory service backed by Vertex AI RAG.
231///
232/// Stores memory as RAG files in a corpus (via `ragFiles:upload`) and uses
233/// semantic search (`:retrieveContexts`) for retrieval.
234///
235/// # Quick start
236///
237/// ```rust,no_run
238/// # use gemini_adk_rs::memory::{VertexAiRagMemoryConfig, VertexAiRagMemoryService, MemoryService, MemoryEntry};
239/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
240/// let svc = VertexAiRagMemoryService::new(
241///     VertexAiRagMemoryConfig::from_corpus(
242///         "projects/my-proj/locations/us-central1/ragCorpora/my-corpus",
243///     ),
244/// )
245/// .with_token("ya29.my-access-token");
246///
247/// let hits = svc.search("session-1", "what did the user order?").await?;
248/// # Ok(())
249/// # }
250/// ```
251pub struct VertexAiRagMemoryService {
252    config: VertexAiRagMemoryConfig,
253    client: reqwest::Client,
254    token_provider: TokenProvider,
255}
256
257impl VertexAiRagMemoryService {
258    /// Create a new Vertex AI RAG memory service.
259    ///
260    /// No auth token is configured yet — call [`with_token`](Self::with_token)
261    /// or [`with_token_refresher`](Self::with_token_refresher) before issuing
262    /// requests.
263    pub fn new(config: VertexAiRagMemoryConfig) -> Self {
264        Self {
265            config,
266            client: reqwest::Client::new(),
267            token_provider: TokenProvider::None,
268        }
269    }
270
271    /// Set a static bearer token for all requests.
272    pub fn with_token(mut self, token: impl Into<String>) -> Self {
273        self.token_provider = TokenProvider::Static(token.into());
274        self
275    }
276
277    /// Set a dynamic token refresher closure (invoked before every request).
278    pub fn with_token_refresher(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self {
279        self.token_provider = TokenProvider::Refresher(Arc::new(f));
280        self
281    }
282
283    /// Returns the configured corpus resource name.
284    pub fn corpus(&self) -> &str {
285        &self.config.corpus
286    }
287
288    /// Upload a text document to the RAG corpus with the given display name.
289    ///
290    /// This mirrors ADK's `rag.upload_file(...)`: the display name carries the
291    /// session info (since the upload API doesn't accept arbitrary metadata).
292    async fn upload_text(&self, contents: &str, display_name: &str) -> Result<(), MemoryError> {
293        let token = self.token_provider.get()?;
294        let url = self.config.upload_rag_file_url();
295
296        // Multipart upload: a metadata part (rag file spec) + the file body.
297        let metadata = json!({
298            "rag_file": { "display_name": display_name },
299        });
300        let form = reqwest::multipart::Form::new()
301            .text("metadata", metadata.to_string())
302            .part(
303                "file",
304                reqwest::multipart::Part::text(contents.to_string())
305                    .file_name(format!("{display_name}.txt"))
306                    .mime_str("text/plain")
307                    .map_err(|e| MemoryError::Storage(format!("invalid mime: {e}")))?,
308            );
309
310        let resp = self
311            .client
312            .post(&url)
313            .header("Authorization", format!("Bearer {token}"))
314            .header("X-Goog-Upload-Protocol", "multipart")
315            .multipart(form)
316            .send()
317            .await
318            .map_err(|e| MemoryError::Storage(format!("HTTP upload failed: {e}")))?;
319
320        let status = resp.status().as_u16();
321        if !(200..300).contains(&status) {
322            let body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
323            return Err(MemoryError::Storage(format!(
324                "Vertex AI RAG upload failed [{status}]: {body}"
325            )));
326        }
327        Ok(())
328    }
329
330    /// Ingest a batch of memory entries for a session, mirroring ADK's
331    /// `add_session_to_memory`. Each entry is serialised as one JSON line and
332    /// the whole batch is uploaded as a single RAG file tagged with an encoded
333    /// display name.
334    ///
335    /// `app_name` / `user_id` / `session_id` are encoded into the file display
336    /// name so that [`search`](Self::search) can later filter by scope.
337    pub async fn add_session_to_memory(
338        &self,
339        app_name: &str,
340        user_id: &str,
341        session_id: &str,
342        entries: &[MemoryEntry],
343    ) -> Result<(), MemoryError> {
344        let mut lines = Vec::new();
345        for entry in entries {
346            let text = match &entry.value {
347                Value::String(s) => s.replace('\n', " "),
348                other => other.to_string(),
349            };
350            lines.push(
351                json!({
352                    "author": entry.key,
353                    "timestamp": entry.updated_at,
354                    "text": text,
355                })
356                .to_string(),
357            );
358        }
359        let contents = lines.join("\n");
360        let display_name = build_source_display_name(app_name, user_id, session_id);
361        self.upload_text(&contents, &display_name).await
362    }
363
364    /// Query the corpus and return the raw retrieved contexts (text +
365    /// source_display_name), mirroring ADK's `search_memory` retrieval step.
366    async fn retrieve_contexts(&self, query: &str) -> Result<Vec<RagContext>, MemoryError> {
367        let token = self.token_provider.get()?;
368        let url = self.config.retrieve_contexts_url();
369        let body = build_retrieve_body(
370            query,
371            &self.config.corpus,
372            self.config.similarity_top_k,
373            self.config.vector_distance_threshold,
374        );
375
376        let resp = self
377            .client
378            .post(&url)
379            .header("Authorization", format!("Bearer {token}"))
380            .header("Content-Type", "application/json")
381            .json(&body)
382            .send()
383            .await
384            .map_err(|e| MemoryError::Storage(format!("HTTP request failed: {e}")))?;
385
386        let status = resp.status().as_u16();
387        if !(200..300).contains(&status) {
388            let err_body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
389            return Err(MemoryError::Storage(format!(
390                "Vertex AI RAG retrieveContexts failed [{status}]: {err_body}"
391            )));
392        }
393
394        let parsed: RetrieveContextsResponse = resp.json().await.map_err(|e| {
395            MemoryError::Storage(format!("failed to parse retrieveContexts response: {e}"))
396        })?;
397        Ok(parsed.contexts.contexts)
398    }
399
400    /// Search memory scoped to a specific `app_name` / `user_id`, mirroring
401    /// ADK's `search_memory`. Filters retrieved contexts by the encoded display
402    /// name and parses the per-event JSON lines back into [`MemoryEntry`]s.
403    pub async fn search_memory(
404        &self,
405        app_name: &str,
406        user_id: &str,
407        query: &str,
408    ) -> Result<Vec<MemoryEntry>, MemoryError> {
409        let contexts = self.retrieve_contexts(query).await?;
410        Ok(parse_scoped_contexts(contexts, Some((app_name, user_id))))
411    }
412}
413
414/// Parse retrieved contexts into memory entries, optionally filtering by
415/// `(app_name, user_id)` encoded in each context's `source_display_name`.
416fn parse_scoped_contexts(
417    contexts: Vec<RagContext>,
418    scope: Option<(&str, &str)>,
419) -> Vec<MemoryEntry> {
420    let mut out = Vec::new();
421    for ctx in contexts {
422        if let Some((app_name, user_id)) = scope {
423            match parse_source_display_name(&ctx.source_display_name) {
424                Some((src_app, src_user, _session)) => {
425                    if src_app != app_name || src_user != user_id {
426                        continue;
427                    }
428                }
429                None => continue,
430            }
431        }
432        for line in ctx.text.split('\n') {
433            let line = line.trim();
434            if line.is_empty() {
435                continue;
436            }
437            if let Ok(event) = serde_json::from_str::<Value>(line) {
438                let author = event
439                    .get("author")
440                    .and_then(Value::as_str)
441                    .unwrap_or("")
442                    .to_string();
443                let timestamp = event
444                    .get("timestamp")
445                    .and_then(|t| t.as_u64().or_else(|| t.as_f64().map(|f| f as u64)))
446                    .unwrap_or(0);
447                let text = event
448                    .get("text")
449                    .and_then(Value::as_str)
450                    .unwrap_or("")
451                    .to_string();
452                out.push(MemoryEntry {
453                    key: author,
454                    value: Value::String(text),
455                    created_at: timestamp,
456                    updated_at: timestamp,
457                });
458            }
459        }
460    }
461    out
462}
463
464#[async_trait]
465impl MemoryService for VertexAiRagMemoryService {
466    /// Ingest a single entry into the corpus. The `session_id` scopes the
467    /// uploaded file's display name (app/user default to the entry key /
468    /// session id when not separately tracked).
469    async fn store(&self, session_id: &str, entry: MemoryEntry) -> Result<(), MemoryError> {
470        // Treat the session_id as both app_name proxy and session id; callers
471        // that need full app/user scoping should use `add_session_to_memory`.
472        self.add_session_to_memory(session_id, session_id, session_id, &[entry])
473            .await
474    }
475
476    /// Vertex AI RAG doesn't support direct key-based retrieval — use
477    /// [`search`](Self::search) instead.
478    async fn get(&self, _session_id: &str, _key: &str) -> Result<Option<MemoryEntry>, MemoryError> {
479        Ok(None)
480    }
481
482    /// Listing is not supported by the RAG retrieval API (semantic only).
483    async fn list(&self, _session_id: &str) -> Result<Vec<MemoryEntry>, MemoryError> {
484        Ok(vec![])
485    }
486
487    /// Semantic search against the corpus via `:retrieveContexts`. Results are
488    /// not scope-filtered here (use [`search_memory`](Self::search_memory) for
489    /// app/user scoping).
490    async fn search(
491        &self,
492        _session_id: &str,
493        query: &str,
494    ) -> Result<Vec<MemoryEntry>, MemoryError> {
495        let contexts = self.retrieve_contexts(query).await?;
496        Ok(parse_scoped_contexts(contexts, None))
497    }
498
499    /// Deleting an individual RAG file by key is not supported by the retrieval
500    /// API — returns [`MemoryError::Unsupported`] rather than silently
501    /// succeeding (which would mask data that was never removed).
502    async fn delete(&self, _session_id: &str, _key: &str) -> Result<(), MemoryError> {
503        Err(MemoryError::Unsupported(
504            "VertexAiRagMemoryService cannot delete individual entries (semantic              retrieval API); manage the corpus via the Vertex AI admin API"
505                .into(),
506        ))
507    }
508
509    /// Clearing a corpus is a destructive admin operation performed out of band —
510    /// returns [`MemoryError::Unsupported`] rather than silently succeeding.
511    async fn clear(&self, _session_id: &str) -> Result<(), MemoryError> {
512        Err(MemoryError::Unsupported(
513            "VertexAiRagMemoryService cannot clear a corpus (destructive admin              operation); use the Vertex AI admin API"
514                .into(),
515        ))
516    }
517}
518
519#[cfg(test)]
520mod tests {
521    use super::*;
522
523    fn test_config() -> VertexAiRagMemoryConfig {
524        VertexAiRagMemoryConfig::from_corpus(
525            "projects/test/locations/us-central1/ragCorpora/test-corpus",
526        )
527    }
528
529    #[test]
530    fn service_metadata() {
531        let svc = VertexAiRagMemoryService::new(test_config());
532        assert!(svc.corpus().contains("test-corpus"));
533    }
534
535    #[test]
536    fn config_parses_project_location() {
537        let cfg = test_config();
538        assert_eq!(cfg.project, "test");
539        assert_eq!(cfg.location, "us-central1");
540        assert!(cfg.retrieve_contexts_url().ends_with(":retrieveContexts"));
541        assert!(cfg.upload_rag_file_url().ends_with("/ragFiles:upload"));
542    }
543
544    #[test]
545    fn display_name_roundtrip() {
546        let dn = build_source_display_name("my-app", "user-1", "sess-42");
547        assert!(dn.starts_with(SOURCE_DISPLAY_NAME_PREFIX));
548        let (a, u, s) = parse_source_display_name(&dn).expect("should parse");
549        assert_eq!(a, "my-app");
550        assert_eq!(u, "user-1");
551        assert_eq!(s, "sess-42");
552    }
553
554    #[test]
555    fn display_name_handles_dots_in_ids() {
556        // Encoded form is dot-safe even when IDs contain dots.
557        let dn = build_source_display_name("a.b", "c.d", "e.f");
558        let (a, u, s) = parse_source_display_name(&dn).expect("should parse");
559        assert_eq!((a.as_str(), u.as_str(), s.as_str()), ("a.b", "c.d", "e.f"));
560    }
561
562    #[test]
563    fn parse_legacy_plain_display_name() {
564        let (a, u, s) = parse_source_display_name("app.user.session").expect("legacy ok");
565        assert_eq!(
566            (a.as_str(), u.as_str(), s.as_str()),
567            ("app", "user", "session")
568        );
569    }
570
571    #[test]
572    fn parse_rejects_malformed_display_name() {
573        assert!(parse_source_display_name("not-a-valid-name").is_none());
574        assert!(parse_source_display_name("adk-memory-v1.onlyonepart").is_none());
575    }
576
577    #[test]
578    fn builds_retrieve_body() {
579        let body = build_retrieve_body("q", "corpus-x", Some(3), Some(10.0));
580        assert_eq!(body["query"]["text"], "q");
581        assert_eq!(body["query"]["ragRetrievalConfig"]["topK"], 3);
582        assert_eq!(
583            body["vertexRagStore"]["ragResources"][0]["ragCorpus"],
584            "corpus-x"
585        );
586        assert_eq!(body["vertexRagStore"]["vectorDistanceThreshold"], 10.0);
587    }
588
589    #[test]
590    fn parse_scoped_contexts_filters_by_scope() {
591        let dn_match = build_source_display_name("app", "alice", "s1");
592        let dn_other = build_source_display_name("app", "bob", "s2");
593        let contexts = vec![
594            RagContext {
595                text: json!({"author": "user", "timestamp": 100, "text": "hello"}).to_string(),
596                source_display_name: dn_match,
597            },
598            RagContext {
599                text: json!({"author": "user", "timestamp": 200, "text": "nope"}).to_string(),
600                source_display_name: dn_other,
601            },
602        ];
603        let entries = parse_scoped_contexts(contexts, Some(("app", "alice")));
604        assert_eq!(entries.len(), 1);
605        assert_eq!(entries[0].key, "user");
606        assert_eq!(entries[0].value, json!("hello"));
607        assert_eq!(entries[0].created_at, 100);
608    }
609
610    #[test]
611    fn parse_scoped_contexts_no_scope_keeps_all() {
612        let contexts = vec![RagContext {
613            text: json!({"author": "model", "timestamp": 5, "text": "hi"}).to_string(),
614            source_display_name: String::new(),
615        }];
616        let entries = parse_scoped_contexts(contexts, None);
617        assert_eq!(entries.len(), 1);
618        assert_eq!(entries[0].key, "model");
619    }
620
621    #[tokio::test]
622    async fn search_without_token_errors() {
623        let svc = VertexAiRagMemoryService::new(test_config());
624        let result = svc.search("s1", "test").await;
625        assert!(result.is_err());
626    }
627}
gemini_adk_rs/memory/vertex_ai_rag.rs

gemini_adk_rs/memory/
vertex_ai_rag.rs
