steeldb/agent/

tools.rs

//! The engine tools the agent may call — pure read primitives over a loaded `Corpus`. The agent
//! plans a question into IKL probes; retrieval is the reasoning (roaring set-algebra), so these three
//! tools (discover schema, discover a facet's vocabulary, run an IKL query) are enough to answer.

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

/// A facet-name string property. When the corpus's real facet names are known, they're injected as an
/// `enum` so a constrained decoder (Needle) can ONLY emit a facet that actually exists — validity by
/// construction, not just post-hoc repair.
fn facet_prop(desc: &str, facets: &[String]) -> serde_json::Value {
    let mut o = json!({ "type": "string", "description": desc });
    if !facets.is_empty() {
        o["enum"] = json!(facets);
    }
    o
}

/// Build the engine tool specs. `facets` = the corpus's real facet names (empty = none known yet); when
/// present they bound the `run_workflow` facet slots to real values.
pub fn engine_tools(facets: &[String]) -> Vec<ToolSpec> {
    vec![
        ToolSpec {
            name: "list_facets".into(),
            description: "List the corpus facets (token prefixes) and how many distinct tokens each has. Call this first to learn the schema.".into(),
            schema: json!({ "type": "object", "properties": {}, "additionalProperties": false }),
        },
        ToolSpec {
            name: "run_workflow".into(),
            description: "Answer in ONE call. Your MAIN job is `constraints`: the facet/value tokens or predicates that narrow the rows to what the question is about, AND-ed together (e.g. [\"powertrain/electric\", \"(num range_km gt 500)\"]). Leave `program` OUT to auto-profile the narrowed rows (the engine picks the informative breakdowns/crosstab/clusters for you). Only set `program` when the question clearly needs a specific one: `crosstab` (facet_a × facet_b, 'for each A which B'), `breakdown` (partition by `facet`), `rank` (salient values of `facet`), `cooccurs` (neighbours of `token`), `s_path` (how tokens `a`,`b` connect), `s_clusters` (concept groupings). Use exact facet NAMES from list_facets.".into(),
            schema: json!({
                "type": "object",
                "properties": {
                    "constraints": { "type": "array", "items": { "type": "string" }, "description": "PRIMARY: filters to AND, as facet/value tokens or predicates, e.g. [\"powertrain/electric\", \"(num range_km gt 500)\"]" },
                    "program": { "type": "string", "enum": ["crosstab", "breakdown", "rank", "cooccurs", "s_path", "s_clusters", "narrow"], "description": "OPTIONAL specific analysis; omit to auto-profile" },
                    "facet_a": facet_prop("crosstab only: grouping facet NAME (e.g. country)", facets),
                    "facet_b": facet_prop("crosstab only: facet NAME counted within each group (e.g. powertrain)", facets),
                    "facet": facet_prop("breakdown/rank only: the facet NAME (e.g. country)", facets),
                    "token": { "type": "string", "description": "cooccurs only: the focus facet/value token (e.g. aws-service/lambda)" },
                    "a": { "type": "string", "description": "s_path only: first facet/value token" },
                    "b": { "type": "string", "description": "s_path only: second facet/value token" }
                },
                "additionalProperties": false
            }),
        },
        ToolSpec {
            name: "facet_tokens".into(),
            description: "List the most common tokens under a facet (with support counts), so you know the exact tokens you can query. Tokens look like `facet/value`.".into(),
            schema: json!({
                "type": "object",
                "properties": {
                    "facet": { "type": "string", "description": "the facet prefix, e.g. `geo` or `country`" },
                    "limit": { "type": "integer", "description": "max tokens to return (default 40)" }
                },
                "required": ["facet"],
                "additionalProperties": false
            }),
        },
        ToolSpec {
            name: "ikl_query".into(),
            description: "Run an IKL s-expression over the corpus and return matching situations. IKL is set-algebra over tokens: `(and A B)`, `(or A B)`, `(not A)`, a bare token, a glob like `geo/*`, or a numeric range `(num <field> <op> <value>)` with op ge|gt|le|lt|eq|ne over a numeric_field from list_facets. Example: `(and powertrain/electric (num range_km ge 500))`.".into(),
            schema: json!({
                "type": "object",
                "properties": {
                    "ikl": { "type": "string", "description": "the IKL s-expression" },
                    "limit": { "type": "integer", "description": "max situations to return (default 20)" }
                },
                "required": ["ikl"],
                "additionalProperties": false
            }),
        },
        ToolSpec {
            name: "search".into(),
            description: "Free-text search across the corpus documents/rows: give any keywords or a natural-language phrase, get back matching situations (with their text). Use this when the user asks WHAT the documents say about something, WHICH files talk about X, or wants to READ evidence — as opposed to counting/comparing (use breakdown/crosstab for those).".into(),
            schema: json!({
                "type": "object",
                "properties": {
                    "query": { "type": "string", "description": "keywords or natural-language phrase" },
                    "limit": { "type": "integer", "description": "max hits (default 8)" }
                },
                "required": ["query"],
                "additionalProperties": false
            }),
        },
        ToolSpec {
            name: "entity_link".into(),
            description: "Resolve a question's mentions (acronyms, entities, phrases) to the corpus's PRECISE existing tokens and retrieve on them. Use for entity-specific questions: the exact tokens docs are indexed under are often more specific than what a plain query guesses.".into(),
            schema: json!({
                "type": "object",
                "properties": { "question": { "type": "string", "description": "the question or sub-question to link" } },
                "required": ["question"],
                "additionalProperties": false
            }),
        },
    ]
}

fn str_arg<'a>(input: &'a Value, key: &str) -> &'a str {
    input.get(key).and_then(|v| v.as_str()).unwrap_or("")
}
fn usize_arg(input: &Value, key: &str, default: usize) -> usize {
    input.get(key).and_then(|v| v.as_u64()).map(|n| n as usize).unwrap_or(default)
}
fn str_list(input: &Value, key: &str) -> Vec<String> {
    input
        .get(key)
        .and_then(|v| v.as_array())
        .map(|a| a.iter().filter_map(|x| x.as_str().map(String::from)).collect())
        .unwrap_or_default()
}

/// Execute a tool call against the corpus. Returns (json_string, is_error).
pub fn exec_tool(corpus: &Corpus, name: &str, input: &Value) -> (String, bool) {
    match name {
        "list_facets" => {
            let s = corpus.stats();
            // Ambient facets that carry structural noise (open-ended text-derived relations, per-file
            // src tags). Kept on the corpus but hidden from list_facets so the model doesn't fixate on
            // them; still queryable if the model explicitly asks facet_tokens on them.
            let ambient = ["rel", "src"];
            let facets: Vec<Value> = s
                .facets
                .iter()
                .filter(|(f, _)| !ambient.contains(&f.as_str()))
                .map(|(f, n)| json!({ "facet": f, "distinct_tokens": n }))
                .collect();
            (json!({ "situations": s.situations, "facets": facets, "numeric_fields": s.numeric_fields }).to_string(), false)
        }
        "facet_tokens" => {
            let facet = input.get("facet").and_then(|v| v.as_str()).unwrap_or("");
            if facet.is_empty() {
                return (json!({ "error": "missing `facet`" }).to_string(), true);
            }
            // Cap responses so a facet with thousands of long-tail values (e.g. CJK `rel/*`) can't
            // overwhelm the model's context. min_support filters singleton noise unless overridden.
            let limit = (input.get("limit").and_then(|v| v.as_u64()).unwrap_or(40) as usize).min(60);
            let min_support = input.get("min_support").and_then(|v| v.as_u64()).unwrap_or(1) as usize;
            let all = corpus.facet_tokens(facet, 1000);
            let total_distinct = all.len();
            let filtered: Vec<(String, usize)> = all.into_iter().filter(|(_, n)| *n >= min_support).collect();
            let shown = filtered.iter().take(limit).cloned().collect::<Vec<_>>();
            let truncated = filtered.len() > shown.len();
            let toks: Vec<Value> = shown.iter().map(|(t, n)| json!({ "token": t, "support": n })).collect();
            let mut out = json!({ "facet": facet, "tokens": toks, "distinct_tokens": total_distinct });
            if truncated {
                out["truncated"] = Value::Bool(true);
                out["note"] = Value::String(format!("showing top {} by support; {} more not shown — raise min_support or query specific tokens with ikl_query", shown.len(), filtered.len() - shown.len()));
            }
            (out.to_string(), false)
        }
        "ikl_query" => {
            let ikl = input.get("ikl").and_then(|v| v.as_str()).unwrap_or("");
            if ikl.is_empty() {
                return (json!({ "error": "missing `ikl`" }).to_string(), true);
            }
            let limit = input.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as usize;
            let out = corpus.query(ikl, limit);
            let hits: Vec<Value> = out
                .hits
                .iter()
                .map(|h| {
                    // truncate long cells so the context stays lean
                    let cells: Vec<String> = h
                        .cells
                        .iter()
                        .map(|c| if c.len() > 300 { format!("{}…", &c[..300]) } else { c.clone() })
                        .collect();
                    json!({ "sid": h.sid, "cells": cells })
                })
                .collect();
            (
                json!({ "ikl": ikl, "count": out.count, "micros": out.micros, "returned": hits.len(), "hits": hits }).to_string(),
                false,
            )
        }
        "search" => {
            let q = str_arg(input, "query");
            if q.is_empty() {
                return (json!({ "error": "missing `query`" }).to_string(), true);
            }
            let limit = usize_arg(input, "limit", 8);
            let linked = corpus.entity_link(q);
            if linked.is_empty() {
                return (json!({ "query": q, "count": 0, "note": "no matching tokens found in the corpus vocabulary" }).to_string(), false);
            }
            // relevance-ranked: most query-token coverage first (not document order)
            let ranked = corpus.search_ranked(&linked, limit);
            let hits: Vec<Value> = ranked.iter().map(|(sid, cov, cells)| {
                let cells: Vec<String> = cells.iter().map(|c| if c.chars().count() > 400 { format!("{}…", c.chars().take(400).collect::<String>()) } else { c.clone() }).collect();
                json!({ "sid": sid, "match_score": cov, "cells": cells })
            }).collect();
            (json!({ "query": q, "matched_tokens": linked, "returned": hits.len(), "hits": hits }).to_string(), false)
        }
        "entity_link" => {
            let q = str_arg(input, "question");
            if q.is_empty() {
                return (json!({ "error": "missing `question`" }).to_string(), true);
            }
            let linked = corpus.entity_link(q);
            if linked.is_empty() {
                return (json!({ "linked": [], "note": "no precise tokens linked from the question" }).to_string(), false);
            }
            let ikl = if linked.len() > 1 { format!("(or {})", linked.join(" ")) } else { linked[0].clone() };
            let out = corpus.query(&ikl, 12);
            let hits: Vec<Value> = out
                .hits
                .iter()
                .map(|h| {
                    let cells: Vec<String> = h.cells.iter().map(|c| if c.len() > 300 { format!("{}…", &c[..300]) } else { c.clone() }).collect();
                    json!({ "sid": h.sid, "cells": cells })
                })
                .collect();
            (json!({ "linked": linked, "ikl": ikl, "count": out.count, "returned": hits.len(), "hits": hits }).to_string(), false)
        }
        "breakdown" => {
            let anchor = str_arg(input, "anchor");
            let anchor = if anchor.is_empty() || anchor == "all" { "*" } else { anchor }; // default = whole corpus
            let facet = str_arg(input, "facet");
            if facet.is_empty() {
                return (json!({ "error": "need `facet`" }).to_string(), true);
            }
            (corpus.breakdown(anchor, facet, usize_arg(input, "k", 15)).to_string(), false)
        }
        "crosstab" => {
            let anchor = str_arg(input, "anchor");
            let anchor = if anchor.is_empty() || anchor == "all" { "*" } else { anchor }; // default = whole corpus
            let (fa, fb) = (str_arg(input, "facet_a"), str_arg(input, "facet_b"));
            if fa.is_empty() || fb.is_empty() {
                return (json!({ "error": "need `facet_a`, `facet_b`" }).to_string(), true);
            }
            (corpus.crosstab(anchor, fa, fb, usize_arg(input, "k", 6)).to_string(), false)
        }
        "rank" => {
            let facet = str_arg(input, "facet");
            if facet.is_empty() {
                return (json!({ "error": "missing `facet`" }).to_string(), true);
            }
            (corpus.rank(facet, usize_arg(input, "k", 10)).to_string(), false)
        }
        "cooccurs" => {
            let token = str_arg(input, "token");
            if token.is_empty() {
                return (json!({ "error": "missing `token`" }).to_string(), true);
            }
            (corpus.cooccurs(token, usize_arg(input, "k", 12)).to_string(), false)
        }
        "s_path" => {
            let (a, b) = (str_arg(input, "a"), str_arg(input, "b"));
            if a.is_empty() || b.is_empty() {
                return (json!({ "error": "need `a` and `b`" }).to_string(), true);
            }
            (corpus.s_path(a, b, usize_arg(input, "s", 1)).to_string(), false)
        }
        "s_clusters" => (corpus.s_clusters(usize_arg(input, "s", 2), usize_arg(input, "k", 8)).to_string(), false),
        "narrow" => (corpus.narrow(&str_list(input, "scope"), &str_list(input, "filters")).to_string(), false),
        other => (json!({ "error": format!("unknown tool {other}") }).to_string(), true),
    }
}

pub const SYSTEM_PROMPT: &str = "\
You are SteelDB, a retrieval-as-reasoning agent over a corpus of documents and structured rows. Answers \
come ONLY from the corpus, reached through the tools. Pick the right tool for the question:

DOCUMENT / EVIDENCE questions (what do the docs say about X, which files mention Y, quote the passage \
about Z, summarise the sections on W): call `search` with plain-language keywords — it returns matching \
situations with their text. Quote from the hit `cells` in your answer.

ANALYTIC questions on structured data (compare, enumerate, rank, relate, break down, distribute): first \
`list_facets` to see the schema, `facet_tokens` for exact token vocabulary, then use one of the bitmap \
programs — `breakdown` (partition), `crosstab` (relate two facets), `rank` (most salient values), \
`cooccurs` / `s_path` / `s_clusters` (structural graph), or `narrow` (find which constraint the corpus \
cannot satisfy). Do NOT dump facet_tokens on a facet with hundreds of long-tail values — use `search` \
or `ikl_query` instead.

IKL for precise set queries: `ikl_query` runs `(and A B)`, `(or A B)`, `(not A)`, globs like `geo/*`, or \
NUMERIC ranges `(num <field> <op> <value>)` (op ge|gt|le|lt|eq|ne) over the numeric_fields shown by \
list_facets — e.g. `(and powertrain/electric (num range_km ge 500))` for 'electric vehicles with range \
over 500'. `entity_link` has already been run for you up front — reuse those linked tokens as scope. \
Answer concisely, grounded in what the tools returned, and say plainly if the corpus does not support \
an answer. Never invent tokens; only query tokens you have seen.";