Skip to main content

phrona_api/
grounding.rs

1use std::sync::Arc;
2
3use axum::Json;
4use axum::extract::State;
5use axum::response::IntoResponse;
6use serde::{Deserialize, Serialize};
7
8use phrona::SearchOptions;
9use phrona::models::{Category, TimeRange};
10
11use crate::{AppError, AppResult, AppState, HeaderAuth, JsonBody, JsonQuery};
12
13/// POST /v1/grounding - AI grounding request; credentials via headers or
14/// the JSON body.
15#[derive(Deserialize)]
16pub struct GroundingRequest {
17    pub query: String,
18    #[serde(default)]
19    pub api_key: Option<String>,
20    #[serde(default)]
21    pub max_results: Option<usize>,
22    #[serde(default)]
23    pub category: Option<String>,
24    #[serde(default)]
25    pub time_range: Option<String>,
26}
27
28/// GET /v1/grounding?query=... - same feature; auth is header-only.
29#[derive(Deserialize)]
30pub struct GroundingGetParams {
31    pub query: String,
32    #[serde(default)]
33    pub max_results: Option<usize>,
34    #[serde(default)]
35    pub category: Option<String>,
36    #[serde(default)]
37    pub time_range: Option<String>,
38}
39
40#[derive(Serialize)]
41pub struct GroundingResponse {
42    pub query: String,
43    pub answer: String,
44    pub sources: Vec<GroundingSource>,
45    pub response_time: f64,
46}
47
48#[derive(Serialize)]
49pub struct GroundingSource {
50    pub title: String,
51    pub url: String,
52    pub content: String,
53    pub score: f64,
54}
55
56/// Build an extractive answer from the top results. The library answer
57/// (e.g. from grokipedia) takes precedence; otherwise the strongest
58/// snippets are stitched together without claiming LLM generation.
59fn synthesize_answer(resp: &phrona::SearchResponse, sources: &[GroundingSource]) -> String {
60    if let Some(a) = &resp.answer
61        && !a.trim().is_empty()
62    {
63        return a.clone();
64    }
65    if sources.is_empty() {
66        return format!("No results found for \"{}\".", resp.query);
67    }
68    let mut parts = Vec::new();
69    for (i, s) in sources.iter().take(3).enumerate() {
70        let content = s.content.trim();
71        if content.is_empty() {
72            continue;
73        }
74        let excerpt = content.chars().take(400).collect::<String>();
75        parts.push(format!("Source {} ({}): {excerpt}", i + 1, s.url));
76    }
77    if parts.is_empty() {
78        return format!("{} sources found for \"{}\".", sources.len(), resp.query);
79    }
80    format!(
81        "Extractive summary for \"{}\":\n{}",
82        resp.query,
83        parts.join("\n")
84    )
85}
86
87pub async fn get(
88    State(state): State<Arc<AppState>>,
89    auth: HeaderAuth,
90    JsonQuery(p): JsonQuery<GroundingGetParams>,
91) -> AppResult<impl IntoResponse> {
92    if !state.authorized(auth.key()) {
93        return Err(AppError::unauthorized());
94    }
95    run(
96        &state,
97        &p.query,
98        p.max_results,
99        p.category.as_deref(),
100        p.time_range.as_deref(),
101    )
102    .await
103}
104
105pub async fn post(
106    State(state): State<Arc<AppState>>,
107    headers: axum::http::HeaderMap,
108    JsonBody(p): JsonBody<GroundingRequest>,
109) -> AppResult<impl IntoResponse> {
110    if !state.authorized(crate::auth_key(&headers, p.api_key.as_deref()).as_deref()) {
111        return Err(AppError::unauthorized());
112    }
113    run(
114        &state,
115        &p.query,
116        p.max_results,
117        p.category.as_deref(),
118        p.time_range.as_deref(),
119    )
120    .await
121}
122
123async fn run(
124    state: &AppState,
125    query: &str,
126    max_results: Option<usize>,
127    category: Option<&str>,
128    time_range: Option<&str>,
129) -> AppResult<Json<GroundingResponse>> {
130    let mut opts = SearchOptions::new(query.to_string());
131    opts.max_results = max_results.unwrap_or(10).clamp(1, 50);
132    if let Some(c) = category {
133        opts.category = c.parse::<Category>().map_err(|_| {
134            AppError::bad_request(
135                "invalid category, expected one of: web, images, news, videos, books",
136            )
137        })?;
138    }
139    if let Some(t) = time_range {
140        opts.time_range = Some(t.parse::<TimeRange>().map_err(|_| {
141            AppError::bad_request("invalid time_range, expected day|week|month|year")
142        })?);
143    }
144
145    let started = std::time::Instant::now();
146    let resp = state.client.search(opts).await?;
147    let response_time = started.elapsed().as_secs_f64();
148
149    let mut sources: Vec<GroundingSource> = Vec::new();
150    for (i, r) in resp.results.iter().enumerate() {
151        let score = (1.0 - i as f64 * 0.05).max(0.05);
152        let (title, url, content) = match r {
153            phrona::ResultItem::Web(w) => (&w.title, &w.url, &w.description),
154            phrona::ResultItem::News(n) => (&n.title, &n.url, &n.description),
155            phrona::ResultItem::Video(v) => (&v.title, &v.url, &v.description),
156            phrona::ResultItem::Image(i) => (&i.title, &i.url, &i.source),
157            phrona::ResultItem::Book(b) => (&b.title, &b.url, &b.info),
158        };
159        sources.push(GroundingSource {
160            title: title.clone(),
161            url: url.clone(),
162            content: content.clone(),
163            score,
164        });
165    }
166
167    let answer = synthesize_answer(&resp, &sources);
168    Ok(Json(GroundingResponse {
169        query: resp.query.clone(),
170        answer,
171        sources,
172        response_time,
173    }))
174}