steeldb/agent/

run.rs

//! The agent loop — provider-agnostic. Plan → probe (tool calls against the corpus) → synthesise,
//! bounded by `max_steps`. Records a trace of every tool call for transparency.

use crate::agent::provider::LlmProvider;
use crate::agent::tools::{engine_tools, exec_tool, SYSTEM_PROMPT};
use crate::agent::types::{Block, Msg, Role, Stop};
use crate::db::Corpus;
use serde::Serialize;
use serde_json::{json, Value};
use std::collections::HashMap;

#[derive(Debug, Clone, Serialize)]
pub struct ToolCallLog {
    pub name: String,
    pub input: Value,
    pub result: String,
    pub is_error: bool,
}

#[derive(Debug, Clone, Serialize)]
pub struct AgentAnswer {
    pub answer: String,
    pub trace: Vec<ToolCallLog>,
    pub steps: usize,
}

pub async fn run_agent(
    provider: &dyn LlmProvider,
    corpus: &Corpus,
    question: &str,
    max_steps: usize,
) -> Result<AgentAnswer, String> {
    // Bound the workflow DSL's facet slots to the corpus's real facets → constrained-by-construction.
    let tools = engine_tools(&corpus.facet_names());
    let mut trace: Vec<ToolCallLog> = Vec::new();
    // Latest structured program result — used to render a deterministic template answer for providers
    // that don't synthesise (Needle / native), so the numbers can't be hallucinated.
    let mut last_program: Option<Value> = None;

    // ALWAYS-ON entity-linked recall (17d6490): resolve the question to the corpus's precise tokens and
    // retrieve on them BEFORE the model plans — independent of what it chooses to call — so recall for
    // entity-specific questions never depends on the model guessing the right tokens. Injected as
    // grounding context in the opening turn, and recorded in the trace for transparency.
    let mut opening = question.to_string();
    let linked = corpus.entity_link(question);
    if !linked.is_empty() {
        // Evidence passages from the SAME projection space (SPLADE/tagger/gazetteer tokens), ranked by
        // token coverage over the roaring bitmap — symmetric query→corpus recall, not vector retrieval.
        // Lets the answer be grounded even if the model never calls search/ikl itself.
        let ranked = corpus.search_ranked(&linked, 8);
        let sample = ranked
            .iter()
            .enumerate()
            .map(|(i, (_sid, _cov, cells))| {
                let joined: String = cells.join(" · ").chars().take(240).collect();
                format!("  [{}] {}", i + 1, joined)
            })
            .collect::<Vec<_>>()
            .join("\n");
        opening.push_str(&format!(
            "\n\n[Retrieved evidence for your question (corpus tokens: {}; {} matching situations). \
             ANSWER FROM these passages when the question is about what the documents say; for \
             counting/comparing use the analytics programs.{}]",
            linked.join(", "),
            ranked.len(),
            if sample.is_empty() { String::new() } else { format!("\n{sample}") }
        ));
        trace.push(ToolCallLog {
            name: "entity_link".into(),
            input: json!({ "question": question }),
            result: json!({ "linked": linked, "passages": ranked.len() }).to_string(),
            is_error: false,
        });
    }
    let mut msgs: Vec<Msg> = vec![Msg::user_text(opening)];

    // small-model loop guards: cache tool results (so repeats are free) and break out of a model that
    // keeps emitting the same probe instead of answering.
    let mut cache: HashMap<String, (String, bool)> = HashMap::new();
    let mut last_sig = String::new();
    let mut repeat = 0usize;

    for step in 1..=max_steps {
        let turn = provider.chat(SYSTEM_PROMPT, &msgs, &tools).await?;

        // record the assistant turn (text + any tool-use blocks) so the model sees its own calls
        let mut ablocks: Vec<Block> = Vec::new();
        if !turn.text.trim().is_empty() {
            ablocks.push(Block::Text(turn.text.clone()));
        }
        for (id, name, input) in &turn.tool_uses {
            ablocks.push(Block::ToolUse { id: id.clone(), name: name.clone(), input: input.clone() });
        }
        if !ablocks.is_empty() {
            msgs.push(Msg { role: Role::Assistant, blocks: ablocks });
        }

        if turn.tool_uses.is_empty() || turn.stop == Stop::EndTurn {
            // Discovery fallback: an on-device driver may not route a vague "what's in here?" to a tool.
            // Rather than surface its refusal, profile the corpus deterministically over the query's
            // candidate scope — a grounded overview needs no model.
            if !provider.synthesizes() && last_program.is_none() {
                let anchor = crate::agent::workflow::compile_anchor(&linked, &[]);
                let profile = crate::agent::profile::auto_profile(corpus, &anchor);
                if let Some(ans) = crate::agent::synth::render_program(&profile) {
                    trace.push(ToolCallLog { name: "auto_profile".into(), input: json!({ "anchor": anchor }), result: profile.to_string(), is_error: false });
                    return Ok(AgentAnswer { answer: ans, trace, steps: step });
                }
            }
            let answer = finalize(provider, &last_program, &turn.text);
            return Ok(AgentAnswer { answer, trace, steps: step });
        }

        // execute every requested tool (from cache when repeated), feed results back as one user message
        let mut rblocks: Vec<Block> = Vec::new();
        for (id, name, input) in &turn.tool_uses {
            let key = format!("{name}|{input}");
            let (result, is_error) = cache
                .entry(key)
                .or_insert_with(|| {
                    // run_workflow compiles the DSL against the query's SPLADE/entity-link candidate
                    // scope; other tools dispatch normally.
                    if name == "run_workflow" {
                        crate::agent::workflow::execute(corpus, input, &linked)
                    } else {
                        exec_tool(corpus, name, input)
                    }
                })
                .clone();
            if !is_error {
                if let Ok(v) = serde_json::from_str::<Value>(&result) {
                    if v.get("program").is_some() {
                        last_program = Some(v);
                    }
                }
            }
            trace.push(ToolCallLog { name: name.clone(), input: input.clone(), result: result.clone(), is_error });
            rblocks.push(Block::ToolResult { id: id.clone(), content: result, is_error });
        }

        // Short-circuit for on-device tool-callers (Needle): a single valid program result IS the answer
        // — render the template now instead of looping for prose the model won't produce well.
        if !provider.synthesizes() {
            if let Some(p) = &last_program {
                if let Some(ans) = crate::agent::synth::render_program(p) {
                    return Ok(AgentAnswer { answer: ans, trace, steps: step });
                }
            }
        }

        // repeat-breaker: if this turn's calls are identical to the previous turn's, the model is stuck —
        // nudge it to answer, and bail with a best-effort answer if it keeps looping.
        let sig = turn.tool_uses.iter().map(|(_, n, i)| format!("{n}|{i}")).collect::<Vec<_>>().join(";");
        if sig == last_sig {
            repeat += 1;
            if repeat >= 3 {
                let ans = if last_program.is_some() || !provider.synthesizes() {
                    finalize(provider, &last_program, &turn.text)
                } else if turn.text.trim().is_empty() {
                    "I could not converge on an answer from the corpus with the available evidence.".to_string()
                } else {
                    sanitize_answer(&turn.text)
                };
                return Ok(AgentAnswer { answer: ans, trace, steps: step });
            }
            rblocks.push(Block::Text(
                "You already ran these exact queries and the results are unchanged. Do NOT repeat tool \
                 calls — give your final answer now, grounded in the evidence above."
                    .into(),
            ));
        } else {
            repeat = 0;
            last_sig = sig;
        }
        msgs.push(Msg { role: Role::User, blocks: rblocks });
    }

    // Ran out of steps: if a program produced a result, render it rather than failing.
    if let Some(p) = &last_program {
        if let Some(ans) = crate::agent::synth::render_program(p) {
            return Ok(AgentAnswer { answer: ans, trace, steps: max_steps });
        }
    }
    Err(format!("agent exceeded max_steps ({max_steps}) without a final answer"))
}

/// Choose the final answer: for providers that synthesise (large models), keep the model's prose; for
/// small tool-callers, render a deterministic template from the latest program result so the figures are
/// grounded, falling back to the model text only if no program ran.
fn finalize(provider: &dyn LlmProvider, last_program: &Option<Value>, model_text: &str) -> String {
    if !provider.synthesizes() {
        if let Some(p) = last_program {
            if let Some(rendered) = crate::agent::synth::render_program(p) {
                return rendered;
            }
        }
    }
    sanitize_answer(model_text)
}

/// Guardrail for small-model degenerate output: if any short substring (2-10 chars) repeats 6+ times
/// in a row (typical of a stuck decoding loop on `と`, `,`, etc.), truncate to before the run.
fn sanitize_answer(text: &str) -> String {
    for w in 2..=10 {
        let chars: Vec<char> = text.chars().collect();
        let mut run_start = 0usize;
        let mut run_len = 1usize;
        for i in w..chars.len() {
            if chars[i] == chars[i - w] {
                run_len += 1;
                if run_len >= w * 6 {
                    // find the actual byte position of the run start
                    let cutoff: usize = text.chars().take(run_start).map(|c| c.len_utf8()).sum();
                    return text[..cutoff].trim_end_matches(|c: char| c.is_whitespace() || c == ',' || c == ';').to_string();
                }
            } else {
                run_start = i - w + 1;
                run_len = 1;
            }
        }
    }
    text.to_string()
}