diff --git a/deepinfer/structure_output/compiler.py b/deepinfer/structure_output/compiler.py index 4bbcddf..829d4e4 100644 --- a/deepinfer/structure_output/compiler.py +++ b/deepinfer/structure_output/compiler.py @@ -10,6 +10,10 @@ from deepinfer import envs logger = logging.getLogger(__name__) +#: The ``schema_type`` values ``dispatch_and_compile`` knows how to compile. Shared with the +#: matcher builder so a wrong type is rejected at the request boundary, not inside the pool. +SUPPORTED_SCHEMA_TYPES: tuple[str, ...] = ("json", "regex", "ebnf") + class XGrammarCompiler: def __init__(self, tokenizer_info: TokenizerInfo, num_threads: int) -> None: @@ -52,15 +56,18 @@ class XGrammarCompiler: self.close() def dispatch_and_compile(self, schema_type: str, schema: str) -> CompiledGrammar: + # Checked outside the try: the `except ValueError` below is for xgrammar's own input + # errors and used to swallow this one too, re-raising it as a RuntimeError that + # try_result() then reported as a timeout. + if schema_type not in SUPPORTED_SCHEMA_TYPES: + raise ValueError(f"unsupported schema_type {schema_type!r}; expected one of {SUPPORTED_SCHEMA_TYPES}") try: if schema_type == "json": ctx = self.inner_compiler.compile_json_schema(schema, strict_mode=True) elif schema_type == "regex": ctx = self.inner_compiler.compile_regex(schema) - elif schema_type == "ebnf": + else: # "ebnf" — the guard above makes this exhaustive ctx = self.inner_compiler.compile_grammar(schema) - else: - raise ValueError("Unsupported Grammar Type") except RuntimeError as e: raise RuntimeError(f"xgrammar compilation RE: {e}") from e except ValueError as e: diff --git a/deepinfer/structure_output/so_matcher_builder.py b/deepinfer/structure_output/so_matcher_builder.py index ab3b867..2afc2d7 100644 --- a/deepinfer/structure_output/so_matcher_builder.py +++ b/deepinfer/structure_output/so_matcher_builder.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING from xgrammar import TokenizerInfo -from deepinfer.structure_output.compiler import XGrammarCompiler +from deepinfer.structure_output.compiler import SUPPORTED_SCHEMA_TYPES, XGrammarCompiler from deepinfer.structure_output.matcher import XGrammarMatcher, XGrammarMatcherFuture, validate_trigger_bound if TYPE_CHECKING: @@ -65,6 +65,8 @@ class SOMatcherBuilder: if missing: raise ValueError(f"so_args is missing required key(s) {missing}; provide schema_type and schema") schema_type = so_args["schema_type"] + if schema_type not in SUPPORTED_SCHEMA_TYPES: + raise ValueError(f"unsupported schema_type {schema_type!r}; expected one of {SUPPORTED_SCHEMA_TYPES}") schema = so_args["schema"] trigger_tokens = so_args.get("trigger_token_ids")