gemini_adk_rs/tools/
vertex_ai_search.rs

1//! Vertex AI Search tool — enterprise search via Vertex AI.
2//!
3//! Mirrors ADK-Python's `vertex_ai_search_tool` / `discovery_engine_search_tool`.
4//! Provides server-side search using Vertex AI Discovery Engine / Search.
5
6use async_trait::async_trait;
7
8use crate::error::ToolError;
9use crate::tool::ToolFunction;
10
11/// Vertex AI Search tool configuration.
12#[derive(Debug, Clone)]
13pub struct VertexAiSearchConfig {
14    /// The Vertex AI search datastore resource name.
15    /// Format: `projects/{project}/locations/{location}/collections/default_collection/dataStores/{data_store_id}`
16    pub datastore: String,
17    /// Optional search filter expression.
18    pub filter: Option<String>,
19    /// Maximum number of results to return.
20    pub max_results: usize,
21}
22
23/// Tool that searches using Vertex AI Discovery Engine / Search.
24///
25/// This tool calls the Vertex AI Search API to perform enterprise
26/// search over configured data stores.
27#[derive(Debug, Clone)]
28pub struct VertexAiSearchTool {
29    config: VertexAiSearchConfig,
30}
31
32impl VertexAiSearchTool {
33    /// Create a new Vertex AI Search tool.
34    pub fn new(config: VertexAiSearchConfig) -> Self {
35        Self { config }
36    }
37
38    /// Returns the configured datastore resource name.
39    pub fn datastore(&self) -> &str {
40        &self.config.datastore
41    }
42}
43
44#[async_trait]
45impl ToolFunction for VertexAiSearchTool {
46    fn name(&self) -> &str {
47        "vertex_ai_search"
48    }
49
50    fn description(&self) -> &str {
51        "Search enterprise data using Vertex AI Search. Returns relevant documents \
52         and snippets from the configured data store."
53    }
54
55    fn parameters(&self) -> Option<serde_json::Value> {
56        Some(serde_json::json!({
57            "type": "object",
58            "properties": {
59                "query": {
60                    "type": "string",
61                    "description": "The search query."
62                }
63            },
64            "required": ["query"]
65        }))
66    }
67
68    async fn call(&self, args: serde_json::Value) -> Result<serde_json::Value, ToolError> {
69        let query = args
70            .get("query")
71            .and_then(|v| v.as_str())
72            .ok_or_else(|| ToolError::InvalidArgs("Missing query".into()))?;
73
74        // In a real integration, this would call the Vertex AI Search API.
75        // The actual API integration requires the Google Cloud SDK.
76        Ok(serde_json::json!({
77            "status": "search_requested",
78            "query": query,
79            "datastore": self.config.datastore,
80            "results": []
81        }))
82    }
83}
84
85/// Discovery Engine Search tool — alias for Vertex AI Search.
86///
87/// Mirrors ADK-Python's `DiscoveryEngineSearchTool`, which is equivalent
88/// to `VertexAiSearchTool` but specifically for Discovery Engine endpoints.
89pub type DiscoveryEngineSearchTool = VertexAiSearchTool;
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use serde_json::json;
95
96    fn test_config() -> VertexAiSearchConfig {
97        VertexAiSearchConfig {
98            datastore: "projects/my-proj/locations/global/collections/default_collection/dataStores/my-store".into(),
99            filter: None,
100            max_results: 10,
101        }
102    }
103
104    #[test]
105    fn tool_metadata() {
106        let tool = VertexAiSearchTool::new(test_config());
107        assert_eq!(tool.name(), "vertex_ai_search");
108        assert!(tool.parameters().is_some());
109        assert!(tool.datastore().contains("my-store"));
110    }
111
112    #[tokio::test]
113    async fn call_with_query() {
114        let tool = VertexAiSearchTool::new(test_config());
115        let result = tool
116            .call(json!({"query": "machine learning"}))
117            .await
118            .unwrap();
119        assert_eq!(result["query"], "machine learning");
120    }
121
122    #[tokio::test]
123    async fn missing_query() {
124        let tool = VertexAiSearchTool::new(test_config());
125        let result = tool.call(json!({})).await;
126        assert!(result.is_err());
127    }
128}
gemini_adk_rs/tools/vertex_ai_search.rs

gemini_adk_rs/tools/
vertex_ai_search.rs
