//! Smart Search API //! Implements the 5W1H search capability using semantic vectors. use axum::{extract::State, http::StatusCode, response::Json, routing::post, Router}; use serde::{Deserialize, Serialize}; use serde_json; use crate::core::embedding::Embedder; // --- Request / Response Structures --- #[derive(Debug, Deserialize)] pub struct SmartSearchRequest { pub uuid: String, pub query: String, pub page: Option, pub page_size: Option, pub limit: Option, } #[derive(Debug, Serialize)] pub struct SearchResult { pub id: i32, pub parent_id: i32, pub scene_order: Option, // Primary: frame-accurate position (authoritative unit) pub start_frame: i64, pub end_frame: i64, pub fps: f64, // Reference: time derived from frames (subject to FPS variation, not precise) pub start_time: f64, pub end_time: f64, pub raw_text: Option, // Text content of the child chunk pub summary: Option, // Summary from parent context pub metadata: Option, pub similarity: Option, } #[derive(Debug, Serialize)] pub struct SmartSearchResponse { pub query: String, pub results: Vec, pub page: usize, pub page_size: usize, pub strategy: String, } // --- API Handler --- pub async fn smart_search( State(state): State, Json(req): Json, ) -> Result, (StatusCode, Json)> { let db = &state.db; let page = req.page.unwrap_or(1).max(1); // Backward compat: if old `limit` sent without `page_size`, use limit as page_size let page_size = if req.page_size.is_some() { req.page_size.unwrap() } else if req.limit.is_some() && req.page.is_none() { req.limit.unwrap() } else { 5 } .max(1); let hard_limit = req.limit.unwrap_or(usize::MAX); let limit = hard_limit.min(page_size); // 1. Generate Embedding using EmbeddingGemma via MOMENTRY_EMBED_URL let embedder = Embedder::new("embeddinggemma-300m".to_string()); let embedding = embedder.embed_query(&req.query).await.map_err( |e| -> (StatusCode, Json) { tracing::error!("Embedding failed: {}", e); ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": e.to_string() })), ) }, )?; // 2. Search Database (Drill-Down: Find Parents First) let db_parents: Vec = db .search_parent_chunks_semantic(&req.uuid, &embedding, limit) .await .map_err( |e: anyhow::Error| -> (StatusCode, Json) { tracing::error!("DB search failed: {}", e); ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": e.to_string() })), ) }, )?; if db_parents.is_empty() { return Ok(Json(SmartSearchResponse { query: req.query, results: vec![], page, page_size, strategy: "semantic_vector_search".to_string(), })); } // Collect Parent IDs let parent_ids: Vec = db_parents.iter().map(|p| p.id).collect(); // 3. Fetch Children for these Parents (Drill Down) // We fetch all children for these parents (limit can be adjusted) let children: Vec = db .get_children_for_parents(&parent_ids, 10) // Fetch top 10 children per parent .await .map_err( |e: anyhow::Error| -> (StatusCode, Json) { tracing::error!("Fetching children failed: {}", e); ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": e.to_string() })), ) }, )?; // 4. Map Parents to a lookup table let parent_map: std::collections::HashMap< i32, &crate::core::db::postgres_db::SemanticSearchResult, > = db_parents.iter().map(|p| (p.id, p)).collect(); // Map Children to API response struct let results: Vec = children .into_iter() .map(|c| { let parent = parent_map.get(&c.parent_id); SearchResult { id: c.id, parent_id: c.parent_id, scene_order: parent.map(|p| p.scene_order), start_frame: c.start_frame, end_frame: c.end_frame, fps: c.fps, start_time: c.start_time, end_time: c.end_time, raw_text: Some(c.raw_text), summary: parent.map(|p| p.summary.clone()), metadata: parent.map(|p| p.metadata.clone()), similarity: parent.and_then(|p| p.similarity), } }) .collect(); // 6. Sort results by similarity (descending) // Since all children of a parent have the same parent similarity, this groups relevant chunks together let mut results = results; results.sort_by(|a, b| { b.similarity .partial_cmp(&a.similarity) .unwrap_or(std::cmp::Ordering::Equal) }); // 7. Limit the final results (optional, but good for API consistency) let truncate_limit = hard_limit.min(page_size * 5); // Allow more children per parent context results.truncate(truncate_limit); // 8. Format Response let response = SmartSearchResponse { query: req.query, results, page, page_size, strategy: "drill_down_semantic_search".to_string(), }; Ok(Json(response)) } // --- Router Setup --- pub fn search_routes() -> Router { Router::new().route("/api/v1/search/smart", post(smart_search)) }