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#[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#[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
56fn 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}