searchi/tools/relevance-eval/src/main.rs

//! Frozen calibration and policy-gated held-out diagnostics through the shared engine.
use anyhow::{ensure, Context, Result};
use searchi_core::{SearchIndex, SearchOptions, SearchResult};
use searchi_service::HybridSearch;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
    fs,
    io::{Read, Write},
    path::{Path, PathBuf},
    time::Instant,
};

fn digest(path: &Path) -> Result<String> {
    let mut stream = fs::File::open(path)?;
    let mut hash = Sha256::new();
    let mut buffer = [0; 65536];
    loop {
        let count = stream.read(&mut buffer)?;
        if count == 0 {
            break;
        }
        hash.update(&buffer[..count]);
    }
    Ok(format!("{:x}", hash.finalize()))
}

fn calibration_queries(fixture: &Value) -> Result<Vec<&Value>> {
    split_queries(fixture, "calibration")
}

fn split_queries<'a>(fixture: &'a Value, split: &str) -> Result<Vec<&'a Value>> {
    ensure!(
        matches!(split, "calibration" | "held_out"),
        "Unknown requested split"
    );
    let queries = fixture["queries"].as_array().context("Missing queries")?;
    ensure!(
        queries
            .iter()
            .all(|q| matches!(q["split"].as_str(), Some("calibration" | "held_out"))),
        "Unknown split"
    );
    let selected: Vec<_> = queries.iter().filter(|q| q["split"] == split).collect();
    ensure!(
        selected.len() == 24,
        "Expected exactly 24 frozen queries in selected split"
    );
    Ok(selected)
}

fn checkpoint(path: &Path, report: &Value) -> Result<()> {
    let parent = path
        .parent()
        .filter(|p| !p.as_os_str().is_empty())
        .unwrap_or(Path::new("."));
    let mut staged = tempfile::NamedTempFile::new_in(parent)?;
    staged.write_all(&serde_json::to_vec_pretty(report)?)?;
    staged.as_file().sync_all()?;
    staged.persist(path)?;
    Ok(())
}

fn score<'a>(
    hits: impl IntoIterator<Item = &'a SearchResult>,
    query: &Value,
    fixture: &Value,
) -> Result<Value> {
    let expected = query["expected_ids"]
        .as_array()
        .context("Missing expected labels")?;
    let mut ids = Vec::new();
    for hit in hits {
        let asset = fixture["assets"]
            .as_array()
            .context("Missing assets")?
            .iter()
            .find(|asset| {
                Path::new(asset["path"].as_str().unwrap_or_default())
                    .file_name()
                    .is_some_and(|name| name == hit.filename.as_str())
            })
            .context("Result outside fixed corpus")?;
        if !ids.contains(&asset["id"]) {
            ids.push(asset["id"].clone());
        }
    }
    let rank = ids
        .iter()
        .position(|id| expected.contains(id))
        .map(|rank| rank + 1);
    let absent = expected.is_empty();
    Ok(
        json!({"asset_ranking":ids,"first_relevant_rank":rank,"top1":if absent {None} else {Some(rank == Some(1))},"recall_at_5":if absent {None} else {Some(ids.iter().take(5).filter(|id| expected.contains(id)).count() as f64 / expected.len() as f64)},"absent_artifact_returned_neighbors":if absent {Some(!ids.is_empty())} else {None}}),
    )
}

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args_os().skip(1).collect();
    ensure!(
        args.len() == 4 || (args.len() == 6 && args[4] == "--held-out-policy"),
        "Supply FIXTURE MODEL_DIR WORKER OUTPUT_JSON [--held-out-policy FROZEN_POLICY]"
    );
    ensure!(
        std::env::var("SEARCHI_FORCE_CPU").as_deref() == Ok("1"),
        "Set SEARCHI_FORCE_CPU=1"
    );
    let fixture_root = PathBuf::from(&args[0]).canonicalize()?;
    let models = PathBuf::from(&args[1]).canonicalize()?;
    let worker = PathBuf::from(&args[2]).canonicalize()?;
    let output = PathBuf::from(&args[3]);
    let policy = if args.len() == 6 {
        let path = PathBuf::from(&args[5]);
        ensure!(
            digest(&path)? == "df0a887b97dd7fda6aeff1b6c966fe7558521f40dbe1c911e91cfd312dca7461",
            "Policy differs from frozen pre-evaluation hypothesis"
        );
        Some(serde_json::from_slice::<Value>(&fs::read(path)?)?)
    } else {
        None
    };
    let split = if policy.is_some() {
        "held_out"
    } else {
        "calibration"
    };
    fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(&output)
        .context("Cannot reserve output; existing evidence is never overwritten")?;
    let mut report = json!({"complete":false,"phase":"verifying","split":split,"held_out_evaluated":policy.is_some(),"policy":policy,"cpu_threads":1,"rows":[]});
    checkpoint(&output, &report)?;
    let fixture_path = fixture_root.join("fixture.json");
    let fixture: Value = serde_json::from_slice(&fs::read(&fixture_path)?)?;
    ensure!(
        digest(&fixture_path)?
            == "000047fd17accd315b6138093bbdc9f353cd877b4bf1958fbe336e4a6b10cd2e",
        "This runner accepts only the reviewed frozen v2 fixture"
    );
    let lock: Value = serde_json::from_slice(&fs::read(fixture_root.join("fixture.lock.json"))?)?;
    ensure!(
        lock["sha256"]["fixture.json"] == digest(&fixture_path)?,
        "Frozen fixture changed"
    );
    let queries = if policy.is_some() {
        split_queries(&fixture, split)?
    } else {
        calibration_queries(&fixture)?
    };
    let manifest_path =
        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../embeddinggemma2/model-manifest.json");
    let manifest: Value = serde_json::from_slice(&fs::read(&manifest_path)?)?;
    for asset in manifest["files"]
        .as_array()
        .context("Missing model assets")?
    {
        let path = models.join(asset["path"].as_str().context("Missing model path")?);
        ensure!(
            fs::metadata(&path)?.len() == asset["bytes"].as_u64().context("Missing model size")?,
            "Model size mismatch"
        );
        ensure!(digest(&path)? == asset["sha256"], "Model hash mismatch");
    }
    report["fixture_sha256"] = json!(digest(&fixture_path)?);
    report["model_manifest_sha256"] = json!(digest(&manifest_path)?);
    report["model"] = manifest;
    report["assets"] = fixture["assets"].clone();
    report["worker_sha256"] = json!(digest(&worker)?);
    report["runner_sha256"] = json!(digest(&std::env::current_exe()?)?);
    report["phase"] = json!("materializing_and_lexical_index");
    checkpoint(&output, &report)?;
    let temporary = tempfile::tempdir()?;
    let corpus = temporary.path().join("corpus");
    fs::create_dir(&corpus)?;
    // Index both splits as the fixed distractor library, execute only the selected split.
    for asset in fixture["assets"].as_array().context("Missing assets")? {
        let source = fixture_root
            .join(asset["path"].as_str().context("Missing asset path")?)
            .canonicalize()?;
        ensure!(
            source.starts_with(&fixture_root),
            "Source escapes frozen fixture"
        );
        ensure!(
            fs::metadata(&source)?.len()
                == asset["bytes"].as_u64().context("Missing asset size")?
                && digest(&source)? == asset["sha256"],
            "Source changed"
        );
        let destination = corpus.join(source.file_name().context("Missing filename")?);
        ensure!(!destination.exists(), "Duplicate materialized filename");
        fs::copy(source, destination)?;
    }
    let db = temporary.path().join("index.sqlite");
    let mut lexical = SearchIndex::open(&db)?;
    let start = Instant::now();
    let lexical_report = lexical.index_folder(&corpus)?;
    let lexical_ms = start.elapsed().as_secs_f64() * 1000.0;
    report["lexical_report"] = json!(lexical_report);
    report["lexical_index_ms"] = json!(lexical_ms);
    report["phase"] = json!("semantic_index");
    checkpoint(&output, &report)?;
    let start = Instant::now();
    let mut indexing = HybridSearch::open(&db, &models, &worker, 1)?;
    let index_report = indexing.index_embeddings()?;
    let semantic_index_ms = start.elapsed().as_secs_f64() * 1000.0;
    report["index_report"] = json!(index_report);
    report["semantic_index_ms"] = json!(semantic_index_ms);
    report["phase"] = json!("query_warmup");
    checkpoint(&output, &report)?;
    ensure!(index_report.failures.is_empty(), "Semantic indexing failed");
    drop(indexing);
    let start = Instant::now();
    let mut query_session = HybridSearch::open_query(&db, &models, &worker, 1)?;
    query_session.warmup()?;
    let query_startup_ms = start.elapsed().as_secs_f64() * 1000.0;
    report["worker_info"] = json!(query_session.worker_info());
    report["query_startup_ms"] = json!(query_startup_ms);
    report["phase"] = json!(format!("{split}_queries"));
    checkpoint(&output, &report)?;
    for query in queries {
        ensure!(query["split"] == split, "Query outside selected split");
        let text = query["query"].as_str().context("Missing query text")?;
        let code_intent = query["id"]
            .as_str()
            .context("Missing query ID")?
            .starts_with(if split == "held_out" {
                "version-"
            } else {
                "pattern-"
            });
        for code_only in [false, true]
            .into_iter()
            .filter(|code_only| !*code_only || code_intent)
        {
            let options = SearchOptions {
                code_only,
                ..Default::default()
            };
            // Match the production hybrid candidate budget for a final limit of 20.
            let lexical_hits = lexical.search_with_options_measured(text, 80, &options, true)?;
            let start = Instant::now();
            let semantic_hits = query_session.search_semantic(text, 80, &options)?;
            let semantic_ms = start.elapsed().as_secs_f64() * 1000.0;
            let hybrid_hits = query_session.search_measured(text, 20, &options, true)?;
            let retained: Vec<_> = semantic_hits
                .iter()
                .filter(|hit| {
                    policy
                        .as_ref()
                        .is_none_or(|p| hit.score >= p["semantic_cosine_minimum"].as_f64().unwrap())
                })
                .collect();
            let scores = json!({"lexical":score(&lexical_hits.results,query,&fixture)?,"semantic":score(&semantic_hits,query,&fixture)?,"hybrid":score(&hybrid_hits.results,query,&fixture)?,"hypothesis_semantic":score(retained,query,&fixture)?});
            report["rows"].as_array_mut().context("Missing output rows")?.push(json!({"query":query,"scope":if code_only {"code"} else {"all_files"},"lexical":lexical_hits,"semantic_raw":semantic_hits,"semantic_ms":semantic_ms,"hybrid":hybrid_hits,"scores":scores,"order_note":"lexical then semantic then hybrid; hybrid reuses the preceding query-vector cache. Raw semantic scores are cosine, hybrid scores are RRF, not confidence."}));
            checkpoint(&output, &report)?;
        }
    }
    report["complete"] = json!(true);
    report["phase"] = json!("complete");
    checkpoint(&output, &report)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn checkpoint_replaces_complete_json_and_existing_evidence_is_reserved() {
        let temporary = tempfile::tempdir().unwrap();
        let path = temporary.path().join("evidence.json");
        fs::OpenOptions::new()
            .write(true)
            .create_new(true)
            .open(&path)
            .unwrap();
        checkpoint(&path, &json!({"complete":false,"rows":[1]})).unwrap();
        assert!(fs::OpenOptions::new()
            .write(true)
            .create_new(true)
            .open(&path)
            .is_err());
        checkpoint(&path, &json!({"complete":true,"rows":[1,2]})).unwrap();
        let report: Value = serde_json::from_slice(&fs::read(path).unwrap()).unwrap();
        assert_eq!(report["rows"], json!([1, 2]));
        assert_eq!(fs::read_dir(temporary.path()).unwrap().count(), 1);
    }
    #[test]
    fn rejects_unknown_split_and_never_selects_held_out() {
        let mut queries = vec![json!({"split":"calibration"}); 24];
        queries.push(json!({"split":"held_out","query":"DO NOT EXECUTE"}));
        let fixture = json!({"queries":queries});
        assert!(calibration_queries(&fixture)
            .unwrap()
            .iter()
            .all(|q| q["split"] == "calibration"));
        assert!(calibration_queries(&json!({"queries":[{"split":"unknown"}]})).is_err());
        assert!(calibration_queries(&json!({"queries":[{"split":"held_out"}]})).is_err());
    }

    #[test]
    fn held_out_selection_is_disjoint_and_has_six_code_requests() {
        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("../../fixtures/relevance-v2/fixture.json");
        let fixture: Value = serde_json::from_slice(&fs::read(path).unwrap()).unwrap();
        let selected = split_queries(&fixture, "held_out").unwrap();
        assert_eq!(selected.len(), 24);
        assert!(selected.iter().all(|q| q["split"] == "held_out"));
        assert_eq!(
            selected
                .iter()
                .filter(|q| q["id"].as_str().unwrap().starts_with("version-"))
                .count(),
            6
        );
        assert!(split_queries(&fixture, "unknown").is_err());
    }

    #[test]
    fn frozen_calibration_has_six_code_requests_and_disjoint_query_ids() {
        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
            .join("../../fixtures/relevance-v2/fixture.json");
        let fixture: Value = serde_json::from_slice(&fs::read(path).unwrap()).unwrap();
        let selected = calibration_queries(&fixture).unwrap();
        assert_eq!(
            selected
                .iter()
                .filter(|q| q["id"].as_str().unwrap().starts_with("pattern-"))
                .count(),
            6
        );
        for query in selected {
            assert!(!fixture["queries"]
                .as_array()
                .unwrap()
                .iter()
                .any(|held| held["split"] == "held_out" && held["id"] == query["id"]));
        }
    }
}