hypersteeldb 0.5.2

//! Bedrock-teacher trajectory generator. Builds correct tool trajectories from the real engine, then
//! uses Claude (Sonnet on Bedrock) to rewrite each into NATURAL language — a varied question paraphrase
//! and a concise answer stated strictly from the computed facts. The engine owns correctness; the
//! teacher owns fluency. Batched to keep call count low.
//!
//!   `cargo run --features bedrock --bin gen_trajectories_teacher -- <out.jsonl> [n_synth] [corpus.csv ...]`
//!
//! Model: STEELDB_BEDROCK_MODEL (default us.anthropic.claude-sonnet-4-5), region AWS_REGION/us-west-2.

use serde_json::{json, Value};
use steeldb::agent::{Msg, ProviderConfig};
use steeldb::trajectories::{gen_templated, synth_corpora, tools_json, Traj};
use steeldb::Corpus;

const BATCH: usize = 10;

const SYS: &str = "\
You rewrite mechanical analytics Q&A into natural, VARIED language for training a small retrieval agent. \
Each item has an id, a mechanical QUESTION, and FACTS (the exact answer computed from a corpus by a bitmap \
program). For EACH id return: (1) `question` — a natural paraphrase, same meaning, how a real user would \
ask it; (2) `answer` — a concise, natural answer that states ONLY what FACTS contain, citing every number \
verbatim. NEVER invent values. If FACTS indicate nothing is present or the set drops to zero, answer plainly \
that the corpus does not support it.

CRITICAL — VARY THE ANSWER STYLE across items so the agent doesn't overfit any one template. Rotate freely \
between: a lead sentence + bullet list; a compact prose sentence with parentheticals; 'X leads with N, then \
Y with M, and Z with K'; 'The corpus has N X, M Y, and K Z.'; leading with the top item and giving the rest; \
starting with 'Overall,' or 'Looking at the data,' or 'Sorted by …,'; using dashes, or 'while', or 'followed \
by'. Do NOT start every answer with 'Across N situations,' or the same shape; each answer should read like \
a different person wrote it. Keep it factual and grounded.

Respond with ONLY a JSON array: [{\"id\":<n>,\"question\":\"...\",\"answer\":\"...\"}]";

#[tokio::main]
async fn main() {
    let args: Vec<String> = std::env::args().collect();
    let out_path = match args.get(1) {
        Some(p) => p.clone(),
        None => {
            eprintln!("usage: gen_trajectories_teacher <out.jsonl> [n_synth] [corpus.csv ...]");
            std::process::exit(2);
        }
    };
    let n_synth: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(80);

    let model = std::env::var("STEELDB_BEDROCK_MODEL").unwrap_or_else(|_| "us.anthropic.claude-sonnet-4-5-20250929-v1:0".to_string());
    let region = std::env::var("AWS_REGION").ok().or_else(|| Some("us-west-2".to_string()));
    let provider = match (ProviderConfig::Bedrock { model_id: model.clone(), region }).build().await {
        Ok(p) => p,
        Err(e) => {
            eprintln!("bedrock provider: {e}");
            std::process::exit(1);
        }
    };
    eprintln!("teacher: {model}");

    // gather correct trajectories from the engine
    let mut corpora = synth_corpora(n_synth, 0x51ed2701);
    for path in args.iter().skip(3) {
        match Corpus::from_csv(std::path::Path::new(path)) {
            Ok(c) => corpora.push((path.clone(), c)),
            Err(e) => eprintln!("skip {path}: {e}"),
        }
    }
    let mut trajs: Vec<Traj> = Vec::new();
    for (_n, c) in &corpora {
        trajs.extend(gen_templated(c));
    }
    eprintln!("{} corpora → {} trajectories; enriching with {model} in batches of {BATCH}", corpora.len(), trajs.len());

    let tools = tools_json();
    let total_batches = trajs.len().div_ceil(BATCH);
    let mut examples: Vec<Value> = Vec::new();

    for (bi, chunk) in trajs.chunks(BATCH).enumerate() {
        let items: Vec<Value> = chunk.iter().enumerate().map(|(i, t)| json!({ "id": i, "question": t.question, "facts": t.facts() })).collect();
        let user = format!("Items:\n{}", serde_json::to_string(&items).unwrap_or_default());
        let overrides = match provider.chat(SYS, &[Msg::user_text(user)], &[]).await {
            Ok(turn) => parse_overrides(&turn.text, chunk.len()),
            Err(e) => {
                eprintln!("batch {bi} chat error ({e}); falling back to templated");
                Vec::new()
            }
        };
        for (i, t) in chunk.iter().enumerate() {
            let (q, a) = overrides.get(i).cloned().flatten().unwrap_or((None, None));
            examples.push(t.to_example(&tools, q.as_deref(), a.as_deref()));
        }
        if bi % 10 == 0 || bi + 1 == total_batches {
            eprintln!("  batch {}/{} ({} examples)", bi + 1, total_batches, examples.len());
        }
    }

    let body: String = examples.iter().map(|e| format!("{e}\n")).collect();
    if let Err(e) = std::fs::write(&out_path, body) {
        eprintln!("write {out_path}: {e}");
        std::process::exit(1);
    }
    eprintln!("wrote {} trajectories → {out_path}", examples.len());
}

/// Parse the teacher's JSON array into `[(question?, answer?)]` indexed by item id; robust to fences,
/// reasoning preambles, and partial output.
#[allow(clippy::type_complexity)]
fn parse_overrides(text: &str, n: usize) -> Vec<Option<(Option<String>, Option<String>)>> {
    let mut out: Vec<Option<(Option<String>, Option<String>)>> = vec![None; n];
    let cleaned = text.replace("```json", "```");
    let body = match cleaned.split_once("```") {
        Some((_, rest)) => rest.split_once("```").map(|(b, _)| b).unwrap_or(rest).to_string(),
        None => cleaned,
    };
    let (a, b) = match (body.find('['), body.rfind(']')) {
        (Some(a), Some(b)) if b > a => (a, b),
        _ => return out,
    };
    let arr: Vec<Value> = match serde_json::from_str(&body[a..=b]) {
        Ok(v) => v,
        Err(_) => return out,
    };
    for item in arr {
        if let Some(id) = item.get("id").and_then(|x| x.as_u64()) {
            let id = id as usize;
            if id < n {
                let q = item.get("question").and_then(|x| x.as_str()).map(|s| s.to_string());
                let ans = item.get("answer").and_then(|x| x.as_str()).map(|s| s.to_string());
                out[id] = Some((q, ans));
            }
        }
    }
    out
}