magi-code 0.61.0

use super::*;
use serde_json::Value;

impl StreamParser {
    // Gemma 4 / diffusiongemma on vLLM emits tool calls inline in content:
    //   <|tool_call>call:NAME{key:<|"|>value<|"|>}<tool_call|>
    // The <|"|> is an escaped double quote from the tokenizer. Returns None if
    // content does not contain the Gemma inline tool-call marker.
    pub(super) fn parse_gemma_inline_tool_calls(&mut self, text: &str) -> Option<Vec<ToolCall>> {
        if !text.contains("<|tool_call>") {
            return None;
        }
        let mut calls = Vec::new();
        let mut remaining = text.trim_start();
        while !remaining.is_empty() {
            let after_start = remaining.strip_prefix("<|tool_call>")?;
            let end = after_start.find("<tool_call|>")?;
            let body = &after_start[..end];
            let call = self.parse_gemma_single_tool_call(body)?;
            calls.push(call);
            remaining = after_start[end + "<tool_call|>".len()..].trim_start();
        }
        if calls.is_empty() { None } else { Some(calls) }
    }

    pub(super) fn parse_gemma_single_tool_call(&mut self, body: &str) -> Option<ToolCall> {
        // Body format: call:NAME{ARGS} where ARGS uses unquoted keys: {query:"value"}
        // and quotes are escaped as <|"|> by the tokenizer.
        let body = body.strip_prefix("call:")?;
        let brace_start = body.find('{')?;
        let name = body[..brace_start].trim();
        if name.is_empty() {
            return None;
        }
        let args_raw = &body[brace_start..];
        let args_end = args_raw.rfind('}')?;
        let args_with_braces = &args_raw[..=args_end];
        // Unescape Gemma's <|"|> token sequences and normalize duplicated quote
        // delimiters from patterns like <|"|>"value"<|"|>.
        let unescaped = normalize_gemma_inline_argument_text(args_with_braces);
        // Gemma format uses unquoted JSON keys ({query:"value"}); fix up for serde_json.
        let arguments = serde_json::from_str::<Value>(&unescaped)
            .or_else(|_| {
                let fixed = fix_unquoted_json_keys(&unescaped);
                serde_json::from_str::<Value>(&fixed)
            })
            .map(normalize_extra_quoted_tool_arguments)
            .ok()?;
        self.gemma_inline_tool_call_counter += 1;
        Some(ToolCall {
            id: format!("gemma_inline_{}", self.gemma_inline_tool_call_counter),
            name: name.to_string(),
            arguments,
        })
    }
}

pub(super) fn normalize_extra_quoted_tool_arguments(value: Value) -> Value {
    match value {
        Value::String(text) => serde_json::from_str::<String>(&text)
            .map(Value::String)
            .unwrap_or(Value::String(text)),
        Value::Array(items) => Value::Array(
            items
                .into_iter()
                .map(normalize_extra_quoted_tool_arguments)
                .collect(),
        ),
        Value::Object(fields) => Value::Object(
            fields
                .into_iter()
                .map(|(key, value)| (key, normalize_extra_quoted_tool_arguments(value)))
                .collect(),
        ),
        value => value,
    }
}

pub(super) fn normalize_gemma_inline_argument_text(text: &str) -> String {
    text.replace("<|\"|>", "\"")
        .replace(":\"\"", ":\"")
        .replace("[\"\"", "[\"")
        .replace("\"\"]", "\"]")
        .replace("\"\"}", "\"}")
        .replace("\"\",", "\",")
}

// ponytail: Gemma 4 emits unquoted JSON keys ({query:"value"}). This regex-free fix
// quotes bare keys so serde_json can parse them. Ceiling: string values containing
// bare-word patterns could misparse; upgrade to a real tokenizer if that occurs.
pub(super) fn fix_unquoted_json_keys(json: &str) -> String {
    // Insert quotes around bare keys: {query: → {"query": and ,key: → ,"key":
    let mut result = String::with_capacity(json.len() + 32);
    let chars: Vec<char> = json.chars().collect();
    let mut i = 0;
    while i < chars.len() {
        let c = chars[i];
        // Detect key position: after { or , (allowing whitespace)
        if (c == '{' || c == ',') && i + 1 < chars.len() {
            result.push(c);
            i += 1;
            // Skip whitespace
            while i < chars.len() && chars[i].is_whitespace() {
                result.push(chars[i]);
                i += 1;
            }
            // If next char starts a bare key (letter or underscore) and isn't already a quote
            if i < chars.len() && chars[i] != '"' && (chars[i].is_alphabetic() || chars[i] == '_') {
                result.push('"');
                while i < chars.len() && (chars[i].is_alphanumeric() || chars[i] == '_') {
                    result.push(chars[i]);
                    i += 1;
                }
                result.push('"');
            }
        } else {
            result.push(c);
            i += 1;
        }
    }
    result
}