steeldb/agent/

workflow.rs

//! The bitmap-workflow DSL. Needle emits a compact workflow in one shot — WHICH constraints to apply to
//! the SPLADE-expanded candidate rows, and WHICH bitmap program to run — and it compiles here to IKL
//! set-algebra over the roaring index. This is the reasoning contract SteelDB drives on:
//!
//!   query → SPLADE/entity-link expansion (candidate rows = `scope`)
//!         → Needle DSL {constraints, program, args}
//!         → compile: refined = scope ∩ constraint₁ ∩ … (roaring AND)
//!         → bitmap program over `refined`
//!         → deterministic template (crate::agent::synth)
//!
//! Constraints are IKL atoms/expressions (`powertrain/electric`, `(num range_km gt 500)`, `(not …)`),
//! AND-ed together and with the candidate scope. The compiled anchor is a plain IKL string the existing
//! programs already accept, so the DSL "compiles to the bitmap program" with no new engine.

use crate::db::Corpus;
use serde_json::{json, Value};

/// Compile the candidate `scope` (tokens from the query's SPLADE/entity-link expansion) plus the DSL
/// `constraints` into a single IKL anchor: `(and (or scope…) constraint₁ …)`. Empty → whole corpus.
pub fn compile_anchor(scope: &[String], constraints: &[String]) -> String {
    let mut parts: Vec<String> = Vec::new();
    let scope: Vec<&String> = scope.iter().filter(|s| !s.trim().is_empty()).collect();
    match scope.len() {
        0 => {}
        1 => parts.push(scope[0].clone()),
        _ => parts.push(format!("(or {})", scope.iter().map(|s| s.as_str()).collect::<Vec<_>>().join(" "))),
    }
    for c in constraints {
        let c = c.trim();
        if !c.is_empty() {
            parts.push(c.to_string());
        }
    }
    match parts.len() {
        0 => "*".to_string(),
        1 => parts.remove(0),
        _ => format!("(and {})", parts.join(" ")),
    }
}

fn str_field<'a>(v: &'a Value, k: &str) -> &'a str {
    v.get(k).and_then(|x| x.as_str()).unwrap_or("")
}

/// Normalise to lowercase alphanumerics — for tolerant matching of model-supplied names to real ones.
fn norm(s: &str) -> String {
    s.chars().filter(|c| c.is_alphanumeric()).flat_map(|c| c.to_lowercase()).collect()
}

/// Resolve a model-supplied facet name to a REAL corpus facet (exact, else case/punctuation-insensitive).
/// This guards against the model putting a program name or a mis-cased/rephrased word in a facet slot.
fn resolve_facet(facets: &[String], name: &str) -> Option<String> {
    if name.is_empty() {
        return None;
    }
    if facets.iter().any(|f| f == name) {
        return Some(name.to_string());
    }
    let n = norm(name);
    facets.iter().find(|f| norm(f) == n).cloned()
}

/// Guard a constraint against the real extracted tokens. Passes IKL expressions `(…)` through; keeps a
/// token that exists; otherwise resolves `facet/value` to the actual token by matching the value under
/// the resolved facet. Returns a clear error (fed back to the model) when it can't be grounded.
fn resolve_constraint(corpus: &Corpus, facets: &[String], c: &str) -> Result<String, String> {
    let c = c.trim();
    if c.is_empty() {
        return Err("empty constraint".into());
    }
    if c.starts_with('(') {
        return Ok(c.to_string()); // IKL expression (e.g. numeric predicate) — trust as-is
    }
    if corpus.has_token(c) {
        return Ok(c.to_string());
    }
    if let Some((f, v)) = c.split_once('/') {
        let rf = resolve_facet(facets, f).ok_or_else(|| format!("constraint facet '{f}' is not in the corpus"))?;
        let nv = norm(v);
        for (tok, _) in corpus.facet_tokens(&rf, 100_000) {
            let leaf = tok.split_once('/').map(|x| x.1).unwrap_or(&tok);
            if norm(leaf) == nv {
                return Ok(tok);
            }
        }
        Err(format!("no value '{v}' under facet '{rf}'"))
    } else {
        Err(format!("constraint '{c}' must be a facet/value token"))
    }
}

/// Error payload that hands the model the real schema so it can retry with valid names.
fn bad_facet(field: &str, got: &str, facets: &[String]) -> (String, bool) {
    (json!({ "error": format!("{field} '{got}' is not a corpus facet"), "valid_facets": facets }).to_string(), true)
}

fn usize_field(v: &Value, k: &str, default: usize) -> usize {
    v.get(k).and_then(|x| x.as_u64()).map(|n| n as usize).unwrap_or(default)
}

/// Execute a workflow DSL against the corpus. `scope` = candidate tokens from query expansion (may be
/// empty → whole corpus). Returns `(result_json, is_error)`; the result carries a `"program"` key so
/// `crate::agent::synth` renders it deterministically.
pub fn execute(corpus: &Corpus, dsl: &Value, scope: &[String]) -> (String, bool) {
    let raw_constraints: Vec<String> = dsl
        .get("constraints")
        .and_then(|c| c.as_array())
        .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
        .unwrap_or_default();
    let program = str_field(dsl, "program");
    let k = usize_field(dsl, "k", 8);

    // GUARD: validate/resolve all placeholders against the real extracted schema before running. Facet
    // slots must be real facet names; constraints must ground to real tokens. Failures return the schema
    // so the model can retry — this is what stops a program name (e.g. "narrow") landing in a facet slot.
    let facets = corpus.facet_names();
    let constraints: Vec<String> = {
        let mut out = Vec::with_capacity(raw_constraints.len());
        for c in &raw_constraints {
            match resolve_constraint(corpus, &facets, c) {
                Ok(t) => out.push(t),
                Err(e) => return (json!({ "error": format!("constraint {e}"), "valid_facets": facets }).to_string(), true),
            }
        }
        out
    };
    let anchor = compile_anchor(scope, &constraints);

    match program {
        "crosstab" => {
            let (fa, fb) = (str_field(dsl, "facet_a"), str_field(dsl, "facet_b"));
            let Some(fa) = resolve_facet(&facets, fa) else { return bad_facet("facet_a", fa, &facets) };
            let Some(fb) = resolve_facet(&facets, fb) else { return bad_facet("facet_b", fb, &facets) };
            (with_dsl(corpus.crosstab(&anchor, &fa, &fb, k), &anchor, &constraints), false)
        }
        "breakdown" => {
            let facet = str_field(dsl, "facet");
            let Some(facet) = resolve_facet(&facets, facet) else { return bad_facet("facet", facet, &facets) };
            (with_dsl(corpus.breakdown(&anchor, &facet, k.max(15)), &anchor, &constraints), false)
        }
        "rank" => {
            let facet = str_field(dsl, "facet");
            let Some(facet) = resolve_facet(&facets, facet) else { return bad_facet("facet", facet, &facets) };
            // rank is global salience (not scoped); constraints don't apply.
            (with_dsl(corpus.rank(&facet, k), &anchor, &constraints), false)
        }
        "cooccurs" => {
            let token = str_field(dsl, "token");
            if !corpus.has_token(token) {
                return (json!({ "error": format!("cooccurs token '{token}' is not in the corpus"), "valid_facets": facets }).to_string(), true);
            }
            (with_dsl(corpus.cooccurs(token, k), &anchor, &constraints), false)
        }
        "s_path" => {
            let a = match resolve_constraint(corpus, &facets, str_field(dsl, "a")) {
                Ok(t) => t,
                Err(e) => return (json!({ "error": format!("s_path a: {e}") }).to_string(), true),
            };
            let b = match resolve_constraint(corpus, &facets, str_field(dsl, "b")) {
                Ok(t) => t,
                Err(e) => return (json!({ "error": format!("s_path b: {e}") }).to_string(), true),
            };
            let s = usize_field(dsl, "s", 1);
            (with_dsl(corpus.s_path(&a, &b, s), &anchor, &constraints), false)
        }
        "s_clusters" => {
            let s = usize_field(dsl, "s", 2);
            (with_dsl(corpus.s_clusters(s, k.max(6)), &anchor, &constraints), false)
        }
        "narrow" => {
            // Pure refinement: return the constrained set (count + sample) with no aggregation.
            (with_dsl(corpus.narrow(scope, &constraints), &anchor, &constraints), false)
        }
        // No explicit program → sniff the constrained set and auto-route to the high-signal programs.
        "" | "profile" | "auto" => (with_dsl(crate::agent::profile::auto_profile(corpus, &anchor), &anchor, &constraints), false),
        other => (
            json!({ "error": format!("unknown program `{other}`"), "valid_programs": ["crosstab", "breakdown", "rank", "cooccurs", "s_path", "s_clusters", "narrow", "profile"] }).to_string(),
            true,
        ),
    }
}

/// Attach the compiled anchor + constraints to the program result for transparency (the trace/UI shows
/// exactly which bitmap the program ran over).
fn with_dsl(mut result: Value, anchor: &str, constraints: &[String]) -> String {
    if let Value::Object(m) = &mut result {
        m.insert("compiled_anchor".into(), json!(anchor));
        m.insert("constraints".into(), json!(constraints));
    }
    result.to_string()
}

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

    #[test]
    fn anchor_compiles_scope_and_constraints() {
        assert_eq!(compile_anchor(&[], &[]), "*");
        assert_eq!(compile_anchor(&[], &["powertrain/electric".into()]), "powertrain/electric");
        assert_eq!(
            compile_anchor(&["ent/x".into(), "ent/y".into()], &["(num range_km gt 500)".into()]),
            "(and (or ent/x ent/y) (num range_km gt 500))"
        );
        assert_eq!(
            compile_anchor(&[], &["powertrain/electric".into(), "body/suv".into()]),
            "(and powertrain/electric body/suv)"
        );
    }
}