splintr 0.16.1

//! Compile-time agent-token ids, one module per bundled vocabulary.
//!
//! The constants themselves are generated — `include!`d from
//! `agent_tokens_generated.rs`, which
//! `scripts/generate_agent_tokens.py --lang rust` writes from the same table
//! that produces the Python `*_AGENT_TOKENS` classes. That shared table is the
//! point: the two languages drifted for a whole release cycle when Rust's
//! constants were hand-written and Python's were generated, leaving Rust with
//! two vocabularies and Python with seven.
//!
//! This file is the hand-written half, and holds only the tests that pin the
//! generated ids against the tokenizers they name — so regenerating never
//! clobbers them.

include!("agent_tokens_generated.rs");

#[cfg(test)]
mod tests {
    use crate::core::pretrained::{base_vocab_size, from_pretrained, PretrainedVocab};
    use crate::core::Tokenize;

    /// One row per module: the vocabulary it names, and a sample of its
    /// constants spanning the block — the first slot, a middle one, the two
    /// vocabularies most often claim for themselves, and the last.
    #[allow(clippy::type_complexity)]
    fn samples() -> Vec<(&'static str, [(&'static str, u32); 6])> {
        macro_rules! row {
            ($vocab:literal, $m:ident) => {
                (
                    $vocab,
                    [
                        ("<|system|>", super::$m::SYSTEM),
                        ("<|im_start|>", super::$m::IM_START),
                        ("<|im_end|>", super::$m::IM_END),
                        ("<|think|>", super::$m::THINK),
                        ("<|pad|>", super::$m::PAD),
                        ("<|/summary|>", super::$m::SUMMARY_END),
                    ],
                )
            };
        }
        vec![
            row!("cl100k_base", cl100k_agent_tokens),
            row!("o200k_base", o200k_agent_tokens),
            row!("gpt-oss", gpt_oss_agent_tokens),
            row!("llama3", llama3_agent_tokens),
            row!("qwen3", qwen3_agent_tokens),
            row!("glm4", glm4_agent_tokens),
            row!("deepseek_v3", deepseek_v3_agent_tokens),
            row!("mistral_v1", mistral_v1_agent_tokens),
            row!("mistral_v2", mistral_v2_agent_tokens),
            row!("mistral_v3", mistral_v3_agent_tokens),
            row!("kimi_k2", kimi_k2_agent_tokens),
            row!("kimi_k3", kimi_k3_agent_tokens),
        ]
    }

    /// Every generated constant must equal the id its tokenizer resolves that
    /// token to.
    ///
    /// This file and `insert_agent_tokens` in `pretrained.rs` are two
    /// independent statements of the same ids — this one generated by
    /// `scripts/generate_agent_tokens.py --lang rust`, that one written by
    /// hand. A constant is what callers reach for precisely so they need not
    /// think about the id, so a stale one is wrong in the worst way: silently.
    /// The Python half of the same contract is pinned by
    /// `python/tests/test_agent_token_constants.py`.
    #[test]
    fn generated_constants_match_the_tokenizers_they_name() {
        for (vocab, sample) in samples() {
            let tokenizer = from_pretrained(vocab).expect("bundled vocabulary loads");
            for (token, expected) in sample {
                assert_eq!(
                    tokenizer.special_token_id(token),
                    Some(expected),
                    "{vocab}: {token} disagrees with the vocabulary — regenerate with \
                     scripts/generate_agent_tokens.py --lang rust"
                );
            }
        }
    }

    /// Vocabularies that ship an agent-token name themselves keep their own id,
    /// below `base_vocab_size` — the id the checkpoint was trained on. A
    /// constant pointing at a splintr-appended id instead would build chat
    /// templates the model never saw.
    #[test]
    fn constants_defer_to_a_vocabularys_own_ids() {
        let qwen_base = base_vocab_size(PretrainedVocab::Qwen3);
        assert_eq!(super::qwen3_agent_tokens::IM_START, 151644);
        assert_eq!(super::qwen3_agent_tokens::IM_END, 151645);
        assert!(super::qwen3_agent_tokens::IM_START < qwen_base);
        // A name Qwen does not define still comes from the appended block.
        assert_eq!(super::qwen3_agent_tokens::SYSTEM, qwen_base);

        let glm_base = base_vocab_size(PretrainedVocab::Glm4);
        assert_eq!(super::glm4_agent_tokens::SYSTEM, 151335);
        assert_eq!(super::glm4_agent_tokens::IMAGE, 151363);
        assert!(super::glm4_agent_tokens::SYSTEM < glm_base);
        // GLM names only the opening markers, so the closing ones stay in the block.
        assert_eq!(super::glm4_agent_tokens::IMAGE_END, glm_base + 43);
    }

    /// Whisper is the one bundled vocabulary with no module here: it carries no
    /// agent tokens. Pinned so "no module" stays a stated fact rather than an
    /// omission someone later fills in by mistake.
    #[test]
    fn whisper_carries_no_agent_tokens() {
        let tokenizer = from_pretrained("whisper").expect("bundled vocabulary loads");
        assert_eq!(tokenizer.special_token_id("<|think|>"), None);
        assert_eq!(tokenizer.special_token_id("<|pad|>"), None);
        assert_eq!(
            base_vocab_size(PretrainedVocab::WhisperV2),
            tokenizer.vocab_size() as u32,
            "with no agent tokens, the base size is the whole vocabulary"
        );
    }
}