1use std::sync::Arc;
4
5use brain_core::metrics::SubsystemMetrics;
6
7use crate::resilience::{resilient_send, CircuitBreaker};
8
9fn make_cb(
10 name: &str,
11 resilience: &brain_core::config::ResilienceConfig,
12 metrics: Option<Arc<SubsystemMetrics>>,
13) -> CircuitBreaker {
14 let cb = CircuitBreaker::new(
15 name,
16 resilience.circuit_breaker_threshold,
17 resilience.circuit_breaker_cooldown_secs,
18 );
19 if let Some(m) = metrics {
20 cb.with_metrics(m)
21 } else {
22 cb
23 }
24}
25
26fn parse_search_results(
28 candidates: Vec<serde_json::Value>,
29 top_k: usize,
30) -> Vec<cortex::actions::SearchHit> {
31 candidates
32 .into_iter()
33 .filter_map(|entry| {
34 let title = entry
35 .get("title")
36 .and_then(serde_json::Value::as_str)
37 .or_else(|| entry.get("name").and_then(serde_json::Value::as_str))
38 .unwrap_or("untitled")
39 .to_string();
40 let url = entry
41 .get("url")
42 .and_then(serde_json::Value::as_str)
43 .or_else(|| entry.get("link").and_then(serde_json::Value::as_str))
44 .unwrap_or_default()
45 .to_string();
46 if url.is_empty() {
47 return None;
48 }
49 let snippet = entry
50 .get("snippet")
51 .and_then(serde_json::Value::as_str)
52 .or_else(|| entry.get("description").and_then(serde_json::Value::as_str))
53 .or_else(|| entry.get("content").and_then(serde_json::Value::as_str))
54 .unwrap_or_default()
55 .to_string();
56 Some(cortex::actions::SearchHit {
57 title,
58 url,
59 snippet,
60 })
61 })
62 .take(top_k.max(1))
63 .collect()
64}
65
66pub fn build_search_client(timeout_ms: u64) -> anyhow::Result<reqwest::Client> {
67 reqwest::Client::builder()
68 .timeout(std::time::Duration::from_millis(timeout_ms.max(1)))
69 .build()
70 .map_err(|e| anyhow::anyhow!("search client init failed: {e}"))
71}
72
73pub struct SearxngSearchBackend {
75 endpoint: String,
76 client: reqwest::Client,
77 circuit_breaker: Arc<CircuitBreaker>,
78 max_retries: u32,
79 retry_base_ms: u64,
80}
81
82impl SearxngSearchBackend {
83 pub fn new(
84 endpoint: &str,
85 timeout_ms: u64,
86 resilience: &brain_core::config::ResilienceConfig,
87 ) -> anyhow::Result<Self> {
88 Self::new_with_metrics(endpoint, timeout_ms, resilience, None)
89 }
90
91 pub fn new_with_metrics(
92 endpoint: &str,
93 timeout_ms: u64,
94 resilience: &brain_core::config::ResilienceConfig,
95 metrics: Option<Arc<SubsystemMetrics>>,
96 ) -> anyhow::Result<Self> {
97 Ok(Self {
98 endpoint: endpoint.trim_end_matches('/').to_string(),
99 client: build_search_client(timeout_ms)?,
100 circuit_breaker: Arc::new(make_cb("searxng", resilience, metrics)),
101 max_retries: resilience.max_retries,
102 retry_base_ms: resilience.retry_base_ms,
103 })
104 }
105}
106
107#[async_trait::async_trait]
108impl cortex::actions::WebSearchBackend for SearxngSearchBackend {
109 async fn search(
110 &self,
111 query: &str,
112 top_k: usize,
113 ) -> Result<Vec<cortex::actions::SearchHit>, cortex::actions::ActionError> {
114 let url = format!("{}/search", self.endpoint);
115 let client = self.client.clone();
116 let url_clone = url.clone();
117 let query_owned = query.to_string();
118 let response = resilient_send(
119 || {
120 client
121 .get(&url_clone)
122 .query(&[("q", query_owned.as_str()), ("format", "json")])
123 },
124 &self.circuit_breaker,
125 self.max_retries,
126 self.retry_base_ms,
127 )
128 .await?;
129
130 if !response.status().is_success() {
131 return Err(cortex::actions::ActionError::ExecutionFailed(format!(
132 "SearXNG returned HTTP {}",
133 response.status()
134 )));
135 }
136
137 let body = response
138 .json::<serde_json::Value>()
139 .await
140 .map_err(|e| cortex::actions::ActionError::ExecutionFailed(e.to_string()))?;
141
142 let candidates = match body.get("results").and_then(|v| v.as_array()) {
143 Some(arr) => arr.clone(),
144 None => {
145 tracing::warn!(
146 backend = "searxng",
147 "Response missing 'results' array — returning empty"
148 );
149 Vec::new()
150 }
151 };
152
153 Ok(parse_search_results(candidates, top_k))
154 }
155}
156
157pub struct TavilySearchBackend {
159 endpoint: String,
160 api_key: String,
161 client: reqwest::Client,
162 circuit_breaker: Arc<CircuitBreaker>,
163 max_retries: u32,
164 retry_base_ms: u64,
165}
166
167impl TavilySearchBackend {
168 pub fn new(
169 endpoint: &str,
170 api_key: &str,
171 timeout_ms: u64,
172 resilience: &brain_core::config::ResilienceConfig,
173 ) -> anyhow::Result<Self> {
174 Self::new_with_metrics(endpoint, api_key, timeout_ms, resilience, None)
175 }
176
177 pub fn new_with_metrics(
178 endpoint: &str,
179 api_key: &str,
180 timeout_ms: u64,
181 resilience: &brain_core::config::ResilienceConfig,
182 metrics: Option<Arc<SubsystemMetrics>>,
183 ) -> anyhow::Result<Self> {
184 Ok(Self {
185 endpoint: endpoint.trim_end_matches('/').to_string(),
186 api_key: api_key.to_string(),
187 client: build_search_client(timeout_ms)?,
188 circuit_breaker: Arc::new(make_cb("tavily", resilience, metrics)),
189 max_retries: resilience.max_retries,
190 retry_base_ms: resilience.retry_base_ms,
191 })
192 }
193}
194
195#[async_trait::async_trait]
196impl cortex::actions::WebSearchBackend for TavilySearchBackend {
197 async fn search(
198 &self,
199 query: &str,
200 top_k: usize,
201 ) -> Result<Vec<cortex::actions::SearchHit>, cortex::actions::ActionError> {
202 let url = format!("{}/search", self.endpoint);
203 let client = self.client.clone();
204 let url_clone = url.clone();
205 let api_key = self.api_key.clone();
206 let query_owned = query.to_string();
207 let response = resilient_send(
208 || {
209 client
210 .post(&url_clone)
211 .bearer_auth(&api_key)
212 .json(&serde_json::json!({
213 "query": query_owned,
214 "max_results": top_k,
215 "search_depth": "basic",
216 }))
217 },
218 &self.circuit_breaker,
219 self.max_retries,
220 self.retry_base_ms,
221 )
222 .await?;
223
224 if !response.status().is_success() {
225 return Err(cortex::actions::ActionError::ExecutionFailed(format!(
226 "Tavily returned HTTP {}",
227 response.status()
228 )));
229 }
230
231 let body = response
232 .json::<serde_json::Value>()
233 .await
234 .map_err(|e| cortex::actions::ActionError::ExecutionFailed(e.to_string()))?;
235
236 let candidates = match body.get("results").and_then(|v| v.as_array()) {
237 Some(arr) => {
238 if !arr.is_empty() && arr[0].get("url").is_none() {
239 tracing::warn!(
240 backend = "tavily",
241 "Results missing 'url' field — response schema may have changed"
242 );
243 }
244 arr.clone()
245 }
246 None => {
247 tracing::warn!(
248 backend = "tavily",
249 "Response missing 'results' array — returning empty"
250 );
251 Vec::new()
252 }
253 };
254
255 Ok(parse_search_results(candidates, top_k))
256 }
257}
258
259pub struct CustomSearchBackend {
261 endpoint: String,
262 client: reqwest::Client,
263 circuit_breaker: Arc<CircuitBreaker>,
264 max_retries: u32,
265 retry_base_ms: u64,
266}
267
268impl CustomSearchBackend {
269 pub fn new(
270 endpoint: &str,
271 timeout_ms: u64,
272 resilience: &brain_core::config::ResilienceConfig,
273 ) -> anyhow::Result<Self> {
274 Self::new_with_metrics(endpoint, timeout_ms, resilience, None)
275 }
276
277 pub fn new_with_metrics(
278 endpoint: &str,
279 timeout_ms: u64,
280 resilience: &brain_core::config::ResilienceConfig,
281 metrics: Option<Arc<SubsystemMetrics>>,
282 ) -> anyhow::Result<Self> {
283 Ok(Self {
284 endpoint: endpoint.to_string(),
285 client: build_search_client(timeout_ms)?,
286 circuit_breaker: Arc::new(make_cb("custom-search", resilience, metrics)),
287 max_retries: resilience.max_retries,
288 retry_base_ms: resilience.retry_base_ms,
289 })
290 }
291}
292
293#[async_trait::async_trait]
294impl cortex::actions::WebSearchBackend for CustomSearchBackend {
295 async fn search(
296 &self,
297 query: &str,
298 top_k: usize,
299 ) -> Result<Vec<cortex::actions::SearchHit>, cortex::actions::ActionError> {
300 let client = self.client.clone();
301 let endpoint = self.endpoint.clone();
302 let query_owned = query.to_string();
303 let response = resilient_send(
304 || {
305 client.post(&endpoint).json(&serde_json::json!({
306 "query": query_owned,
307 "top_k": top_k,
308 }))
309 },
310 &self.circuit_breaker,
311 self.max_retries,
312 self.retry_base_ms,
313 )
314 .await?;
315
316 if !response.status().is_success() {
317 return Err(cortex::actions::ActionError::ExecutionFailed(format!(
318 "search endpoint returned HTTP {}",
319 response.status()
320 )));
321 }
322
323 let body = response
324 .json::<serde_json::Value>()
325 .await
326 .map_err(|e| cortex::actions::ActionError::ExecutionFailed(e.to_string()))?;
327
328 let candidates: Vec<serde_json::Value> = body
329 .get("hits")
330 .and_then(|v| v.as_array())
331 .cloned()
332 .or_else(|| body.get("results").and_then(|v| v.as_array()).cloned())
333 .or_else(|| body.as_array().cloned())
334 .unwrap_or_default();
335
336 Ok(parse_search_results(candidates, top_k))
337 }
338}
339
340#[cfg(test)]
341mod tests {
342 use super::*;
343 use cortex::actions::WebSearchBackend;
344
345 fn fast_resilience() -> brain_core::config::ResilienceConfig {
346 brain_core::config::ResilienceConfig {
347 max_retries: 0,
348 retry_base_ms: 10,
349 circuit_breaker_threshold: 5,
350 circuit_breaker_cooldown_secs: 60,
351 }
352 }
353
354 #[tokio::test]
355 async fn test_searxng_successful_search() {
356 let mut server = mockito::Server::new_async().await;
357 let _mock = server
358 .mock("GET", "/search")
359 .match_query(mockito::Matcher::UrlEncoded("q".into(), "rust".into()))
360 .with_status(200)
361 .with_header("content-type", "application/json")
362 .with_body(
363 r#"{
364 "results": [
365 {"title": "Rust docs", "url": "https://doc.rust-lang.org", "content": "language docs"},
366 {"title": "Rust book", "url": "https://rust-book.rs", "content": "book"}
367 ]
368 }"#,
369 )
370 .create_async()
371 .await;
372
373 let backend = SearxngSearchBackend::new(&server.url(), 5000, &fast_resilience()).unwrap();
374 let hits = backend.search("rust", 10).await.unwrap();
375 assert_eq!(hits.len(), 2);
376 assert_eq!(hits[0].title, "Rust docs");
377 assert_eq!(hits[0].url, "https://doc.rust-lang.org");
378 assert_eq!(hits[0].snippet, "language docs");
379 }
380
381 #[tokio::test]
382 async fn test_searxng_empty_results() {
383 let mut server = mockito::Server::new_async().await;
384 let _mock = server
385 .mock("GET", "/search")
386 .match_query(mockito::Matcher::Any)
387 .with_status(200)
388 .with_header("content-type", "application/json")
389 .with_body(r#"{"results": []}"#)
390 .create_async()
391 .await;
392
393 let backend = SearxngSearchBackend::new(&server.url(), 5000, &fast_resilience()).unwrap();
394 let hits = backend.search("nothing", 10).await.unwrap();
395 assert!(hits.is_empty());
396 }
397
398 #[tokio::test]
399 async fn test_searxng_5xx_surfaces_as_error() {
400 let mut server = mockito::Server::new_async().await;
401 let _mock = server
402 .mock("GET", "/search")
403 .match_query(mockito::Matcher::Any)
404 .with_status(500)
405 .with_body("internal error")
406 .expect_at_least(1)
407 .create_async()
408 .await;
409
410 let backend = SearxngSearchBackend::new(&server.url(), 5000, &fast_resilience()).unwrap();
411 let result = backend.search("boom", 10).await;
412 assert!(result.is_err(), "expected 5xx to surface as error");
413 }
414
415 #[tokio::test]
416 async fn test_searxng_top_k_limit() {
417 let mut server = mockito::Server::new_async().await;
418 let _mock = server
419 .mock("GET", "/search")
420 .match_query(mockito::Matcher::Any)
421 .with_status(200)
422 .with_header("content-type", "application/json")
423 .with_body(
424 r#"{
425 "results": [
426 {"title": "a", "url": "https://a.com", "content": ""},
427 {"title": "b", "url": "https://b.com", "content": ""},
428 {"title": "c", "url": "https://c.com", "content": ""},
429 {"title": "d", "url": "https://d.com", "content": ""}
430 ]
431 }"#,
432 )
433 .create_async()
434 .await;
435
436 let backend = SearxngSearchBackend::new(&server.url(), 5000, &fast_resilience()).unwrap();
437 let hits = backend.search("q", 2).await.unwrap();
438 assert_eq!(hits.len(), 2);
439 }
440
441 #[tokio::test]
442 async fn test_searxng_missing_results_field_returns_empty() {
443 let mut server = mockito::Server::new_async().await;
444 let _mock = server
445 .mock("GET", "/search")
446 .match_query(mockito::Matcher::Any)
447 .with_status(200)
448 .with_header("content-type", "application/json")
449 .with_body(r#"{"something": "else"}"#)
450 .create_async()
451 .await;
452
453 let backend = SearxngSearchBackend::new(&server.url(), 5000, &fast_resilience()).unwrap();
454 let hits = backend.search("q", 10).await.unwrap();
455 assert!(hits.is_empty());
456 }
457
458 #[tokio::test]
459 async fn test_tavily_successful_search_sends_bearer() {
460 let mut server = mockito::Server::new_async().await;
461 let _mock = server
462 .mock("POST", "/search")
463 .match_header("authorization", "Bearer tvly-test-key")
464 .with_status(200)
465 .with_header("content-type", "application/json")
466 .with_body(
467 r#"{
468 "results": [
469 {"title": "Tavily hit", "url": "https://example.com", "content": "snippet"}
470 ]
471 }"#,
472 )
473 .create_async()
474 .await;
475
476 let backend =
477 TavilySearchBackend::new(&server.url(), "tvly-test-key", 5000, &fast_resilience())
478 .unwrap();
479 let hits = backend.search("question", 5).await.unwrap();
480 assert_eq!(hits.len(), 1);
481 assert_eq!(hits[0].title, "Tavily hit");
482 }
483}