Move router to pile-dataset

crates/pile-dataset/src/serve/lookup.rs

@@ -0,0 +1,80 @@
+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>,
+}
+
+#[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();
+	debug!(
+		message = "Serving /lookup",
+		query = body.query,
+		limit = body.limit.unwrap_or(10)
+	);
+
+	let results: Vec<LookupResult> = match state.fts_lookup(&body.query, body.limit.unwrap_or(10)) {
+		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();
+		}
+	};
+
+	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 })).into_response();
+}