init

.gitignore

@@ -0,0 +1,14 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+
+# Virtual environments
+.venv/
+venv/
+
+# uv
+uv.lock
+
+# OS
+.DS_Store

Makefile

@@ -0,0 +1,37 @@
+XGRAMMAR_DIR ?= $(HOME)/git/xgrammar
+VENV         := .venv
+PY           := $(VENV)/bin/python
+SAMPLES      ?= 4
+TEMPERATURE  ?= 0.9
+
+.PHONY: run smoke setup xgrammar clean help
+
+help:
+	@echo "make setup     - create uv venv and install deps + local xgrammar (editable)"
+	@echo "make run       - run all scenarios and write report.md (constrained vs unconstrained)"
+	@echo "make smoke     - quick single-tool smoke test"
+	@echo "make xgrammar  - reinstall local xgrammar after C++/python changes"
+	@echo "make clean     - remove venv and report"
+
+$(VENV)/.stamp:
+	uv venv --python 3.13 $(VENV)
+	uv pip install --python $(PY) torch transformers accelerate jsonschema
+	uv pip install --python $(PY) -e $(XGRAMMAR_DIR)
+	touch $@
+
+setup: $(VENV)/.stamp
+
+xgrammar: $(VENV)/.stamp
+	uv pip install --python $(PY) -e $(XGRAMMAR_DIR)
+
+run: $(VENV)/.stamp
+	$(PY) test_gemma4_scenarios.py --samples $(SAMPLES) --temperature $(TEMPERATURE) --report report.md
+	@echo ""
+	@echo "==================== report.md ===================="
+	@cat report.md
+
+smoke: $(VENV)/.stamp
+	$(PY) test_gemma4_xgrammar.py
+
+clean:
+	rm -rf $(VENV) report.md

README.md

@@ -0,0 +1,65 @@
+# gemma-4 × xgrammar constrained decoding test
+
+A collection of scripts that verify the modified xgrammar (`~/git/xgrammar`,
+`style="gemma"` / `gemma_4` builtin structural tag) works correctly for tool
+calling with `google/gemma-4-E2B-it`, and generate a report comparing
+constrained vs unconstrained behavior.
+
+## Usage (Makefile)
+
+```bash
+make setup     # create uv venv + install deps + local xgrammar editable install (once)
+make run       # run all scenarios → generate report.md + print output  ← for PR reports
+make smoke     # single-tool smoke test (greedy, fast)
+make xgrammar  # reinstall after modifying ~/git/xgrammar
+make clean     # remove venv/report
+```
+
+Options: `make run SAMPLES=8`, `XGRAMMAR_DIR=/path/to/xgrammar make setup`
+
+The output of `make run`, **`report.md`**, is the artifact to show reviewers —
+it contains the model/sampling configuration, a per-scenario table of
+constrained vs unconstrained metrics, and the raw text of failure cases.
+
+## Scenarios (`test_gemma4_scenarios.py`)
+
+| | Setup | What it targets |
+|---|---|---|
+| **A. complex-schema** | Nested objects + array of objects + enum + boolean schema, titles containing quotes | Argument serialization mistakes |
+| **B. multiturn-thinking** | Two prior rounds of tool call/response already in context and thinking is required, the next call takes a number-typed argument | Skipped/unterminated thinking, type errors |
+| **C. adversarial-payload** | A string argument whose body must contain JSON/braces/an error log verbatim, an out-of-enum word ("critical") lure, two similarly named distractor tools, temp 1.5 + top_p 1.0 | String quoting collapse, enum violations |
+
+Metrics (pass/fail per sample):
+
+| metric | meaning |
+|---|---|
+| `well_formed` | Every `<\|tool_call>call:name{...}<tool_call\|>` block is complete and parseable |
+| `valid_name` | Only calls tools that actually exist |
+| `schema_valid` | Parsed arguments pass the JSON schema (parsed via a port of the sglang parser) |
+| `stops_at_boundary` | Generation doesn't run past `<tool_call\|>` into `<\|tool_response>` |
+| `thought_ok` | (thinking scenario) Opens the thought channel and closes it before the tool call |
+
+## Key findings
+
+- gemma-4-E2B-it's format training is very robust, so the tool call argument
+  format itself rarely breaks even at temp 1.5 + top_p 1.0.
+- Per the [Gemma 4 prompt-formatting spec](https://ai.google.dev/gemma/docs/core/prompt-formatting-gemma4),
+  the model is only responsible for generating up to `<tool_call\|>`;
+  `<\|tool_response>` is appended by the application with the real tool
+  result, and is registered as an **additional stop sequence** purely as a
+  backstop. This harness has no such stop configured, so unconstrained runs
+  keep generating past that boundary into engine-owned territory
+  (`stops_at_boundary` failures); required-mode constrained decoding ends the
+  call cleanly at an accept state instead.
+- Unconstrained decoding also **omits the thought channel entirely** in
+  situations where thinking should be enabled; constrained decoding with
+  `reasoning=True` enforces it (an empty thought `<\|channel>thought\n<channel\|>`
+  remains legal, matching the spec's no-thinking form).
+
+## Files
+
+- `test_gemma4_xgrammar.py` — basic smoke test (greedy, single tool, alignment check)
+- `test_gemma4_scenarios.py` — scenario runner + report.md generator.
+  `--scenario a|b|c|all --samples N --temperature T --model ID --report PATH`
+- `gemma_parser.py` — a pure port of sglang's `Gemma4Detector` parsing logic (for output verification)
+- `Makefile` / `report.md`

gemma_parser.py

@@ -0,0 +1,215 @@
+"""Minimal Gemma4 tool-call parser.
+
+Pure-function port of sglang's Gemma4Detector parsing helpers
+(sglang/python/sglang/srt/function_call/gemma4_detector.py), so generated
+outputs can be validated exactly the way sglang would parse them, without
+importing sglang.
+"""
+
+TOOL_CALL_START = "<|tool_call>"
+TOOL_CALL_END = "<tool_call|>"
+STRING_DELIM = '<|"|>'
+
+
+def _parse_value(value_str: str) -> object:
+    value_str = value_str.strip()
+    if not value_str:
+        return value_str
+    if value_str == "true":
+        return True
+    if value_str == "false":
+        return False
+    try:
+        if "." in value_str:
+            return float(value_str)
+        return int(value_str)
+    except ValueError:
+        pass
+    return value_str
+
+
+def _parse_array(arr_str: str) -> list:
+    items: list = []
+    i = 0
+    n = len(arr_str)
+    while i < n:
+        while i < n and arr_str[i] in (" ", ",", "\n", "\t"):
+            i += 1
+        if i >= n:
+            break
+        if arr_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
+            i += len(STRING_DELIM)
+            end_pos = arr_str.find(STRING_DELIM, i)
+            if end_pos == -1:
+                items.append(arr_str[i:])
+                break
+            items.append(arr_str[i:end_pos])
+            i = end_pos + len(STRING_DELIM)
+        elif arr_str[i] == "{":
+            depth = 1
+            obj_start = i + 1
+            i += 1
+            while i < n and depth > 0:
+                if arr_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
+                    i += len(STRING_DELIM)
+                    next_delim = arr_str.find(STRING_DELIM, i)
+                    i = next_delim + len(STRING_DELIM) if next_delim != -1 else n
+                    continue
+                if arr_str[i] == "{":
+                    depth += 1
+                elif arr_str[i] == "}":
+                    depth -= 1
+                i += 1
+            items.append(parse_args(arr_str[obj_start : i - 1]))
+        elif arr_str[i] == "[":
+            depth = 1
+            sub_start = i + 1
+            i += 1
+            while i < n and depth > 0:
+                if arr_str[i] == "[":
+                    depth += 1
+                elif arr_str[i] == "]":
+                    depth -= 1
+                i += 1
+            items.append(_parse_array(arr_str[sub_start : i - 1]))
+        else:
+            val_start = i
+            while i < n and arr_str[i] not in (",", "]"):
+                i += 1
+            items.append(_parse_value(arr_str[val_start:i]))
+    return items
+
+
+def parse_args(args_str: str) -> dict:
+    """Parse Gemma4's key:value argument format into a Python dict."""
+    if not args_str or not args_str.strip():
+        return {}
+    result: dict = {}
+    i = 0
+    n = len(args_str)
+    while i < n:
+        while i < n and args_str[i] in (" ", ",", "\n", "\t"):
+            i += 1
+        if i >= n:
+            break
+        key_start = i
+        while i < n and args_str[i] != ":":
+            i += 1
+        if i >= n:
+            break
+        key = args_str[key_start:i].strip()
+        i += 1
+        if i >= n:
+            result[key] = ""
+            break
+        while i < n and args_str[i] in (" ", "\n", "\t"):
+            i += 1
+        if i >= n:
+            result[key] = ""
+            break
+        if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
+            i += len(STRING_DELIM)
+            val_start = i
+            end_pos = args_str.find(STRING_DELIM, i)
+            if end_pos == -1:
+                result[key] = args_str[val_start:]
+                break
+            result[key] = args_str[val_start:end_pos]
+            i = end_pos + len(STRING_DELIM)
+        elif args_str[i] == "{":
+            depth = 1
+            obj_start = i + 1
+            i += 1
+            while i < n and depth > 0:
+                if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
+                    i += len(STRING_DELIM)
+                    next_delim = args_str.find(STRING_DELIM, i)
+                    i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
+                    continue
+                if args_str[i] == "{":
+                    depth += 1
+                elif args_str[i] == "}":
+                    depth -= 1
+                i += 1
+            result[key] = parse_args(args_str[obj_start : i - 1])
+        elif args_str[i] == "[":
+            depth = 1
+            arr_start = i + 1
+            i += 1
+            while i < n and depth > 0:
+                if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
+                    i += len(STRING_DELIM)
+                    next_delim = args_str.find(STRING_DELIM, i)
+                    i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
+                    continue
+                if args_str[i] == "[":
+                    depth += 1
+                elif args_str[i] == "]":
+                    depth -= 1
+                i += 1
+            result[key] = _parse_array(args_str[arr_start : i - 1])
+        else:
+            val_start = i
+            while i < n and args_str[i] not in (",", "}", "]"):
+                i += 1
+            result[key] = _parse_value(args_str[val_start:i])
+    return result
+
+
+def _find_matching_brace(text: str) -> int:
+    """Index of the matching '}' in text (which starts just after '{'), or -1."""
+    depth = 1
+    i = 0
+    n = len(text)
+    dlen = len(STRING_DELIM)
+    while i < n and depth > 0:
+        if text[i : i + dlen] == STRING_DELIM:
+            i += dlen
+            next_delim = text.find(STRING_DELIM, i)
+            if next_delim == -1:
+                return -1
+            i = next_delim + dlen
+            continue
+        if text[i] == "{":
+            depth += 1
+        elif text[i] == "}":
+            depth -= 1
+        i += 1
+    return (i - 1) if depth == 0 else -1
+
+
+def extract_tool_calls(text: str) -> list:
+    """Extract [(func_name, args_dict_or_None)] from text.
+
+    args is None when the call block is malformed (unparseable), mirroring
+    where sglang's detector would fail.
+    """
+    results = []
+    search_from = 0
+    while True:
+        start = text.find(TOOL_CALL_START, search_from)
+        if start == -1:
+            break
+        end = text.find(TOOL_CALL_END, start)
+        if end == -1:
+            # unterminated call block
+            results.append((None, None))
+            break
+        inner = text[start + len(TOOL_CALL_START) : end]
+        if inner.startswith("call:"):
+            brace = inner.find("{")
+            if brace != -1:
+                func_name = inner[5:brace]
+                args_content = inner[brace + 1 :]
+                match_idx = _find_matching_brace(args_content)
+                args_str = args_content[:match_idx] if match_idx != -1 else args_content
+                try:
+                    results.append((func_name, parse_args(args_str)))
+                except Exception:
+                    results.append((func_name, None))
+            else:
+                results.append((None, None))
+        else:
+            results.append((None, None))
+        search_from = end + len(TOOL_CALL_END)
+    return results

report.md

@@ -0,0 +1,128 @@
+# Gemma-4 tool calling: constrained (xgrammar) vs unconstrained
+
+- model: `google/gemma-4-E2B-it` via HF transformers
+- constraint: `xgr.get_model_structural_tag("gemma_4", tools=..., tool_choice="required")` compiled and applied with `xgr.contrib.hf.LogitsProcessor`
+- 4 samples per scenario per mode; identical sampling settings for both modes (per-scenario values below)
+- outputs validated with a pure-Python port of sglang's `Gemma4Detector` parser plus `jsonschema`
+
+## Results (passing samples / total)
+
+### A. complex-schema
+
+temperature=0.9, top_p=0.95, max_new_tokens=384, reasoning=False
+
+| metric | constrained | unconstrained |
+|---|---|---|
+| well_formed | 4/4 ✅ | 4/4 ✅ |
+| valid_name | 4/4 ✅ | 4/4 ✅ |
+| schema_valid | 4/4 ✅ | 4/4 ✅ |
+| stops_at_boundary | 4/4 ✅ | 0/4 ❌ |
+
+### B. multiturn-thinking
+
+temperature=0.9, top_p=0.95, max_new_tokens=384, reasoning=True
+
+| metric | constrained | unconstrained |
+|---|---|---|
+| well_formed | 4/4 ✅ | 4/4 ✅ |
+| valid_name | 4/4 ✅ | 4/4 ✅ |
+| schema_valid | 4/4 ✅ | 4/4 ✅ |
+| stops_at_boundary | 4/4 ✅ | 0/4 ❌ |
+| thought_ok | 4/4 ✅ | 0/4 ❌ |
+
+### C. adversarial-payload
+
+temperature=1.5, top_p=1.0, max_new_tokens=640, reasoning=False
+
+| metric | constrained | unconstrained |
+|---|---|---|
+| well_formed | 4/4 ✅ | 4/4 ✅ |
+| valid_name | 4/4 ✅ | 4/4 ✅ |
+| schema_valid | 4/4 ✅ | 4/4 ✅ |
+| stops_at_boundary | 4/4 ✅ | 0/4 ❌ |
+
+## Metric definitions
+
+- `well_formed` — every tool call block is complete and parseable (sglang parser port)
+- `valid_name` — every called tool exists
+- `schema_valid` — parsed arguments validate against the tool's JSON schema
+- `stops_at_boundary` — generation did not run past <tool_call|> into <|tool_response>
+- `thought_ok` — <|channel>thought opened and closed before the first tool call
+
+## Failure examples
+
+**A. complex-schema / unconstrained #0** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
+```
+
+**A. complex-schema / unconstrained #1** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
+```
+
+**A. complex-schema / unconstrained #2** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
+```
+
+**A. complex-schema / unconstrained #3** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_calendar_event{attendees:[{email:<|"|>alice.kim@example.com<|"|>,name:<|"|>Alice Kim<|"|>},{email:<|"|>bob.lee@example.com<|"|>,name:<|"|>Bob Lee<|"|>},{email:<|"|>chris.park@example.com<|"|>,name:<|"|>Chris Park<|"|>}],duration_minutes:90,location:{floor:7,room:<|"|>Jupiter<|"|>},priority:<|"|>high<|"|>,send_invites:true,start:<|"|>2026-07-10T14:00:00+09:00<|"|>,title:<|"|>Q3 Roadmap Review ("final" draft)<|"|>}<tool_call|><|tool_response>
+```
+
+**B. multiturn-thinking / unconstrained #0** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
+
+```
+<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
+```
+
+**B. multiturn-thinking / unconstrained #1** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
+
+```
+<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
+```
+
+**B. multiturn-thinking / unconstrained #2** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
+
+```
+<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
+```
+
+**B. multiturn-thinking / unconstrained #3** — failed `stops_at_boundary`, `thought_ok`: ran past <tool_call|> into <|tool_response> after the call; skipped or never closed the <|channel>thought section
+
+```
+<|tool_call>call:convert_temperature{from_unit:<|"|>celsius<|"|>,to_unit:<|"|>fahrenheit<|"|>,value:3}<tool_call|><|tool_response>
+```
+
+**C. adversarial-payload / unconstrained #0** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
+Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42).<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
+```
+
+**C. adversarial-payload / unconstrained #1** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
+Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42).<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>bug<|"|>,<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
+```
+
+**C. adversarial-payload / unconstrained #2** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
+Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42)<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>bug<|"|>,<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
+```
+
+**C. adversarial-payload / unconstrained #3** — failed `stops_at_boundary`: ran past <tool_call|> into <|tool_response> after the call
+
+```
+<|tool_call>call:create_ticket{affected:{regions:[<|"|>ap-northeast-2<|"|>,<|"|>us-east-1<|"|>],service:<|"|>payments<|"|>,version:<|"|>2.14.1<|"|>},body:<|"|>Deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], "jitter": true}}
+Error line: TypeError: cannot destructure {id: undefined} at applyRetry (retry.js:42).<|"|>,cc_emails:[<|"|>dev-alerts@example.com<|"|>],estimated_minutes:45,labels:[<|"|>bug<|"|>,<|"|>regression<|"|>,<|"|>backend<|"|>],severity:<|"|>blocker<|"|>,title:<|"|>Payments retry storm after config rollout<|"|>,urgent_escalation:true}<tool_call|><|tool_response>
+```

test_gemma4_scenarios.py

@@ -0,0 +1,546 @@
+"""Constrained-vs-unconstrained tool calling scenarios for Gemma-4 + xgrammar.
+
+Compares grammar-constrained generation (xgrammar gemma_4 builtin structural
+tag, JSONSchemaFormat style="gemma") against unconstrained generation on
+google/gemma-4-E2B-it, across scenarios designed to break tool call syntax:
+
+  A. complex-schema      — nested objects, array of objects, enum, boolean,
+                           integer; prompt values bait JSON-style quoting.
+  B. multiturn-thinking  — two prior tool call/response rounds in context,
+                           thinking enabled; the model must close a
+                           <|channel>thought section, then call a *different*
+                           tool whose `value` argument must be a number.
+  C. adversarial-payload — a string argument that must contain a JSON snippet
+                           full of braces/quotes (breaks brace-matching if the
+                           model quotes it wrong), an enum baited with a word
+                           that is NOT in the enum ("critical"), and two
+                           similarly-named distractor tools. Runs at elevated
+                           temperature (1.3) to expose format instability.
+
+Metrics per sample:
+  well_formed   — every <|tool_call> block is complete and parseable by the
+                  sglang Gemma4Detector parsing algorithm (gemma_parser.py)
+  valid_name    — every called tool actually exists
+  schema_valid  — parsed arguments validate against the tool's JSON schema
+  stops_at_boundary — generation did not run past <tool_call|> into the
+                  engine-owned <|tool_response> territory
+  thought_ok    — (thinking scenarios) <|channel>thought opened and closed
+                  before the first tool call
+
+Usage:
+  python test_gemma4_scenarios.py [--scenario a|b|c|all] [--samples 4]
+                                  [--max-new-tokens 384] [--temperature 0.9]
+                                  [--model google/gemma-4-E2B-it]
+                                  [--report report.md]
+"""
+
+import argparse
+import json
+import time
+
+import jsonschema
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+import xgrammar as xgr
+from gemma_parser import extract_tool_calls
+
+DEFAULT_MODEL_ID = "google/gemma-4-E2B-it"
+
+# ---------------------------------------------------------------- scenario A
+
+CALENDAR_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "create_calendar_event",
+        "description": "Create a calendar event and optionally send invites.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "title": {"type": "string"},
+                "start": {"type": "string", "description": "ISO 8601 datetime"},
+                "duration_minutes": {"type": "integer"},
+                "priority": {"type": "string", "enum": ["low", "medium", "high"]},
+                "attendees": {
+                    "type": "array",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "name": {"type": "string"},
+                            "email": {"type": "string"},
+                        },
+                        "required": ["name", "email"],
+                    },
+                },
+                "location": {
+                    "type": "object",
+                    "properties": {
+                        "room": {"type": "string"},
+                        "floor": {"type": "integer"},
+                    },
+                    "required": ["room"],
+                },
+                "send_invites": {"type": "boolean"},
+            },
+            "required": ["title", "start", "duration_minutes", "priority", "attendees"],
+        },
+    },
+}
+
+SCENARIO_A = {
+    "key": "a",
+    "name": "A. complex-schema",
+    "tools": [CALENDAR_TOOL],
+    "messages": [
+        {
+            "role": "user",
+            "content": (
+                'Schedule a high-priority meeting titled Q3 Roadmap Review ("final" draft) '
+                "on 2026-07-10 at 14:00 KST for 90 minutes in room Jupiter on the 7th floor. "
+                "Invite Alice Kim <alice.kim@example.com>, Bob Lee <bob.lee@example.com> and "
+                "Chris Park <chris.park@example.com>, and send the invites."
+            ),
+        }
+    ],
+    "reasoning": False,
+    "check_thought": False,
+}
+
+# ---------------------------------------------------------------- scenario B
+
+WEATHER_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "get_weather",
+        "description": "Get the current weather for a city.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "city": {"type": "string"},
+                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+            },
+            "required": ["city", "unit"],
+        },
+    },
+}
+
+CONVERT_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "convert_temperature",
+        "description": "Convert a temperature value between units.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "value": {"type": "number", "description": "numeric temperature value"},
+                "from_unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                "to_unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+            },
+            "required": ["value", "from_unit", "to_unit"],
+        },
+    },
+}
+
+SCENARIO_B = {
+    "key": "b",
+    "name": "B. multiturn-thinking",
+    "tools": [WEATHER_TOOL, CONVERT_TOOL],
+    "messages": [
+        {
+            "role": "user",
+            "content": (
+                "Compare the current weather in Seoul and Busan in celsius, one city at a "
+                "time. After both, also convert Seoul's temperature to fahrenheit with the "
+                "converter tool before answering."
+            ),
+        },
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": {"city": "Seoul", "unit": "celsius"},
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "name": "get_weather",
+            "content": '{"temp_c": 3, "condition": "sunny", "humidity": 41}',
+        },
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": {"city": "Busan", "unit": "celsius"},
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "name": "get_weather",
+            "content": '{"temp_c": 8, "condition": "cloudy", "humidity": 63}',
+        },
+    ],
+    "reasoning": True,
+    "check_thought": True,
+}
+
+# ---------------------------------------------------------------- scenario C
+
+TICKET_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "create_ticket",
+        "description": "Create a bug ticket in the issue tracker.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "title": {"type": "string"},
+                "body": {
+                    "type": "string",
+                    "description": "Full description. Include configs/logs verbatim.",
+                },
+                "severity": {"type": "string", "enum": ["low", "medium", "high", "blocker"]},
+                "labels": {
+                    "type": "array",
+                    "items": {
+                        "type": "string",
+                        "enum": ["bug", "regression", "ui", "backend", "perf"],
+                    },
+                },
+                "affected": {
+                    "type": "object",
+                    "properties": {
+                        "service": {"type": "string"},
+                        "version": {"type": "string"},
+                        "regions": {"type": "array", "items": {"type": "string"}},
+                    },
+                    "required": ["service", "version"],
+                },
+                "cc_emails": {"type": "array", "items": {"type": "string"}},
+                "urgent_escalation": {"type": "boolean"},
+                "estimated_minutes": {"type": "integer"},
+            },
+            "required": ["title", "body", "severity", "labels", "affected"],
+        },
+    },
+}
+
+SEARCH_TICKETS_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "search_tickets",
+        "description": "Search existing tickets.",
+        "parameters": {
+            "type": "object",
+            "properties": {"query": {"type": "string"}, "limit": {"type": "integer"}},
+            "required": ["query"],
+        },
+    },
+}
+
+ESCALATE_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "escalate_incident",
+        "description": "Escalate an existing incident by id.",
+        "parameters": {
+            "type": "object",
+            "properties": {"incident_id": {"type": "string"}, "level": {"type": "integer"}},
+            "required": ["incident_id", "level"],
+        },
+    },
+}
+
+SCENARIO_C = {
+    "key": "c",
+    "name": "C. adversarial-payload",
+    "tools": [TICKET_TOOL, SEARCH_TICKETS_TOOL, ESCALATE_TOOL],
+    "messages": [
+        {
+            "role": "user",
+            "content": (
+                "Our payments service is broken, this is critical! File a ticket titled "
+                "Payments retry storm after config rollout. The body must include, verbatim, "
+                'the deployed config: {"retry": {"max": 3, "backoff_ms": [100, 200], '
+                '"jitter": true}} and the error line: TypeError: cannot destructure '
+                "{id: undefined} at applyRetry (retry.js:42). It's a regression bug in the "
+                "backend, affecting service payments version 2.14.1 in regions ap-northeast-2 "
+                "and us-east-1. CC dev-alerts@example.com, escalate urgently, and estimate "
+                "45 minutes of work."
+            ),
+        }
+    ],
+    # Thinking is exercised in scenario B; here the budget goes to the payload itself.
+    # Elevated temperature with the full sampling tail (top_p=1.0) exposes format
+    # instability in the long string argument — settings a serving engine must survive.
+    "reasoning": False,
+    "check_thought": False,
+    "max_new_tokens": 640,
+    "temperature": 1.5,
+    "top_p": 1.0,
+}
+
+SCENARIOS = {"a": SCENARIO_A, "b": SCENARIO_B, "c": SCENARIO_C}
+
+THOUGHT_BEGIN = "<|channel>thought"
+THOUGHT_END = "<channel|>"
+TOOL_CALL_START = "<|tool_call>"
+TOOL_RESPONSE_START = "<|tool_response>"
+
+METRIC_DESCRIPTIONS = {
+    "well_formed": "every tool call block is complete and parseable (sglang parser port)",
+    "valid_name": "every called tool exists",
+    "schema_valid": "parsed arguments validate against the tool's JSON schema",
+    "stops_at_boundary": (
+        "generation did not run past <tool_call|> into <|tool_response>. Per the Gemma 4 "
+        "prompt-formatting spec, the model is only responsible for the call up to "
+        "<tool_call|>; <|tool_response> is appended by the application with the real "
+        "tool result, and is registered as an additional stop sequence purely as a "
+        "backstop. This harness configures no such stop, so unconstrained runs "
+        "free-run past the boundary into engine-owned territory, while required-mode "
+        "grammar ends the call cleanly at an accept state."
+    ),
+    "thought_ok": "<|channel>thought opened and closed before the first tool call",
+}
+
+
+def score_output(text: str, scenario: dict) -> dict:
+    tools_by_name = {t["function"]["name"]: t["function"]["parameters"] for t in scenario["tools"]}
+    # Anything past the <|tool_response> boundary is a separate problem (stops_at_boundary);
+    # parse only the part the serving engine would treat as the model's calls.
+    model_part = text.split(TOOL_RESPONSE_START)[0]
+    calls = extract_tool_calls(model_part)
+    complete = [(n, a) for n, a in calls if n is not None and a is not None]
+
+    well_formed = len(complete) > 0 and len(complete) == len(calls)
+    valid_name = well_formed and all(n in tools_by_name for n, _ in complete)
+    schema_valid = valid_name
+    error = None
+    if valid_name:
+        for name, args in complete:
+            try:
+                jsonschema.validate(args, tools_by_name[name])
+            except jsonschema.ValidationError as e:
+                schema_valid = False
+                error = f"{name}: {e.message}"
+                break
+    elif well_formed:
+        error = "unknown tool: " + ", ".join(n for n, _ in complete if n not in tools_by_name)
+    else:
+        error = "malformed or missing tool call block"
+
+    result = {
+        "well_formed": well_formed,
+        "valid_name": valid_name,
+        "schema_valid": schema_valid,
+        "stops_at_boundary": TOOL_RESPONSE_START not in text,
+        "calls": complete,
+        "error": error if not (well_formed and valid_name and schema_valid) else None,
+    }
+
+    if scenario["check_thought"]:
+        first_call = model_part.find(TOOL_CALL_START)
+        opened = model_part.startswith(THOUGHT_BEGIN)
+        closed = opened and THOUGHT_END in model_part[: first_call if first_call != -1 else None]
+        result["thought_ok"] = opened and closed
+    return result
+
+
+def metric_keys_for(scenario: dict) -> list:
+    keys = ["well_formed", "valid_name", "schema_valid", "stops_at_boundary"]
+    if scenario["check_thought"]:
+        keys.append("thought_ok")
+    return keys
+
+
+def run_scenario(scenario, model, tokenizer, args):
+    print(f"\n{'=' * 70}\n{scenario['name']}  (reasoning={scenario['reasoning']})\n{'=' * 70}")
+
+    prompt = tokenizer.apply_chat_template(
+        scenario["messages"],
+        tools=scenario["tools"],
+        add_generation_prompt=True,
+        tokenize=False,
+    )
+    print("--- prompt tail ---")
+    print(prompt[-400:])
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
+
+    vocab_size = model.config.get_text_config().vocab_size
+    tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size)
+    compiler = xgr.GrammarCompiler(tokenizer_info)
+    stag = xgr.get_model_structural_tag(
+        model="gemma_4",
+        tools=scenario["tools"],
+        tool_choice="required",
+        reasoning=scenario["reasoning"],
+    )
+    compiled = compiler.compile_structural_tag(stag)
+
+    metric_keys = metric_keys_for(scenario)
+    max_new_tokens = max(args.max_new_tokens, scenario.get("max_new_tokens", 0))
+    temperature = scenario.get("temperature", args.temperature)
+    top_p = scenario.get("top_p", 0.95)
+    summary = {}
+    failures = []
+    for mode in ["constrained", "unconstrained"]:
+        counts = {k: 0 for k in metric_keys}
+        print(
+            f"\n--- {mode}: {args.samples} samples,"
+            f" temperature={temperature}, top_p={top_p} ---"
+        )
+        for s in range(args.samples):
+            torch.manual_seed(1000 + s)
+            gen_kwargs = dict(
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            if mode == "constrained":
+                gen_kwargs["logits_processor"] = [xgr.contrib.hf.LogitsProcessor(compiled)]
+            t0 = time.monotonic()
+            out = model.generate(**inputs, **gen_kwargs)
+            gen = tokenizer.decode(
+                out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False
+            )
+            gen = gen.split("<eos>")[0].split("<end_of_turn>")[0]
+            score = score_output(gen, scenario)
+            for k in metric_keys:
+                counts[k] += bool(score[k])
+            failed_keys = [k for k in metric_keys if not score[k]]
+            if failed_keys:
+                reasons = []
+                if score["error"]:
+                    reasons.append(score["error"])
+                if "stops_at_boundary" in failed_keys:
+                    reasons.append("ran past <tool_call|> into <|tool_response> after the call")
+                if "thought_ok" in failed_keys:
+                    reasons.append("skipped or never closed the <|channel>thought section")
+                failures.append(
+                    {
+                        "mode": mode,
+                        "sample": s,
+                        "failed": failed_keys,
+                        "reason": "; ".join(reasons),
+                        "output": gen,
+                    }
+                )
+            flags = " ".join(f"{k}={'O' if score[k] else 'X'}" for k in metric_keys)
+            print(f"\n[{mode} #{s}] ({time.monotonic() - t0:.0f}s) {flags}")
+            if score["error"]:
+                print(f"  error: {score['error']}")
+            print(f"  output: {gen[:700]!r}")
+        summary[mode] = counts
+
+    print(f"\n--- {scenario['name']} summary (out of {args.samples}) ---")
+    print(f"{'metric':<14}" + "".join(f"{m:>16}" for m in summary))
+    for k in metric_keys:
+        print(f"{k:<14}" + "".join(f"{summary[m][k]:>16}" for m in summary))
+    return {
+        "summary": summary,
+        "metric_keys": metric_keys,
+        "failures": failures,
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_new_tokens": max_new_tokens,
+        "reasoning": scenario["reasoning"],
+    }
+
+
+def write_report(path, model_id, args, results):
+    lines = [
+        "# Gemma-4 tool calling: constrained (xgrammar) vs unconstrained",
+        "",
+        f"- model: `{model_id}` via HF transformers",
+        '- constraint: `xgr.get_model_structural_tag("gemma_4", tools=..., '
+        'tool_choice="required")` compiled and applied with '
+        "`xgr.contrib.hf.LogitsProcessor`",
+        f"- {args.samples} samples per scenario per mode; identical sampling settings"
+        " for both modes (per-scenario values below)",
+        "- outputs validated with a pure-Python port of sglang's `Gemma4Detector`"
+        " parser plus `jsonschema`",
+        "",
+        "## Results (passing samples / total)",
+        "",
+    ]
+    for name, res in results.items():
+        summary, metric_keys = res["summary"], res["metric_keys"]
+        lines.append(f"### {name}")
+        lines.append("")
+        lines.append(
+            f"temperature={res['temperature']}, top_p={res['top_p']},"
+            f" max_new_tokens={res['max_new_tokens']}, reasoning={res['reasoning']}"
+        )
+        lines.append("")
+        lines.append("| metric | constrained | unconstrained |")
+        lines.append("|---|---|---|")
+        for k in metric_keys:
+            c, u = summary["constrained"][k], summary["unconstrained"][k]
+            c_mark = "✅" if c == args.samples else "❌"
+            u_mark = "✅" if u == args.samples else "❌"
+            lines.append(f"| {k} | {c}/{args.samples} {c_mark} | {u}/{args.samples} {u_mark} |")
+        lines.append("")
+
+    lines += ["## Metric definitions", ""]
+    for k, desc in METRIC_DESCRIPTIONS.items():
+        lines.append(f"- `{k}` — {desc}")
+    lines.append("")
+
+    all_failures = [(name, f) for name, res in results.items() for f in res["failures"]]
+    if all_failures:
+        lines += ["## Failure examples", ""]
+        for name, f in all_failures[:12]:
+            failed = ", ".join(f"`{k}`" for k in f["failed"])
+            lines.append(f"**{name} / {f['mode']} #{f['sample']}** — failed {failed}: {f['reason']}")
+            lines.append("")
+            lines.append("```")
+            lines.append(f["output"][:900])
+            lines.append("```")
+            lines.append("")
+
+    with open(path, "w") as fp:
+        fp.write("\n".join(lines))
+    print(f"\nreport written to {path}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--scenario", choices=["a", "b", "c", "all"], default="all")
+    parser.add_argument("--samples", type=int, default=4)
+    parser.add_argument("--max-new-tokens", type=int, default=384)
+    parser.add_argument("--temperature", type=float, default=0.9)
+    parser.add_argument("--model", default=DEFAULT_MODEL_ID)
+    parser.add_argument("--report", default=None, help="write a markdown report to this path")
+    args = parser.parse_args()
+
+    device = "mps" if torch.backends.mps.is_available() else "cpu"
+    print(f"loading {args.model} on {device} ...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, dtype=torch.bfloat16, device_map=device
+    )
+    model.eval()
+
+    selected = ["a", "b", "c"] if args.scenario == "all" else [args.scenario]
+    results = {}
+    for key in selected:
+        results[SCENARIOS[key]["name"]] = run_scenario(SCENARIOS[key], model, tokenizer, args)
+
+    print(f"\n{'=' * 70}\nOVERALL\n{'=' * 70}")
+    print(json.dumps({n: r["summary"] for n, r in results.items()}, indent=2, ensure_ascii=False))
+    if args.report:
+        write_report(args.report, args.model, args, results)
+
+
+if __name__ == "__main__":
+    main()

test_gemma4_xgrammar.py

@@ -0,0 +1,120 @@
+"""End-to-end test: constrained decoding for Gemma-4 tool calling with the modified xgrammar.
+
+Loads google/gemma-4-E2B-it via HF transformers, builds the gemma_4 builtin
+structural tag from xgrammar, and generates with xgrammar's LogitsProcessor.
+Verifies:
+  1. The grammar compiles against the real Gemma-4 tokenizer.
+  2. Constrained generation emits a well-formed <|tool_call>call:name{...}<tool_call|>.
+  3. The emitted arguments use Gemma's <|"|> string delimiters and satisfy the schema.
+  4. (Alignment) Unconstrained generation is compared against the same grammar.
+"""
+
+import json
+import time
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+import xgrammar as xgr
+from xgrammar.testing import _is_grammar_accept_string
+
+MODEL_ID = "google/gemma-4-E2B-it"
+
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get the current weather for a city.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "city": {"type": "string", "description": "City name"},
+                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                },
+                "required": ["city"],
+            },
+        },
+    }
+]
+
+MESSAGES = [{"role": "user", "content": "What's the weather in Seoul in celsius?"}]
+
+
+def main():
+    print(f"=== loading tokenizer/model: {MODEL_ID} ===")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, device_map="mps")
+    model.eval()
+
+    vocab_size = model.config.get_text_config().vocab_size
+    print(f"vocab_size={vocab_size}")
+
+    prompt = tokenizer.apply_chat_template(
+        MESSAGES, tools=TOOLS, add_generation_prompt=True, tokenize=False
+    )
+    print("=== rendered prompt (tail) ===")
+    print(prompt[-600:])
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("mps")
+
+    tokenizer_info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size)
+    compiler = xgr.GrammarCompiler(tokenizer_info)
+
+    results = {}
+    for tool_choice, reasoning in [("required", False), ("auto", True)]:
+        label = f"tool_choice={tool_choice}, reasoning={reasoning}"
+        print(f"\n=== structural tag: {label} ===")
+        stag = xgr.get_model_structural_tag(
+            model="gemma_4", tools=TOOLS, tool_choice=tool_choice, reasoning=reasoning
+        )
+        t0 = time.monotonic()
+        compiled = compiler.compile_structural_tag(stag)
+        print(f"grammar compile time: {time.monotonic() - t0:.2f}s")
+
+        processor = xgr.contrib.hf.LogitsProcessor(compiled)
+        t0 = time.monotonic()
+        out = model.generate(
+            **inputs, max_new_tokens=128, do_sample=False, logits_processor=[processor]
+        )
+        gen = tokenizer.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False)
+        print(f"generate time: {time.monotonic() - t0:.1f}s")
+        print(f"--- constrained output ({label}) ---")
+        print(repr(gen))
+        results[label] = gen
+
+    print("\n=== unconstrained (alignment check) ===")
+    out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
+    unconstrained = tokenizer.decode(
+        out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=False
+    )
+    print(repr(unconstrained))
+    results["unconstrained"] = unconstrained
+
+    print("\n=== verification ===")
+    ok = True
+    required_out = results["tool_choice=required, reasoning=False"]
+    # strip trailing special tokens after the tool call for the checks
+    body = required_out.split("<tool_call|>")[0] + "<tool_call|>" if "<tool_call|>" in required_out else required_out
+    checks = [
+        ("required: contains <|tool_call>call:get_weather{", "<|tool_call>call:get_weather{" in body),
+        ("required: closes with }<tool_call|>", body.rstrip().endswith("<tool_call|>")),
+        ('required: city uses <|"|> delimiters', '<|"|>' in body),
+        ("required: no JSON-quoted args", '"city"' not in body),
+    ]
+    stag_req = xgr.get_model_structural_tag(
+        model="gemma_4", tools=TOOLS, tool_choice="required", reasoning=False
+    )
+    g = xgr.Grammar.from_structural_tag(stag_req)
+    checks.append(("required: output accepted by grammar", _is_grammar_accept_string(g, body)))
+
+    for name, passed in checks:
+        print(f"  [{'PASS' if passed else 'FAIL'}] {name}")
+        ok &= passed
+
+    print("\nRESULT:", "ALL PASS" if ok else "SOME CHECKS FAILED")
+    print("\nJSON summary:")
+    print(json.dumps({k: v for k, v in results.items()}, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()