Skip to main content

oxi_agent/tools/
web_search.rs

1use super::search_cache::{SearchCache, SearchResult};
2/// Web search tool — searches via a3s-search library (DuckDuckGo, Wikipedia, Bing, Brave).
3///
4/// Uses a3s-search as a Rust library (not CLI), so no external binary is needed.
5/// Results are structured types — no text parsing required.
6///
7/// Features:
8/// - Multiple search engines (ddg, wiki, bing, brave)
9/// - Result caching with search IDs for later retrieval via `get_search_results`
10/// - Configurable engine selection and result count
11/// - Zero-config: no API keys, no external binary needed
12use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
13use async_trait::async_trait;
14use serde_json::{json, Value};
15use std::sync::Arc;
16use tokio::sync::oneshot;
17
18/// Maximum number of results to return by default.
19const DEFAULT_MAX_RESULTS: usize = 10;
20
21/// Maximum number of results allowed.
22const MAX_RESULTS: usize = 30;
23
24/// Default search engines.
25const DEFAULT_ENGINES: &str = "ddg,wiki";
26
27// ── Engine shortcut → a3s engine mapping ──────────────────────────
28
29/// Build a3s-search engine instances from shortcuts and add to Search.
30fn add_engines(search: &mut a3s_search::Search, shortcuts: &str) {
31    for shortcut in shortcuts.split(',') {
32        match shortcut.trim() {
33            "ddg" => search.add_engine(a3s_search::engines::DuckDuckGo::new()),
34            "wiki" => search.add_engine(a3s_search::engines::Wikipedia::new()),
35            "bing" => search.add_engine(a3s_search::engines::Bing::new()),
36            "brave" => search.add_engine(a3s_search::engines::Brave::new()),
37            s if !s.is_empty() => tracing::warn!("Unknown search engine: {}", s),
38            _ => {}
39        }
40    }
41}
42
43// ── WebSearchTool ─────────────────────────────────────────────────
44
45/// Multi-engine web search tool using a3s-search library.
46pub struct WebSearchTool {
47    cache: Arc<SearchCache>,
48}
49
50impl WebSearchTool {
51    /// Create a new WebSearchTool with the given search cache.
52    pub fn new(cache: Arc<SearchCache>) -> Self {
53        Self { cache }
54    }
55
56    /// Execute search using a3s-search library.
57    async fn do_search(
58        &self,
59        query: &str,
60        engines: &str,
61        limit: usize,
62    ) -> Result<Vec<SearchResult>, ToolError> {
63        let mut search = a3s_search::Search::new();
64        add_engines(&mut search, engines);
65
66        if search.engine_count() == 0 {
67            return Err(
68                "No valid engines specified. Available: ddg, wiki, bing, brave".to_string(),
69            );
70        }
71
72        search.set_timeout(std::time::Duration::from_secs(15));
73
74        let a3s_query = a3s_search::SearchQuery::new(query);
75        let results = search
76            .search(a3s_query)
77            .await
78            .map_err(|e| format!("Search failed: {}", e))?;
79
80        let formatted: Vec<SearchResult> = results
81            .items()
82            .iter()
83            .take(limit)
84            .map(|r| SearchResult {
85                title: r.title.clone(),
86                url: r.url.clone(),
87                snippet: r.content.clone(),
88                engines: r.engines.iter().cloned().collect(),
89                score: r.score,
90            })
91            .collect();
92
93        // Log engine errors if any
94        for (engine, error) in results.errors() {
95            tracing::warn!("Search engine {} error: {}", engine, error);
96        }
97
98        Ok(formatted)
99    }
100}
101
102// ── Formatting ────────────────────────────────────────────────────
103
104/// Format search results for display.
105fn format_results(results: &[SearchResult]) -> String {
106    if results.is_empty() {
107        return "No results found.".to_string();
108    }
109    results
110        .iter()
111        .enumerate()
112        .map(|(i, r)| {
113            let snippet = if r.snippet.chars().count() > 200 {
114                let truncated: String = r.snippet.chars().take(200).collect();
115                format!("{}...", truncated)
116            } else {
117                r.snippet.clone()
118            };
119            format!("{}. **{}**\n   {}\n   {}", i + 1, r.title, r.url, snippet)
120        })
121        .collect::<Vec<_>>()
122        .join("\n\n")
123}
124
125// ── AgentTool impl ────────────────────────────────────────────────
126
127#[async_trait]
128impl AgentTool for WebSearchTool {
129    fn name(&self) -> &str {
130        "web_search"
131    }
132
133    fn label(&self) -> &str {
134        "Web Search"
135    }
136
137    fn description(&self) -> &str {
138        "Search the web using multiple engines (DuckDuckGo, Wikipedia, Bing, Brave). No server or API key needed. Returns results with titles, URLs, and snippets."
139    }
140
141    fn parameters_schema(&self) -> Value {
142        json!({
143            "type": "object",
144            "properties": {
145                "query": {
146                    "type": "string",
147                    "description": "Search query string"
148                },
149                "engines": {
150                    "type": "string",
151                    "description": "Comma-separated engines (ddg,wiki,bing,brave). Default: ddg,wiki",
152                    "default": "ddg,wiki"
153                },
154                "limit": {
155                    "type": "integer",
156                    "description": "Maximum number of results to return (default: 10, max: 30)",
157                    "default": 10
158                }
159            },
160            "required": ["query"]
161        })
162    }
163
164    async fn execute(
165        &self,
166        _tool_call_id: &str,
167        params: Value,
168        _signal: Option<oneshot::Receiver<()>>,
169        _ctx: &ToolContext,
170    ) -> Result<AgentToolResult, ToolError> {
171        let query = params["query"]
172            .as_str()
173            .ok_or_else(|| "Missing required parameter: query".to_string())?;
174
175        let engines = params["engines"].as_str().unwrap_or(DEFAULT_ENGINES);
176
177        let limit = params["limit"]
178            .as_u64()
179            .unwrap_or(DEFAULT_MAX_RESULTS as u64)
180            .min(MAX_RESULTS as u64) as usize;
181
182        let results = self.do_search(query, engines, limit).await?;
183
184        if results.is_empty() {
185            return Ok(AgentToolResult::success(format!(
186                "No results found for: {}",
187                query
188            )));
189        }
190
191        // Cache results and generate a search ID
192        let search_id = self.cache.insert(query, results.clone());
193
194        let output = format_results(&results);
195
196        let results_json: Vec<Value> = results
197            .iter()
198            .map(|r| {
199                json!({
200                    "title": r.title,
201                    "url": r.url,
202                    "snippet": r.snippet,
203                    "engines": r.engines,
204                    "score": r.score
205                })
206            })
207            .collect();
208
209        Ok(AgentToolResult::success(output).with_metadata(json!({
210            "results": results_json,
211            "query": query,
212            "searchId": search_id,
213            "resultCount": results.len()
214        })))
215    }
216}
217
218// ── Tests ─────────────────────────────────────────────────────────
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn test_add_engines_ddg() {
226        let mut search = a3s_search::Search::new();
227        add_engines(&mut search, "ddg");
228        assert_eq!(search.engine_count(), 1);
229    }
230
231    #[test]
232    fn test_add_engines_multiple() {
233        let mut search = a3s_search::Search::new();
234        add_engines(&mut search, "ddg,wiki,brave");
235        assert_eq!(search.engine_count(), 3);
236    }
237
238    #[test]
239    fn test_add_engines_unknown() {
240        let mut search = a3s_search::Search::new();
241        add_engines(&mut search, "ddg,unknown,wiki");
242        assert_eq!(search.engine_count(), 2);
243    }
244
245    #[test]
246    fn test_add_engines_empty() {
247        let mut search = a3s_search::Search::new();
248        add_engines(&mut search, "");
249        assert_eq!(search.engine_count(), 0);
250    }
251
252    #[test]
253    fn test_format_results_empty() {
254        assert_eq!(format_results(&[]), "No results found.");
255    }
256
257    #[test]
258    fn test_format_results() {
259        let results = vec![SearchResult {
260            title: "Test".to_string(),
261            url: "https://example.com".to_string(),
262            snippet: "A snippet".to_string(),
263            engines: vec!["DuckDuckGo".to_string()],
264            score: 1.0,
265        }];
266        let formatted = format_results(&results);
267        assert!(formatted.contains("**Test**"));
268        assert!(formatted.contains("https://example.com"));
269    }
270
271    #[test]
272    fn test_schema() {
273        let cache = Arc::new(SearchCache::new());
274        let tool = WebSearchTool::new(cache);
275        let schema = tool.parameters_schema();
276        assert_eq!(schema["type"], "object");
277        assert!(schema["properties"]["query"].is_object());
278        assert!(schema["properties"]["engines"].is_object());
279        assert!(schema["properties"]["limit"].is_object());
280        assert!(schema["required"]
281            .as_array()
282            .unwrap()
283            .contains(&json!("query")));
284    }
285}