Files
pile/crates/pile-dataset/src/serve/lookup.rs
rm-dr f7ea25f059
Some checks failed
CI / Typos (push) Successful in 31s
CI / Clippy (push) Successful in 1m17s
CI / Build and test (all features) (push) Failing after 5m32s
CI / Build and test (push) Failing after 7m49s
Return len
2026-03-24 03:43:26 -07:00

81 lines
1.9 KiB
Rust

use axum::{
Json,
extract::State,
http::StatusCode,
response::{IntoResponse, Response},
};
use serde::{Deserialize, Serialize};
use std::{sync::Arc, time::Instant};
use tracing::debug;
use utoipa::ToSchema;
use crate::Datasets;
#[derive(Serialize, Deserialize, ToSchema, Debug)]
pub struct LookupRequest {
pub query: String,
#[serde(default)]
pub limit: Option<usize>,
}
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct LookupResponse {
pub results: Vec<LookupResult>,
pub total: u64,
}
#[derive(Debug, Serialize, Deserialize, ToSchema)]
pub struct LookupResult {
pub score: f32,
pub source: String,
pub key: String,
}
/// Search for an item in this dataset
#[utoipa::path(
post,
path = "/lookup",
responses(
(status = 200, description = "Search results", body = Vec<LookupResponse>),
(status = 400, description = "Invalid request"),
(status = 401, description = "Unauthorized"),
(status = 404, description = "URL not found"),
(status = 500, description = "Internal server error"),
)
)]
pub async fn lookup(
State(state): State<Arc<Datasets>>,
Json(body): Json<LookupRequest>,
) -> Response {
let start = Instant::now();
let limit = body.limit.unwrap_or(128).min(1024);
debug!(message = "Serving /lookup", query = body.query, limit);
let results: Vec<LookupResult> = match state.fts_lookup(&body.query, limit) {
Ok(x) => x
.into_iter()
.map(|x| LookupResult {
key: x.key,
score: x.score,
source: x.source.into(),
})
.collect(),
Err(error) => {
return (StatusCode::INTERNAL_SERVER_ERROR, format!("{error:?}")).into_response();
}
};
let total: u64 = state.sources.iter().map(|x| x.1.len() as u64).sum();
debug!(
message = "Served /lookup",
query = body.query,
limit = body.limit.unwrap_or(10),
time_ms = start.elapsed().as_millis()
);
return (StatusCode::OK, Json(LookupResponse { results, total })).into_response();
}