inferencelayer 0.2.13

//! Architecture-based chat templating for `lfm2-serve` — the FALLBACK path.
//!
//! This module used to claim that HF Jinja "needs Python-compatible filters that no Rust Jinja
//! engine provides". That was false, and it cost us the checkpoint's own prompt format. See
//! [`super::template`]: `minijinja` plus a handful of CPython-fidelity shims renders every family
//! we serve byte-for-byte against `transformers.apply_chat_template`. When the checkpoint ships a
//! `chat_template`, that path wins.
//!
//! What remains here is the fallback, and it is still needed: checkpoints that ship no template
//! at all (converted dirs, GGUF-only), a template that fails to compile, and the
//! `OSFKB_CHAT_TEMPLATE=arch` A/B arm. It mirrors the engine-side `chat_wrap`: Gemma speaks
//! `<start_of_turn>` turns with the assistant addressed as `model`; every other arch uses the
//! ChatML `<|im_start|>` convention. The function is pure, so the exact rendered prompt is
//! unit-testable without a GPU or tokenizer.
//!
//! One difference the caller must respect: this renderer emits NO BOS and relies on the
//! tokenizer's post-processor (`encode(prompt, add_special_tokens = true)`), whereas a Jinja
//! render is complete and must be tokenized with `add_special_tokens = false`.
//!
//! Tool calling rides on the same arch match. When `tools_block` is supplied (see
//! [`super::tools::tools_system_block`]) it is folded into the system turn, prior `assistant`
//! tool calls are replayed as `<tool_call>` blocks, and `tool`-role results are wrapped in
//! `<tool_response>` — the Hermes/Qwen convention, shared across the ChatML families.

use crate::weights::Arch;

/// One tool call an assistant turn made, replayed from history. `arguments` is a raw JSON string
/// (an object literal like `{"city":"Paris"}`), inserted verbatim into the rendered `<tool_call>`.
pub struct ToolCall<'a> {
    pub name: &'a str,
    pub arguments: &'a str,
}

/// One chat turn as received on the wire. `tool_calls` is non-empty only for an `assistant` turn
/// replayed from a prior tool round; `content` carries the tool result for a `tool` turn.
#[derive(Default)]
pub struct ChatTurn<'a> {
    pub role: &'a str,
    pub content: &'a str,
    pub tool_calls: Vec<ToolCall<'a>>,
}

/// Render `turns` into the flat prompt string for `arch`, ending with the assistant generation
/// primer (`<start_of_turn>model\n` for Gemma, `<|im_start|>assistant\n` otherwise) so the model
/// continues in the assistant's voice.
///
/// `tools_block`, when `Some`, is the pre-rendered `# Tools` advertisement: it is appended to the
/// first system turn, or, when the conversation has no system turn, emitted as a synthetic leading
/// system turn — matching how the Qwen template surfaces tools.
pub fn render_chat_prompt(arch: Arch, turns: &[ChatTurn], tools_block: Option<&str>) -> String {
    let mut prompt = String::new();
    match arch {
        Arch::Gemma3 | Arch::Gemma4 => {
            // Gemma-4 differs from Gemma-3 at the template level in TWO ways, both verified
            // against the real checkpoint (tokenizer.json + chat_template.jinja + a
            // transformers apply_chat_template render):
            //  - <bos> moved OUT of the tokenizer: the gemma-4 post-processor injects NOTHING
            //    (gemma-3's injects <bos>) while its canonical jinja opens with
            //    `{{- bos_token -}}`. Render it here for Gemma4 ONLY — without it the real E2B
            //    degenerates into fluent-but-semantically-dead output (the classic missing-BOS
            //    Gemma failure); doubling it on Gemma3 would be the symmetric mistake.
            //  - The turn markers were RENAMED: `<|turn>role\n…<turn|>\n` (token 106 =
            //    <turn|>, matching the checkpoint's eos list [1, 106, 50]). Gemma-3 keeps
            //    <start_of_turn>/<end_of_turn>.
            let (open, close) = if arch == Arch::Gemma4 {
                prompt.push_str("<bos>");
                ("<|turn>", "<turn|>")
            } else {
                ("<start_of_turn>", "<end_of_turn>")
            };
            // Gemma has no standardized tool grammar; surface the block best-effort as a leading
            // user preamble so the same Hermes call convention still reaches the model.
            if let Some(block) = tools_block {
                prompt.push_str(&format!("{open}user\n{block}{close}\n"));
            }
            for t in turns {
                let role = if t.role == "assistant" {
                    "model"
                } else {
                    t.role
                };
                prompt.push_str(&format!("{open}{role}\n{}{close}\n", t.content));
            }
            prompt.push_str(&format!("{open}model\n"));
        }
        _ => {
            let first_is_system = turns.first().is_some_and(|t| t.role == "system");
            // No system turn to fold into → synthesize one carrying just the tools block.
            if let Some(block) = tools_block
                && !first_is_system
            {
                prompt.push_str(&format!("<|im_start|>system\n{block}<|im_end|>\n"));
            }
            for (i, t) in turns.iter().enumerate() {
                match t.role {
                    "system" => {
                        let mut content = t.content.to_string();
                        if i == 0
                            && let Some(block) = tools_block
                        {
                            if !content.is_empty() {
                                content.push_str("\n\n");
                            }
                            content.push_str(block);
                        }
                        prompt.push_str(&format!("<|im_start|>system\n{content}<|im_end|>\n"));
                    }
                    "assistant" => {
                        prompt.push_str(&format!("<|im_start|>assistant\n{}", t.content));
                        for tc in &t.tool_calls {
                            let args = if tc.arguments.trim().is_empty() {
                                "{}"
                            } else {
                                tc.arguments
                            };
                            prompt.push_str(&format!(
                                "\n<tool_call>\n{{\"name\": \"{}\", \"arguments\": {}}}\n</tool_call>",
                                tc.name, args
                            ));
                        }
                        prompt.push_str("<|im_end|>\n");
                    }
                    "tool" => {
                        // Consecutive tool results share one `user` turn (Qwen convention).
                        let prev_is_tool = i > 0 && turns[i - 1].role == "tool";
                        if !prev_is_tool {
                            prompt.push_str("<|im_start|>user\n");
                        }
                        prompt.push_str(&format!(
                            "<tool_response>\n{}\n</tool_response>\n",
                            t.content
                        ));
                        let next_is_tool = turns.get(i + 1).is_some_and(|n| n.role == "tool");
                        if !next_is_tool {
                            prompt.push_str("<|im_end|>\n");
                        }
                    }
                    _ => {
                        prompt.push_str(&format!(
                            "<|im_start|>{}\n{}<|im_end|>\n",
                            t.role, t.content
                        ));
                    }
                }
            }
            prompt.push_str("<|im_start|>assistant\n");
            // Qwen3.5 checkpoints ship a NO-THINK-default chat template: the assistant turn is
            // primed with an empty think block, and the model was tuned on that framing (the
            // production claim-extractor-4B's own chat_template.jinja renders exactly this;
            // `generate.rs --chat` has always added it manually). Omitting it serves the model
            // off-distribution — it may open with its own `<think>` prose, which downstream
            // grammar-constrained decoding then fights. Found by scripts/template_audit.py;
            // gated by tests/chat_template_audit.rs against the HF-jinja render.
            if arch == Arch::Qwen35 {
                prompt.push_str("<think>\n\n</think>\n\n");
            }
        }
    }
    prompt
}

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

    /// A plain (no-tool-call) turn.
    fn turn<'a>(role: &'a str, content: &'a str) -> ChatTurn<'a> {
        ChatTurn {
            role,
            content,
            ..Default::default()
        }
    }

    #[test]
    fn should_render_chatml_for_qwen_family() {
        let turns = [
            turn("user", "hi"),
            turn("assistant", "hello"),
            turn("user", "bye"),
        ];
        let p = render_chat_prompt(Arch::Qwen3, &turns, None);
        assert_eq!(
            p,
            "<|im_start|>user\nhi<|im_end|>\n\
             <|im_start|>assistant\nhello<|im_end|>\n\
             <|im_start|>user\nbye<|im_end|>\n\
             <|im_start|>assistant\n"
        );
    }

    #[test]
    fn should_render_gemma_turns_and_map_assistant_to_model() {
        let turns = [turn("user", "hi"), turn("assistant", "hello")];
        let p = render_chat_prompt(Arch::Gemma3, &turns, None);
        assert_eq!(
            p,
            "<start_of_turn>user\nhi<end_of_turn>\n\
             <start_of_turn>model\nhello<end_of_turn>\n\
             <start_of_turn>model\n"
        );
        // The Gemma assistant is addressed as `model`, never `assistant`.
        assert!(!p.contains("assistant"));
    }

    #[test]
    fn should_prime_assistant_turn_for_empty_history() {
        assert_eq!(
            render_chat_prompt(Arch::Lfm2, &[], None),
            "<|im_start|>assistant\n"
        );
        // Gemma-4's canonical shape (verified against the real E2B checkpoint's jinja via
        // transformers apply_chat_template): <bos> rendered by the TEMPLATE (the gemma-4
        // tokenizer post-processor injects nothing), and RENAMED turn markers <|turn>/<turn|>.
        assert_eq!(
            render_chat_prompt(Arch::Gemma4, &[], None),
            "<bos><|turn>model\n"
        );
    }

    #[test]
    fn tools_block_folds_into_an_existing_system_turn() {
        let turns = [turn("system", "You are helpful."), turn("user", "hi")];
        let p = render_chat_prompt(Arch::Qwen3, &turns, Some("# Tools\n<tools></tools>"));
        assert_eq!(
            p,
            "<|im_start|>system\nYou are helpful.\n\n# Tools\n<tools></tools><|im_end|>\n\
             <|im_start|>user\nhi<|im_end|>\n\
             <|im_start|>assistant\n"
        );
    }

    #[test]
    fn tools_block_synthesizes_a_system_turn_when_absent() {
        let turns = [turn("user", "hi")];
        let p = render_chat_prompt(Arch::Qwen3, &turns, Some("# Tools"));
        assert_eq!(
            p,
            "<|im_start|>system\n# Tools<|im_end|>\n\
             <|im_start|>user\nhi<|im_end|>\n\
             <|im_start|>assistant\n"
        );
    }

    #[test]
    fn replays_assistant_tool_call_and_grouped_tool_results() {
        let turns = [
            turn("user", "weather in Paris and Rome?"),
            ChatTurn {
                role: "assistant",
                content: "",
                tool_calls: vec![
                    ToolCall {
                        name: "get_weather",
                        arguments: "{\"city\":\"Paris\"}",
                    },
                    ToolCall {
                        name: "get_weather",
                        arguments: "{\"city\":\"Rome\"}",
                    },
                ],
            },
            turn("tool", "18C"),
            turn("tool", "24C"),
        ];
        let p = render_chat_prompt(Arch::Qwen3, &turns, None);
        assert_eq!(
            p,
            "<|im_start|>user\nweather in Paris and Rome?<|im_end|>\n\
             <|im_start|>assistant\n\
             \n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\":\"Paris\"}}\n</tool_call>\
             \n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\":\"Rome\"}}\n</tool_call>\
             <|im_end|>\n\
             <|im_start|>user\n\
             <tool_response>\n18C\n</tool_response>\n\
             <tool_response>\n24C\n</tool_response>\n\
             <|im_end|>\n\
             <|im_start|>assistant\n"
        );
    }

    #[test]
    fn empty_tool_call_arguments_default_to_empty_object() {
        let turns = [ChatTurn {
            role: "assistant",
            content: "",
            tool_calls: vec![ToolCall {
                name: "noop",
                arguments: "",
            }],
        }];
        let p = render_chat_prompt(Arch::Qwen3, &turns, None);
        assert!(
            p.contains("{\"name\": \"noop\", \"arguments\": {}}"),
            "empty args render as {{}}: {p}"
        );
    }
}