gemini_adk_rs/tools/retrieval/
vertex_ai_rag.rs

1//! Vertex AI RAG retrieval tool — retrieve context via Vertex AI RAG API.
2//!
3//! Mirrors ADK-Python's `vertex_ai_rag_retrieval` tool. Calls the Vertex AI
4//! RAG `:retrieveContexts` endpoint for the configured rag corpora / resources
5//! and returns the retrieved contexts as the tool result.
6
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use serde::Deserialize;
11use serde_json::{Value, json};
12
13use super::base::{BaseRetrievalTool, RetrievalResult};
14use crate::error::ToolError;
15
16/// Configuration for Vertex AI RAG retrieval.
17///
18/// Mirrors ADK-Python's `VertexRagStore` retrieval parameters. A retrieval may
19/// target one or more rag corpora (resource names) and constrain the result set
20/// via `similarity_top_k` and `vector_distance_threshold`.
21#[derive(Debug, Clone)]
22pub struct VertexAiRagConfig {
23    /// Google Cloud project ID.
24    pub project: String,
25    /// Google Cloud location (e.g. `us-central1`).
26    pub location: String,
27    /// The RAG corpus resource name(s).
28    /// Format: `projects/{project}/locations/{location}/ragCorpora/{corpus_id}`
29    pub rag_corpora: Vec<String>,
30    /// Number of contexts to retrieve (maps to `similarity_top_k`).
31    pub similarity_top_k: Option<u32>,
32    /// Only return contexts with a vector distance smaller than this threshold.
33    pub vector_distance_threshold: Option<f64>,
34}
35
36impl VertexAiRagConfig {
37    /// Create a new config from a single corpus resource name.
38    ///
39    /// The project and location are parsed from the corpus resource name when
40    /// it is fully-qualified (`projects/{project}/locations/{location}/...`).
41    pub fn from_corpus(corpus: impl Into<String>) -> Self {
42        let corpus = corpus.into();
43        let (project, location) = parse_project_location(&corpus);
44        Self {
45            project,
46            location,
47            rag_corpora: vec![corpus],
48            similarity_top_k: None,
49            vector_distance_threshold: None,
50        }
51    }
52
53    /// Set the explicit project / location used to build the endpoint URL.
54    pub fn with_project_location(
55        mut self,
56        project: impl Into<String>,
57        location: impl Into<String>,
58    ) -> Self {
59        self.project = project.into();
60        self.location = location.into();
61        self
62    }
63
64    /// Set the number of contexts to retrieve.
65    pub fn with_similarity_top_k(mut self, top_k: u32) -> Self {
66        self.similarity_top_k = Some(top_k);
67        self
68    }
69
70    /// Set the vector distance threshold.
71    pub fn with_vector_distance_threshold(mut self, threshold: f64) -> Self {
72        self.vector_distance_threshold = Some(threshold);
73        self
74    }
75
76    /// Build the `:retrieveContexts` endpoint URL.
77    ///
78    /// Format:
79    /// `https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}:retrieveContexts`
80    fn retrieve_contexts_url(&self) -> String {
81        format!(
82            "https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project}/locations/{location}:retrieveContexts",
83            project = self.project,
84            location = self.location,
85        )
86    }
87}
88
89/// Best-effort extraction of `{project}` / `{location}` from a rag corpus
90/// resource name of the form
91/// `projects/{project}/locations/{location}/ragCorpora/{id}`.
92fn parse_project_location(corpus: &str) -> (String, String) {
93    let mut project = String::new();
94    let mut location = String::new();
95    let mut parts = corpus.split('/');
96    while let Some(seg) = parts.next() {
97        match seg {
98            "projects" => project = parts.next().unwrap_or_default().to_string(),
99            "locations" => location = parts.next().unwrap_or_default().to_string(),
100            _ => {}
101        }
102    }
103    (project, location)
104}
105
106// ──────────────────────────────────────────────────────────────────────────────
107// Token provider (mirrors session/vertex_ai.rs)
108// ──────────────────────────────────────────────────────────────────────────────
109
110/// How to supply a bearer token for Vertex AI requests.
111enum TokenProvider {
112    /// No token configured — requests will fail with a clear message.
113    None,
114    /// A static, pre-fetched bearer token string.
115    Static(String),
116    /// A dynamic refresher: called before every request.
117    Refresher(Arc<dyn Fn() -> String + Send + Sync>),
118}
119
120impl TokenProvider {
121    fn get(&self) -> Result<String, ToolError> {
122        match self {
123            TokenProvider::None => Err(ToolError::ExecutionFailed(
124                "missing auth token: call .with_token() or .with_token_refresher()".into(),
125            )),
126            TokenProvider::Static(t) => Ok(t.clone()),
127            TokenProvider::Refresher(f) => Ok(f()),
128        }
129    }
130}
131
132// ──────────────────────────────────────────────────────────────────────────────
133// Wire DTOs — shapes returned by `:retrieveContexts`
134// ──────────────────────────────────────────────────────────────────────────────
135
136/// Top-level response envelope for `:retrieveContexts`.
137#[derive(Debug, Default, Deserialize)]
138#[serde(rename_all = "camelCase")]
139pub(crate) struct RetrieveContextsResponse {
140    #[serde(default)]
141    pub contexts: RagContexts,
142}
143
144/// The `contexts` object inside the response (itself wrapping a list).
145#[derive(Debug, Default, Deserialize)]
146#[serde(rename_all = "camelCase")]
147pub(crate) struct RagContexts {
148    #[serde(default)]
149    pub contexts: Vec<RagContext>,
150}
151
152/// A single retrieved context chunk.
153#[derive(Debug, Default, Deserialize)]
154#[serde(rename_all = "camelCase")]
155pub(crate) struct RagContext {
156    #[serde(default)]
157    pub text: String,
158    /// Source resource URI (e.g. the GCS path of the ingested file).
159    #[serde(default)]
160    pub source_uri: String,
161    /// Display name supplied at ingest time (carries session info for memory).
162    #[serde(default)]
163    pub source_display_name: String,
164    /// Vector distance / relevance score (lower distance = more relevant).
165    #[serde(default)]
166    pub distance: Option<f64>,
167    /// Some API revisions return a `score` field instead of `distance`.
168    #[serde(default)]
169    pub score: Option<f64>,
170}
171
172/// Build the JSON request body for a `:retrieveContexts` call.
173pub(crate) fn build_retrieve_body(
174    query: &str,
175    rag_corpora: &[String],
176    similarity_top_k: Option<u32>,
177    vector_distance_threshold: Option<f64>,
178) -> Value {
179    let resources: Vec<Value> = rag_corpora
180        .iter()
181        .map(|c| json!({ "ragCorpus": c }))
182        .collect();
183
184    let mut vertex_rag_store = json!({ "ragResources": resources });
185    if let Some(threshold) = vector_distance_threshold {
186        vertex_rag_store["vectorDistanceThreshold"] = json!(threshold);
187    }
188
189    let mut retrieval_config = json!({});
190    if let Some(top_k) = similarity_top_k {
191        retrieval_config["topK"] = json!(top_k);
192    }
193    let mut query_obj = json!({ "text": query });
194    if retrieval_config
195        .as_object()
196        .map(|o| !o.is_empty())
197        .unwrap_or(false)
198    {
199        query_obj["ragRetrievalConfig"] = retrieval_config;
200    }
201
202    json!({
203        "vertexRagStore": vertex_rag_store,
204        "query": query_obj,
205    })
206}
207
208/// Map a wire response into the crate's `RetrievalResult` list.
209pub(crate) fn map_contexts(resp: RetrieveContextsResponse) -> Vec<RetrievalResult> {
210    resp.contexts
211        .contexts
212        .into_iter()
213        .map(|c| {
214            // Prefer an explicit score; otherwise derive one from distance
215            // (distance is "smaller is better", so invert into a 0..1 score).
216            let score = c.score.unwrap_or_else(|| match c.distance {
217                Some(d) => 1.0 / (1.0 + d.max(0.0)),
218                None => 0.0,
219            });
220            let source = if !c.source_uri.is_empty() {
221                c.source_uri
222            } else {
223                c.source_display_name.clone()
224            };
225            RetrievalResult {
226                content: c.text,
227                source,
228                score,
229                metadata: json!({ "sourceDisplayName": c.source_display_name }),
230            }
231        })
232        .collect()
233}
234
235/// Retrieval tool that searches via the Vertex AI RAG API.
236///
237/// Calls the Vertex AI RAG `:retrieveContexts` endpoint to retrieve relevant
238/// document chunks from the configured corpora.
239///
240/// # Quick start
241///
242/// ```rust,no_run
243/// # use gemini_adk_rs::tools::retrieval::{VertexAiRagConfig, VertexAiRagRetrievalTool, BaseRetrievalTool};
244/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
245/// let tool = VertexAiRagRetrievalTool::new(
246///     VertexAiRagConfig::from_corpus(
247///         "projects/my-proj/locations/us-central1/ragCorpora/my-corpus",
248///     )
249///     .with_similarity_top_k(5),
250/// )
251/// .with_token("ya29.my-access-token");
252///
253/// let results = tool.retrieve("what is the refund policy?", 5).await?;
254/// # Ok(())
255/// # }
256/// ```
257pub struct VertexAiRagRetrievalTool {
258    config: VertexAiRagConfig,
259    client: reqwest::Client,
260    token_provider: TokenProvider,
261}
262
263impl VertexAiRagRetrievalTool {
264    /// Create a new Vertex AI RAG retrieval tool.
265    ///
266    /// No auth token is configured yet — call [`with_token`](Self::with_token)
267    /// or [`with_token_refresher`](Self::with_token_refresher) before issuing
268    /// requests.
269    pub fn new(config: VertexAiRagConfig) -> Self {
270        Self {
271            config,
272            client: reqwest::Client::new(),
273            token_provider: TokenProvider::None,
274        }
275    }
276
277    /// Set a static bearer token for all requests.
278    pub fn with_token(mut self, token: impl Into<String>) -> Self {
279        self.token_provider = TokenProvider::Static(token.into());
280        self
281    }
282
283    /// Set a dynamic token refresher closure (invoked before every request).
284    pub fn with_token_refresher(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self {
285        self.token_provider = TokenProvider::Refresher(Arc::new(f));
286        self
287    }
288
289    /// Returns the configured corpora resource names.
290    pub fn corpora(&self) -> &[String] {
291        &self.config.rag_corpora
292    }
293
294    /// Returns the first configured corpus (convenience accessor).
295    pub fn corpus(&self) -> &str {
296        self.config
297            .rag_corpora
298            .first()
299            .map(String::as_str)
300            .unwrap_or("")
301    }
302}
303
304#[async_trait]
305impl BaseRetrievalTool for VertexAiRagRetrievalTool {
306    fn name(&self) -> &str {
307        "vertex_ai_rag_retrieval"
308    }
309
310    async fn retrieve(&self, query: &str, top_k: usize) -> Result<Vec<RetrievalResult>, ToolError> {
311        // Per-call top_k overrides the configured default when set.
312        let top_k = self
313            .config
314            .similarity_top_k
315            .or_else(|| (top_k > 0).then_some(top_k as u32));
316
317        let body = build_retrieve_body(
318            query,
319            &self.config.rag_corpora,
320            top_k,
321            self.config.vector_distance_threshold,
322        );
323
324        let token = self.token_provider.get()?;
325        let url = self.config.retrieve_contexts_url();
326
327        let resp = self
328            .client
329            .post(&url)
330            .header("Authorization", format!("Bearer {token}"))
331            .header("Content-Type", "application/json")
332            .json(&body)
333            .send()
334            .await
335            .map_err(|e| ToolError::ExecutionFailed(format!("HTTP request failed: {e}")))?;
336
337        let status = resp.status().as_u16();
338        if !(200..300).contains(&status) {
339            let err_body = resp.text().await.unwrap_or_else(|_| "<unreadable>".into());
340            return Err(ToolError::ExecutionFailed(format!(
341                "Vertex AI RAG retrieveContexts failed [{status}]: {err_body}"
342            )));
343        }
344
345        let parsed: RetrieveContextsResponse = resp.json().await.map_err(|e| {
346            ToolError::ExecutionFailed(format!("failed to parse retrieveContexts response: {e}"))
347        })?;
348
349        Ok(map_contexts(parsed))
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356
357    fn test_config() -> VertexAiRagConfig {
358        VertexAiRagConfig::from_corpus(
359            "projects/my-proj/locations/us-central1/ragCorpora/my-corpus",
360        )
361    }
362
363    #[test]
364    fn tool_metadata() {
365        let tool = VertexAiRagRetrievalTool::new(test_config());
366        assert_eq!(tool.name(), "vertex_ai_rag_retrieval");
367        assert!(tool.corpus().contains("my-corpus"));
368    }
369
370    #[test]
371    fn parses_project_location_from_corpus() {
372        let cfg = test_config();
373        assert_eq!(cfg.project, "my-proj");
374        assert_eq!(cfg.location, "us-central1");
375        assert!(
376            cfg.retrieve_contexts_url()
377                .contains("us-central1-aiplatform.googleapis.com")
378        );
379        assert!(cfg.retrieve_contexts_url().ends_with(":retrieveContexts"));
380    }
381
382    #[test]
383    fn builds_request_body() {
384        let body = build_retrieve_body(
385            "hello",
386            &["projects/p/locations/l/ragCorpora/c".into()],
387            Some(7),
388            Some(0.5),
389        );
390        assert_eq!(body["query"]["text"], "hello");
391        assert_eq!(body["query"]["ragRetrievalConfig"]["topK"], 7);
392        assert_eq!(
393            body["vertexRagStore"]["ragResources"][0]["ragCorpus"],
394            "projects/p/locations/l/ragCorpora/c"
395        );
396        assert_eq!(body["vertexRagStore"]["vectorDistanceThreshold"], 0.5);
397    }
398
399    #[test]
400    fn builds_request_body_minimal() {
401        let body = build_retrieve_body("q", &["c".into()], None, None);
402        assert_eq!(body["query"]["text"], "q");
403        // No top_k → no ragRetrievalConfig.
404        assert!(body["query"].get("ragRetrievalConfig").is_none());
405        assert!(
406            body["vertexRagStore"]
407                .get("vectorDistanceThreshold")
408                .is_none()
409        );
410    }
411
412    #[test]
413    fn maps_contexts_with_distance() {
414        let resp = RetrieveContextsResponse {
415            contexts: RagContexts {
416                contexts: vec![RagContext {
417                    text: "chunk text".into(),
418                    source_uri: "gs://bucket/file.txt".into(),
419                    source_display_name: "adk-memory-v1.abc".into(),
420                    distance: Some(0.0),
421                    score: None,
422                }],
423            },
424        };
425        let results = map_contexts(resp);
426        assert_eq!(results.len(), 1);
427        assert_eq!(results[0].content, "chunk text");
428        assert_eq!(results[0].source, "gs://bucket/file.txt");
429        assert!((results[0].score - 1.0).abs() < f64::EPSILON);
430    }
431
432    #[test]
433    fn missing_token_errors() {
434        let tool = VertexAiRagRetrievalTool::new(test_config());
435        let err = tool.token_provider.get().unwrap_err();
436        assert!(matches!(err, ToolError::ExecutionFailed(_)));
437        assert!(err.to_string().contains("missing auth token"));
438    }
439
440    #[tokio::test]
441    async fn retrieve_without_token_errors() {
442        let tool = VertexAiRagRetrievalTool::new(test_config());
443        let result = tool.retrieve("test query", 5).await;
444        assert!(result.is_err());
445    }
446}
gemini_adk_rs/tools/retrieval/vertex_ai_rag.rs

gemini_adk_rs/tools/retrieval/
vertex_ai_rag.rs
