| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Lark grammars in the xgrammar backend. |
| 4 | |
| 5 | `structured_outputs.grammar` accepts either EBNF or Lark. The xgrammar backend |
| 6 | parses Lark natively, both when validating a request and when compiling it, so |
| 7 | the request's grammar is never rewritten and frontends that skip Python |
| 8 | request validation get the same behavior. |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | from mistral_common.protocol.instruct.tool_calls import Tool, ToolChoiceEnum |
| 13 | from transformers import AutoTokenizer |
| 14 | |
| 15 | from vllm.config import StructuredOutputsConfig, VllmConfig |
| 16 | from vllm.config.structured_outputs import StructuredOutputsBackend |
| 17 | from vllm.exceptions import VLLMValidationError |
| 18 | from vllm.sampling_params import SamplingParams, StructuredOutputsParams |
| 19 | from vllm.tokenizers.mistral import MistralTokenizer |
| 20 | from vllm.v1.structured_output.backend_types import StructuredOutputOptions |
| 21 | from vllm.v1.structured_output.backend_xgrammar import ( |
| 22 | XgrammarBackend, |
| 23 | validate_xgrammar_grammar, |
| 24 | ) |
| 25 | |
| 26 | pytestmark = pytest.mark.cpu_test |
| 27 | |
| 28 | TOKENIZER = "openai-community/gpt2" |
| 29 | VOCAB_SIZE = 50257 |
| 30 | MISTRAL_TOKENIZER = "mistralai/Ministral-3-3B-Instruct-2512" |
| 31 | |
| 32 | # Uses Lark features beyond plain rules: a regex terminal, `%import common`, |
| 33 | # and a rule alias. The Codex `apply_patch` tool grammar has the same shape. |
| 34 | LARK_GRAMMAR = """ |
| 35 | start: "id=" NUMBER tag? LF? |
| 36 | tag: " #" /[a-z]+/ -> label |
| 37 | NUMBER: /[0-9]+/ |
| 38 | %import common.LF |
| 39 | """ |
| 40 | |
| 41 | EBNF_GRAMMAR = 'root ::= "id=" [0-9]+' |
| 42 | |
| 43 | |
| 44 | def _validate( |
| 45 | grammar: str, |
| 46 | backend: StructuredOutputsBackend = "xgrammar", |
| 47 | tokenizer: object | None = None, |
| 48 | ) -> StructuredOutputsParams: |
| 49 | params = SamplingParams(structured_outputs=StructuredOutputsParams(grammar=grammar)) |
| 50 | params._validate_structured_outputs( |
| 51 | _StubModelConfig(), |
| 52 | StructuredOutputsConfig(backend=backend), |
| 53 | tokenizer=tokenizer or object(), |
| 54 | ) |
| 55 | assert params.structured_outputs is not None |
| 56 | return params.structured_outputs |
| 57 | |
| 58 | |
| 59 | class _StubModelConfig: |
| 60 | is_diffusion = False |
| 61 | |
| 62 | |
| 63 | GRAMMARS = [ |
| 64 | pytest.param(LARK_GRAMMAR, id="lark"), |
| 65 | pytest.param(EBNF_GRAMMAR, id="ebnf"), |
| 66 | ] |
| 67 | |
| 68 | |
| 69 | @pytest.mark.parametrize("grammar", GRAMMARS) |
| 70 | def test_validation_accepts_grammar_without_rewriting_it(grammar: str): |
| 71 | assert _validate(grammar).grammar == grammar |
| 72 | |
| 73 | |
| 74 | def test_auto_backend_selects_xgrammar_for_lark(): |
| 75 | assert _validate(LARK_GRAMMAR, backend="auto")._backend == "xgrammar" |
| 76 | |
| 77 | |
| 78 | @pytest.mark.parametrize("mode", [ToolChoiceEnum.required, ToolChoiceEnum.none]) |
| 79 | def test_auto_backend_selects_guidance_for_mistral_lark(mode: ToolChoiceEnum): |
| 80 | """The Mistral tool parser generates Lark grammars for llguidance, which |
| 81 | keeps special tokens such as `[TOOL_CALLS]` out of regex matches.""" |
| 82 | tokenizer = MistralTokenizer.from_pretrained(MISTRAL_TOKENIZER) |
| 83 | factory = tokenizer.grammar_factory |
| 84 | tool = { |
| 85 | "type": "function", |
| 86 | "function": { |
| 87 | "name": "get_weather", |
| 88 | "parameters": {"type": "object", "properties": {}}, |
| 89 | }, |
| 90 | } |
| 91 | grammar = factory.get_lark_from_jinja( |
| 92 | template=factory.select_jinja_template(), |
| 93 | mode=mode, |
| 94 | tools=[Tool.from_openai(tool)], |
| 95 | json_schema=None, |
| 96 | parallel_tool_calls=True, |
| 97 | json_only=False, |
| 98 | ) |
| 99 | assert _validate(grammar, "auto", tokenizer)._backend == "guidance" |
| 100 | |
| 101 | |
| 102 | @pytest.mark.parametrize( |
| 103 | "grammar", |
| 104 | [ |
| 105 | pytest.param('rule: "hello"', id="missing-start-rule"), |
| 106 | pytest.param("start: 'hello'", id="single-quoted-literal"), |
| 107 | ], |
| 108 | ) |
| 109 | def test_validation_rejects_invalid_lark(grammar: str): |
| 110 | params = SamplingParams(structured_outputs=StructuredOutputsParams(grammar=grammar)) |
| 111 | with pytest.raises(VLLMValidationError, match="Invalid grammar specification"): |
| 112 | validate_xgrammar_grammar(params) |
| 113 | |
| 114 | |
| 115 | @pytest.fixture(scope="module") |
| 116 | def backend() -> XgrammarBackend: |
| 117 | vllm_config = VllmConfig( |
| 118 | structured_outputs_config=StructuredOutputsConfig(backend="xgrammar") |
| 119 | ) |
| 120 | return XgrammarBackend( |
| 121 | vllm_config, |
| 122 | tokenizer=AutoTokenizer.from_pretrained(TOKENIZER), |
| 123 | vocab_size=VOCAB_SIZE, |
| 124 | ) |
| 125 | |
| 126 | |
| 127 | @pytest.mark.parametrize("grammar", GRAMMARS) |
| 128 | @pytest.mark.parametrize( |
| 129 | "text, accepted", |
| 130 | [("id=42", True), ("id=x", False)], |
| 131 | ) |
| 132 | def test_compile_grammar_matches_lark_and_ebnf( |
| 133 | backend: XgrammarBackend, grammar: str, text: str, accepted: bool |
| 134 | ): |
| 135 | compiled = backend.compile_grammar(StructuredOutputOptions.GRAMMAR, grammar) |
| 136 | token_ids = backend.tokenizer.encode(text, add_special_tokens=False) |
| 137 | assert compiled.accept_tokens("req", token_ids) is accepted |
| 138 | |
| 139 | |
| 140 | def test_compile_grammar_uses_lark_semantics(backend: XgrammarBackend): |
| 141 | compiled = backend.compile_grammar(StructuredOutputOptions.GRAMMAR, LARK_GRAMMAR) |
| 142 | token_ids = backend.tokenizer.encode("id=7 #ok\n", add_special_tokens=False) |
| 143 | assert compiled.accept_tokens("req", token_ids) |
| 144 | assert compiled.is_terminated() or compiled.accept_tokens( |
| 145 | "req", [backend.tokenizer.eos_token_id] |
| 146 | ) |