hypersteeldb 0.2.3

//! DSL finetuning-trajectory generator — produce `(question → run_workflow DSL)` pairs for finetuning
//! Needle to drive SteelDB. Every pair is instantiated from a corpus's REAL facets/values and then
//! VERIFIED by executing the DSL against that corpus (the engine is the oracle): a pair is kept only if
//! its workflow compiles and returns a non-degenerate result. Output is Needle's training format —
//! `{query, tools, target}` where `target` is the tool-call JSON the model must emit.
//!
//! This is the tractable, high-value half of the finetuning harness: pure Rust, reuses the engine, and
//! the emitted data is exactly what the candle trainer (next) consumes.

use crate::agent::tools::engine_tools;
use crate::agent::workflow;
use crate::db::Corpus;
use serde_json::{json, Value};

/// One finetuning example: a natural-language query and the gold DSL tool call it should produce.
#[derive(Clone, Debug)]
pub struct DslTraj {
    pub query: String,
    pub dsl: Value, // the run_workflow `arguments` object
    pub kind: &'static str,
}

impl DslTraj {
    /// Needle training line: `{query, tools, target}`. `tools` is the run_workflow schema (with this
    /// corpus's facet enums); `target` is the tool-call array the decoder must emit.
    pub fn to_needle_example(&self, tools_json: &str) -> Value {
        let target = json!([{ "name": "run_workflow", "arguments": self.dsl }]);
        json!({ "query": self.query, "tools": tools_json, "target": target.to_string(), "kind": self.kind })
    }
}

/// The run_workflow tool schema for a corpus (facet enums bound to its real facets), serialised as the
/// JSON-string array Needle expects.
pub fn needle_tools_json(corpus: &Corpus) -> String {
    let facets = corpus.facet_names();
    let rw = engine_tools(&facets).into_iter().find(|t| t.name == "run_workflow");
    match rw {
        Some(t) => Value::Array(vec![json!({ "name": t.name, "description": t.description, "parameters": t.schema })]).to_string(),
        None => "[]".to_string(),
    }
}

/// True if a program result is worth training on (compiled + non-empty).
fn nondegenerate(result: &Value, kind: &str) -> bool {
    match kind {
        "crosstab" => result.get("matrix").and_then(|m| m.as_array()).map(|a| a.len() >= 2).unwrap_or(false),
        "breakdown" | "constrained_breakdown" => {
            result.get("partition").and_then(|p| p.as_array()).map(|a| !a.is_empty()).unwrap_or(false)
        }
        "rank" => result.get("ranked").and_then(|r| r.as_array()).map(|a| !a.is_empty()).unwrap_or(false),
        "cooccurs" => result.get("cooccurs").and_then(|c| c.as_array()).map(|a| !a.is_empty()).unwrap_or(false),
        "profile" => result.get("sections").and_then(|s| s.as_array()).map(|a| !a.is_empty()).unwrap_or(false),
        _ => result.get("error").is_none(),
    }
}

/// Verify a candidate DSL against the corpus; keep it only if it runs cleanly and returns signal.
fn verify(corpus: &Corpus, dsl: &Value, kind: &'static str) -> bool {
    let (result, is_error) = workflow::execute(corpus, dsl, &[]);
    if is_error {
        return false;
    }
    serde_json::from_str::<Value>(&result).map(|v| nondegenerate(&v, kind)).unwrap_or(false)
}

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

/// Generate verified DSL trajectories from one corpus, over its real facets/values, with varied phrasing.
pub fn gen(corpus: &Corpus) -> Vec<DslTraj> {
    let facets: Vec<String> = corpus.facet_names().into_iter().filter(|f| f != "src").collect();
    let mut out = Vec::new();
    let mut push = |query: String, dsl: Value, kind: &'static str| {
        if verify(corpus, &dsl, kind) {
            out.push(DslTraj { query, dsl, kind });
        }
    };

    // representative value(s) per facet
    let vals = |f: &str| -> Vec<String> {
        corpus.facet_tokens(f, 3).into_iter().map(|(t, _)| leaf(&t).to_string()).collect()
    };

    for f in &facets {
        // breakdown
        for q in [format!("What is the distribution of {f}?"), format!("Break the records down by {f}."), format!("How many records are there per {f}?")] {
            push(q, json!({ "program": "breakdown", "facet": f }), "breakdown");
        }
        // rank
        push(format!("Which {f} values are most common?"), json!({ "program": "rank", "facet": f }), "rank");

        // constraints-only (auto-profile) + cooccurs + constrained breakdown, per a sample value
        for v in vals(f).into_iter().take(2) {
            let tok = format!("{f}/{v}");
            for q in [format!("Tell me about the {v} records."), format!("Summarise the {v} {f}.")] {
                push(q, json!({ "constraints": [tok] }), "profile");
            }
            push(format!("What is associated with {v}?"), json!({ "program": "cooccurs", "token": tok }), "cooccurs");
            // constrained breakdown by another facet
            if let Some(g) = facets.iter().find(|g| *g != f) {
                for q in [format!("Among {v} records, break down by {g}."), format!("For {v}, what is the {g} distribution?")] {
                    push(q, json!({ "constraints": [tok.clone()], "program": "breakdown", "facet": g }), "constrained_breakdown");
                }
            }
        }
    }

    // crosstab over ordered facet pairs
    for a in &facets {
        for b in &facets {
            if a == b {
                continue;
            }
            for q in [format!("For each {a}, which {b} is most common?"), format!("How does {b} vary across {a}?")] {
                push(q, json!({ "program": "crosstab", "facet_a": a, "facet_b": b }), "crosstab");
            }
        }
    }

    out
}