phrona_api/
tavily.rs

1use std::sync::Arc;
2
3use axum::extract::State;
4use axum::response::IntoResponse;
5use axum::{Json, http::HeaderMap};
6use serde::{Deserialize, Serialize};
7
8use phrona::SearchOptions;
9use phrona::models::{Category, ResultItem, TimeRange};
10
11use crate::{AppError, AppResult, AppState, JsonBody};
12
13/// Tavily-compatible request body.
14///
15/// The Tavily API (https://docs.tavily.com) is the de-facto standard for
16/// AI search. Clients such as `tavily-python` can target this server by
17/// setting `base_url` to it and calling `/search`.
18#[derive(Deserialize)]
19pub struct TavilyRequest {
20    pub query: String,
21    #[serde(default)]
22    pub api_key: Option<String>,
23    #[serde(default)]
24    pub search_depth: Option<String>,
25    #[serde(default)]
26    pub topic: Option<String>,
27    #[serde(default)]
28    pub days: Option<u32>,
29    #[serde(default)]
30    pub max_results: Option<usize>,
31    #[serde(default)]
32    pub include_images: bool,
33    #[serde(default)]
34    pub include_answer: bool,
35    #[serde(default)]
36    pub include_raw_content: bool,
37    /// Accepted for compatibility with Tavily clients; only meaningful for
38    /// Tavily's image-search endpoint, which this server does not expose.
39    #[serde(default)]
40    pub include_image_descriptions: bool,
41    #[serde(default)]
42    pub include_domains: Option<Vec<String>>,
43    #[serde(default)]
44    pub exclude_domains: Option<Vec<String>>,
45}
46
47#[derive(Serialize)]
48pub struct TavilyResponse {
49    pub query: String,
50    pub follow_up_questions: Vec<String>,
51    pub response_time: f64,
52    pub answer: Option<String>,
53    pub images: Option<Vec<String>>,
54    pub results: Vec<TavilyResult>,
55}
56
57#[derive(Serialize)]
58pub struct TavilyResult {
59    pub title: String,
60    pub url: String,
61    pub content: String,
62    pub score: f64,
63    pub raw_content: Option<String>,
64}
65
66fn days_to_range(days: u32) -> Option<TimeRange> {
67    Some(match days {
68        0..=1 => TimeRange::Day,
69        2..=7 => TimeRange::Week,
70        8..=30 => TimeRange::Month,
71        _ => TimeRange::Year,
72    })
73}
74
75fn apply_domains(query: &mut String, include: &[String], exclude: &[String]) {
76    if !include.is_empty() {
77        let sites: Vec<String> = include.iter().map(|d| format!("site:{d}")).collect();
78        query.push_str(&format!(" ({})", sites.join(" OR ")));
79    }
80    for d in exclude {
81        query.push_str(&format!(" -site:{d}"));
82    }
83}
84
85fn to_tavily_result(r: &ResultItem, pos: usize) -> (String, String, String, f64) {
86    let score = (1.0 - pos as f64 * 0.05).max(0.05);
87    match r {
88        ResultItem::Web(w) => (w.title.clone(), w.url.clone(), w.description.clone(), score),
89        ResultItem::News(n) => (n.title.clone(), n.url.clone(), n.description.clone(), score),
90        ResultItem::Video(v) => (v.title.clone(), v.url.clone(), v.description.clone(), score),
91        ResultItem::Image(i) => (i.title.clone(), i.url.clone(), i.source.clone(), score),
92        ResultItem::Book(b) => (b.title.clone(), b.url.clone(), b.info.clone(), score),
93    }
94}
95
96pub async fn search(
97    State(state): State<Arc<AppState>>,
98    headers: HeaderMap,
99    JsonBody(req): JsonBody<TavilyRequest>,
100) -> AppResult<impl IntoResponse> {
101    // Tavily SDKs pass credentials as `api_key` in the JSON body
102    // (langchain-tavily, llama-index) or as headers; both are honored.
103    let key = crate::auth_key(&headers, req.api_key.as_deref());
104    if !state.authorized(key.as_deref()) {
105        return Err(AppError::unauthorized());
106    }
107
108    let mut opts = SearchOptions::new(req.query.clone());
109    let topic_is_news = matches!(req.topic.as_deref(), Some("news"));
110    opts.category = if topic_is_news {
111        Category::News
112    } else {
113        Category::Web
114    };
115    let depth = req.search_depth.as_deref().unwrap_or("basic");
116    // "advanced" is honored by querying every engine in the category;
117    // anything else is rejected loudly instead of silently coerced.
118    if !matches!(depth, "basic" | "advanced") {
119        return Err(AppError::bad_request(format!(
120            "invalid search_depth '{depth}', expected 'basic' or 'advanced'"
121        )));
122    }
123    if depth == "basic" {
124        let mut engines = match opts.category {
125            Category::News => vec!["bing_news".into(), "duckduckgo_news".into()],
126            _ => vec!["bing".into(), "duckduckgo".into()],
127        };
128        if req.include_answer {
129            engines.push("grokipedia".into());
130        }
131        opts.engines = engines;
132    }
133    if let Some(days) = req.days {
134        opts.time_range = days_to_range(days);
135    } else if topic_is_news {
136        // news topic without an explicit window: last week, like Tavily
137        opts.time_range = Some(TimeRange::Week);
138    }
139    opts.max_results = req.max_results.unwrap_or(5).clamp(1, 20);
140    if let Some(include) = &req.include_domains {
141        apply_domains(&mut opts.query, include, &[]);
142    }
143    if let Some(exclude) = &req.exclude_domains {
144        apply_domains(&mut opts.query, &[], exclude);
145    }
146
147    let started = std::time::Instant::now();
148    let resp = state.client.search(opts).await?;
149    let response_time = started.elapsed().as_secs_f64();
150
151    let limit = resp.total.min(req.max_results.unwrap_or(5).clamp(1, 20));
152    let mut results: Vec<TavilyResult> = Vec::with_capacity(limit);
153    let mut images: Vec<String> = Vec::new();
154    for (i, r) in resp.results.iter().take(limit).enumerate() {
155        let (title, url, content, score) = to_tavily_result(r, i);
156        if let ResultItem::Image(img) = r
157            && !img.image_url.is_empty()
158        {
159            images.push(img.image_url.clone());
160        }
161        results.push(TavilyResult {
162            title,
163            url,
164            content,
165            score,
166            raw_content: None,
167        });
168    }
169
170    if req.include_images {
171        // Tavily's `images` field lists image results alongside the web
172        // hits; run a dedicated image search to populate it honestly.
173        let mut img_opts = SearchOptions::new(req.query.clone());
174        img_opts.category = Category::Images;
175        img_opts.max_results = limit.clamp(1, 8);
176        if let Ok(img_resp) = state.client.search(img_opts).await {
177            for r in img_resp.results.iter().take(8) {
178                if let ResultItem::Image(img) = r
179                    && !img.image_url.is_empty()
180                {
181                    images.push(img.image_url.clone());
182                }
183            }
184        }
185    }
186
187    if req.include_raw_content {
188        let client = state.client.http();
189        let urls: Vec<String> = results.iter().map(|r| r.url.clone()).collect();
190        let pages = phrona::extract_many(client, &urls, 8000, Some(&req.query)).await;
191        for (r, page) in results.iter_mut().zip(pages) {
192            r.raw_content = Some(match page {
193                Ok(p) => p.text,
194                Err(e) => format!("extract failed: {e}"),
195            });
196        }
197    }
198
199    Ok(Json(TavilyResponse {
200        query: req.query.clone(),
201        follow_up_questions: Vec::new(),
202        response_time,
203        answer: req.include_answer.then_some(resp.answer.clone()).flatten(),
204        images: (req.include_images && !images.is_empty()).then_some(images),
205        results,
206    }))
207}
phrona_api/tavily.rs

phrona_api/
tavily.rs
