// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
​
//! Chat-template 渲染模块，用于缓存感知路由。
//!
//! 引擎对应用聊天模板后的 token 序列进行 KV 缓存。路由器必须在哈希前
//! 渲染相同的模板，否则查询哈希与引擎缓存不匹配，导致重叠块始终为零。
​
use anyhow::{Context, Result};
use minijinja::{Environment, UndefinedBehavior};
use std::collections::BTreeMap;
​
/// 在环境注册的固定模板名称。
const TEMPLATE_NAME: &str = "chat";
​
/// HuggingFace 注入模板上下文中的特殊 token 键列表。
const SPECIAL_TOKEN_KEYS: [&str; 7] = [
    "bos_token", "eos_token", "unk_token", "sep_token",
    "pad_token", "cls_token", "mask_token",
];
​
/// 编译后的聊天模板及所引用的特殊 token 字符串。
pub struct ChatTemplate {
    env: Environment<'static>,
    /// `(名称, token 字符串)` 对，缺失的 token 为空字符串。
    special_tokens: Vec<(&'static str, String)>,
}
​
impl ChatTemplate {
    /// 从 `tokenizer_config.json` 构建聊天模板。
    ///
    /// 若配置中无 `chat_template` 则返回 `Ok(None)`，模型将走原始文本路由。
    pub fn from_tokenizer_config(cfg: &serde_json::Value) -> Result<Option<Self>> {
        // 提取模板源码。仅支持 String 和 Array 变体；
        // 若为 Object（多命名模板）则返回 None，存在兼容性风险。
        let Some(template_src) = extract_chat_template(cfg) else {
            return Ok(None);
        };
​
        // 收集特殊 token，缺失时用空字符串替代（匹配 jinja2 `none`→`""` 行为）
        let special_tokens = SPECIAL_TOKEN_KEYS
            .iter()
            .map(|&key| (key, extract_token_str(cfg, key).unwrap_or_default()))
            .collect();
​
        let mut env = Environment::new();
​
        // 对齐 HuggingFace 的 trim_blocks + lstrip_blocks，否则空白字符差异
        // 将导致每个块的哈希值永久偏离引擎的缓存。
        env.set_trim_blocks(true);
        env.set_lstrip_blocks(true);
​
        // 使用半严格未定义行为：渲染未提供的变量将产生错误，
        // 从而使调用者降级为原始文本哈希，避免产生看似合理但实际不匹配的 prompt。
        env.set_undefined_behavior(UndefinedBehavior::SemiStrict);
​
        // 添加 Python 字符串 / 字典方法的兼容回调，用于渲染真实聊天模板。
        env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
​
        // 暴露 `raise_exception` 和 `strftime_now` 函数
        env.add_function("raise_exception", raise_exception);
        env.add_function("strftime_now", strftime_now);
​
        // 编译模板
        env.add_template_owned(TEMPLATE_NAME, template_src)
            .context("compile chat template from tokenizer_config.json")?;
​
        Ok(Some(Self { env, special_tokens }))
    }
}

// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
​
//! DeepSeek-V4 prompt 编码器，用于缓存感知路由。
//!
//! DSV4 没有 Jinja 模板；引擎在 Python 中构建 prompt。此模块复现了其核心逻辑，
//! 使路由器哈希的 token 序列与引擎缓存一致。
​
/// 以下常量与引擎中定义的 marker token 对应的字面量一致。
const BOS: &str = "<｜begin▁of▁sentence｜>"; // token id 0
const EOS: &str = "<｜end▁of▁sentence｜>"; // token id 1
const USER: &str = "<｜User｜>"; // token id 128803
const ASSISTANT: &str = "<｜Assistant｜>"; // token id 128804
const THINK_END: &str = "</think>"; // token id 128822
​
/// 将 `messages` 数组渲染为 DSV4 聊天 prompt 字符串。
///
/// 与引擎的 `encoding_dsv4.encode_messages` 对齐（仅聊天模式、纯文本、无 tools/tasks）。
pub fn render_messages(messages: &serde_json::Value) -> String {
    let mut msgs: Vec<(String, String)> = messages
        .as_array()
        .map(|arr| {
            arr.iter()
                .map(|m| {
                    let role = m.get("role")
                        .and_then(|r| r.as_str())
                        .unwrap_or("")
                        .to_string();
                    (role, content_to_string(m.get("content")))
                })
                .collect()
        })
        .unwrap_or_default();
​
    // 引擎会隐式插入空 system 消息（若非 system 开头）
    if msgs.first().map(|(r, _)| r != "system").unwrap_or(true) {
        msgs.insert(0, ("system".to_string(), String::new()));
    }
​
    // 合并连续 user 轮次（用 \n\n 连接）
    merge_consecutive_user_turns(&mut msgs);
​
    let mut out = String::from(BOS);
    for i in 0..msgs.len() {
        render_one(i, &msgs, &mut out);
    }
    out
}
​
/// 合并连续 user 轮次，匹配引擎预处理。
fn merge_consecutive_user_turns(msgs: &mut Vec<(String, String)>) {
    let mut merged: Vec<(String, String)> = Vec::with_capacity(msgs.len());
    for (role, content) in msgs.drain(..) {
        match merged.last_mut() {
            Some((last_role, last_content)) if last_role == "user" && role == "user" => {
                last_content.push_str("\n\n");
                last_content.push_str(&content);
            }
            _ => merged.push((role, content)),
        }
    }
    *msgs = merged;
}
​
/// 渲染单个轮次到 out 字符串。
fn render_one(i: usize, msgs: &[(String, String)], out: &mut String) {
    let (role, content) = &msgs[i];
    match role.as_str() {
        "system" => out.push_str(content),
        "user" | "developer" => {
            out.push_str(USER);
            out.push_str(content);
        }
        "assistant" => {
            // 聊天模式无推理块，直接添加标记
            out.push_str(ASSISTANT);
            out.push_str(content);
            out.push_str(EOS);
        }
        _ => {} // 忽略未知 role
    }
}