phrona_api/
tavily.rs

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

phrona_api/
tavily.rs
