1use crate::utils::{RetryConfig, classify_host, retry_async_if, truncate_content};
2use anyhow::{Result, anyhow};
3use async_trait::async_trait;
4use reqwest::Client;
5use serde::{Deserialize, Serialize};
6use std::time::Duration;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct SearchResult {
11 pub title: String,
12 pub url: String,
13 pub snippet: String,
14 pub full_content: String,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct WebFetchResult {
20 pub title: String,
21 pub content: String,
22}
23
24#[async_trait]
27pub trait SearchProvider: Send + Sync {
28 async fn search(&self, query: &str, count: usize) -> Result<Vec<SearchResult>>;
29}
30
31#[async_trait]
35pub trait FetchProvider: Send + Sync {
36 async fn fetch(&self, url: &str) -> Result<WebFetchResult>;
37}
38
39#[derive(Debug, Deserialize)]
41struct OllamaSearchResponse {
42 results: Vec<OllamaSearchResult>,
43}
44
45#[derive(Debug, Deserialize)]
46struct OllamaSearchResult {
47 title: String,
48 url: String,
49 content: String,
50}
51
52#[derive(Debug, Deserialize)]
54struct OllamaFetchResponse {
55 title: Option<String>,
56 content: Option<String>,
57}
58
59#[derive(Debug, Deserialize)]
62struct SearxngResponse {
63 #[serde(default)]
64 results: Vec<SearxngResult>,
65}
66
67#[derive(Debug, Deserialize)]
68struct SearxngResult {
69 #[serde(default)]
70 title: String,
71 url: String,
72 #[serde(default)]
73 content: String,
74}
75
76const OLLAMA_API_BASE: &str = "https://ollama.com/api";
77
78const NATIVE_FETCH_UA: &str =
81 "Mozilla/5.0 (compatible; MermaidBot/1.0; +https://github.com/noahsabaj/mermaid-cli)";
82
83#[derive(Debug)]
87struct HttpStatusError {
88 status: u16,
89}
90
91impl std::fmt::Display for HttpStatusError {
92 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
93 write!(f, "HTTP {}", self.status)
94 }
95}
96
97impl std::error::Error for HttpStatusError {}
98
99fn web_error_is_retryable(e: &anyhow::Error) -> bool {
104 if let Some(re) = e.downcast_ref::<reqwest::Error>() {
105 return re.is_timeout() || re.is_connect();
106 }
107 if let Some(h) = e.downcast_ref::<HttpStatusError>() {
108 return h.status == 429 || (500..600).contains(&h.status);
109 }
110 false
111}
112
113#[derive(Clone)]
116pub struct OllamaWebClient {
117 client: Client,
118 api_key: String,
119}
120
121impl OllamaWebClient {
122 pub fn new(api_key: String) -> Self {
123 Self {
124 client: Client::new(),
125 api_key,
126 }
127 }
128
129 async fn search_impl(&self, query: &str, count: usize) -> Result<Vec<SearchResult>> {
135 if count == 0 || count > 10 {
136 return Err(anyhow!(
137 "Result count must be between 1 and 10, got {}",
138 count
139 ));
140 }
141
142 let retry_config = RetryConfig {
143 max_attempts: 3,
144 initial_delay_ms: 500,
145 max_delay_ms: 5000,
146 backoff_multiplier: 2.0,
147 };
148
149 let client = self.client.clone();
150 let api_key = self.api_key.clone();
151 let query_owned = query.to_string();
152 let ollama_response: OllamaSearchResponse = retry_async_if(
154 || {
155 let client = client.clone();
156 let api_key = api_key.clone();
157 let query = query_owned.clone();
158 async move {
159 let response = client
160 .post(format!("{}/web_search", OLLAMA_API_BASE))
161 .header("Authorization", format!("Bearer {}", api_key))
162 .json(&serde_json::json!({
163 "query": query,
164 "max_results": count,
165 }))
166 .timeout(Duration::from_secs(30))
167 .send()
168 .await
169 .map_err(|e| {
170 anyhow::Error::new(e).context("Failed to reach Ollama web search API")
171 })?;
172
173 if !response.status().is_success() {
174 let status = response.status();
175 let body = response.text().await.unwrap_or_default();
176 return Err(anyhow::Error::new(HttpStatusError {
177 status: status.as_u16(),
178 })
179 .context(format!(
180 "Ollama web search API returned error {}: {}",
181 status, body
182 )));
183 }
184
185 let body =
186 read_body_capped(response, crate::constants::MAX_WEB_BODY_BYTES).await?;
187 serde_json::from_slice::<OllamaSearchResponse>(&body)
188 .map_err(|e| anyhow!("Failed to parse Ollama search response: {}", e))
189 }
190 },
191 &retry_config,
192 web_error_is_retryable,
193 )
194 .await?;
195
196 let search_results = map_search_results(
197 ollama_response
198 .results
199 .into_iter()
200 .map(|r| (r.title, r.url, r.content)),
201 count,
202 );
203
204 Ok(search_results)
206 }
207
208 async fn fetch_impl(&self, url: &str) -> Result<WebFetchResult> {
210 let retry_config = RetryConfig {
211 max_attempts: 2,
212 initial_delay_ms: 200,
213 max_delay_ms: 2000,
214 backoff_multiplier: 2.0,
215 };
216
217 let client = self.client.clone();
218 let api_key = self.api_key.clone();
219 let url_owned = url.to_string();
220 let response: OllamaFetchResponse = retry_async_if(
221 || {
222 let client = client.clone();
223 let api_key = api_key.clone();
224 let url = url_owned.clone();
225 async move {
226 let response = client
227 .post(format!("{}/web_fetch", OLLAMA_API_BASE))
228 .header("Authorization", format!("Bearer {}", api_key))
229 .json(&serde_json::json!({ "url": url }))
230 .timeout(Duration::from_secs(15))
231 .send()
232 .await
233 .map_err(|e| {
234 anyhow::Error::new(e).context(format!("Failed to fetch {}", url))
235 })?;
236
237 if !response.status().is_success() {
238 let status = response.status();
239 return Err(anyhow::Error::new(HttpStatusError {
240 status: status.as_u16(),
241 })
242 .context(format!("Failed to fetch {}", url)));
243 }
244
245 let body =
246 read_body_capped(response, crate::constants::MAX_WEB_BODY_BYTES).await?;
247 serde_json::from_slice::<OllamaFetchResponse>(&body)
248 .map_err(|e| anyhow!("Failed to parse fetch response: {}", e))
249 }
250 },
251 &retry_config,
252 web_error_is_retryable,
253 )
254 .await?;
255
256 Ok(WebFetchResult {
257 title: response.title.unwrap_or_default(),
258 content: response.content.unwrap_or_default(),
259 })
260 }
261}
262
263#[async_trait]
264impl SearchProvider for OllamaWebClient {
265 async fn search(&self, query: &str, count: usize) -> Result<Vec<SearchResult>> {
266 self.search_impl(query, count).await
267 }
268}
269
270#[async_trait]
271impl FetchProvider for OllamaWebClient {
272 async fn fetch(&self, url: &str) -> Result<WebFetchResult> {
273 self.fetch_impl(url).await
274 }
275}
276
277#[derive(Clone)]
281pub struct SearxngClient {
282 client: Client,
283 base_url: String,
284}
285
286impl SearxngClient {
287 pub fn new(base_url: String) -> Self {
288 Self {
289 client: Client::new(),
290 base_url: base_url.trim_end_matches('/').to_string(),
291 }
292 }
293}
294
295#[async_trait]
296impl SearchProvider for SearxngClient {
297 async fn search(&self, query: &str, count: usize) -> Result<Vec<SearchResult>> {
298 let request_url = reqwest::Url::parse_with_params(
299 &format!("{}/search", self.base_url),
300 &[("q", query), ("format", "json")],
301 )
302 .map_err(|e| anyhow!("invalid SearXNG URL {}: {e}", self.base_url))?;
303
304 let response = self
305 .client
306 .get(request_url)
307 .timeout(Duration::from_secs(30))
308 .send()
309 .await
310 .map_err(|e| {
311 anyhow::Error::new(e).context(format!(
312 "Failed to reach SearXNG at {} — is it running?",
313 self.base_url
314 ))
315 })?;
316
317 if !response.status().is_success() {
318 let status = response.status();
319 return Err(anyhow!(
320 "SearXNG at {} returned {status}. A 403 usually means the JSON format is \
321 disabled — add `json` to `search.formats` in its settings.yml.",
322 self.base_url
323 ));
324 }
325
326 let body = read_body_capped(response, crate::constants::MAX_WEB_BODY_BYTES).await?;
327 let parsed: SearxngResponse = serde_json::from_slice(&body).map_err(|e| {
328 anyhow!("Failed to parse SearXNG response (is `format=json` enabled?): {e}")
329 })?;
330
331 let results = map_search_results(
332 parsed
333 .results
334 .into_iter()
335 .map(|r| (r.title, r.url, r.content)),
336 count,
337 );
338
339 Ok(results)
342 }
343}
344
345pub struct ManagedSearxngBackend;
350
351#[async_trait]
352impl SearchProvider for ManagedSearxngBackend {
353 async fn search(&self, query: &str, count: usize) -> Result<Vec<SearchResult>> {
354 let base_url = crate::searxng::manager().ensure_running().await?;
355 SearxngClient::new(base_url).search(query, count).await
356 }
357}
358
359pub struct NativeFetchClient {
365 client: Client,
366}
367
368impl Default for NativeFetchClient {
369 fn default() -> Self {
370 Self::new()
371 }
372}
373
374impl NativeFetchClient {
375 pub fn new() -> Self {
376 let client = Client::builder()
377 .user_agent(NATIVE_FETCH_UA)
378 .timeout(Duration::from_secs(20))
379 .redirect(reqwest::redirect::Policy::limited(5))
380 .dns_resolver(std::sync::Arc::new(VettingResolver))
386 .build()
387 .unwrap_or_else(|_| Client::new());
388 Self { client }
389 }
390}
391
392#[async_trait]
393impl FetchProvider for NativeFetchClient {
394 async fn fetch(&self, url: &str) -> Result<WebFetchResult> {
395 let response = self
400 .client
401 .get(url)
402 .send()
403 .await
404 .map_err(|e| anyhow::Error::new(e).context(format!("Failed to fetch {url}")))?;
405
406 if !response.status().is_success() {
407 let status = response.status();
408 return Err(anyhow::Error::new(HttpStatusError {
409 status: status.as_u16(),
410 })
411 .context(format!("Failed to fetch {url}")));
412 }
413
414 let content_type = response
415 .headers()
416 .get(reqwest::header::CONTENT_TYPE)
417 .and_then(|v| v.to_str().ok())
418 .unwrap_or("")
419 .to_ascii_lowercase();
420
421 let body = read_body_capped(response, crate::constants::MAX_WEB_BODY_BYTES).await?;
422
423 let looks_textual = content_type.is_empty()
427 || content_type.contains("html")
428 || content_type.contains("xml")
429 || content_type.contains("json")
430 || content_type.starts_with("text/");
431 if !looks_textual {
432 return Err(anyhow!(
433 "web_fetch: unsupported content-type '{content_type}' (only html/text pages)"
434 ));
435 }
436
437 let html = String::from_utf8_lossy(&body).into_owned();
438 let url_owned = url.to_string();
439 let (title, content) =
442 tokio::task::spawn_blocking(move || extract_readable(&html, &url_owned))
443 .await
444 .map_err(|e| anyhow!("content extraction failed: {e}"))?;
445
446 Ok(WebFetchResult { title, content })
447 }
448}
449
450struct VettingResolver;
457
458impl reqwest::dns::Resolve for VettingResolver {
459 fn resolve(&self, name: reqwest::dns::Name) -> reqwest::dns::Resolving {
460 Box::pin(async move {
461 let host = name.as_str().to_string();
462 let addrs: Vec<std::net::SocketAddr> =
465 tokio::net::lookup_host((host.as_str(), 0)).await?.collect();
466 for addr in &addrs {
467 if classify_host(&addr.ip().to_string()).is_internal() {
468 return Err(format!(
469 "refusing to connect to '{host}' — it resolves to an internal address"
470 )
471 .into());
472 }
473 }
474 Ok(Box::new(addrs.into_iter()) as reqwest::dns::Addrs)
475 })
476 }
477}
478
479fn extract_readable(html: &str, url: &str) -> (String, String) {
484 use dom_smoothie::Readability;
485
486 if let Ok(mut readability) = Readability::new(html, Some(url), None)
487 && let Ok(article) = readability.parse()
488 {
489 let content_html = article.content.to_string();
490 let markdown = htmd::convert(&content_html).unwrap_or_default();
491 let markdown = markdown.trim();
492 if !markdown.is_empty() {
496 let title = if article.title.trim().is_empty() {
497 fallback_title(html)
498 } else {
499 article.title
500 };
501 return (title, markdown.to_string());
502 }
503 }
504
505 let markdown = htmd::convert(html).unwrap_or_default();
506 (fallback_title(html), markdown.trim().to_string())
507}
508
509fn fallback_title(html: &str) -> String {
512 let lower = html.to_ascii_lowercase();
513 let Some(open) = lower.find("<title") else {
514 return String::new();
515 };
516 let after_tag = match html[open..].find('>') {
517 Some(gt) => &html[open + gt + 1..],
518 None => return String::new(),
519 };
520 match after_tag.to_ascii_lowercase().find("</title>") {
521 Some(end) => after_tag[..end].trim().to_string(),
522 None => String::new(),
523 }
524}
525
526fn map_search_results(
529 hits: impl Iterator<Item = (String, String, String)>,
530 count: usize,
531) -> Vec<SearchResult> {
532 hits.take(count)
533 .map(|(title, url, content)| {
534 let full_content = truncate_content(&content, crate::constants::WEB_CONTENT_MAX_CHARS);
535 let snippet = content.chars().take(200).collect();
536 SearchResult {
537 title,
538 url,
539 snippet,
540 full_content,
541 }
542 })
543 .collect()
544}
545
546pub fn format_results(results: &[SearchResult]) -> String {
551 let mut formatted = String::from("[SEARCH_RESULTS]\n");
552
553 for (i, result) in results.iter().enumerate() {
554 formatted.push_str(&format!(
555 "[{}] Title: {}\nURL: {}\nContent:\n{}\n---\n",
556 i + 1,
557 result.title,
558 result.url,
559 result.full_content
560 ));
561 }
562
563 formatted.push_str("[/SEARCH_RESULTS]\n\n");
564
565 formatted.push_str("Sources:\n");
567 for (i, result) in results.iter().enumerate() {
568 formatted.push_str(&format!("{}. {} - {}\n", i + 1, result.title, result.url));
569 }
570
571 formatted
572}
573
574async fn read_body_capped(response: reqwest::Response, max_bytes: usize) -> Result<Vec<u8>> {
581 use futures::StreamExt;
582 if let Some(len) = response.content_length()
583 && len as usize > max_bytes
584 {
585 return Err(anyhow!(
586 "response body too large: {len} bytes exceeds {max_bytes} cap"
587 ));
588 }
589 let mut stream = response.bytes_stream();
590 let mut buf = Vec::new();
591 while let Some(chunk) = stream.next().await {
592 let chunk = chunk.map_err(|e| anyhow!("error reading response body: {e}"))?;
593 if buf.len() + chunk.len() > max_bytes {
594 return Err(anyhow!("response body exceeded {max_bytes} byte cap"));
595 }
596 buf.extend_from_slice(&chunk);
597 }
598 Ok(buf)
599}
600
601#[cfg(test)]
602mod tests {
603 use super::*;
604
605 #[test]
606 fn test_ollama_web_client_creation() {
607 let client = OllamaWebClient::new("test-key".to_string());
608 assert_eq!(client.api_key, "test-key");
609 }
610
611 #[test]
612 fn test_format_results() {
613 let results = vec![SearchResult {
614 title: "Test Article".to_string(),
615 url: "https://example.com".to_string(),
616 snippet: "This is a test".to_string(),
617 full_content: "Full content here".to_string(),
618 }];
619
620 let formatted = format_results(&results);
621 assert!(formatted.contains("[SEARCH_RESULTS]"));
622 assert!(formatted.contains("Test Article"));
623 assert!(formatted.contains("https://example.com"));
624 assert!(formatted.contains("[/SEARCH_RESULTS]"));
625 }
626
627 #[test]
628 fn map_search_results_truncates_and_caps_count() {
629 let hits = (0..5).map(|i| {
630 (
631 format!("t{i}"),
632 format!("https://e{i}.com"),
633 "x".repeat(crate::constants::WEB_CONTENT_MAX_CHARS * 2),
634 )
635 });
636 let out = map_search_results(hits, 3);
637 assert_eq!(out.len(), 3, "count cap applied");
638 assert!(
639 out[0].full_content.len() <= crate::constants::WEB_CONTENT_MAX_CHARS + 64,
640 "content truncated"
641 );
642 assert!(out[0].snippet.chars().count() <= 200);
643 }
644
645 #[test]
646 fn searxng_response_parses_results() {
647 let json = serde_json::json!({
648 "results": [
649 {"title": "A", "url": "https://a.com", "content": "alpha"},
650 {"url": "https://b.com"},
651 ]
652 })
653 .to_string();
654 let parsed: SearxngResponse = serde_json::from_str(&json).unwrap();
655 assert_eq!(parsed.results.len(), 2);
656 assert_eq!(parsed.results[0].url, "https://a.com");
657 assert_eq!(parsed.results[1].title, "");
659 assert_eq!(parsed.results[1].content, "");
660 }
661
662 #[test]
663 fn extract_readable_produces_markdown() {
664 let html = r#"<html><head><title>My Page</title></head>
665 <body><article><h1>Heading</h1><p>Hello <a href="https://x.com">link</a>.</p>
666 <p>More text to satisfy the readability length heuristic so this block is
667 treated as the main article content rather than boilerplate chrome.</p>
668 </article></body></html>"#;
669 let (title, md) = extract_readable(html, "https://example.com/page");
670 assert!(!title.is_empty(), "title extracted, got {title:?}");
671 assert!(
672 md.contains("Hello"),
673 "content converted to markdown: {md:?}"
674 );
675 assert!(md.contains("](https://x.com)"), "links preserved: {md:?}");
676 }
677
678 #[test]
679 fn extract_readable_fallback_title_on_unparseable() {
680 let html = "<title>Bare</title><p>just a snippet</p>";
682 let (title, md) = extract_readable(html, "https://example.com");
683 assert_eq!(title, "Bare");
684 assert!(md.contains("just a snippet"));
685 }
686
687 #[tokio::test]
688 async fn vetting_resolver_rejects_a_name_resolving_to_loopback() {
689 use reqwest::dns::Resolve;
690 use std::str::FromStr;
691 let name = reqwest::dns::Name::from_str("localhost").expect("valid host name");
696 assert!(
697 VettingResolver.resolve(name).await.is_err(),
698 "a name resolving to a loopback address must be rejected"
699 );
700 }
701
702 #[test]
703 fn web_error_is_retryable_classifies_status() {
704 assert!(web_error_is_retryable(&anyhow::Error::new(
706 HttpStatusError { status: 500 }
707 )));
708 assert!(web_error_is_retryable(&anyhow::Error::new(
709 HttpStatusError { status: 429 }
710 )));
711 assert!(!web_error_is_retryable(&anyhow::Error::new(
712 HttpStatusError { status: 404 }
713 )));
714 assert!(!web_error_is_retryable(&anyhow::Error::new(
715 HttpStatusError { status: 401 }
716 )));
717 assert!(!web_error_is_retryable(&anyhow!("parse failed")));
718 let wrapped = anyhow::Error::new(HttpStatusError { status: 503 }).context("upstream");
721 assert!(web_error_is_retryable(&wrapped));
722 }
723}