import argparse
from pathlib import Path

from transformers import AutoTokenizer


parser = argparse.ArgumentParser()
parser.add_argument(
    "--chat-template",
    type=Path,
    help="Path to a Jinja chat template override.",
)
args = parser.parse_args()

apply_chat_template_kwargs = {}
if args.chat_template is not None:
    apply_chat_template_kwargs["chat_template"] = (
        args.chat_template.expanduser().read_text(encoding="utf-8"))

tok = AutoTokenizer.from_pretrained(
    "google/gemma-4-31B-it",
)
messages = [
    {"role": "user", "content": "hi"},
    {
        "role": "assistant",
        "tool_calls": [
            {
                "id": "call_B",
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "arguments": {"city": "SF"},
                },
            }
        ],
    },
    {
        "role": "tool",
        "tool_call_id": "call_A",  # missing or incorrect tool_call_id
        "content": "sunny",
    },
]
print(
    tok.apply_chat_template(
        messages,
        tokenize=False,
        **apply_chat_template_kwargs,
    ))

pytest tests/renderers/test_gemma4_chat_template.py
pytest tests/tool_parsers/test_gemma4_tool_parser.py
pytest tests/tool_use/test_gemma4_responses_adjust_request.py

repro.py --chat-template examples/tool_chat_template_gemma4.jinja