hypersteeldb 0.2.3

//! Qwen harness — a reliability wrapper around a small local model so it completes the tool-loop.
//!
//! Small models (Qwen 0.8–1.7B via a local OpenAI-compatible server) are erratic tool callers: they
//! frequently emit the call as Hermes-style `<tool_call>{…}</tool_call>` text (or a fenced ```json
//! block, or a bare object) in the assistant CONTENT instead of the structured `tool_calls` field, and
//! they repeat identical calls. This wrapper post-processes every provider turn to (a) recover
//! tool calls embedded in text and promote them to real tool-uses, and (b) strip that markup from the
//! visible answer. The loop-level guards (dedupe, grounding) live in `run.rs`.

use crate::agent::provider::LlmProvider;
use crate::agent::types::{Msg, Stop, ToolSpec, Turn};
use async_trait::async_trait;
use serde_json::{json, Value};

/// Fuzzy-match a hallucinated tool name against the real tool list — small models routinely emit
/// close variants (`list_facet`, `list_facet_s`, `list-facets`). Returns the real name if a prefix /
/// contains / edit-1 match is unambiguous.
fn resolve_tool_name(name: &str, tools: &[ToolSpec]) -> Option<String> {
    if tools.iter().any(|t| t.name == name) {
        return Some(name.to_string());
    }
    let norm = |s: &str| s.chars().filter(|c| c.is_alphanumeric()).flat_map(|c| c.to_lowercase()).collect::<String>();
    let target = norm(name);
    if target.is_empty() {
        return None;
    }
    let mut hits: Vec<&ToolSpec> = tools.iter().filter(|t| { let n = norm(&t.name); n == target || n.starts_with(&target) || target.starts_with(&n) || n.contains(&target) || target.contains(&n) }).collect();
    if hits.len() == 1 {
        return Some(hits.remove(0).name.clone());
    }
    None
}

/// Extra system guidance for a small local model: how to drive the tools and do MECE query expansion.
/// Appended to the base system prompt only when the harness is active (capable models don't need it).
pub const HARNESS_GUIDANCE: &str = "\
TOOL-USE HINTS (you are a small model — be systematic and literal):
- First decide the QUESTION TYPE:
  · DOCUMENT / EVIDENCE (\"what do the docs say about X\", \"which files mention Y\", \"quote/summarise\") → \
call `search` with plain keywords from the user's question. Quote from the `cells` field of the hits.
  · ANALYTIC on structured data (\"break down by\", \"compare\", \"how many\", \"which is most\", \"relate X to Y\") → \
call list_facets, facet_tokens for 1-3 relevant facets, then breakdown / crosstab / rank / cooccurs / narrow.
- MECE QUERY EXPANSION for the analytic path: decompose into 2-4 orthogonal probes; combine winning \
tokens with (and …)/(or …). Do NOT dump facet_tokens on a facet with hundreds of long-tail values — \
use `search` for text-heavy facets.
- entity_link has already been run for you and its precise tokens are in the opening message — reuse \
them as scope tokens for ikl_query.
- FILTER-THEN-AGGREGATE: for 'which G has the most X where F=v' (a constraint + an aggregation), put the \
FILTER in the `anchor` and aggregate by the OTHER facet — e.g. 'which country makes the most electric \
vehicles' → breakdown(anchor='powertrain/electric', facet='country'), NOT breakdown(anchor='*', \
facet='make'). Use `(and a b)` in the anchor for multiple filters. Aggregate by the facet the question \
asks to rank/compare, not the one in the filter.
- After 2-4 informative probes, STOP calling tools and write the grounded answer. Never repeat an \
identical call.

WORKED EXAMPLE
Q: 'Which Japanese standards cover EV battery safety?'
  1. list_facets  →  facets: geo, ent, capability, …
  2. facet_tokens('geo') → geo/japan ;  facet_tokens('capability') → capability/battery-safety, capability/ev
  3. MECE probes — place: geo/japan ; topic: (or capability/battery-safety ent/battery) ; \
combined: (and geo/japan (or capability/battery-safety ent/battery))
  4. breakdown('(and geo/japan (or capability/battery-safety ent/battery))', 'ent', 10) to list the standards
  5. Answer, citing the retrieved situations.";

/// Wrap any provider with small-model tool-call recovery + MECE/tool guidance. Native tool-callers get
/// the recovery for free; the guidance is appended to the system prompt on every turn.
pub struct QwenHarness {
    inner: Box<dyn LlmProvider>,
}

impl QwenHarness {
    pub fn wrap(inner: Box<dyn LlmProvider>) -> Box<dyn LlmProvider> {
        Box::new(QwenHarness { inner })
    }
}

#[async_trait]
impl LlmProvider for QwenHarness {
    fn name(&self) -> &str {
        self.inner.name()
    }

    async fn chat(&self, system: &str, msgs: &[Msg], tools: &[ToolSpec]) -> Result<Turn, String> {
        // The MECE/tool guidance only helps when tools are on the table; for a tool-less structured
        // request (e.g. "return this JSON") it just derails a small model, so leave the prompt alone.
        let system = if tools.is_empty() { system.to_string() } else { format!("{system}\n\n{HARNESS_GUIDANCE}") };
        let mut turn = self.inner.chat(&system, msgs, tools).await?;
        if turn.tool_uses.is_empty() {
            let (calls, cleaned) = recover_tool_calls(&turn.text);
            if !calls.is_empty() {
                turn.tool_uses = calls
                    .into_iter()
                    .enumerate()
                    .map(|(i, (name, input))| (format!("recovered_{i}"), name, input))
                    .collect();
                turn.text = cleaned;
                turn.stop = Stop::ToolUse;
            }
        }
        // fuzzy-map hallucinated tool names to the real ones (`list_facet_s` → `list_facets`).
        for (_, name, _) in turn.tool_uses.iter_mut() {
            if let Some(real) = resolve_tool_name(name, tools) {
                if real != *name {
                    *name = real;
                }
            }
        }
        Ok(turn)
    }
}

/// Extract tool calls embedded in assistant text and return them plus the text with that markup
/// removed. Handles Hermes `<tool_call>…</tool_call>` blocks, fenced ```json blocks, and a bare
/// top-level `{"name":…,"arguments":…}` object.
pub fn recover_tool_calls(text: &str) -> (Vec<(String, Value)>, String) {
    let mut calls = Vec::new();
    let mut cleaned = String::new();
    let mut rest = text;

    // 1) <tool_call> … </tool_call> (Qwen/Hermes)
    while let Some(start) = rest.find("<tool_call>") {
        cleaned.push_str(&rest[..start]);
        let after = &rest[start + "<tool_call>".len()..];
        match after.find("</tool_call>") {
            Some(end) => {
                if let Some(c) = parse_call(after[..end].trim()) {
                    calls.push(c);
                }
                rest = &after[end + "</tool_call>".len()..];
            }
            None => {
                if let Some(c) = parse_call(after.trim()) {
                    calls.push(c);
                }
                rest = "";
                break;
            }
        }
    }
    cleaned.push_str(rest);

    // 2) fenced ```json { … } ``` or 3) a bare top-level object, only if nothing recovered yet
    if calls.is_empty() {
        if let Some((c, span)) = fenced_or_bare(text) {
            calls.push(c);
            cleaned = text.replacen(&span, "", 1);
        }
    }

    (calls, cleaned.trim().to_string())
}

/// Parse one call object: `{"name": "...", "arguments": {…}|"json-string"}`. `arguments` may be an
/// object or a JSON-encoded string (or absent → `{}`). Tolerates typographic quotes small models emit
/// (Gemma 3n loves `“…”`).
fn parse_call(s: &str) -> Option<(String, Value)> {
    let normalized = s
        .replace('\u{201C}', "\"").replace('\u{201D}', "\"")   // smart double quotes
        .replace('\u{2018}', "\"").replace('\u{2019}', "\"")   // smart single quotes
        .replace('\u{FF1A}', ":").replace('\u{FF0C}', ",")     // full-width : and ,
        .replace('\u{FF08}', "(").replace('\u{FF09}', ")");    // full-width ( and )
    let v: Value = serde_json::from_str(&normalized).ok()?;
    let name = v.get("name").and_then(|n| n.as_str())?.to_string();
    let args = match v.get("arguments") {
        Some(Value::Object(o)) => Value::Object(o.clone()),
        Some(Value::String(s)) => serde_json::from_str(s).unwrap_or_else(|_| json!({})),
        _ => json!({}),
    };
    Some((name, args))
}

/// Find a fenced ```json … ``` block or a bare `{…}` object carrying a call; return the call and the
/// exact source span to strip.
fn fenced_or_bare(text: &str) -> Option<((String, Value), String)> {
    // fenced block
    if let Some(open) = text.find("```") {
        let after = &text[open + 3..];
        // skip an optional language tag on the fence line
        let body_start = after.find('\n').map(|i| i + 1).unwrap_or(0);
        if let Some(close) = after[body_start..].find("```") {
            let inner = after[body_start..body_start + close].trim();
            if let Some(c) = parse_call(inner) {
                let span = &text[open..open + 3 + body_start + close + 3];
                return Some((c, span.to_string()));
            }
        }
    }
    // bare object spanning first '{' to last '}'
    let (a, b) = (text.find('{')?, text.rfind('}')?);
    if b > a {
        let inner = &text[a..=b];
        if let Some(c) = parse_call(inner) {
            return Some((c, inner.to_string()));
        }
    }
    None
}

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

    #[test]
    fn recovers_hermes_tool_call() {
        let t = "Let me look.\n<tool_call>{\"name\": \"list_facets\", \"arguments\": {}}</tool_call>";
        let (calls, cleaned) = recover_tool_calls(t);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].0, "list_facets");
        assert_eq!(cleaned, "Let me look.");
    }

    #[test]
    fn recovers_string_arguments() {
        let t = "<tool_call>{\"name\": \"facet_tokens\", \"arguments\": \"{\\\"facet\\\":\\\"geo\\\"}\"}</tool_call>";
        let (calls, _) = recover_tool_calls(t);
        assert_eq!(calls[0].1["facet"], "geo");
    }

    #[test]
    fn recovers_fenced_block() {
        let t = "Here:\n```json\n{\"name\": \"rank\", \"arguments\": {\"facet\": \"vendor\"}}\n```";
        let (calls, _) = recover_tool_calls(t);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].0, "rank");
    }

    #[test]
    fn tool_name_fuzz_resolve() {
        let specs = vec![
            ToolSpec { name: "list_facets".into(), description: "".into(), schema: json!({}) },
            ToolSpec { name: "facet_tokens".into(), description: "".into(), schema: json!({}) },
        ];
        assert_eq!(resolve_tool_name("list_facet", &specs).as_deref(), Some("list_facets"));
        assert_eq!(resolve_tool_name("list_facet_s", &specs).as_deref(), Some("list_facets"));
        assert_eq!(resolve_tool_name("list-facets", &specs).as_deref(), Some("list_facets"));
        assert_eq!(resolve_tool_name("nonsense_xyz", &specs), None);
    }

    #[test]
    fn plain_text_untouched() {
        let (calls, cleaned) = recover_tool_calls("The answer is 42.");
        assert!(calls.is_empty());
        assert_eq!(cleaned, "The answer is 42.");
    }
}