Skip to main content

cersei_tools/
web_search.rs

1//! WebSearch tool: search the web via a configurable search API.
2
3use super::*;
4use serde::Deserialize;
5
6/// Environment variable for search API key.
7const SEARCH_API_KEY_ENV: &str = "CERSEI_SEARCH_API_KEY";
8/// Environment variable for search API endpoint.
9const SEARCH_API_URL_ENV: &str = "CERSEI_SEARCH_API_URL";
10/// Default search endpoint (Brave Search API).
11const DEFAULT_SEARCH_URL: &str = "https://api.search.brave.com/res/v1/web/search";
12
13pub struct WebSearchTool;
14
15#[async_trait]
16impl Tool for WebSearchTool {
17    fn name(&self) -> &str {
18        "WebSearch"
19    }
20    fn description(&self) -> &str {
21        "Search the web and return relevant results. Requires CERSEI_SEARCH_API_KEY environment variable."
22    }
23    fn permission_level(&self) -> PermissionLevel {
24        PermissionLevel::ReadOnly
25    }
26    fn category(&self) -> ToolCategory {
27        ToolCategory::Web
28    }
29
30    fn input_schema(&self) -> Value {
31        serde_json::json!({
32            "type": "object",
33            "properties": {
34                "query": { "type": "string", "description": "Search query" },
35                "num_results": { "type": "integer", "description": "Number of results (default 8, max 20)" }
36            },
37            "required": ["query"]
38        })
39    }
40
41    async fn execute(&self, input: Value, _ctx: &ToolContext) -> ToolResult {
42        #[derive(Deserialize)]
43        struct Input {
44            query: String,
45            num_results: Option<usize>,
46        }
47
48        let input: Input = match serde_json::from_value(input) {
49            Ok(i) => i,
50            Err(e) => return ToolResult::error(format!("Invalid input: {}", e)),
51        };
52
53        let api_key = match std::env::var(SEARCH_API_KEY_ENV) {
54            Ok(k) if !k.is_empty() => k,
55            _ => {
56                return ToolResult::error(format!(
57                    "Web search requires {}. Set it to your Brave Search API key.",
58                    SEARCH_API_KEY_ENV
59                ))
60            }
61        };
62
63        let search_url =
64            std::env::var(SEARCH_API_URL_ENV).unwrap_or_else(|_| DEFAULT_SEARCH_URL.to_string());
65        let num_results = input.num_results.unwrap_or(8).min(20);
66
67        let client = match reqwest::Client::builder()
68            .timeout(std::time::Duration::from_secs(15))
69            .build()
70        {
71            Ok(c) => c,
72            Err(e) => return ToolResult::error(format!("HTTP client error: {}", e)),
73        };
74
75        let response = match client
76            .get(&search_url)
77            .header("X-Subscription-Token", &api_key)
78            .header("Accept", "application/json")
79            .query(&[
80                ("q", input.query.as_str()),
81                ("count", &num_results.to_string()),
82            ])
83            .send()
84            .await
85        {
86            Ok(r) => r,
87            Err(e) => return ToolResult::error(format!("Search request failed: {}", e)),
88        };
89
90        if !response.status().is_success() {
91            let status = response.status();
92            let body = response.text().await.unwrap_or_default();
93            return ToolResult::error(format!("Search API error ({}): {}", status, body));
94        }
95
96        let json: Value = match response.json().await {
97            Ok(j) => j,
98            Err(e) => return ToolResult::error(format!("Failed to parse response: {}", e)),
99        };
100
101        // Format results
102        let mut output = String::new();
103        if let Some(results) = json["web"]["results"].as_array() {
104            for (i, result) in results.iter().enumerate().take(num_results) {
105                let title = result["title"].as_str().unwrap_or("(no title)");
106                let url = result["url"].as_str().unwrap_or("");
107                let desc = result["description"].as_str().unwrap_or("");
108                output.push_str(&format!(
109                    "{}. **{}**\n   {}\n   {}\n\n",
110                    i + 1,
111                    title,
112                    url,
113                    desc
114                ));
115            }
116        }
117
118        if output.is_empty() {
119            ToolResult::success(format!("No results found for: {}", input.query))
120        } else {
121            ToolResult::success(output)
122        }
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn test_schema() {
132        let tool = WebSearchTool;
133        assert!(tool.input_schema()["properties"]["query"].is_object());
134        assert_eq!(tool.category(), ToolCategory::Web);
135    }
136}