fix xgrammar_backend crash with malformed inputs (#13752)

This commit is contained in:
gongwei-130
2025-11-24 13:54:19 -08:00
committed by GitHub
parent fafaa2ccea
commit 9dc15d8569
@@ -16,7 +16,7 @@
import dataclasses import dataclasses
import json import json
import logging import logging
from typing import List, Optional, Tuple, Union from typing import Dict, List, Optional, Tuple, Union
import torch import torch
from xgrammar import ( from xgrammar import (
@@ -196,6 +196,34 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
self.override_stop_tokens = override_stop_tokens self.override_stop_tokens = override_stop_tokens
self.any_whitespace = any_whitespace self.any_whitespace = any_whitespace
@staticmethod
def _sanitize_structural_format(structural_format):
"""Recursively replace missing json_schema fields with an empty schema."""
if not isinstance(structural_format, dict):
return
fmt_type = structural_format.get("type")
if fmt_type in {"json_schema", "qwen_xml_parameter"}:
if structural_format.get("json_schema") is None:
structural_format["json_schema"] = {}
if fmt_type == "tag":
XGrammarGrammarBackend._sanitize_structural_format(
structural_format.get("content")
)
elif fmt_type in {"sequence", "or"}:
for element in structural_format.get("elements", []):
XGrammarGrammarBackend._sanitize_structural_format(element)
elif fmt_type in {"triggered_tags", "tags_with_separator"}:
for tag in structural_format.get("tags", []):
XGrammarGrammarBackend._sanitize_structural_format(tag)
@staticmethod
def _sanitize_structural_tag_structures(structural_tag: Dict) -> None:
for structure in structural_tag.get("structures", []):
if structure.get("schema") is None:
structure["schema"] = {}
def _from_context( def _from_context(
self, ctx: CompiledGrammar, key_string: str, grammar_stats: GrammarStats self, ctx: CompiledGrammar, key_string: str, grammar_stats: GrammarStats
) -> XGrammarGrammar: ) -> XGrammarGrammar:
@@ -249,6 +277,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
# TODO(dark): it's REALLY stupid to construct object from string and decode it again # TODO(dark): it's REALLY stupid to construct object from string and decode it again
structural_tag = json.loads(key_string) structural_tag = json.loads(key_string)
if is_legacy_structural_tag(structural_tag): if is_legacy_structural_tag(structural_tag):
self._sanitize_structural_tag_structures(structural_tag)
tags = [ tags = [
StructuralTagItem( StructuralTagItem(
begin=structure["begin"], begin=structure["begin"],
@@ -261,6 +290,11 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
tags, structural_tag["triggers"] tags, structural_tag["triggers"]
) )
else: else:
format_dict = structural_tag.get("format")
if isinstance(format_dict, dict):
self._sanitize_structural_format(format_dict)
structural_tag["format"] = format_dict
key_string = json.dumps(structural_tag)
ctx = self.grammar_compiler.compile_structural_tag(key_string) ctx = self.grammar_compiler.compile_structural_tag(key_string)
except (RuntimeError, json.decoder.JSONDecodeError) as e: except (RuntimeError, json.decoder.JSONDecodeError) as e:
logging.error(f"Hit invalid structural_tag: {key_string=}, {e=}") logging.error(f"Hit invalid structural_tag: {key_string=}, {e=}")