hypersteeldb 0.5.5

//! **Prose synthesis** — rewrite the benchmark corpus's templated documents as natural paragraphs using
//! Bedrock, without changing a single fact.
//!
//!   `cargo run --release --bin synth_prose -- <corpus_dir> [--workers N] [--limit N] [--dry-run]`
//!
//! Why this exists: a templated corpus teaches an extraction pipeline the template. Every battle report
//! shares one sentence skeleton, so a tagger can hit the gold spans by matching position rather than
//! language. Rewriting the same records as varied prose removes that crutch while leaving the benchmark
//! measuring the same thing.
//!
//! **The facts must not move.** The gold answers in `questions.jsonl` are computed from the structured
//! records, not from the document text — so if a paraphrase drops a duration or renames a venue, the gold
//! answer silently becomes wrong and the whole benchmark loses its meaning. Two defences:
//!
//! 1. Each record carries the literals that must survive (`must_contain`, emitted by `gen_benchmark_corpus`).
//!    A rewrite is **verified** against them, not trusted. Numbers are matched on digit boundaries, so
//!    "seven minutes" does not pass for `7` and `2025` does not satisfy a required `7`.
//! 2. A rewrite that fails verification is retried once under a stricter instruction; if it fails again the
//!    original deterministic document is **kept**. The corpus degrades to templated prose, never to wrong
//!    prose.
//!
//! Resumable: completed doc ids are appended to `.synth_done`, so an interrupted run continues where it
//! stopped instead of paying for the same 9,000 calls twice.

use std::collections::HashSet;
use std::io::Write as _;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;

const DEFAULT_MODEL: &str = "us.anthropic.claude-haiku-4-5-20251001-v1:0";
const DEFAULT_REGION: &str = "us-east-1";

struct Record {
    doc: usize,
    kind: String,
    fields: serde_json::Value,
    must_contain: Vec<String>,
    /// phrases that must NOT appear — polarity guards. "permitted" is a substring of "not permitted", so a
    /// required-literal check alone cannot tell a permit from a ban.
    must_not_contain: Vec<String>,
}

#[derive(Default)]
struct Stats {
    ok: AtomicUsize,
    retried: AtomicUsize,
    fallback: AtomicUsize,
    failed_call: AtomicUsize,
    in_tokens: AtomicUsize,
    out_tokens: AtomicUsize,
}

/// Is `needle` present in `hay`? Numbers must match on a digit boundary — a bare substring test lets the
/// `7` inside `2025` satisfy a required duration of 7 minutes, which would pass a document that never
/// states the duration at all.
fn contains_literal(hay: &str, needle: &str) -> bool {
    if needle.is_empty() {
        return true;
    }
    if needle.chars().all(|c| c.is_ascii_digit()) {
        let bytes = hay.as_bytes();
        let mut from = 0;
        while let Some(pos) = hay[from..].find(needle).map(|p| p + from) {
            let before_ok = pos == 0 || !bytes[pos - 1].is_ascii_digit();
            let after = pos + needle.len();
            let after_ok = after >= bytes.len() || !bytes[after].is_ascii_digit();
            if before_ok && after_ok {
                return true;
            }
            from = pos + 1;
        }
        return false;
    }
    hay.to_lowercase().contains(&needle.to_lowercase())
}

fn missing_literals(text: &str, must: &[String]) -> Vec<String> {
    must.iter().filter(|m| !contains_literal(text, m)).cloned().collect()
}

/// Forbidden phrases that turned up anyway — a rewrite containing one has flipped the record's polarity or
/// its epistemic status, which is worse than an omission because it reads as a confident fact.
fn forbidden_present(text: &str, must_not: &[String]) -> Vec<String> {
    must_not.iter().filter(|m| contains_literal(text, m)).cloned().collect()
}

fn verification_failures(text: &str, r: &Record) -> Vec<String> {
    let mut f = missing_literals(text, &r.must_contain);
    f.extend(forbidden_present(text, &r.must_not_contain).into_iter().map(|p| format!("forbidden: {p}")));
    f
}

fn field<'a>(f: &'a serde_json::Value, k: &str) -> String {
    match &f[k] {
        serde_json::Value::String(s) => s.clone(),
        serde_json::Value::Number(n) => n.to_string(),
        serde_json::Value::Bool(b) => b.to_string(),
        _ => String::new(),
    }
}

/// The record, stated plainly. The model's job is only to vary the *language*, so it is handed the facts
/// rather than the templated document — paraphrasing a template tends to preserve its skeleton.
fn describe(r: &Record) -> String {
    let f = &r.fields;
    match r.kind.as_str() {
        "battle" => format!(
            "Battle report. Winner: {}. Loser: {}. Venue: {}. Region: {}. Tournament: {}. Year: {}. \
             Duration: {} minutes. {} used {}. {} used {}. Deciding move: {}.",
            field(f, "winner"),
            if field(f, "winner") == field(f, "trainer_a") { field(f, "trainer_b") } else { field(f, "trainer_a") },
            field(f, "venue"), field(f, "region"), field(f, "tournament"), field(f, "year"),
            field(f, "minutes"),
            field(f, "trainer_a"), field(f, "species_a"),
            field(f, "trainer_b"), field(f, "species_b"),
            field(f, "move")
        ),
        "regulation" => format!(
            "Competition regulation clause. Series: {}. Season: {}. Species ruling: {} is {} in this series.{}",
            field(f, "series"), field(f, "year"), field(f, "species"), field(f, "verdict"),
            if f["hedged"].as_bool().unwrap_or(false) {
                " This ruling is PROVISIONAL: use the exact phrase \"under review\" and present the status \
                 as not yet settled."
            } else {
                " This ruling is SETTLED: state it as final. Do not describe it as provisional, under review, \
                 or subject to change."
            }
        ),
        "survey" => format!(
            "Habitat survey. Species observed: {}. Location: {}. Region: {}. Elevation: {} m. \
             Mean temperature: {} °C.",
            field(f, "species"), field(f, "venue"), field(f, "region"),
            field(f, "elevation_m"), field(f, "temp_c")
        ),
        other => format!("{other} record: {f}"),
    }
}

fn system_prompt(strict: bool) -> String {
    let base = "You rewrite structured Pokémon-league records as natural prose for a document corpus.\n\n\
        Rules:\n\
        - Write ONE paragraph of 55-90 words. No heading, no bullet list, no preamble, no closing remark.\n\
        - Include EVERY fact from the record. Do not omit any name, place, number or move.\n\
        - Write all numbers as digits (7, not seven; 2025, not twenty twenty-five).\n\
        - Keep proper names and region names spelled exactly as given.\n\
        - Use the exact permission wording given (\"permitted\" or \"not permitted\"); do not substitute            \"banned\", \"prohibited\", \"disallowed\" or \"allowed\".\n\
        - Invent NO facts: no scores, no crowd sizes, no commentary on skill, no extra Pokémon or people.\n\
        - Vary sentence structure and vocabulary between documents; avoid a fixed formula.\n\
        - Output only the paragraph.";
    if strict {
        format!("{base}\n\nThe previous attempt omitted required details. Be exhaustive: every number and \
                 every proper name from the record must appear literally in your paragraph.")
    } else {
        base.to_string()
    }
}

/// One Bedrock call. Returns the paragraph text plus (input, output) token counts.
fn invoke(model: &str, region: &str, system: &str, user: &str) -> Result<(String, usize, usize), String> {
    let body = serde_json::json!({
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 400,
        "temperature": 1.0,
        "system": system,
        "messages": [{ "role": "user", "content": user }]
    })
    .to_string();

    let tmp = std::env::temp_dir().join(format!("synth_prose_{}_{:?}.json", std::process::id(), std::thread::current().id()));
    std::fs::write(&tmp, &body).map_err(|e| e.to_string())?;
    let outp = tmp.with_extension("out.json");

    let status = Command::new("aws")
        .args([
            "bedrock-runtime", "invoke-model",
            "--region", region,
            "--model-id", model,
            "--cli-binary-format", "raw-in-base64-out",
            "--body",
        ])
        .arg(format!("fileb://{}", tmp.display()))
        .arg(&outp)
        .output()
        .map_err(|e| format!("spawn aws: {e}"))?;

    let _ = std::fs::remove_file(&tmp);
    if !status.status.success() {
        let err = String::from_utf8_lossy(&status.stderr).trim().to_string();
        let _ = std::fs::remove_file(&outp);
        return Err(err.chars().take(200).collect());
    }
    let raw = std::fs::read_to_string(&outp).map_err(|e| e.to_string())?;
    let _ = std::fs::remove_file(&outp);
    let v: serde_json::Value = serde_json::from_str(&raw).map_err(|e| e.to_string())?;
    let text = v["content"][0]["text"].as_str().unwrap_or_default().trim().to_string();
    if text.is_empty() {
        return Err("empty completion".into());
    }
    let it = v["usage"]["input_tokens"].as_u64().unwrap_or(0) as usize;
    let ot = v["usage"]["output_tokens"].as_u64().unwrap_or(0) as usize;
    Ok((text, it, ot))
}

/// Strip any heading or list the model added despite instructions, keeping the prose body.
fn clean(text: &str) -> String {
    let body: Vec<&str> = text
        .lines()
        .map(|l| l.trim())
        .filter(|l| !l.is_empty() && !l.starts_with('#') && !l.starts_with('-') && !l.starts_with('*'))
        .collect();
    if body.is_empty() { text.trim().to_string() } else { body.join(" ") }
}

fn doc_path(dir: &Path, doc: usize, kind: &str) -> PathBuf {
    dir.join(format!("{doc:06}_{kind}.md"))
}

fn title(r: &Record) -> String {
    let f = &r.fields;
    match r.kind.as_str() {
        "battle" => format!("Battle Report — {} vs {}", field(f, "trainer_a"), field(f, "trainer_b")),
        "regulation" => format!("Competition Regulations — Series {} {}", field(f, "series"), field(f, "year")),
        "survey" => format!("Habitat Survey — {} ({})", field(f, "venue"), field(f, "region")),
        other => other.to_string(),
    }
}

fn main() -> Result<(), String> {
    let args: Vec<String> = std::env::args().collect();
    let dir = PathBuf::from(args.get(1).cloned().unwrap_or_else(|| "benchmark_corpus".to_string()));
    let flag = |name: &str| -> Option<String> {
        args.iter().position(|a| a == name).and_then(|i| args.get(i + 1)).cloned()
    };
    let workers: usize = flag("--workers").and_then(|s| s.parse().ok()).unwrap_or(16);
    let limit: usize = flag("--limit").and_then(|s| s.parse().ok()).unwrap_or(usize::MAX);
    let model = flag("--model").unwrap_or_else(|| DEFAULT_MODEL.to_string());
    let region = flag("--region").unwrap_or_else(|| DEFAULT_REGION.to_string());
    let dry = args.iter().any(|a| a == "--dry-run");

    let recs_path = dir.join("records.jsonl");
    let body = std::fs::read_to_string(&recs_path)
        .map_err(|e| format!("{}: {e} — run gen_benchmark_corpus first", recs_path.display()))?;

    // resume: never pay twice for the same document
    let done_path = dir.join(".synth_done");
    let done: HashSet<usize> = std::fs::read_to_string(&done_path)
        .unwrap_or_default()
        .lines()
        .filter_map(|l| l.trim().parse().ok())
        .collect();

    let mut records: Vec<Record> = Vec::new();
    for line in body.lines().filter(|l| !l.trim().is_empty()) {
        let v: serde_json::Value = match serde_json::from_str(line) {
            Ok(v) => v,
            Err(_) => continue,
        };
        let doc = v["doc"].as_u64().unwrap_or(0) as usize;
        if done.contains(&doc) {
            continue;
        }
        records.push(Record {
            doc,
            kind: v["kind"].as_str().unwrap_or("").to_string(),
            fields: v["fields"].clone(),
            must_contain: v["must_contain"]
                .as_array()
                .map(|a| a.iter().filter_map(|x| x.as_str().map(String::from)).collect())
                .unwrap_or_default(),
            must_not_contain: v["must_not_contain"]
                .as_array()
                .map(|a| a.iter().filter_map(|x| x.as_str().map(String::from)).collect())
                .unwrap_or_default(),
        });
    }
    records.truncate(limit);

    println!(
        "synth_prose: {} records pending ({} already done), {workers} workers, model {model}",
        records.len(),
        done.len()
    );
    if records.is_empty() {
        println!("nothing to do");
        return Ok(());
    }
    if dry {
        for r in records.iter().take(3) {
            println!("\n--- doc {} ({}) ---\n{}\nmust_contain: {:?}", r.doc, r.kind, describe(r), r.must_contain);
        }
        println!("\n(dry run: no Bedrock calls made)");
        return Ok(());
    }

    let total = records.len();
    let stats = Stats::default();
    let next = AtomicUsize::new(0);
    let done_log = Mutex::new(
        std::fs::OpenOptions::new().create(true).append(true).open(&done_path).map_err(|e| e.to_string())?,
    );
    let started = std::time::Instant::now();

    std::thread::scope(|scope| {
        for _ in 0..workers.max(1) {
            scope.spawn(|| loop {
                let i = next.fetch_add(1, Ordering::SeqCst);
                if i >= total {
                    return;
                }
                let r = &records[i];
                let user = describe(r);

                let mut written: Option<String> = None;
                for attempt in 0..2 {
                    match invoke(&model, &region, &system_prompt(attempt > 0), &user) {
                        Ok((text, it, ot)) => {
                            stats.in_tokens.fetch_add(it, Ordering::Relaxed);
                            stats.out_tokens.fetch_add(ot, Ordering::Relaxed);
                            let text = clean(&text);
                            if verification_failures(&text, r).is_empty() {
                                written = Some(text);
                                break;
                            }
                            if attempt == 0 {
                                stats.retried.fetch_add(1, Ordering::Relaxed);
                            }
                        }
                        Err(e) => {
                            if attempt == 1 {
                                stats.failed_call.fetch_add(1, Ordering::Relaxed);
                                eprintln!("  doc {}: bedrock error: {e}", r.doc);
                            }
                            std::thread::sleep(std::time::Duration::from_millis(500 * (attempt as u64 + 1)));
                        }
                    }
                }

                match written {
                    Some(text) => {
                        let doc_body = format!("# {}\n\n{}\n", title(r), text);
                        if std::fs::write(doc_path(&dir, r.doc, &r.kind), doc_body).is_ok() {
                            stats.ok.fetch_add(1, Ordering::Relaxed);
                        }
                    }
                    // keep the deterministic document: templated prose beats wrong prose
                    None => {
                        stats.fallback.fetch_add(1, Ordering::Relaxed);
                    }
                }

                if let Ok(mut f) = done_log.lock() {
                    let _ = writeln!(f, "{}", r.doc);
                }
                let n = stats.ok.load(Ordering::Relaxed) + stats.fallback.load(Ordering::Relaxed);
                if n % 250 == 0 {
                    let rate = n as f64 / started.elapsed().as_secs_f64().max(0.001);
                    let eta = (total - n.min(total)) as f64 / rate.max(0.001);
                    println!("  {n}/{total}  {rate:.1} docs/s  eta {:.0}m", eta / 60.0);
                }
            });
        }
    });

    let (ok, fb, rt, fc) = (
        stats.ok.load(Ordering::Relaxed),
        stats.fallback.load(Ordering::Relaxed),
        stats.retried.load(Ordering::Relaxed),
        stats.failed_call.load(Ordering::Relaxed),
    );
    let (it, ot) = (stats.in_tokens.load(Ordering::Relaxed), stats.out_tokens.load(Ordering::Relaxed));
    println!("\nsynthesised {ok}/{total} documents in {:.1} min", started.elapsed().as_secs_f64() / 60.0);
    println!("  retried after a missing fact: {rt}");
    println!("  kept the template (verification never passed): {fb}");
    println!("  bedrock call failures: {fc}");
    println!("  tokens: {it} in / {ot} out");
    if fb > 0 {
        println!("\n{fb} document(s) stayed templated. The gold answers remain exact either way.");
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn numbers_must_match_on_digit_boundaries() {
        // the failure this guards: "2025" contains "2", and "seven minutes" states no digit at all
        assert!(!contains_literal("the 2025 season", "2"));
        assert!(!contains_literal("it ran seven minutes", "7"));
        assert!(contains_literal("it ran 7 minutes", "7"));
        assert!(contains_literal("2369 m elevation", "2369"));
        assert!(!contains_literal("at 12369 m", "2369"));
        assert!(contains_literal("held in 2025", "2025"));
    }

    #[test]
    fn names_match_case_insensitively() {
        // models routinely re-capitalise region names; that is a surface change, not a lost fact
        assert!(contains_literal("in the Johto region", "johto"));
        assert!(contains_literal("Morty Shade won", "Morty Shade"));
        assert!(!contains_literal("Morty Shade won", "Falkner Gale"));
    }

    #[test]
    fn missing_literals_are_reported() {
        let must: Vec<String> = ["Registeel", "7", "Ecruteak City"].iter().map(|s| s.to_string()).collect();
        let text = "Registeel fought at Ecruteak City for seven minutes.";
        assert_eq!(missing_literals(text, &must), vec!["7".to_string()]);
        assert!(missing_literals("Registeel at Ecruteak City, 7 minutes", &must).is_empty());
    }

    #[test]
    fn a_permit_rendered_as_a_ban_is_rejected() {
        // the substring trap: "not permitted" also contains "permitted", so required literals alone pass
        let r = Record {
            doc: 1,
            kind: "regulation".into(),
            fields: serde_json::json!({ "species": "Metagross", "series": 4, "year": 2025 }),
            must_contain: vec!["Metagross".into(), "4".into(), "2025".into(), "permitted".into()],
            must_not_contain: vec!["not permitted".into()],
        };
        let flipped = "In Series 4 of 2025, Metagross is not permitted.";
        assert!(missing_literals(flipped, &r.must_contain).is_empty(), "required literals alone are fooled");
        assert!(!verification_failures(flipped, &r).is_empty(), "polarity flip must be caught");

        let correct = "In Series 4 of 2025, Metagross is permitted for competition.";
        assert!(verification_failures(correct, &r).is_empty());
    }

    #[test]
    fn a_settled_ruling_may_not_read_as_provisional() {
        let r = Record {
            doc: 2,
            kind: "regulation".into(),
            fields: serde_json::json!({}),
            must_contain: vec!["Regice".into(), "not permitted".into()],
            must_not_contain: vec!["under review".into(), "provisional".into()],
        };
        assert!(!verification_failures("Regice is not permitted, though it remains under review.", &r).is_empty());
        assert!(verification_failures("Regice is not permitted in this series.", &r).is_empty());
    }

    #[test]
    fn clean_strips_headings_and_bullets() {
        let t = "# Battle Report\n\nMorty Shade won in 7 minutes.\n- extra bullet";
        assert_eq!(clean(t), "Morty Shade won in 7 minutes.");
    }

    #[test]
    fn describe_names_the_loser_correctly() {
        let r = Record {
            doc: 1,
            kind: "battle".into(),
            fields: serde_json::json!({
                "trainer_a": "A One", "trainer_b": "B Two", "winner": "B Two", "venue": "V",
                "region": "kanto", "tournament": "T", "year": 2025, "minutes": 9,
                "species_a": "Registeel", "species_b": "Gengar", "move": "Surf"
            }),
            must_contain: vec![],
            must_not_contain: vec![],
        };
        let d = describe(&r);
        assert!(d.contains("Winner: B Two"), "{d}");
        assert!(d.contains("Loser: A One"), "{d}");
    }
}