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

View File

@@ -16,7 +16,7 @@
import dataclasses
import json
import logging
from typing import List, Optional, Tuple, Union
from typing import Dict, List, Optional, Tuple, Union
import torch
from xgrammar import (
@@ -196,6 +196,34 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
self.override_stop_tokens = override_stop_tokens
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(
self, ctx: CompiledGrammar, key_string: str, grammar_stats: GrammarStats
) -> XGrammarGrammar:
@@ -249,6 +277,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
# TODO(dark): it's REALLY stupid to construct object from string and decode it again
structural_tag = json.loads(key_string)
if is_legacy_structural_tag(structural_tag):
self._sanitize_structural_tag_structures(structural_tag)
tags = [
StructuralTagItem(
begin=structure["begin"],
@@ -261,6 +290,11 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
tags, structural_tag["triggers"]
)
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)
except (RuntimeError, json.decoder.JSONDecodeError) as e:
logging.error(f"Hit invalid structural_tag: {key_string=}, {e=}")