1use super::search_cache::{SearchCache, SearchResult};
2use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
13use async_trait::async_trait;
14use serde_json::{json, Value};
15use std::sync::Arc;
16use tokio::sync::oneshot;
17
18const DEFAULT_MAX_RESULTS: usize = 10;
20
21const MAX_RESULTS: usize = 30;
23
24const DEFAULT_ENGINES: &str = "ddg,wiki";
26
27fn 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
43pub struct WebSearchTool {
47 cache: Arc<SearchCache>,
48}
49
50impl WebSearchTool {
51 pub fn new(cache: Arc<SearchCache>) -> Self {
53 Self { cache }
54 }
55
56 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 for (engine, error) in results.errors() {
95 tracing::warn!("Search engine {} error: {}", engine, error);
96 }
97
98 Ok(formatted)
99 }
100}
101
102fn 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#[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 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#[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}