hypersteeldb 0.5.2

//! **S-expression trajectories** — training data that teaches a model to emit *compound IKL* rather than
//! a flat JSON tool call.
//!
//! The paper's contract is that the agent emits an **S-expression AST in IKL** which the engine lints,
//! compiles and executes (§2.2, §6). Earlier training here targeted a JSON `run_workflow` object, which
//! cannot express what IKL can: nested conjunction/disjunction, negation, subtree wildcards and numeric
//! range predicates in one term. This module generates `(question → IKL)` pairs over a corpus's real
//! [`VocabularySpace`], covering exactly those constructs.
//!
//! Every candidate is **verified by the engine before it is kept**: the expression must lint clean against
//! the corpus vocabulary *and* evaluate to a non-empty row set (or, for the deliberate refusal cases, to
//! the expected emptiness). A query that looks plausible but matches nothing teaches the model to
//! hallucinate confidently, so the corpus itself is the oracle.

use crate::db::Corpus;
use crate::linter::Linter;
use crate::vocabulary::VocabularySpace;
use serde::{Deserialize, Serialize};
use serde_json::json;

/// Which IKL construct an example exercises — tracked so coverage is planned rather than hoped for.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Construct {
    /// a single grounded atom
    Atom,
    /// `(and A B)` — the intersection similarity search cannot express
    Conjunction,
    /// `(or A B)`
    Disjunction,
    /// `(and A (not B))` — the set difference
    Negation,
    /// a subtree wildcard, e.g. `org/*`
    Wildcard,
    /// `(num field op value)`
    NumericRange,
    /// nested combination of the above
    Compound,
    /// epistemic filter: exclude refuted assertions
    Epistemic,
    /// a question whose terms are NOT in the vocabulary — the model must refuse, not invent
    Unanswerable,
}

/// One training pair: a natural question and the IKL the engine verified.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IklTraj {
    pub query: String,
    /// the S-expression, or empty for [`Construct::Unanswerable`]
    pub ikl: String,
    pub construct: Construct,
    /// rows the expression matched at generation time (0 for unanswerable)
    pub matched: usize,
    /// for unanswerable cases: the in-vocabulary alternatives to offer back
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub alternatives: Vec<String>,
}

/// The tool a model calls to run IKL — the single entry point, so the model's whole job is writing the
/// expression. Presented to the finetuner in Needle's bare `{name, description, parameters}` shape.
pub fn ikl_tool(spec: &VocabularySpace) -> serde_json::Value {
    let stems = spec.valid_prefixes();
    json!({
        "name": "ikl_query",
        "description": format!(
            "Run an IKL s-expression over the corpus. Grammar: a bare token `facet/value`, a subtree wildcard `facet/*`, `(and …)`, `(or …)`, `(not …)`, and `(num <field> <op> <value>)` with op ge|gt|le|lt|eq|ne. Facet stems: {}.",
            stems.join(", ")
        ),
        "parameters": {
            "type": "object",
            "properties": { "ikl": { "type": "string" } },
            "required": ["ikl"]
        }
    })
}

/// Verify an expression against the corpus: it must lint clean and match at least `min_rows` rows.
/// Returns the matched row count when acceptable.
pub fn verify(corpus: &Corpus, linter: &Linter, ikl: &str, min_rows: usize) -> Option<usize> {
    let report = linter.lint(ikl);
    if !report.ok {
        return None;
    }
    let out = corpus.query(ikl, 1);
    (out.count >= min_rows).then_some(out.count)
}

/// Pick tokens worth building questions from: reasonably-supported, non-noise, grouped by facet.
fn facet_tokens(corpus: &Corpus, facet: &str, k: usize) -> Vec<String> {
    corpus.facet_tokens(facet, k).into_iter().map(|(t, _)| t).collect()
}

fn leaf(t: &str) -> &str {
    t.rsplit('/').next().unwrap_or(t)
}

/// Generate verified `(question → IKL)` pairs over a real corpus. Question phrasing is templated — the
/// point of this data is the *expression*, and a teacher model can paraphrase later without touching the
/// verified target.
pub fn generate(corpus: &Corpus, _spec: &VocabularySpace, per_construct: usize) -> Vec<IklTraj> {
    let linter = corpus.linter();
    let stats = corpus.stats();
    let facets: Vec<String> = stats
        .facets
        .iter()
        .map(|(f, _)| f.clone())
        .filter(|f| f != "src" && f != "state")
        .collect();
    let mut out: Vec<IklTraj> = Vec::new();

    let mut push = |query: String, ikl: String, construct: Construct, alternatives: Vec<String>| {
        let min = if construct == Construct::Unanswerable { 0 } else { 1 };
        if construct == Construct::Unanswerable {
            out.push(IklTraj { query, ikl, construct, matched: 0, alternatives });
            return;
        }
        if let Some(n) = verify(corpus, &linter, &ikl, min) {
            out.push(IklTraj { query, ikl, construct, matched: n, alternatives });
        }
    };

    // ATOM + WILDCARD
    for f in facets.iter().take(per_construct.max(2)) {
        for t in facet_tokens(corpus, f, 2) {
            push(format!("Which records mention {}?", leaf(&t)), t.clone(), Construct::Atom, vec![]);
        }
        push(format!("Show everything with a {f}."), format!("{f}/*"), Construct::Wildcard, vec![]);
    }

    // CONJUNCTION / DISJUNCTION / NEGATION across two facets
    for a in facets.iter().take(3) {
        for b in facets.iter().take(3) {
            if a == b {
                continue;
            }
            let (ta, tb) = (facet_tokens(corpus, a, 1), facet_tokens(corpus, b, 1));
            let (Some(ta), Some(tb)) = (ta.first(), tb.first()) else { continue };
            push(
                format!("Which records have both {} and {}?", leaf(ta), leaf(tb)),
                format!("(and {ta} {tb})"),
                Construct::Conjunction,
                vec![],
            );
            push(
                format!("Show records with either {} or {}.", leaf(ta), leaf(tb)),
                format!("(or {ta} {tb})"),
                Construct::Disjunction,
                vec![],
            );
            push(
                format!("Which {} records are not {}?", leaf(ta), leaf(tb)),
                format!("(and {ta} (not {tb}))"),
                Construct::Negation,
                vec![],
            );
            // COMPOUND: wildcard ∩ token, minus another
            push(
                format!("Any {a} that is {} but not {}?", leaf(ta), leaf(tb)),
                format!("(and {a}/* {ta} (not {tb}))"),
                Construct::Compound,
                vec![],
            );
        }
    }

    // NUMERIC RANGE over the corpus's real numeric fields
    for field in stats.numeric_fields.iter().take(3) {
        for (op, word) in [("gt", "more than"), ("lt", "less than")] {
            push(
                format!("Which records have {field} {word} 100?"),
                format!("(num {field} {op} 100)"),
                Construct::NumericRange,
                vec![],
            );
        }
        if let Some(f) = facets.first() {
            push(
                format!("Show {f} records with {field} above 50."),
                format!("(and {f}/* (num {field} gt 50))"),
                Construct::Compound,
                vec![],
            );
        }
    }

    // EPISTEMIC: exclude refuted assertions — only meaningful if the corpus carries state tokens
    if corpus.has_token("state/negated") {
        for f in facets.iter().take(2) {
            push(
                format!("Which {f} facts are asserted, excluding anything negated?"),
                format!("(and {f}/* (not state/negated))"),
                Construct::Epistemic,
                vec![],
            );
        }
    }

    // UNANSWERABLE: terms deliberately outside the vocabulary. The model must route back with the real
    // alternatives rather than inventing an expression — the behaviour the paper's linter enables.
    let real_facets = linter.facet_names();
    for bogus in ["gene/brca1", "platform/uav", "artifact/power_cube"] {
        let bogus_facet = bogus.split('/').next().unwrap_or(bogus);
        if real_facets.iter().any(|f| f == bogus_facet) {
            continue; // it would actually be answerable here
        }
        push(
            format!("Which records involve {}?", leaf(bogus)),
            String::new(),
            Construct::Unanswerable,
            real_facets.iter().take(6).cloned().collect(),
        );
    }

    out
}

/// Render as Needle finetune rows: `{query, tools, answers}`. An unanswerable question maps to an **empty**
/// `answers` list, which is how Needle represents "no tool applies" — so refusal is learned the same way
/// tool selection is.
pub fn to_needle_jsonl(spec: &VocabularySpace, trajs: &[IklTraj]) -> String {
    let tool = ikl_tool(spec);
    trajs
        .iter()
        .filter_map(|t| {
            let answers = if t.construct == Construct::Unanswerable {
                json!([])
            } else {
                json!([{ "name": "ikl_query", "arguments": { "ikl": t.ikl } }])
            };
            serde_json::to_string(&json!({
                "query": t.query,
                "tools": [tool.clone()],
                "answers": answers,
            }))
            .ok()
            .map(|l| l + "\n")
        })
        .collect()
}

/// Coverage summary for a generated set.
pub fn coverage(trajs: &[IklTraj]) -> std::collections::BTreeMap<String, usize> {
    let mut m = std::collections::BTreeMap::new();
    for t in trajs {
        // use the serde (snake_case) name so keys match the wire format, not Debug's CamelCase
        let key = serde_json::to_value(t.construct).ok().and_then(|v| v.as_str().map(String::from)).unwrap_or_default();
        *m.entry(key).or_insert(0) += 1;
    }
    m
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::projector::CorpusKind;
    use crate::vocabulary::{EntityFacet, VocabularySpace};

    fn spec() -> VocabularySpace {
        VocabularySpace {
            version: 1,
            corpus: "t".into(),
            entity_facets: vec![
                EntityFacet { name: "country".into(), parent: None, description: "".into(), examples: vec![], structural: false },
                EntityFacet { name: "powertrain".into(), parent: None, description: "".into(), examples: vec![], structural: false },
            ],
            relation_facets: vec![],
            gazetteer: vec![],
            metrics: None,
        }
    }

    fn corpus() -> Corpus {
        let mut c = Corpus::new_incremental("t", vec!["row".into()], CorpusKind::Csv);
        let rows = [
            (vec!["country/japan", "powertrain/electric", "state/negated"], 600.0),
            (vec!["country/japan", "powertrain/diesel"], 1100.0),
            (vec!["country/usa", "powertrain/electric"], 650.0),
            (vec!["country/usa", "powertrain/diesel"], 1300.0),
        ];
        for (toks, km) in rows {
            let t: Vec<String> = toks.iter().map(|s| s.to_string()).collect();
            c.add_situation_num(t, vec!["r".into()], vec![("range_km".into(), km)]);
        }
        c
    }

    #[test]
    fn every_generated_expression_is_engine_verified() {
        let (c, s) = (corpus(), spec());
        let trajs = generate(&c, &s, 3);
        assert!(!trajs.is_empty());
        let linter = c.linter();
        for t in trajs.iter().filter(|t| t.construct != Construct::Unanswerable) {
            assert!(linter.lint(&t.ikl).ok, "must lint clean: {}", t.ikl);
            assert!(t.matched >= 1, "must match rows: {} → {}", t.ikl, t.matched);
            assert_eq!(c.query(&t.ikl, 1).count, t.matched, "recorded count must be reproducible: {}", t.ikl);
        }
    }

    #[test]
    fn covers_the_constructs_similarity_search_cannot_express() {
        let (c, s) = (corpus(), spec());
        let cov = coverage(&generate(&c, &s, 3));
        for needed in ["conjunction", "negation", "wildcard", "numeric_range", "compound", "unanswerable"] {
            assert!(cov.get(needed).copied().unwrap_or(0) > 0, "missing {needed}: {cov:?}");
        }
    }

    #[test]
    fn negation_actually_subtracts() {
        let (c, s) = (corpus(), spec());
        let trajs = generate(&c, &s, 3);
        let neg = trajs.iter().find(|t| t.construct == Construct::Negation).expect("a negation example");
        // an AND-NOT must match strictly fewer rows than the bare positive atom it narrows
        let positive = neg.ikl.split_whitespace().nth(1).unwrap().to_string();
        let bare = c.query(&positive, 1).count;
        assert!(neg.matched < bare, "{} ({}) should be narrower than {positive} ({bare})", neg.ikl, neg.matched);
    }

    #[test]
    fn unanswerable_questions_carry_alternatives_and_no_expression() {
        let (c, s) = (corpus(), spec());
        let trajs = generate(&c, &s, 3);
        let u = trajs.iter().find(|t| t.construct == Construct::Unanswerable).expect("an unanswerable example");
        assert!(u.ikl.is_empty(), "no expression may be invented");
        assert!(!u.alternatives.is_empty(), "must offer what DOES exist");
        assert!(u.alternatives.iter().any(|a| a == "country" || a == "powertrain"));
        // and it must serialise as an empty answers list (Needle's "no tool applies")
        let jsonl = to_needle_jsonl(&s, std::slice::from_ref(u));
        let v: serde_json::Value = serde_json::from_str(jsonl.trim()).unwrap();
        assert_eq!(v["answers"].as_array().unwrap().len(), 0);
    }

    #[test]
    fn needle_rows_expose_the_ikl_tool_with_real_stems() {
        let (c, s) = (corpus(), spec());
        let trajs = generate(&c, &s, 2);
        let jsonl = to_needle_jsonl(&s, &trajs);
        let first: serde_json::Value = serde_json::from_str(jsonl.lines().next().unwrap()).unwrap();
        assert_eq!(first["tools"][0]["name"], "ikl_query");
        let desc = first["tools"][0]["description"].as_str().unwrap();
        assert!(desc.contains("country"), "facet stems must be advertised: {desc}");
        assert!(desc.contains("(and"), "grammar must be advertised");
    }
}