hypersteeldb 0.5.2

//! 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)
    }

    /// Forwarded untouched.
    ///
    /// A decorator that implements a trait by overriding only the methods it cares about silently inherits the
    /// DEFAULT for everything else, and a default is not the same as the inner provider's answer. Two bugs came
    /// from exactly that here, so every method the trait grows must be forwarded unless there is a stated reason
    /// not to — `harness_forwards_every_capability_of_the_inner_provider` fails if one is forgotten.
    ///
    /// Structured output needs no harness help: the grammar constrains the decoder directly, and the recovery
    /// this wrapper performs is for tool calls embedded in prose, which cannot occur when the output is a schema.
    async fn chat_json(
        &self,
        system: &str,
        msgs: &[Msg],
        schema: &Value,
        name: &str,
    ) -> Result<Option<Value>, String> {
        self.inner.chat_json(system, msgs, schema, name).await
    }

    /// Forwarded, and it must be: `NativeProvider` and `NeedleProvider` both return `false` here precisely
    /// because a small model cannot be trusted to phrase numbers, and this wrapper wraps both. Inheriting the
    /// default `true` meant the agent rendered a 1.7B model's prose instead of the deterministic template —
    /// silently undoing the guard, for the models that need it most.
    fn synthesizes(&self) -> bool {
        self.inner.synthesizes()
    }
}

/// 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.");
    }


    /// An inner provider whose answers differ from every trait default, so anything the wrapper fails to
    /// forward shows up as the default rather than as this.
    struct Stub;

    #[async_trait]
    impl LlmProvider for Stub {
        fn name(&self) -> &str {
            "stub"
        }
        async fn chat(&self, _: &str, _: &[Msg], _: &[ToolSpec]) -> Result<Turn, String> {
            Ok(Turn { text: "plain".into(), tool_uses: Vec::new(), stop: Stop::EndTurn })
        }
        async fn chat_json(&self, _: &str, _: &[Msg], _: &Value, _: &str) -> Result<Option<Value>, String> {
            Ok(Some(json!({ "reached": "inner" })))
        }
        fn synthesizes(&self) -> bool {
            false
        }
    }

    /// The bug this guards against cost two real failures.
    ///
    /// `chat_json` was added to the trait with a default of `Ok(None)`. The harness did not forward it, so every
    /// locally-served model silently lost structured output and `learn` fell through to the prose path — the
    /// capability was implemented, wired up, and dead. `synthesizes` had the same hole already: the harness wraps
    /// `NativeProvider`, which returns `false` so the agent will not trust a 1.7B model to phrase numbers, and
    /// the wrapper reported the default `true` instead.
    ///
    /// A defaulted trait method plus a decorator is a silent-failure machine. This test is the interlock.
    #[test]
    fn harness_forwards_every_capability_of_the_inner_provider() {
        let wrapped = QwenHarness::wrap(Box::new(Stub));

        assert_eq!(wrapped.name(), "stub", "name must come from the inner provider");
        assert!(
            !wrapped.synthesizes(),
            "synthesizes() returned the trait default instead of the inner provider's answer"
        );

        let got = block_on(wrapped.chat_json("sys", &[], &json!({}), "n"))
            .expect("forwarded call should not error");
        assert_eq!(
            got,
            Some(json!({ "reached": "inner" })),
            "chat_json returned the trait default; the harness is not forwarding it"
        );
    }

    /// Drive a future with no runtime. Sound here because the stub never yields.
    fn block_on<F: std::future::Future>(mut fut: F) -> F::Output {
        use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
        fn noop(_: *const ()) {}
        fn clone(p: *const ()) -> RawWaker {
            RawWaker::new(p, &VTABLE)
        }
        static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, noop, noop, noop);
        let waker = unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) };
        let mut cx = Context::from_waker(&waker);
        let mut fut = unsafe { std::pin::Pin::new_unchecked(&mut fut) };
        loop {
            match fut.as_mut().poll(&mut cx) {
                Poll::Ready(v) => return v,
                Poll::Pending => panic!("the stub must not yield"),
            }
        }
    }
}