use axum::Json;
use axum::extract::{Extension, Path, RawQuery, State};
use axum::http::{HeaderMap, header};
use axum::response::{IntoResponse, Response};
use crate::limits::{Caller, PageHold};
use crate::server::AppState;
use crate::types::{Document, SearchRequest};
const JINA_TOP_RESULTS: usize = 5;
pub(crate) async fn handle_jina_search(
State(s): State<AppState>,
caller: Option<Extension<Caller>>,
headers: HeaderMap,
Path(raw): Path<String>,
RawQuery(query): RawQuery,
) -> Response {
let mut q = raw.trim().to_string();
let sites: Vec<String> = query
.as_deref()
.unwrap_or_default()
.split('&')
.filter_map(|kv| kv.strip_prefix("site="))
.map(url_decode)
.filter(|v| !v.is_empty())
.collect();
if q.is_empty() {
return crate::error::ApiError::invalid("the search query is empty").into_response();
}
if !sites.is_empty() {
let any: Vec<String> = sites.iter().map(|d| format!("site:{d}")).collect();
q = format!("{q} ({})", any.join(" OR "));
}
let (opts, mut params) = match super::headers::parse(&headers) {
Ok(v) => v,
Err(e) => return e.into_response(),
};
params.token_budget = None;
if let Err(e) = crate::validate::check_scrape_options(&opts, s.engine.as_ref()) {
return e.into_response();
}
if let Some(Extension(c)) = &caller
&& let Err(e) = s.limiter.check_quota(c)
{
return e.into_response();
}
let req: SearchRequest = match serde_json::from_value(serde_json::json!({
"query": q, "limit": JINA_TOP_RESULTS,
})) {
Ok(r) => r,
Err(e) => return crate::error::ApiError::invalid(e.to_string()).into_response(),
};
let hits = match s.engine.search(&req).await {
Ok(v) => v["web"].as_array().cloned().unwrap_or_default(),
Err(e) => return e.into_response(),
};
let mut handles = Vec::new();
let mut pages: Vec<(
String,
Result<Document, crate::error::ApiError>,
Option<PageHold>,
)> = Vec::with_capacity(hits.len().min(JINA_TOP_RESULTS));
for h in hits.iter().take(JINA_TOP_RESULTS) {
let url = h["url"].as_str().unwrap_or_default().to_string();
let hold = match &caller {
Some(Extension(c)) => match s.limiter.reserve_page(c).await {
Ok(hold) => Some(hold),
Err(e) => {
pages.push((url, Err(e), None));
continue;
}
},
None => None,
};
let (engine, opts) = (s.engine.clone(), opts.clone());
handles.push(tokio::spawn(async move {
let r = engine.scrape(&url, &opts).await;
(url, r, hold)
}));
}
for h in handles {
match h.await {
Ok(p) => pages.push(p),
Err(e) => pages.push((
String::new(),
Err(crate::error::ApiError::new(
crate::error::ErrorCode::Internal,
e.to_string(),
)),
None,
)),
}
}
let mut out = Vec::new();
for ((url, r, hold), hit) in pages.into_iter().zip(&hits) {
let title = hit["title"].as_str().unwrap_or_default().to_string();
let description = hit["description"].as_str().unwrap_or_default().to_string();
let (content, chunks, warnings) =
match r.and_then(|doc| super::respond::render(&doc, &url, ¶ms).map(|x| (x, doc))) {
Ok((x, doc)) => {
if let Some(hold) = hold {
hold.commit().await;
}
let mut w = doc
.metadata
.pagesift
.map(|m| m.warnings)
.unwrap_or_default();
w.extend(x.warnings);
(x.text, x.chunks, w)
}
Err(e) => (
String::new(),
None,
vec![format!("{}: {}", e.code.as_str(), e.message)],
),
};
out.push((title, url, description, content, chunks, warnings));
}
let accept = headers
.get(header::ACCEPT)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if accept.contains("application/json") {
let data: Vec<serde_json::Value> = out
.iter()
.map(|(t, u, d, c, chunks, w)| {
let mut one = serde_json::json!({ "title": t, "url": u, "description": d, "content": c, "warnings": w });
if let Some(chunks) = chunks {
one["chunks"] = serde_json::json!(chunks);
}
one
})
.collect();
return Json(serde_json::json!({ "code": 200, "status": 20000, "data": data }))
.into_response();
}
let mut text = String::new();
for (i, (t, u, d, c, _, w)) in out.iter().enumerate() {
let n = i + 1;
text.push_str(&format!(
"[{n}] Title: {t}\n[{n}] URL Source: {u}\n[{n}] Description: {d}\n"
));
if !w.is_empty() {
text.push_str(&format!("[{n}] Warnings: {}\n", w.join("; ")));
}
text.push_str(&format!("[{n}] Markdown Content:\n{c}\n\n"));
}
([(header::CONTENT_TYPE, "text/plain; charset=utf-8")], text).into_response()
}
fn url_decode(s: &str) -> String {
let bytes = s.replace('+', " ").into_bytes();
let mut out = Vec::with_capacity(bytes.len());
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'%'
&& let Some(b) = std::str::from_utf8(&bytes[i + 1..(i + 3).min(bytes.len())])
.ok()
.and_then(|h| u8::from_str_radix(h, 16).ok())
{
out.push(b);
i += 3; continue;
}
out.push(bytes[i]);
i += 1;
}
String::from_utf8_lossy(&out).into_owned()
}