"""Audit chat-template prefix preservation for tool messages.

The blog's §6 shows a simplified 12-line ``is_prefix_preserving`` for readability.
This script uses the production-hardened helper that ships in TRL —
``trl.chat_template_utils.is_chat_template_prefix_preserving`` — which adds the
edge-case handling the illustration omits (DeepSeek-V3 dict-arg fallback, VLM
processors, real tool-name in the probe, EOS-boundary alignment). The two agree on
every family below; the TRL helper is the one to actually call.

    uv pip install trl transformers
    python scripts/audit_prefix.py
"""
from transformers import AutoTokenizer
from trl.chat_template_utils import is_chat_template_prefix_preserving

MODELS = [
    ("Qwen2.5",      "Qwen/Qwen2.5-0.5B-Instruct"),
    ("Qwen3",        "Qwen/Qwen3-0.6B"),
    ("Qwen3-VL",     "Qwen/Qwen3-VL-2B-Instruct"),
    ("Llama 3.1",    "hf-internal-testing/Llama-3.1-8B-Instruct"),
    ("Llama 3.2",    "meta-llama/Llama-3.2-1B-Instruct"),
    ("DeepSeek-V3",  "deepseek-ai/DeepSeek-V3"),
    ("GLM-4.5",      "zai-org/GLM-4.5"),
    ("GPT-OSS",      "hf-internal-testing/gpt-oss-20b"),
    ("SmolLM3",      "HuggingFaceTB/SmolLM3-3B"),
]

print(f"{'family':14s} prefix-preserving for tool messages?")
print("-" * 52)
for label, mid in MODELS:
    try:
        tok = AutoTokenizer.from_pretrained(mid, trust_remote_code=True)
    except Exception as e:
        print(f"{label:14s} SKIP — load failed: {str(e).splitlines()[0][:30]}")
        continue
    try:
        ok = is_chat_template_prefix_preserving(tok)
    except Exception as e:
        print(f"{label:14s} ERROR — {type(e).__name__}: {str(e)[:30]}")
        continue
    print(f"{label:14s} {'✅ PASS' if ok else '❌ FAIL (one-line Jinja patch — see §6)'}")