inferencelayer 0.2.3

//! Portable tool-calling glue for `lfm2-serve`: the prompt-side tool advertisement and the
//! output-side tool-call extraction, both pure and GPU-free so they are unit-testable in isolation.
//!
//! The convention is the Hermes / Qwen `<tools>` + `<tool_call>` format. It is the de-facto
//! standard shared across the ChatML open-model families this engine serves (Qwen3, Qwen2,
//! Llama-as-ChatML, Lfm2), so tool calling is driven by that shared convention rather than a
//! per-checkpoint Jinja `chat_template` (which would need a Python-compatible Jinja engine — see
//! [`super::chat`]). Model-family divergence is confined to [`super::chat`]'s arch match; the tool
//! block text and the parser here are format-level and reused across arches.
//!
//! Prompt side: [`tools_system_block`] renders the function signatures into the text that
//! [`super::chat::render_chat_prompt`] injects into the system turn. Output side:
//! [`parse_tool_calls`] pulls `<tool_call>…</tool_call>` blocks back out of the model's reply into
//! structured calls plus whatever prose sat outside them.

/// A tool call recovered from the model's output: the function name and its arguments as a compact
/// JSON string (the OpenAI wire shape — `function.arguments` is a string, never an object).
#[derive(Debug, Clone, PartialEq)]
pub struct ParsedToolCall {
    pub name: String,
    pub arguments: String,
}

/// Render the `# Tools` system-prompt block advertising `tools` (each a full OpenAI tool object,
/// e.g. `{"type":"function","function":{"name":…,"parameters":{…}}}`) in the Hermes/Qwen
/// convention. The caller injects the returned text into the system turn; the model is told to emit
/// calls as `<tool_call>{"name":…,"arguments":{…}}</tool_call>`, which [`parse_tool_calls`] reads
/// back. Callers only invoke this when `tools` is non-empty.
pub fn tools_system_block(tools: &[serde_json::Value]) -> String {
    let mut s = String::from(
        "# Tools\n\nYou may call one or more functions to assist with the user query.\n\n\
         You are provided with function signatures within <tools></tools> XML tags:\n<tools>",
    );
    for t in tools {
        s.push('\n');
        // Compact one-line JSON per signature, matching the Qwen template's `| tojson`.
        s.push_str(&serde_json::to_string(t).unwrap_or_default());
    }
    s.push_str(
        "\n</tools>\n\nFor each function call, return a json object with function name and \
         arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n\
         {\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>",
    );
    s
}

/// Split a model reply into (leftover prose, tool calls). Every `<tool_call>…</tool_call>` block is
/// parsed into a [`ParsedToolCall`]; text outside the blocks is concatenated and trimmed as the
/// assistant's `content`. Blocks whose body is not the expected `{"name":…,"arguments":…}` JSON are
/// dropped (their raw text is not leaked into `content` — a malformed call is a non-call, not prose)
/// so a garbled block can never masquerade as a message. An unterminated trailing `<tool_call>`
/// (generation hit the token cap mid-call) is parsed best-effort from what arrived.
pub fn parse_tool_calls(text: &str) -> (String, Vec<ParsedToolCall>) {
    const OPEN: &str = "<tool_call>";
    const CLOSE: &str = "</tool_call>";
    let mut calls = Vec::new();
    let mut content = String::new();
    let mut rest = text;
    while let Some(start) = rest.find(OPEN) {
        content.push_str(&rest[..start]);
        let after = &rest[start + OPEN.len()..];
        match after.find(CLOSE) {
            Some(end) => {
                if let Some(c) = parse_one(after[..end].trim()) {
                    calls.push(c);
                }
                rest = &after[end + CLOSE.len()..];
            }
            None => {
                // Unterminated: best-effort parse of the tail, then stop.
                if let Some(c) = parse_one(after.trim()) {
                    calls.push(c);
                }
                rest = "";
                break;
            }
        }
    }
    content.push_str(rest);
    (content.trim().to_string(), calls)
}

/// Parse one `<tool_call>` body (`{"name":…,"arguments":…}`) into a [`ParsedToolCall`], normalizing
/// `arguments` to a compact JSON string whether the model emitted it as an object (the norm) or,
/// defensively, as an already-stringified value. Returns `None` if it is not a name-bearing object.
fn parse_one(body: &str) -> Option<ParsedToolCall> {
    let v: serde_json::Value = serde_json::from_str(body).ok()?;
    let name = v.get("name")?.as_str()?.to_string();
    let arguments = match v.get("arguments") {
        Some(serde_json::Value::String(s)) => s.clone(),
        Some(other) => serde_json::to_string(other).ok()?,
        None => "{}".to_string(),
    };
    Some(ParsedToolCall { name, arguments })
}

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

    fn call(name: &str, args: &str) -> ParsedToolCall {
        ParsedToolCall {
            name: name.to_string(),
            arguments: args.to_string(),
        }
    }

    #[test]
    fn system_block_lists_each_signature_and_the_call_convention() {
        let tools = vec![serde_json::json!({
            "type": "function",
            "function": {"name": "get_weather", "parameters": {"type": "object"}}
        })];
        let s = tools_system_block(&tools);
        assert!(
            s.contains("<tools>\n{"),
            "signatures live inside <tools>: {s}"
        );
        assert!(
            s.contains("\"name\":\"get_weather\""),
            "compact json line: {s}"
        );
        assert!(s.contains("</tools>"));
        assert!(
            s.contains("<tool_call>"),
            "call convention is spelled out: {s}"
        );
    }

    #[test]
    fn parses_a_single_tool_call_and_stringifies_object_arguments() {
        let (content, calls) = parse_tool_calls(
            "<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n</tool_call>",
        );
        assert_eq!(content, "");
        assert_eq!(calls, vec![call("get_weather", "{\"city\":\"Paris\"}")]);
    }

    #[test]
    fn keeps_prose_outside_the_blocks_as_content() {
        let (content, calls) = parse_tool_calls(
            "Let me check.\n<tool_call>\n{\"name\": \"f\", \"arguments\": {}}\n</tool_call>",
        );
        assert_eq!(content, "Let me check.");
        assert_eq!(calls, vec![call("f", "{}")]);
    }

    #[test]
    fn parses_multiple_calls() {
        let (_, calls) = parse_tool_calls(
            "<tool_call>\n{\"name\": \"a\", \"arguments\": {\"x\": 1}}\n</tool_call>\
             <tool_call>\n{\"name\": \"b\", \"arguments\": {}}\n</tool_call>",
        );
        assert_eq!(calls, vec![call("a", "{\"x\":1}"), call("b", "{}")]);
    }

    #[test]
    fn plain_reply_yields_no_calls() {
        let (content, calls) = parse_tool_calls("The capital of France is Paris.");
        assert_eq!(content, "The capital of France is Paris.");
        assert!(calls.is_empty());
    }

    #[test]
    fn missing_arguments_defaults_to_empty_object() {
        let (_, calls) = parse_tool_calls("<tool_call>{\"name\": \"noop\"}</tool_call>");
        assert_eq!(calls, vec![call("noop", "{}")]);
    }

    #[test]
    fn malformed_block_is_dropped_not_leaked_as_content() {
        // A block whose body is not name-bearing JSON is neither a call nor prose.
        let (content, calls) = parse_tool_calls("<tool_call>not json</tool_call>");
        assert_eq!(content, "");
        assert!(calls.is_empty());
    }

    #[test]
    fn unterminated_trailing_block_is_parsed_best_effort() {
        // Generation hit the token cap after the JSON but before </tool_call>.
        let (_, calls) =
            parse_tool_calls("<tool_call>\n{\"name\": \"f\", \"arguments\": {\"a\": 1}}");
        assert_eq!(calls, vec![call("f", "{\"a\":1}")]);
    }
}