hypersteeldb 0.3.2

//! Replay held-out trajectories against any OpenAI-compatible provider (stock ollama, tuned server,
//! Bedrock) and score two things:
//!   1. **Tool-selection accuracy** — Jaccard over the multiset of programs the model chose vs the gold
//!      program sequence in the trajectory.
//!   2. **Answer grounding** — fraction of gold numeric facts (integers + string tokens) that appear
//!      verbatim in the model's final answer. Punishes "corpus does not support this" hedging when the
//!      numbers were actually retrieved.
//!
//!   `cargo run --features paddock,bedrock --bin eval_agent -- <trajectories.jsonl> [n_eval] [seed]`
//!
//! Provider selection via the usual env: STEELDB_LLM, STEELDB_PADDOCK_URL, STEELDB_PADDOCK_MODEL,
//! STEELDB_BEDROCK_MODEL, AWS_REGION.

use serde_json::Value;
use std::collections::BTreeSet;
use steeldb::agent::{run_agent, ProviderConfig};
use steeldb::trajectories::{gen_templated, synth_corpora};
use steeldb::Corpus;

#[tokio::main]
async fn main() {
    let args: Vec<String> = std::env::args().collect();
    let path = args.get(1).cloned().unwrap_or_else(|| "/tmp/traj_1k.jsonl".to_string());
    let n_eval: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(20);
    let seed: u64 = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(0);

    // Load the full file in order (the teacher wrote trajectories in the same order it generated them),
    // then hold out a random slice deterministically for eval.
    let lines_all: Vec<String> = match std::fs::read_to_string(&path) {
        Ok(s) => s.lines().filter(|l| !l.trim().is_empty()).map(String::from).collect(),
        Err(e) => {
            eprintln!("read {path}: {e}");
            std::process::exit(1);
        }
    };

    // Rebuild the synth corpora at the same seed the teacher used, so trajectory index → (corpus,
    // templated-trajectory) is aligned. Any extras (vehicles.csv etc.) sit at the end; we cap eval to
    // the synth region so alignment is guaranteed.
    let mut corpora: Vec<Corpus> = Vec::new();
    let mut trajs_by_idx: Vec<usize> = Vec::new();
    for (_n, c) in synth_corpora(95, 0x51ed2701) {
        let n = gen_templated(&c).len();
        let ci = corpora.len();
        corpora.push(c);
        for _ in 0..n {
            trajs_by_idx.push(ci);
        }
    }
    let usable = trajs_by_idx.len().min(lines_all.len());

    // Random hold-out from the usable prefix
    let mut order: Vec<usize> = (0..usable).collect();
    let mut s = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
    let mut rnd = |m: usize| {
        s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
        (s >> 33) as usize % m.max(1)
    };
    for i in (1..order.len()).rev() {
        order.swap(i, rnd(i + 1));
    }
    order.truncate(n_eval);
    let lines: Vec<(usize, &String)> = order.iter().map(|&i| (i, &lines_all[i])).collect();

    // Provider under test.
    let cfg = ProviderConfig::from_env();
    eprintln!("provider config: {cfg:?}");
    let provider = match cfg.build().await {
        Ok(p) => p,
        Err(e) => {
            eprintln!("provider error: {e}");
            std::process::exit(1);
        }
    };
    eprintln!("provider: {}\nevaluating {} held-out trajectories from {path}\n", provider.name(), lines.len());

    let mut n = 0usize;
    let mut sum_tool_j = 0.0f64;
    let mut sum_ground = 0.0f64;
    let mut correct_top_program = 0usize;
    let mut skipped = 0usize;
    let mut per_program: std::collections::BTreeMap<String, (f64, f64, usize)> = std::collections::BTreeMap::new();

    for (i, (traj_idx, l)) in lines.iter().enumerate() {
        let ex: Value = match serde_json::from_str(l) {
            Ok(v) => v,
            Err(_) => {
                skipped += 1;
                continue;
            }
        };
        let question = ex["messages"][1]["content"].as_str().unwrap_or("");
        let (gold_programs, gold_facts) = extract_gold(&ex);
        let ci = trajs_by_idx[*traj_idx];
        let corpus = &corpora[ci];

        let ans = match run_agent(provider.as_ref(), corpus, question, 12).await {
            Ok(a) => a,
            Err(e) => {
                eprintln!("[{i}] agent error: {e}");
                skipped += 1;
                continue;
            }
        };
        let called: Vec<String> = ans.trace.iter().map(|t| t.name.clone()).collect();
        let j = jaccard_multiset(&called, &gold_programs);
        let g = grounding(&ans.answer, &gold_facts);
        let top_ok = ans.trace.iter().any(|t| gold_programs.last().is_some_and(|g| *g == t.name));

        n += 1;
        sum_tool_j += j;
        sum_ground += g;
        if top_ok {
            correct_top_program += 1;
        }
        if let Some(gp) = gold_programs.last() {
            let e = per_program.entry(gp.clone()).or_insert((0.0, 0.0, 0));
            e.0 += j;
            e.1 += g;
            e.2 += 1;
        }

        let short = if question.len() > 60 { format!("{}…", &question[..60]) } else { question.to_string() };
        eprintln!("[{i:>3}] tool_j={j:.2} ground={g:.2} top={} · {short}", if top_ok { "✓" } else { "✗" });
    }

    if n == 0 {
        eprintln!("\nno examples scored ({skipped} skipped)");
        std::process::exit(1);
    }
    println!("\n════════════ eval summary ════════════");
    println!("provider:        {}", provider.name());
    println!("scored:          {} (skipped {})", n, skipped);
    println!("tool jaccard:    {:.3}  (calls-multiset overlap with gold)", sum_tool_j / n as f64);
    println!("answer grounded: {:.3}  (fraction of gold facts stated verbatim in the answer)", sum_ground / n as f64);
    println!("top-program hit: {:.3}  (final gold program appeared in tool trace)", correct_top_program as f64 / n as f64);
    if !per_program.is_empty() {
        println!("\nby target program:");
        for (p, (tj, gr, k)) in &per_program {
            println!("  {p:<10} n={k:<3}  tool_j={:.2}  ground={:.2}", tj / *k as f64, gr / *k as f64);
        }
    }
}

/// Pull the gold program sequence + the numeric/token facts from the trajectory.
fn extract_gold(ex: &Value) -> (Vec<String>, BTreeSet<String>) {
    let mut programs = Vec::new();
    let mut facts = BTreeSet::new();
    let empty = Vec::new();
    let msgs = ex["messages"].as_array().unwrap_or(&empty);
    for m in msgs {
        // gold tool calls (assistant with tool_calls) → program sequence
        if let Some(tcs) = m.get("tool_calls").and_then(|v| v.as_array()) {
            for tc in tcs {
                if let Some(name) = tc["function"]["name"].as_str() {
                    programs.push(name.to_string());
                }
            }
        }
        // tool results → mine facts (counts, tokens, values, situation totals)
        if m["role"].as_str() == Some("tool") {
            if let Some(content) = m["content"].as_str() {
                if let Ok(v) = serde_json::from_str::<Value>(content) {
                    walk_facts(&v, &mut facts);
                }
            }
        }
    }
    (programs, facts)
}

/// Collect the numeric strings ≥2 chars and the leaf portion of `facet/value` tokens found anywhere in
/// a tool-result JSON value. These are the concrete facts an answer must cite to be grounded.
fn walk_facts(v: &Value, out: &mut BTreeSet<String>) {
    match v {
        Value::Number(n) => {
            let s = n.to_string();
            if s.len() >= 2 && n.as_u64().unwrap_or(0) > 0 {
                out.insert(s);
            }
        }
        Value::String(s) => {
            if let Some((_, tail)) = s.split_once('/') {
                if tail.len() >= 2 {
                    out.insert(tail.to_string());
                }
            } else if s.len() >= 3 && s.chars().any(|c| c.is_alphanumeric()) {
                // short strings only if they look like facet-value leaves already
                out.insert(s.clone());
            }
        }
        Value::Array(a) => a.iter().for_each(|x| walk_facts(x, out)),
        Value::Object(o) => o.values().for_each(|x| walk_facts(x, out)),
        _ => {}
    }
}

/// Jaccard over multisets of program names — counts matter (calling `breakdown` twice ≠ once).
fn jaccard_multiset(a: &[String], b: &[String]) -> f64 {
    let mut ma: std::collections::HashMap<&String, u32> = std::collections::HashMap::new();
    for x in a {
        *ma.entry(x).or_default() += 1;
    }
    let mut mb: std::collections::HashMap<&String, u32> = std::collections::HashMap::new();
    for x in b {
        *mb.entry(x).or_default() += 1;
    }
    let mut inter = 0u32;
    let mut union = 0u32;
    for k in ma.keys().chain(mb.keys()).collect::<BTreeSet<_>>() {
        let (x, y) = (ma.get(k).copied().unwrap_or(0), mb.get(k).copied().unwrap_or(0));
        inter += x.min(y);
        union += x.max(y);
    }
    if union == 0 {
        1.0
    } else {
        inter as f64 / union as f64
    }
}

/// Fraction of gold facts that appear verbatim (case-insensitive substring) in the answer text.
fn grounding(answer: &str, facts: &BTreeSet<String>) -> f64 {
    if facts.is_empty() {
        return 1.0;
    }
    let low = answer.to_lowercase();
    let hit = facts.iter().filter(|f| low.contains(&f.to_lowercase())).count();
    hit as f64 / facts.len() as f64
}