feat: score-based search, LLM re-ranking endpoint, video title search, pipeline module

Core search changes: - Replace RRF with score-based merge (max of semantic/keyword/identity) - Add video title ILIKE search for brand/name queries (score 0.9) - Add /api/v1/search/llm-smart endpoint with Gemma 4 re-ranking - Fix LLM JSON parsing (markdown fences, empty responses) Infrastructure: - Rebuild Qdrant collection (clear 347K contaminated points) - Add dotenv loading to main.rs for config parity - Implement store_pre_chunk in postgres_db.rs Pipeline module (WordPress): - store-asrx, rule1, vectorize, phase1, complete endpoints - CLI commands for pipeline operations Docs: - SEARCH_SCORE_IMPROVEMENT.md (score-based merge proposal)
2026-06-04 07:40:41 +08:00
parent e1572907ae
commit 834b0d4865
14 changed files with 835 additions and 31 deletions
@@ -3308,10 +3308,38 @@ impl PostgresDb {

    pub async fn store_pre_chunk(
        &self,
-        _uuid: &str,
-        _chunk_type: &str,
-        _data: serde_json::Value,
+        uuid: &str,
+        processor_type: &str,
+        data: serde_json::Value,
    ) -> Result<()> {
+        let table = schema::table_name("pre_chunks");
+        let pre_chunk: PreChunk = serde_json::from_value(data)?;
+        let start_time = pre_chunk.start_frame as f64 / pre_chunk.fps;
+        let end_time = pre_chunk.end_frame as f64 / pre_chunk.fps;
+        sqlx::query(&format!(
+            "INSERT INTO {} (file_uuid, file_id, source_type, source_file, chunk_type, \
+             start_frame, end_frame, start_time, end_time, fps, data, text_content, \
+             processed, chunk_id, processor_type) \
+             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)",
+            table
+        ))
+        .bind(uuid)
+        .bind(pre_chunk.file_id)
+        .bind(&pre_chunk.source_type)
+        .bind(&pre_chunk.source_file)
+        .bind(&pre_chunk.chunk_type)
+        .bind(pre_chunk.start_frame)
+        .bind(pre_chunk.end_frame)
+        .bind(start_time)
+        .bind(end_time)
+        .bind(pre_chunk.fps)
+        .bind(&pre_chunk.raw_json)
+        .bind(&pre_chunk.text_content)
+        .bind(pre_chunk.processed)
+        .bind(&pre_chunk.chunk_id)
+        .bind(processor_type)
+        .execute(&self.pool)
+        .await?;
        Ok(())
    }

@@ -1,2 +1,3 @@
 pub mod client;
 pub mod function_calling;
+pub mod rerank;
@@ -0,0 +1,168 @@
+use std::collections::HashSet;
+
+use anyhow::Result;
+use reqwest::Client;
+use serde::{Deserialize, Serialize};
+use std::time::Duration;
+use tracing::{debug, warn};
+
+use crate::core::config;
+
+#[derive(Debug, Serialize)]
+struct ChatRequest {
+    model: String,
+    messages: Vec<ChatMessage>,
+    temperature: f32,
+    max_tokens: u32,
+    stream: bool,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+struct ChatMessage {
+    role: String,
+    content: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct ChatResponse {
+    choices: Vec<Choice>,
+}
+
+#[derive(Debug, Deserialize)]
+struct Choice {
+    message: ChatMessage,
+}
+
+#[derive(Debug, Deserialize)]
+struct RerankResponse {
+    ranked: Vec<usize>,
+}
+
+pub async fn rerank_search_results(query: &str, candidates: &[(usize, &str)]) -> Result<Vec<usize>> {
+    if candidates.is_empty() {
+        return Ok(vec![]);
+    }
+
+    let mut chunks_text = String::new();
+    for (i, (_, text)) in candidates.iter().enumerate() {
+        let display = if text.len() > 100 {
+            format!("{}...", &text[..100])
+        } else {
+            text.to_string()
+        };
+        chunks_text.push_str(&format!("[{}] {}\n", i + 1, display));
+    }
+
+    let prompt = format!(
+        r#"You are a search relevance judge. Rank ALL chunks by relevance to the query.
+
+Query: "{}"
+
+Chunks:
+{}
+
+Return a JSON object with ALL chunk numbers in order of relevance (most relevant first).
+Example: {{"ranked": [5, 1, 3, 2, 4, 6, 7, 8, 9, 10]}}
+Include every chunk number exactly once. Only respond with the JSON."#,
+        query, chunks_text
+    );
+
+    let client = Client::builder()
+        .timeout(Duration::from_secs(15))
+        .build()?;
+
+    let req = ChatRequest {
+        model: config::llm::CHAT_MODEL.clone(),
+        messages: vec![
+            ChatMessage {
+                role: "system".to_string(),
+                content: "You are a precise search relevance judge.".to_string(),
+            },
+            ChatMessage {
+                role: "user".to_string(),
+                content: prompt,
+            },
+        ],
+        temperature: 0.1,
+        max_tokens: 512,
+        stream: false,
+    };
+
+    debug!("LLM rerank: {} candidates for query '{}'", candidates.len(), query);
+
+    let res = client
+        .post(&*config::llm::CHAT_URL)
+        .json(&req)
+        .send()
+        .await?;
+
+    if !res.status().is_success() {
+        let status = res.status();
+        let body = res.text().await.unwrap_or_default();
+        warn!("LLM rerank API error: {} — body: {}", status, body);
+        return Ok(candidates.iter().map(|(idx, _)| *idx).collect());
+    }
+
+    let chat_res: ChatResponse = res.json().await?;
+    let content = chat_res
+        .choices
+        .into_iter()
+        .next()
+        .map(|c| c.message.content)
+        .unwrap_or_default();
+
+    let content = content.trim();
+
+    // Strip markdown code fences if present
+    let content = if content.starts_with("```") {
+        let lines: Vec<&str> = content.lines().collect();
+        let start = if lines.first().map(|l| l.contains("```")).unwrap_or(false) { 1 } else { 0 };
+        let end = if lines.last().map(|l| l.contains("```")).unwrap_or(false) {
+            lines.len().saturating_sub(1)
+        } else {
+            lines.len()
+        };
+        lines[start..end].join("\n").trim().to_string()
+    } else {
+        content.to_string()
+    };
+
+    let json_start = content.find('{');
+    let json_end = content.rfind('}');
+
+    if let (Some(start), Some(end)) = (json_start, json_end) {
+        let json_str = &content[start..=end];
+        match serde_json::from_str::<RerankResponse>(json_str) {
+            Ok(parsed) => {
+                let mut ranked: Vec<usize> = parsed
+                    .ranked
+                    .into_iter()
+                    .filter_map(|i| {
+                        if i > 0 && i <= candidates.len() {
+                            Some(candidates[i - 1].0)
+                        } else {
+                            None
+                        }
+                    })
+                    .collect();
+
+                if !ranked.is_empty() {
+                    let seen: HashSet<usize> = ranked.iter().cloned().collect();
+                    for (orig_idx, _) in candidates {
+                        if !seen.contains(orig_idx) {
+                            ranked.push(*orig_idx);
+                        }
+                    }
+                    return Ok(ranked);
+                }
+                warn!("LLM rerank returned empty ranked list");
+            }
+            Err(e) => {
+                warn!("Failed to parse LLM rerank JSON: {}", e);
+            }
+        }
+    }
+
+    warn!("LLM rerank: could not parse response — content: {}", &content[..content.len().min(200)]);
+    Ok(candidates.iter().map(|(idx, _)| *idx).collect())
+}
@@ -12,6 +12,7 @@ pub mod ingestion;
 pub mod llm;
 pub mod overlay;
 pub mod person_identity;
+pub mod pipeline;
 pub mod probe;
 pub mod processor;
 pub mod storage;
@@ -0,0 +1,172 @@
+use anyhow::{Context, Result};
+
+use crate::core::chunk::rule1_ingest;
+use crate::core::config;
+use crate::core::db::postgres_db::PostgresDb;
+use crate::core::db::qdrant_db::QdrantDb;
+use crate::core::db::schema;
+use crate::core::db::VectorPayload;
+use crate::core::embedding::Embedder;
+use crate::core::processor::asrx::AsrxResult;
+use crate::core::processor::PythonExecutor;
+use crate::core::storage::output_dir::OutputDir;
+
+pub async fn store_asrx_chunks(db: &PostgresDb, uuid: &str) -> Result<()> {
+    let output_dir = OutputDir::new();
+    let asrx_path = output_dir.get_output_path(uuid, "asrx.json");
+
+    let json_str = std::fs::read_to_string(&asrx_path)
+        .with_context(|| format!("ASRX file not found: {:?}", asrx_path))?;
+    let result: AsrxResult = serde_json::from_str(&json_str)
+        .context("Failed to parse ASRX JSON")?;
+
+    let segments_count = result.segments.len();
+    let mut pre_chunks = Vec::new();
+    let mut speaker_detections = Vec::new();
+
+    for (i, segment) in result.segments.iter().enumerate() {
+        let data = serde_json::json!({
+            "text": segment.text,
+            "speaker_id": segment.speaker_id,
+            "timestamp": segment.start_time,
+        });
+        pre_chunks.push((i as i64, Some(segment.start_time), data, None, None));
+        speaker_detections.push((
+            segment.speaker_id.clone().unwrap_or_default(),
+            segment.start_time,
+            segment.end_time,
+            segment.text.clone(),
+            None::<String>,
+            0.0,
+        ));
+    }
+
+    db.store_raw_pre_chunks_batch(uuid, "asrx", &pre_chunks).await?;
+    db.store_raw_pre_chunks_batch(uuid, "asr", &pre_chunks).await?;
+    db.store_speaker_detections_batch(uuid, &speaker_detections).await?;
+
+    println!("Stored {} ASRX pre-chunks for {}", segments_count, uuid);
+    Ok(())
+}
+
+pub async fn execute_rule1(db: &PostgresDb, uuid: &str) -> Result<usize> {
+    let video = db.get_video_by_uuid(uuid)
+        .await?
+        .context("Video not found")?;
+    let fps = video.fps;
+
+    let count = rule1_ingest::execute_rule1(db, uuid, fps).await
+        .context("Rule 1 ingestion failed")?;
+
+    println!("Rule 1 completed: {} chunks inserted for {}", count, uuid);
+    Ok(count)
+}
+
+pub async fn vectorize_chunks(uuid: &str) -> Result<()> {
+    let db = PostgresDb::new(&config::DATABASE_URL).await?;
+    let qdrant = QdrantDb::new();
+    let embedder = Embedder::new("embeddinggemma-300m".to_string());
+
+    let chunk_table = schema::table_name("chunk");
+    let rows = sqlx::query_as::<_, (String, String, String, i64, i64, f64, f64, String)>(
+        &format!(
+            "SELECT chunk_id, chunk_type, text_content, start_frame, end_frame, \
+             start_time, end_time, content::text \
+             FROM {} WHERE file_uuid = $1 AND chunk_type = 'sentence' \
+             AND embedding IS NULL \
+             AND (text_content IS NOT NULL AND text_content != '') \
+             ORDER BY id",
+            chunk_table
+        ),
+    )
+    .bind(uuid)
+    .fetch_all(db.pool())
+    .await?;
+
+    if rows.is_empty() {
+        println!("No sentence chunks to vectorize for {}", uuid);
+        return Ok(());
+    }
+
+    let total = rows.len();
+    let mut stored = 0usize;
+
+    for (chunk_id, _chunk_type, text, start_frame, end_frame, start_time, end_time, _content_str) in &rows {
+        if text.is_empty() {
+            continue;
+        }
+
+        match embedder.embed_document(text).await {
+            Ok(vector) => {
+                if let Err(e) = db.store_vector(chunk_id, &vector, uuid).await {
+                    eprintln!("PG store failed for {}: {}", chunk_id, e);
+                    continue;
+                }
+                let payload = VectorPayload {
+                    file_uuid: uuid.to_string(),
+                    chunk_id: chunk_id.clone(),
+                    chunk_type: "sentence".to_string(),
+                    start_frame: *start_frame,
+                    end_frame: *end_frame,
+                    start_time: *start_time,
+                    end_time: *end_time,
+                    text: Some(text.clone()),
+                };
+                if let Err(e) = qdrant.upsert_vector(chunk_id, &vector, payload).await {
+                    eprintln!("Qdrant upsert failed for {}: {}", chunk_id, e);
+                    continue;
+                }
+                stored += 1;
+                if stored % 50 == 0 {
+                    println!("Vectorized {}/{} chunks for {}", stored, total, uuid);
+                }
+            }
+            Err(e) => {
+                eprintln!("Embedding failed for {}: {}", chunk_id, e);
+            }
+        }
+    }
+
+    println!("Vectorization complete: {}/{} vectors for {}", stored, total, uuid);
+    Ok(())
+}
+
+pub async fn run_phase1(uuid: &str) -> Result<()> {
+    let executor = PythonExecutor::new()
+        .context("Failed to create PythonExecutor")?;
+
+    executor
+        .run(
+            "release_pack.py",
+            &["--phase", "1", "--file-uuid", uuid],
+            None,
+            "RELEASE_P1",
+            Some(std::time::Duration::from_secs(120)),
+        )
+        .await
+        .context("Phase 1 release pack failed")?;
+
+    println!("Phase 1 release packaged for {}", uuid);
+    Ok(())
+}
+
+pub async fn mark_complete(db: &PostgresDb, uuid: &str) -> Result<()> {
+    use crate::core::db::MonitorJobStatus;
+    use crate::core::db::VideoStatus;
+
+    let job_id = sqlx::query_scalar::<_, i32>(
+        &format!("SELECT id FROM {} WHERE uuid = $1 LIMIT 1", schema::table_name("monitor_jobs")),
+    )
+    .bind(uuid)
+    .fetch_optional(db.pool())
+    .await?;
+
+    if let Some(job_id) = job_id {
+        db.update_job_status(job_id, MonitorJobStatus::Completed).await?;
+        println!("Job {} marked as completed", job_id);
+    }
+
+    db.update_video_status(uuid, VideoStatus::Completed).await?;
+    println!("Video {} marked as completed", uuid);
+    Ok(())
+}