cascade_agent/tools/
search.rs

1//! Web search tools: Tavily and Brave.
2
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use serde_json::{json, Value};
6
7use super::{Tool, ToolResult};
8
9// ---------------------------------------------------------------------------
10// Shared types
11// ---------------------------------------------------------------------------
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct SearchResult {
15    pub title: String,
16    pub url: String,
17    pub snippet: String,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub score: Option<f64>,
20}
21
22// ---------------------------------------------------------------------------
23// TavilySearchTool
24// ---------------------------------------------------------------------------
25
26/// Calls the [Tavily Search API](https://tavily.com).
27///
28/// Parameters:
29///   - `query`       (string, required): the search query.
30///   - `max_results` (int, optional):    max number of results (default 5).
31#[derive(Debug)]
32pub struct TavilySearchTool {
33    api_key: String,
34    client: reqwest::Client,
35    default_max_results: usize,
36}
37
38impl TavilySearchTool {
39    /// Create from a raw API key.
40    pub fn new(api_key: String) -> Self {
41        Self {
42            api_key,
43            client: reqwest::Client::new(),
44            default_max_results: 5,
45        }
46    }
47
48    /// Create by reading the API key from an environment variable.
49    pub fn from_env(env_var: &str, default_max_results: usize) -> Option<Self> {
50        let api_key = std::env::var(env_var).ok()?;
51        Some(Self {
52            api_key,
53            client: reqwest::Client::new(),
54            default_max_results,
55        })
56    }
57
58    async fn do_search(&self, query: &str, max_results: usize) -> ToolResult {
59        let body = json!({
60            "api_key": self.api_key,
61            "query": query,
62            "max_results": max_results,
63            "include_answer": true,
64        });
65
66        let resp = match self
67            .client
68            .post("https://api.tavily.com/search")
69            .header("Content-Type", "application/json")
70            .json(&body)
71            .send()
72            .await
73        {
74            Ok(r) => r,
75            Err(e) => return ToolResult::err(format!("Tavily request failed: {}", e)),
76        };
77
78        let status = resp.status();
79        let text = match resp.text().await {
80            Ok(t) => t,
81            Err(e) => return ToolResult::err(format!("Failed to read Tavily response: {}", e)),
82        };
83
84        if !status.is_success() {
85            return ToolResult::err(format!("Tavily API returned HTTP {}: {}", status, text));
86        }
87
88        // Parse the response
89        let parsed: Value = match serde_json::from_str(&text) {
90            Ok(v) => v,
91            Err(e) => return ToolResult::err(format!("Failed to parse Tavily response: {}", e)),
92        };
93
94        // Extract results into a uniform format
95        let results = parse_tavily_results(&parsed);
96        let answer = parsed.get("answer").and_then(|v| v.as_str()).unwrap_or("");
97
98        ToolResult::ok(json!({
99            "query": query,
100            "answer": answer,
101            "results": results,
102        }))
103    }
104}
105
106fn parse_tavily_results(parsed: &Value) -> Vec<Value> {
107    let mut out = Vec::new();
108    if let Some(arr) = parsed.get("results").and_then(|v| v.as_array()) {
109        for item in arr {
110            out.push(json!({
111                "title": item.get("title").and_then(|v| v.as_str()).unwrap_or(""),
112                "url": item.get("url").and_then(|v| v.as_str()).unwrap_or(""),
113                "snippet": item.get("content").and_then(|v| v.as_str()).unwrap_or(""),
114                "score": item.get("score").and_then(|v| v.as_f64()),
115            }));
116        }
117    }
118    out
119}
120
121#[async_trait]
122impl Tool for TavilySearchTool {
123    fn name(&self) -> &str {
124        "tavily_search"
125    }
126
127    fn description(&self) -> &str {
128        "Search the web using the Tavily Search API. Returns relevant results with titles, URLs, and snippets."
129    }
130
131    fn parameters_schema(&self) -> Value {
132        json!({
133            "type": "object",
134            "properties": {
135                "query": {
136                    "type": "string",
137                    "description": "The search query."
138                },
139                "max_results": {
140                    "type": "integer",
141                    "description": "Maximum number of results to return (default: 5).",
142                    "default": self.default_max_results
143                }
144            },
145            "required": ["query"]
146        })
147    }
148
149    async fn execute(&self, args: Value) -> ToolResult {
150        let query = match args.get("query").and_then(|v| v.as_str()) {
151            Some(q) => q,
152            None => return ToolResult::err("Missing required parameter 'query'"),
153        };
154        let max_results = args
155            .get("max_results")
156            .and_then(|v| v.as_u64())
157            .unwrap_or(self.default_max_results as u64) as usize;
158
159        self.do_search(query, max_results).await
160    }
161}
162
163// ---------------------------------------------------------------------------
164// BraveSearchTool
165// ---------------------------------------------------------------------------
166
167/// Calls the [Brave Web Search API](https://brave.com/search/api/).
168///
169/// Parameters:
170///   - `query`       (string, required): the search query.
171///   - `max_results` (int, optional):    max number of results (default 5).
172#[derive(Debug)]
173pub struct BraveSearchTool {
174    api_key: String,
175    client: reqwest::Client,
176    default_max_results: usize,
177}
178
179impl BraveSearchTool {
180    /// Create from a raw API key.
181    pub fn new(api_key: String) -> Self {
182        Self {
183            api_key,
184            client: reqwest::Client::new(),
185            default_max_results: 5,
186        }
187    }
188
189    /// Create by reading the API key from an environment variable.
190    pub fn from_env(env_var: &str, default_max_results: usize) -> Option<Self> {
191        let api_key = std::env::var(env_var).ok()?;
192        Some(Self {
193            api_key,
194            client: reqwest::Client::new(),
195            default_max_results,
196        })
197    }
198
199    async fn do_search(&self, query: &str, max_results: usize) -> ToolResult {
200        let url = format!(
201            "https://api.search.brave.com/res/v1/web/search?q={}&count={}",
202            urlencoding::encode(query),
203            max_results
204        );
205
206        let resp = match self
207            .client
208            .get(&url)
209            .header("Accept", "application/json")
210            .header("Accept-Encoding", "gzip")
211            .header("X-Subscription-Token", &self.api_key)
212            .send()
213            .await
214        {
215            Ok(r) => r,
216            Err(e) => return ToolResult::err(format!("Brave request failed: {}", e)),
217        };
218
219        let status = resp.status();
220        let text = match resp.text().await {
221            Ok(t) => t,
222            Err(e) => return ToolResult::err(format!("Failed to read Brave response: {}", e)),
223        };
224
225        if !status.is_success() {
226            return ToolResult::err(format!("Brave API returned HTTP {}: {}", status, text));
227        }
228
229        let parsed: Value = match serde_json::from_str(&text) {
230            Ok(v) => v,
231            Err(e) => return ToolResult::err(format!("Failed to parse Brave response: {}", e)),
232        };
233
234        let results = parse_brave_results(&parsed);
235
236        ToolResult::ok(json!({
237            "query": query,
238            "results": results,
239        }))
240    }
241}
242
243fn parse_brave_results(parsed: &Value) -> Vec<Value> {
244    let mut out = Vec::new();
245    let web = parsed.get("web").and_then(|v| v.get("results"));
246    if let Some(arr) = web.and_then(|v| v.as_array()) {
247        for item in arr {
248            out.push(json!({
249                "title": item.get("title").and_then(|v| v.as_str()).unwrap_or(""),
250                "url": item.get("url").and_then(|v| v.as_str()).unwrap_or(""),
251                "snippet": item.get("description").and_then(|v| v.as_str()).unwrap_or(""),
252            }));
253        }
254    }
255    out
256}
257
258#[async_trait]
259impl Tool for BraveSearchTool {
260    fn name(&self) -> &str {
261        "brave_search"
262    }
263
264    fn description(&self) -> &str {
265        "Search the web using the Brave Web Search API. Returns relevant results with titles, URLs, and snippets."
266    }
267
268    fn parameters_schema(&self) -> Value {
269        json!({
270            "type": "object",
271            "properties": {
272                "query": {
273                    "type": "string",
274                    "description": "The search query."
275                },
276                "max_results": {
277                    "type": "integer",
278                    "description": "Maximum number of results to return (default: 5).",
279                    "default": self.default_max_results
280                }
281            },
282            "required": ["query"]
283        })
284    }
285
286    async fn execute(&self, args: Value) -> ToolResult {
287        let query = match args.get("query").and_then(|v| v.as_str()) {
288            Some(q) => q,
289            None => return ToolResult::err("Missing required parameter 'query'"),
290        };
291        let max_results = args
292            .get("max_results")
293            .and_then(|v| v.as_u64())
294            .unwrap_or(self.default_max_results as u64) as usize;
295
296        self.do_search(query, max_results).await
297    }
298}
299
300// ---------------------------------------------------------------------------
301// Minimal URL-encoding helper (avoid adding a dependency just for this)
302// ---------------------------------------------------------------------------
303mod urlencoding {
304    pub fn encode(s: &str) -> String {
305        let mut out = String::with_capacity(s.len() * 2);
306        for byte in s.as_bytes() {
307            match *byte {
308                b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
309                    out.push(*byte as char);
310                }
311                _ => {
312                    out.push_str(&format!("%{:02X}", byte));
313                }
314            }
315        }
316        out
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    #[test]
325    fn tavily_schema_is_valid() {
326        let tool = TavilySearchTool::new("test-key".into());
327        let schema = tool.parameters_schema();
328        assert_eq!(schema["type"], "object");
329        assert!(schema["required"]
330            .as_array()
331            .unwrap()
332            .contains(&json!("query")));
333    }
334
335    #[test]
336    fn brave_schema_is_valid() {
337        let tool = BraveSearchTool::new("test-key".into());
338        let schema = tool.parameters_schema();
339        assert_eq!(schema["type"], "object");
340        assert!(schema["required"]
341            .as_array()
342            .unwrap()
343            .contains(&json!("query")));
344    }
345
346    #[test]
347    fn urlencode_works() {
348        assert_eq!(urlencoding::encode("hello world"), "hello%20world");
349        assert_eq!(urlencoding::encode("a+b=c"), "a%2Bb%3Dc");
350    }
351}
cascade_agent/tools/search.rs

cascade_agent/tools/
search.rs
