batuta/agent/driver/

realizar.rs

//! RealizarDriver — sovereign local inference via GGUF/APR models.
//!
//! Uses the `realizar` crate for local LLM inference. All data
//! stays on-device (Sovereign privacy tier, Genchi Genbutsu).
//!
//! Tool call parsing: local models output `<tool_call>` JSON blocks
//! in their text. The driver extracts these into `ToolCall` structs.
//!
//! Feature-gated behind `inference`.

use async_trait::async_trait;
use std::path::PathBuf;
use tracing::info;

use super::chat_template::{format_prompt_with_template, ChatTemplate};
use super::validate::validate_model_file;
use super::{CompletionRequest, CompletionResponse, LlmDriver, ToolCall};
use crate::agent::result::{AgentError, DriverError, StopReason, TokenUsage};
use crate::serve::backends::PrivacyTier;

/// Local inference driver using realizar (GGUF/APR/SafeTensors).
pub struct RealizarDriver {
    /// Path to model file.
    model_path: PathBuf,
    /// Context window size.
    context_window_size: usize,
    /// Auto-detected chat template.
    template: ChatTemplate,
}

impl RealizarDriver {
    /// Create a new RealizarDriver from a model path.
    ///
    /// **Contract: `apr_model_validity` (apr-code-v1.yaml)**
    ///
    /// Preconditions enforced at the load boundary (Jidoka):
    /// - File must exist
    /// - APR files: must have embedded tokenizer (checked via header)
    /// - GGUF files: must have valid magic bytes
    ///
    /// Violation → actionable error with re-conversion instructions.
    /// No broken model ever reaches the inference loop.
    pub fn new(model_path: PathBuf, context_window: Option<usize>) -> Result<Self, AgentError> {
        if !model_path.exists() {
            return Err(AgentError::Driver(DriverError::InferenceFailed(format!(
                "model not found: {}",
                model_path.display()
            ))));
        }

        // ═══ CONTRACT: apr_model_validity — Jidoka boundary check ═══
        validate_model_file(&model_path)?;

        let context_window_size = context_window.unwrap_or(4096);
        let template = ChatTemplate::from_model_path(&model_path);
        Ok(Self { model_path, context_window_size, template })
    }
}

#[async_trait]
impl LlmDriver for RealizarDriver {
    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, AgentError> {
        // Format messages using auto-detected chat template
        let prompt = format_prompt_with_template(&request, self.template);

        // Build inference config (explicit fields — no Default impl)
        let config = realizar::infer::InferenceConfig {
            model_path: self.model_path.clone(),
            prompt: Some(prompt),
            input_tokens: None,
            max_tokens: request.max_tokens as usize,
            temperature: request.temperature,
            top_k: 0,
            // PMAT-823: new sampling fields default to greedy/disabled so this
            // agent driver's behavior is unchanged (it only exposes temperature).
            top_p: None,
            seed: 42,
            repeat_penalty: 1.0,
            repeat_last_n: 64,
            // PMAT-156/158: Disable GPU only for APR models (wgpu shader bug).
            // GGUF models work fine with CUDA — keep GPU enabled for them.
            no_gpu: self.model_path.extension().is_some_and(|e| e == "apr"),
            trace: false,
            trace_verbose: false,
            trace_output: None,
            trace_steps: None,
            verbose: false,
            use_mock_backend: false,
            stop_tokens: vec![],
        };

        // Run inference in blocking thread (realizar is sync)
        let result = tokio::task::spawn_blocking(move || realizar::infer::run_inference(&config))
            .await
            .map_err(|e| {
                AgentError::Driver(DriverError::InferenceFailed(format!("spawn_blocking: {e}")))
            })?
            .map_err(|e| AgentError::Driver(DriverError::InferenceFailed(e.to_string())))?;

        // Parse tool calls from text output
        let (raw_text, tool_calls) = parse_tool_calls(&result.text);

        // Sanitize output: strip echoed system prompt and chat template markers
        let text = sanitize_output(&raw_text, request.system.as_deref());

        let stop_reason =
            if tool_calls.is_empty() { StopReason::EndTurn } else { StopReason::ToolUse };

        Ok(CompletionResponse {
            text,
            stop_reason,
            tool_calls,
            usage: TokenUsage {
                input_tokens: result.input_token_count as u64,
                output_tokens: result.generated_token_count as u64,
            },
        })
    }

    fn context_window(&self) -> usize {
        self.context_window_size
    }

    fn privacy_tier(&self) -> PrivacyTier {
        PrivacyTier::Sovereign
    }
}

/// Parse tool calls from model output text.
///
/// Supports multiple formats (PMAT-158):
/// 1. `<tool_call>{"name":...}</tool_call>` — custom XML tags
/// 2. `<tool_call>{"name":...}` — unclosed XML (small model fallback)
/// 3. `` ```json\n{"name":...}\n``` `` — markdown code block (Qwen native)
///
/// CCPA-m296 SALVAGE: if the envelope parser above finds NOTHING but the
/// text is recoverably tool-call-shaped — a generically-fenced block
/// (`` ```...{"name":..,"input":..}... ``` ``, any language tag or none) or a
/// bare top-level `{"name":..,"input":..}` JSON object — [`salvage_tool_calls`]
/// recovers it. This converts "the model almost emitted a tool_call" near-misses
/// into real tool calls instead of letting the raw Markdown re-prime prose mode
/// across turns (the self-reinforcing text loop that defeats tool-calling).
///
/// Returns the remaining text (with tool call blocks removed)
/// and the extracted tool calls.
/// Public wrapper for tool call parsing (used by AprServeDriver).
pub fn parse_tool_calls_pub(text: &str) -> (String, Vec<ToolCall>) {
    parse_tool_calls(text)
}

fn parse_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
    let (remaining, calls) = parse_tool_calls_envelope(text);
    if !calls.is_empty() {
        return (remaining, calls);
    }
    // CCPA-m296: envelope parser found no tool call — try the conservative
    // salvage parser before scoring this turn as inert prose text.
    let (salvaged_remaining, salvaged) = salvage_tool_calls(&remaining);
    if !salvaged.is_empty() {
        info!(count = salvaged.len(), "salvaged tool call(s) from non-envelope output (CCPA-m296)");
        return (salvaged_remaining, salvaged);
    }
    (remaining, calls)
}

fn parse_tool_calls_envelope(text: &str) -> (String, Vec<ToolCall>) {
    let mut tool_calls = Vec::new();
    let mut remaining = String::new();
    let mut call_counter = 0u32;

    let mut cursor = text;
    loop {
        // Find next tool call start — try <tool_call> first, then ```json
        let xml_pos = cursor.find("<tool_call>");
        let md_pos = cursor.find("```json");

        let (start, tag_len, is_markdown) = match (xml_pos, md_pos) {
            (Some(x), Some(m)) if x <= m => (x, "<tool_call>".len(), false),
            (Some(x), None) => (x, "<tool_call>".len(), false),
            (_, Some(m)) => (m, "```json".len(), true),
            (None, None) => {
                remaining.push_str(cursor);
                break;
            }
        };

        remaining.push_str(&cursor[..start]);
        let after_tag = &cursor[start + tag_len..];

        // Find closing tag and extract JSON
        let (json_str, advance_past) = if is_markdown {
            // Markdown: ```json\n...\n```
            if let Some(end) = after_tag.find("```") {
                (&after_tag[..end], &after_tag[end + "```".len()..])
            } else {
                (after_tag, "")
            }
        } else if let Some(end) = after_tag.find("</tool_call>") {
            (&after_tag[..end], &after_tag[end + "</tool_call>".len()..])
        } else {
            // PMAT-158: No closing tag — try parsing to end-of-string
            (after_tag, "")
        };
        let json_str = json_str.trim();

        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str) {
            // Must have "name" field to be a tool call (not just any JSON)
            if let Some(name) = parsed.get("name").and_then(|n| n.as_str()) {
                let name = name.to_string();
                let input = parsed.get("input").cloned().unwrap_or(serde_json::json!({}));
                call_counter += 1;
                tool_calls.push(ToolCall { id: format!("local-{call_counter}"), name, input });
            } else {
                remaining.push_str(&cursor[start..]);
                break;
            }
        } else {
            remaining.push_str(&cursor[start..]);
            break;
        }

        cursor = advance_past;
        if cursor.is_empty() {
            break;
        }
    }

    (remaining.trim().to_string(), tool_calls)
}

/// CCPA-m296 salvage parser: recover a tool call the model emitted OUTSIDE the
/// exact `<tool_call>` / ```json envelope, but in an unambiguous, recoverable
/// shape. Two recoverable shapes are accepted, in priority order:
///
/// 1. A generically-fenced code block — `` ```<anylang>\n{...}\n``` `` — whose
///    inner content parses as a tool-call-shaped JSON object. (The envelope
///    parser only recognises the exact `` ```json `` tag; coder-finetuned models
///    routinely emit `` ```tool_call ``, `` ```rust ``, or a bare `` ``` ``.)
/// 2. A bare top-level `{"name": "...", "input": {...}}` JSON object embedded in
///    prose (no fence, no tags).
///
/// CONSERVATIVE BY DESIGN: only JSON objects that (a) parse cleanly and (b) have
/// a string `name` field AND an `input` field are salvaged. Plain JSON examples
/// (e.g. `{"key": "value"}`) and prose are never mistaken for tool calls. This
/// directly recovers the "model almost emitted a tool_call" near-misses that
/// would otherwise be scored as inert text and re-prime prose mode next turn.
///
/// Returns the remaining text (with the salvaged span removed) and the
/// recovered calls (`salvage-{n}` ids so salvage events stay traceable).
fn salvage_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
    // Shape 1: a generic fenced block ```<tag>\n ... \n```
    if let Some((before, inner, after)) = extract_first_fenced_block(text) {
        if let Some(call) = tool_call_from_json_str(inner.trim(), 1) {
            let remaining = format!("{before}{after}");
            return (remaining.trim().to_string(), vec![call]);
        }
    }

    // Shape 2: a bare top-level {"name":..,"input":..} object embedded in prose.
    if let Some((start, end)) = find_balanced_json_object(text) {
        if let Some(call) = tool_call_from_json_str(text[start..end].trim(), 1) {
            let remaining = format!("{}{}", &text[..start], &text[end..]);
            return (remaining.trim().to_string(), vec![call]);
        }
    }

    (text.trim().to_string(), Vec::new())
}

/// Parse a JSON string into a tool call iff it is unambiguously tool-call-shaped:
/// a JSON object with a string `name` field AND an `input` field. Returns `None`
/// otherwise (plain JSON, arrays, scalars, prose).
fn tool_call_from_json_str(json_str: &str, idx: u32) -> Option<ToolCall> {
    let parsed = serde_json::from_str::<serde_json::Value>(json_str).ok()?;
    let obj = parsed.as_object()?;
    // Require BOTH name (string) and an explicit input field — stricter than the
    // envelope parser (which defaults input to {}) so prose/JSON examples that
    // merely contain a "name" key are never salvaged.
    let name = obj.get("name")?.as_str()?.to_string();
    if name.is_empty() {
        return None;
    }
    let input = obj.get("input")?.clone();
    Some(ToolCall { id: format!("salvage-{idx}"), name, input })
}

/// Extract the first ```...``` fenced block: returns (text-before, inner, text-after).
/// Accepts any language tag (or none); the inner content is everything between the
/// opening fence's newline and the closing fence.
fn extract_first_fenced_block(text: &str) -> Option<(&str, &str, &str)> {
    let open = text.find("```")?;
    let before = &text[..open];
    let rest = &text[open + 3..];
    // Skip the optional language tag up to (and including) the first newline.
    let inner_start = rest.find('\n').map(|i| i + 1)?;
    let body = &rest[inner_start..];
    let close = body.find("```")?;
    let inner = &body[..close];
    let after = &body[close + 3..];
    Some((before, inner, after))
}

/// Find the first balanced top-level `{...}` JSON object span in `text`.
/// Returns `(start, end)` byte indices (end exclusive) of the object including
/// braces, tracking string literals + escapes so braces inside strings don't
/// unbalance the scan. Returns `None` if no balanced object is found.
fn find_balanced_json_object(text: &str) -> Option<(usize, usize)> {
    let bytes = text.as_bytes();
    let start = text.find('{')?;
    let mut depth = 0i32;
    let mut in_str = false;
    let mut escaped = false;
    let mut i = start;
    while i < bytes.len() {
        let c = bytes[i];
        if in_str {
            if escaped {
                escaped = false;
            } else if c == b'\\' {
                escaped = true;
            } else if c == b'"' {
                in_str = false;
            }
        } else {
            match c {
                b'"' => in_str = true,
                b'{' => depth += 1,
                b'}' => {
                    depth -= 1;
                    if depth == 0 {
                        return Some((start, i + 1));
                    }
                }
                _ => {}
            }
        }
        i += 1;
    }
    None
}

/// Sanitize model output: strip echoed system prompt and chat template markers.
///
/// Small models (<3B) often echo the system prompt or leak chat template
/// tokens into their response. This strips those artifacts so the agent
/// loop sees clean assistant text.
fn sanitize_output(text: &str, system_prompt: Option<&str>) -> String {
    let mut cleaned = text.to_string();

    // Strip echoed system prompt (common with small models)
    if let Some(sys) = system_prompt {
        // Check if output starts with a significant prefix of the system prompt
        let sys_prefix = &sys[..sys.len().min(80)];
        if cleaned.starts_with(sys_prefix) {
            // The model regurgitated the system prompt — strip it
            cleaned = cleaned[sys.len().min(cleaned.len())..].to_string();
        }
    }

    // Strip leaked chat template markers
    for marker in &[
        "<|im_start|>",
        "<|im_end|>",
        "<|start_header_id|>",
        "<|end_header_id|>",
        "<|eot_id|>",
        "<|system|>",
        "<|user|>",
        "<|assistant|>",
        "<|end|>",
    ] {
        cleaned = cleaned.replace(marker, "");
    }

    // Strip leading/trailing whitespace and role labels
    let cleaned = cleaned.trim();
    let cleaned = cleaned.strip_prefix("system\n").unwrap_or(cleaned);
    let cleaned = cleaned.strip_prefix("assistant\n").unwrap_or(cleaned);
    cleaned.trim().to_string()
}

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

    #[test]
    fn test_parse_no_tool_calls() {
        let (text, calls) = parse_tool_calls("Hello world");
        assert_eq!(text, "Hello world");
        assert!(calls.is_empty());
    }

    #[test]
    fn test_parse_single_tool_call() {
        let input = r#"Before text
<tool_call>
{"name": "rag", "input": {"query": "SIMD"}}
</tool_call>
After text"#;
        let (text, calls) = parse_tool_calls(input);
        assert_eq!(text, "Before text\n\nAfter text");
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].name, "rag");
        assert_eq!(calls[0].id, "local-1");
        assert_eq!(calls[0].input, serde_json::json!({"query": "SIMD"}));
    }

    #[test]
    fn test_parse_multiple_tool_calls() {
        let input = r#"<tool_call>
{"name": "rag", "input": {"query": "a"}}
</tool_call>
Middle
<tool_call>
{"name": "memory", "input": {"action": "recall", "query": "b"}}
</tool_call>"#;
        let (text, calls) = parse_tool_calls(input);
        assert_eq!(text, "Middle");
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].name, "rag");
        assert_eq!(calls[0].id, "local-1");
        assert_eq!(calls[1].name, "memory");
        assert_eq!(calls[1].id, "local-2");
    }

    #[test]
    fn test_parse_malformed_json() {
        let input = r#"<tool_call>
not valid json
</tool_call>"#;
        let (_text, calls) = parse_tool_calls(input);
        assert!(calls.is_empty());
    }

    #[test]
    fn test_parse_missing_close_tag_with_valid_json() {
        // PMAT-158: Small models omit </tool_call>. Parser should still extract.
        let input =
            "<tool_call>\n{\"name\": \"file_read\", \"input\": {\"path\": \"src/main.rs\"}}";
        let (text, calls) = parse_tool_calls(input);
        assert_eq!(calls.len(), 1, "should extract tool call without closing tag");
        assert_eq!(calls[0].name, "file_read");
        assert!(text.is_empty(), "no remaining text expected");
    }

    #[test]
    fn test_parse_missing_close_tag_with_trailing_text() {
        // Unclosed tag with text before it — text preserved, tool call extracted
        let input =
            "Let me read that.\n<tool_call> {\"name\": \"file_read\", \"input\": {\"path\": \"foo.rs\"}}";
        let (text, calls) = parse_tool_calls(input);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].name, "file_read");
        assert!(text.contains("Let me read that"));
    }

    #[test]
    fn test_parse_missing_close_tag_invalid_json() {
        // Unclosed tag with invalid JSON — treated as plain text
        let input = "<tool_call>\nnot valid json at all";
        let (text, calls) = parse_tool_calls(input);
        assert!(calls.is_empty(), "invalid JSON should not produce tool call");
        assert!(text.contains("<tool_call>"));
    }

    #[test]
    fn test_parse_markdown_code_block() {
        // PMAT-158: Qwen2.5-Coder native format — ```json blocks
        let input = "Let me read that file.\n```json\n{\"name\": \"file_read\", \"input\": {\"path\": \"src/main.rs\"}}\n```";
        let (text, calls) = parse_tool_calls(input);
        assert_eq!(calls.len(), 1, "should extract tool call from markdown block");
        assert_eq!(calls[0].name, "file_read");
        assert_eq!(calls[0].input["path"], "src/main.rs");
        assert!(text.contains("Let me read that"));
    }

    #[test]
    fn test_parse_markdown_code_block_not_tool_call() {
        // JSON in code block without "name" field — not a tool call
        let input = "Here's an example:\n```json\n{\"key\": \"value\"}\n```";
        let (text, calls) = parse_tool_calls(input);
        assert!(calls.is_empty(), "JSON without name field should not be a tool call");
        assert!(text.contains("example"));
    }

    #[test]
    fn test_parse_missing_name() {
        let input = r#"<tool_call>
{"input": {"query": "test"}}
</tool_call>"#;
        let (_, calls) = parse_tool_calls(input);
        assert!(calls.is_empty(), "JSON without name should not be extracted");
    }

    #[test]
    fn test_privacy_tier_always_sovereign() {
        assert_eq!(PrivacyTier::Sovereign, PrivacyTier::Sovereign);
    }

    // ── Output sanitization tests ──

    #[test]
    fn test_sanitize_strips_echoed_system_prompt() {
        let sys = "You are apr code, a sovereign AI coding assistant.";
        let output = format!("{sys} And then the model continues here.");
        let cleaned = sanitize_output(&output, Some(sys));
        assert!(!cleaned.contains("sovereign AI coding assistant"));
        assert!(cleaned.contains("continues here"));
    }

    #[test]
    fn test_sanitize_strips_chat_markers() {
        let output = "<|im_start|>assistant\nHello world<|im_end|>";
        let cleaned = sanitize_output(output, None);
        assert_eq!(cleaned, "Hello world");
    }

    #[test]
    fn test_sanitize_preserves_clean_output() {
        let output = "The answer is 42.";
        let cleaned = sanitize_output(output, Some("You are helpful."));
        assert_eq!(cleaned, "The answer is 42.");
    }

    #[test]
    fn test_sanitize_strips_role_prefix() {
        let output = "assistant\nHere is my response.";
        let cleaned = sanitize_output(output, None);
        assert_eq!(cleaned, "Here is my response.");
    }

    // ── CCPA-m296 salvage parser tests ──

    #[test]
    fn test_salvage_bare_top_level_json_tool_call() {
        // Model emitted a bare {"name","input"} object with NO envelope/fence.
        // Without salvage this scores as inert prose and re-primes prose mode.
        let input =
            "Sure, I'll read it.\n{\"name\": \"file_read\", \"input\": {\"path\": \"src/lib.rs\"}}";
        let (text, calls) = parse_tool_calls(input);
        assert_eq!(calls.len(), 1, "salvage must recover a bare tool-call JSON object");
        assert_eq!(calls[0].name, "file_read");
        assert_eq!(calls[0].input["path"], "src/lib.rs");
        assert!(calls[0].id.starts_with("salvage-"), "salvaged calls carry a traceable id");
        assert!(text.contains("Sure, I'll read it"), "prose around the call is preserved");
        assert!(!text.contains("file_read"), "the salvaged JSON span is removed from text");
    }

    #[test]
    fn test_salvage_generic_fenced_block_non_json_tag() {
        // Coder models fence tool calls with ```tool_call / ```rust, not ```json.
        // The envelope parser only knows ```json; salvage must catch the rest.
        let input =
            "```tool_call\n{\"name\": \"shell\", \"input\": {\"command\": \"cargo test\"}}\n```";
        let (_text, calls) = parse_tool_calls(input);
        assert_eq!(calls.len(), 1, "salvage must recover a generically-fenced tool call");
        assert_eq!(calls[0].name, "shell");
        assert_eq!(calls[0].input["command"], "cargo test");
    }

    #[test]
    fn test_salvage_conservative_rejects_plain_json() {
        // A bare JSON object WITHOUT name+input is NOT a tool call — never salvage it.
        let input = "Here is some config:\n{\"key\": \"value\", \"count\": 3}";
        let (text, calls) = parse_tool_calls(input);
        assert!(calls.is_empty(), "plain JSON (no name+input) must not be salvaged");
        assert!(text.contains("config"));
    }

    #[test]
    fn test_salvage_conservative_rejects_name_without_input() {
        // Stricter than the envelope parser: salvage requires an explicit `input`.
        let input = "{\"name\": \"file_read\"}";
        let (_text, calls) = parse_tool_calls(input);
        assert!(calls.is_empty(), "name without input is too ambiguous to salvage");
    }

    #[test]
    fn test_salvage_handles_braces_inside_strings() {
        // The balanced-object scanner must not unbalance on braces inside strings.
        let input = "{\"name\": \"shell\", \"input\": {\"command\": \"echo ${HOME} and }{\"}}";
        let (_text, calls) = parse_tool_calls(input);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].name, "shell");
        assert_eq!(calls[0].input["command"], "echo ${HOME} and }{");
    }

    #[test]
    fn test_envelope_takes_precedence_over_salvage() {
        // A proper <tool_call> envelope must be parsed by the envelope path
        // (id "local-1"), never falling through to salvage.
        let input =
            "<tool_call>\n{\"name\": \"glob\", \"input\": {\"pattern\": \"*.rs\"}}\n</tool_call>";
        let (_text, calls) = parse_tool_calls(input);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].id, "local-1", "envelope parser owns this, not salvage");
    }
}