searchi/tools/media-eval/src/bin/pipeline.rs

//! End-to-end corpus indexing and querying through actual Searchi APIs.
use anyhow::{ensure, Result};
use searchi_core::{SearchIndex, SearchOptions};
use searchi_media::{FfmpegTools, MediaLimits};
use searchi_service::HybridSearch;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{fs, path::PathBuf, time::Instant};

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args_os().skip(1).collect();
    ensure!(
        args.len() == 6 || args.len() == 7,
        "Supply CORPUS MODEL CPU_WORKER FFMPEG FFPROBE OUTPUT_JSON [AUDIO_WINDOW_SECONDS]"
    );
    ensure!(
        std::env::var("SEARCHI_FORCE_CPU").as_deref() == Ok("1"),
        "CPU only"
    );
    let root = PathBuf::from(&args[0]).canonicalize()?;
    let bytes = fs::read(root.join("fixture.json"))?;
    let hash = format!("{:x}", Sha256::digest(&bytes));
    let lock: Value = serde_json::from_slice(&fs::read(root.join("labels.lock.json"))?)?;
    ensure!(lock["fixture_sha256"] == hash, "Frozen labels changed");
    let fixture: Value = serde_json::from_slice(&bytes)?;
    for asset in fixture["assets"].as_array().unwrap() {
        let content = fs::read(root.join(asset["path"].as_str().unwrap()))?;
        ensure!(
            asset["sha256"] == format!("{:x}", Sha256::digest(content)),
            "Asset changed"
        );
    }
    let temporary = tempfile::tempdir()?;
    let db = temporary.path().join("index.sqlite");
    let mut lexical = SearchIndex::open(&db)?;
    let filename_index = lexical.index_folder(root.join("assets"))?;
    let mut hybrid = HybridSearch::open(&db, &args[1], &args[2], 1)?;
    let tools = FfmpegTools::new(PathBuf::from(&args[3]), PathBuf::from(&args[4]))?;
    let start = Instant::now();
    let mut limits = MediaLimits::default();
    if let Some(seconds) = args.get(6) {
        let seconds: u32 = seconds.to_string_lossy().parse()?;
        ensure!(
            (5..=30).contains(&seconds),
            "Audio window must be 5–30 seconds"
        );
        limits.audio_window_ms = seconds * 1000;
    }
    let index_report = hybrid.index_embeddings_with_media(&tools, &limits)?;
    let encoding_ms = start.elapsed().as_secs_f64() * 1000.0;
    println!("Actual pipeline indexed {} chunks", index_report.chunks);
    hybrid.warmup()?;
    let mut rows = Vec::new();
    for query in fixture["queries"].as_array().unwrap() {
        if query["filters"].get("time_range_s").is_some() {
            rows.push(json!({"query":query,"unsupported":"Production query API has no temporal-range filter yet"}));
            continue;
        }
        let mut options = SearchOptions::default();
        if let Some(extension) = query["filters"]["extension"].as_str() {
            options.extension = Some(extension.into());
        }
        let start = Instant::now();
        let hits = hybrid.search(query["text"].as_str().unwrap(), 20, &options)?;
        let query_ms = start.elapsed().as_secs_f64() * 1000.0;
        let file_rankings: Vec<_> = hits.iter().map(|hit| {
            let asset = fixture["assets"].as_array().unwrap().iter().find(|asset| PathBuf::from(asset["path"].as_str().unwrap()).file_name().unwrap().to_string_lossy() == hit.filename).unwrap();
            let modality = if hit.match_kind.contains("audio") { "audio" } else if hit.match_kind == "media_metadata" { "metadata" } else { "image" };
            let mut candidate = json!({"asset":asset["id"],"modality":modality,"match_kind":hit.match_kind,"source":hit.source});
            if let Some(source) = &hit.source { if source.kind == "timestamp_ms" {
                let stamp = f64::from(source.number)/1000.0;
                candidate["start_s"] = json!(stamp);
                candidate["end_s"] = json!((stamp + if modality == "audio" {f64::from(limits.audio_window_ms)/1000.0} else {2.0}).min(asset["duration_s"].as_f64().unwrap_or(stamp)));
                candidate["interval_note"] = json!("End inferred from configured pipeline window policy, actual API exposes only start timestamp");
            }}
            json!({"candidate":candidate,"score":hit.score})
        }).collect();
        rows.push(json!({"query":query,"segment_top5":file_rankings.iter().take(5).cloned().collect::<Vec<_>>(),"file_rankings":file_rankings,"query_ms":query_ms}));
    }
    let report = json!({"fixture_sha256":hash,"worker":hybrid.worker_info(),"index_report":index_report,"filename_index":filename_index,"encoding_ms":encoding_ms,"audio_window_ms":limits.audio_window_ms,"audio_overlap_ms":limits.audio_overlap_ms,"encoding_failures":index_report.failures,"rows":rows,"scope":"Actual Searchi hybrid pipeline and native extension filters; temporal-range filter unsupported; segment ends inferred from configured audio/frame policy, exact point timestamps come from actual result API"});
    fs::write(&args[5], serde_json::to_vec_pretty(&report)?)?;
    Ok(())
}