feat(core): JSON-syntax tool-argument grammar for Qwen/Llama/DeepSeek (#376)

[pekkah]

Summary

Extends schema/grammar-constrained tool-call argument decoding (#374, PR #375) to the families that emit standard JSON — Qwen (<tool_call>{json}</tool_call>), Llama-3 (<|python_tag|>{json}), and DeepSeek (NAME<|tool_sep|>{json}) — the JSON sibling of GemmaToolArgumentConstraint. Closes #376.

What changed

• JsonToolArgumentConstraint — a byte-level pushdown automaton over standard JSON: {"key":"value",...} with double-quoted keys/strings (\-escape aware), bare numbers/booleans/null, [...] arrays, nested {...} objects, and enum restriction. Matching is fully byte-level because a real BPE merges structural bytes with content/whitespace — one token carries {", another ":, another "}}, and the empty object is a single {} token (verified by probing the real Qwen vocab).
• Early-engage at the argument key's colon ("arguments": / "parameters":) pushes the object frame in OExpectOpenBrace so the next token is masked — a merged {} token can't drop the required arguments (the feat(engine): schema/grammar-constrained decoding for tool-call arguments #374 failure mode, in JSON form).
• Two envelopes, parameterized per family:
  • NameValueObject (Qwen/Llama): a preamble parser walks {"name":"X","argskey":{ to capture the tool name and find the args object; argskey is arguments (Qwen) or parameters (Llama).
  • NameThenSeparator (DeepSeek): bare name up to <|tool_sep|>, then the args object directly.
    Wired via QwenToolCallAdapter / LlamaToolCallAdapter / DeepSeekToolCallAdapter.BuildArgumentConstraint.
• Shared compile logic extracted — CompiledObject/CompiledNode + TryCompile* moved out of GemmaToolArgumentConstraint into ToolSchemaCompiler, used by both constraints (DRY, and sets up Tool-arg grammar: constrain partially-typed tools (Any/open-object values) (#374 follow-up) #378 to modify one place). The Gemma constraint now calls it; all 72 prior Gemma grammar/adapter tests stay green — behavior-preserving.
• EOG forbidden inside the region — empty-byte pruning plus an explicit EOG mask pass, so an EOS that decodes to ordinary text can't slip through a free string and truncate the call.
• Default-off byte-identical — inert when no tool is constrainable, when the open marker isn't a vocabulary token, or when the model emits a non-JSON shape (Qwen3.6's XML <function=> never engages).

Tests

• CI mock (JsonToolGrammarMockTests, 14) — FakeJsonTokenizer provides single-byte tokens plus realistic merged structural pieces ({", ":, {}, …). Covers: required key can't close early, foreign-key rejection, enum, string with \" escape, number, nested object, array, the merged-{} early-engage, and the Qwen/Llama/DeepSeek envelopes.
• Model-gated (JsonToolGrammarConstraintTests, 4) — drives the constraint over a real Qwen3.6 BPE vocabulary (production merges: {}, {", ":), asserting it engages at the colon, rejects merged {}, allows {", enforces declared keys, restricts the enum, and stays inert on the XML form.

90 grammar/tool tests green; full solution builds clean under TreatWarningsAsErrors + trim/AOT analyzers.

Note on GPU e2e

A live generation A/B (like #377's) isn't included: the on-disk Qwen3.6 emits the XML tool shape (the JSON constraint correctly stays inert on it), and no JSON-emitting Qwen/Llama that fits VRAM was available. The model-gated test instead exercises the byte-level matching against the real production tokenizer.

Follow-up

Qwen3-Coder's <parameter=…> XML wrapper isn't JSON and is left for a separate matcher (its adapter has no override → arguments generate unconstrained, never wrong) — the issue flagged this as possibly separate.

Follow-up to #374 / #375.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request extracts the tool schema compilation logic into a shared ToolSchemaCompiler class and introduces JsonToolArgumentConstraint to support constrained decoding of standard JSON tool-call arguments for Qwen, Llama-3, and DeepSeek. Corresponding adapters and tests have been added or updated. Feedback highlights a critical performance bottleneck in ComputeMask during free string content generation, where calling SimulateToken for the entire vocabulary on every step will severely degrade performance; a fast-path optimization is suggested to resolve this.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Performance Bottleneck in ComputeMask during String Content Generation

When the state machine is in SContent (generating free string content), CollectStr sets all 256 bytes in _firstByteOk to true. This causes candidate to be true for almost the entire vocabulary (typically 150,000+ tokens for Qwen/Llama).

Without an optimization, SimulateToken(id) is called for every single token in the vocabulary on every generated token step inside a string. SimulateToken copies the stack of frames and walks the token's bytes, which introduces a massive $O(V)$ overhead (where $V$ is the vocabulary size) and will severely degrade generation speed (potentially taking seconds per token).

Since the vast majority of tokens do not contain the string delimiter " or the escape character \\, we can fast-path them when !escaped and skip the expensive SimulateToken call entirely.

        int kept = 0;
        ref var top = ref _stack[_depth - 1];
        bool isStrContent = top.Kind == FK.Str && top.State == SContent;
        bool escaped = top.Escaped;

        for (int id = 0; id < buf.Length; id++)
        {
            var bytes = _vocab.TokenBytes(id);
            // Empty-byte tokens (EOG / control) never advance the structure — forbidding them keeps
            // an end-of-generation token from truncating the call mid-object.
            bool candidate = bytes.Length != 0 && _firstByteOk[bytes[0]];
            if (candidate)
            {
                bool ok = isStrContent && !escaped && !ContainsQuoteOrSlash(bytes)
                    ? true
                    : SimulateToken(id);

                if (ok)
                    kept++;
                else
                    buf[id] = float.NegativeInfinity;
            }
            else
            {
                buf[id] = float.NegativeInfinity;
            }
        }

        // Belt-and-suspenders: forbid every EOG id regardless of its bytes (a tokenizer whose EOS
        // decodes to ordinary text could otherwise pass first-byte pruning inside a free string).
        foreach (int id in _forbidden)
            if ((uint)id < (uint)buf.Length && !float.IsNegativeInfinity(buf[id]))
            { buf[id] = float.NegativeInfinity; kept--; }

        return kept;

        static bool ContainsQuoteOrSlash(ReadOnlySpan<byte> span)
        {
            for (int i = 0; i < span.Length; i++)
            {
                byte b = span[i];
                if (b is (byte)'"' or (byte)'\\') return true;
            }
            return false;
        }

[pekkah]

Addressed the ComputeMask perf bottleneck in f5bd249: added the fast-path you suggested — inside SContent (and when not mid-escape), a token containing neither " nor \ is pure content and is admitted without SimulateToken; only tokens that could close/escape get the full replay. This is the byte-level analogue of the Gemma sibling's token-level free-content shortcut. All 18 JSON tests stay green.