| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import copy |
| 5 | import json |
| 6 | import os |
| 7 | from dataclasses import dataclass |
| 8 | from typing import TYPE_CHECKING, Any |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm.exceptions import VLLMValidationError |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.sampling_params import SamplingParams |
| 15 | from vllm.utils.import_utils import LazyLoader |
| 16 | from vllm.utils.mistral import is_mistral_tokenizer |
| 17 | from vllm.v1.structured_output.backend_types import ( |
| 18 | StructuredOutputBackend, |
| 19 | StructuredOutputGrammar, |
| 20 | StructuredOutputOptions, |
| 21 | ) |
| 22 | from vllm.v1.structured_output.request import get_structured_output_key |
| 23 | from vllm.v1.structured_output.utils import strip_speculative_padding |
| 24 | |
| 25 | if TYPE_CHECKING: |
| 26 | import llguidance |
| 27 | import llguidance.hf as llguidance_hf |
| 28 | import llguidance.torch as llguidance_torch |
| 29 | else: |
| 30 | llguidance = LazyLoader("llguidance", globals(), "llguidance") |
| 31 | llguidance_hf = LazyLoader("llguidance.hf", globals(), "llguidance.hf") |
| 32 | llguidance_torch = LazyLoader("llguidance.torch", globals(), "llguidance.torch") |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | def _walk_json_for_additional_properties(data: object): |
| 38 | if isinstance(data, dict): |
| 39 | for value in data.values(): |
| 40 | _walk_json_for_additional_properties(value) |
| 41 | if "additionalProperties" not in data and ( |
| 42 | "properties" in data or "patternProperties" in data |
| 43 | ): |
| 44 | data["additionalProperties"] = False |
| 45 | elif isinstance(data, list): |
| 46 | for item in data: |
| 47 | _walk_json_for_additional_properties(item) |
| 48 | |
| 49 | |
| 50 | def has_guidance_unsupported_json_features(schema: dict[str, Any]) -> bool: |
| 51 | """Check if JSON schema contains features unsupported by guidance/llguidance.""" |
| 52 | |
| 53 | def check_object(obj: dict[str, Any]) -> bool: |
| 54 | if not isinstance(obj, dict): |
| 55 | return False |
| 56 | |
| 57 | # patternProperties is not supported by llguidance |
| 58 | if "patternProperties" in obj: |
| 59 | return True |
| 60 | |
| 61 | # Recursively check all nested objects and arrays |
| 62 | for value in obj.values(): |
| 63 | if isinstance(value, dict): |
| 64 | if check_object(value): |
| 65 | return True |
| 66 | elif isinstance(value, list): |
| 67 | for item in value: |
| 68 | if isinstance(item, dict) and check_object(item): |
| 69 | return True |
| 70 | |
| 71 | return False |
| 72 | |
| 73 | return check_object(schema) |
| 74 | |
| 75 | |
| 76 | def process_for_additional_properties( |
| 77 | guide_json: str | dict[str, Any], |
| 78 | ) -> dict[str, Any]: |
| 79 | if isinstance(guide_json, str): |
| 80 | guide_json_obj = json.loads(guide_json) |
| 81 | else: |
| 82 | # copy for modifications |
| 83 | guide_json_obj = copy.deepcopy(guide_json) |
| 84 | _walk_json_for_additional_properties(guide_json_obj) |
| 85 | return guide_json_obj |
| 86 | |
| 87 | |
| 88 | @dataclass |
| 89 | class GuidanceBackend(StructuredOutputBackend): |
| 90 | def __post_init__(self): |
| 91 | self.disable_any_whitespace = ( |
| 92 | self.vllm_config.structured_outputs_config.disable_any_whitespace |
| 93 | ) |
| 94 | self.disable_additional_properties = ( |
| 95 | self.vllm_config.structured_outputs_config.disable_additional_properties |
| 96 | ) |
| 97 | |
| 98 | if is_mistral_tokenizer(self.tokenizer): |
| 99 | self.ll_tokenizer = self.tokenizer.llg_tokenizer |
| 100 | else: |
| 101 | self.ll_tokenizer = llguidance_hf.from_tokenizer( |
| 102 | self.tokenizer, max(self.vocab_size, len(self.tokenizer)) |
| 103 | ) |
| 104 | |
| 105 | def compile_grammar( |
| 106 | self, |
| 107 | request_type: StructuredOutputOptions, |
| 108 | grammar_spec: str, |
| 109 | stop_token_ids: set[int] | None = None, |
| 110 | ) -> StructuredOutputGrammar: |
| 111 | self.serialized_grammar = serialize_guidance_grammar( |
| 112 | request_type, |
| 113 | grammar_spec, |
| 114 | self.disable_any_whitespace, |
| 115 | self.disable_additional_properties, |
| 116 | ) |
| 117 | |
| 118 | ll_matcher = llguidance.LLMatcher( |
| 119 | self.ll_tokenizer, |
| 120 | self.serialized_grammar, |
| 121 | log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")), |
| 122 | ) |
| 123 | |
| 124 | r = GuidanceGrammar( |
| 125 | ll_matcher=ll_matcher, |
| 126 | ll_tokenizer=self.ll_tokenizer, |
| 127 | vocab_size=self.vocab_size, |
| 128 | ) |
| 129 | |
| 130 | r.check_error() |
| 131 | return r |
| 132 | |
| 133 | def allocate_token_bitmask(self, max_num_seqs: int): |
| 134 | return llguidance_torch.allocate_token_bitmask( |
| 135 | max_num_seqs, self.ll_tokenizer.vocab_size |
| 136 | ) |
| 137 | |
| 138 | def destroy(self): |
| 139 | pass |
| 140 | |
| 141 | |
| 142 | @dataclass |
| 143 | class GuidanceGrammar(StructuredOutputGrammar): |
| 144 | ll_matcher: llguidance.LLMatcher |
| 145 | ll_tokenizer: llguidance.LLTokenizer |
| 146 | vocab_size: int |
| 147 | printed_error: bool = False |
| 148 | terminated: bool = False |
| 149 | rollback_lag: int = 0 |
| 150 | |
| 151 | def check_error(self): |
| 152 | if not self.printed_error: |
| 153 | err = self.ll_matcher.get_error() |
| 154 | if err: |
| 155 | self.printed_error = True |
| 156 | logger.warning("LLMatcher error: %s", err) |
| 157 | |
| 158 | def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: |
| 159 | """Accepts a list of tokens and advances the parser. |
| 160 | |
| 161 | Returns True if the parser was advanced successfully. |
| 162 | Returns False if the parser failed to advance. |
| 163 | """ |
| 164 | if self.ll_tokenizer.eos_token in tokens: |
| 165 | if self.ll_matcher.is_stopped() and not self.terminated: |
| 166 | self.rollback_lag = 1 |
| 167 | self.terminated = True |
| 168 | |
| 169 | if self.ll_matcher.is_stopped(): |
| 170 | return True |
| 171 | |
| 172 | # TODO - Add jump decoding support in the future: |
| 173 | # self.ll_matcher.compute_ff_bytes() - this should always work |
| 174 | # self.ll_matcher.compute_ff_tokens() - this only works for |
| 175 | # "canonical" tokenizers |
| 176 | # For conversion between the two, see |
| 177 | # https://github.com/guidance-ai/llguidance/blob/main/docs/fast_forward.md |
| 178 | |
| 179 | r = self.ll_matcher.consume_tokens(tokens) |
| 180 | |
| 181 | self.check_error() |
| 182 | |
| 183 | return r |
| 184 | |
| 185 | def validate_tokens(self, tokens: list[int]) -> list[int]: |
| 186 | """Checks if the list of tokens are accepted by the parser in sequence. |
| 187 | Will not advance the parser. |
| 188 | |
| 189 | Returns the prefix list of tokens that are accepted by the parser. |
| 190 | """ |
| 191 | if len(tokens) == 0: |
| 192 | return [] |
| 193 | if self.ll_matcher.is_stopped(): |
| 194 | return [] |
| 195 | |
| 196 | tokens = strip_speculative_padding(tokens) |
| 197 | if len(tokens) == 0: |
| 198 | return [] |
| 199 | |
| 200 | num_tokens = self.ll_matcher.validate_tokens(tokens) |
| 201 | |
| 202 | self.check_error() |
| 203 | |
| 204 | return tokens[:num_tokens] |
| 205 | |
| 206 | def rollback(self, num_tokens: int) -> None: |
| 207 | if num_tokens > 0: |
| 208 | self.ll_matcher.rollback(num_tokens - self.rollback_lag) |
| 209 | self.terminated = False |
| 210 | self.rollback_lag = 0 |
| 211 | self.check_error() |
| 212 | |
| 213 | def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None: |
| 214 | # this will automatically return [EOS] mask if the matcher is stopped |
| 215 | # or otherwise in an error state |
| 216 | llguidance_torch.fill_next_token_bitmask(self.ll_matcher, bitmask, idx) |
| 217 | self.check_error() |
| 218 | |
| 219 | def is_terminated(self) -> bool: |
| 220 | return self.terminated |
| 221 | |
| 222 | def reset(self): |
| 223 | # This method may be not needed anymore? TODO |
| 224 | self.ll_matcher.reset() |
| 225 | |
| 226 | |
| 227 | def serialize_guidance_grammar( |
| 228 | request_type: StructuredOutputOptions, |
| 229 | grammar_spec: str | dict[str, Any], |
| 230 | disable_any_whitespace: bool = False, |
| 231 | disable_additional_properties: bool = False, |
| 232 | ) -> str: |
| 233 | def _process_schema( |
| 234 | grammar_spec: str | dict[str, Any], |
| 235 | ) -> str: |
| 236 | if disable_additional_properties: |
| 237 | grammar_spec = process_for_additional_properties(grammar_spec) |
| 238 | return llguidance.LLMatcher.grammar_from_json_schema( |
| 239 | grammar_spec, |
| 240 | defaults={ |
| 241 | "whitespace_flexible": not disable_any_whitespace, |
| 242 | }, |
| 243 | ) |
| 244 | |
| 245 | if request_type == StructuredOutputOptions.JSON: |
| 246 | return _process_schema(grammar_spec) |
| 247 | elif request_type == StructuredOutputOptions.JSON_OBJECT: |
| 248 | return llguidance.LLMatcher.grammar_from_json_schema( |
| 249 | '{"type": "object"}', |
| 250 | defaults={ |
| 251 | "whitespace_flexible": not disable_any_whitespace, |
| 252 | }, |
| 253 | ) |
| 254 | else: |
| 255 | if request_type == StructuredOutputOptions.REGEX: |
| 256 | tp = "regex" |
| 257 | elif request_type == StructuredOutputOptions.GRAMMAR: |
| 258 | tp = "grammar" |
| 259 | elif request_type == StructuredOutputOptions.CHOICE: |
| 260 | tp = "choice" |
| 261 | elif request_type == StructuredOutputOptions.STRUCTURAL_TAG: |
| 262 | if isinstance(grammar_spec, str): |
| 263 | s_tag = json.loads(grammar_spec) |
| 264 | else: |
| 265 | s_tag = grammar_spec |
| 266 | triggers: list[str] = s_tag["triggers"] |
| 267 | tags: list[llguidance.StructTag] = [] |
| 268 | for s in s_tag["structures"]: |
| 269 | begin: str = s["begin"] |
| 270 | trig = next((t for t in triggers if begin.startswith(t)), None) |
| 271 | if trig is None: |
| 272 | raise VLLMValidationError( |
| 273 | f"Trigger {begin} not found in triggers {triggers}" |
| 274 | ) |
| 275 | tags.append( |
| 276 | llguidance.StructTag( |
| 277 | trigger=trig, |
| 278 | begin=s["begin"], |
| 279 | grammar=_process_schema(s["schema"]), |
| 280 | end=s["end"], |
| 281 | ) |
| 282 | ) |
| 283 | if not tags: |
| 284 | raise VLLMValidationError( |
| 285 | "No structural tags found in the grammar spec." |
| 286 | ) |
| 287 | return llguidance.StructTag.to_grammar(tags) |
| 288 | else: |
| 289 | logger.error( |
| 290 | "Validation should have already occurred. Please file an issue." |
| 291 | ) |
| 292 | raise ValueError( |
| 293 | f"grammar is not of valid supported types. ({request_type!s})" |
| 294 | ) |
| 295 | return llguidance.grammar_from(tp, grammar_spec) |
| 296 | |
| 297 | |
| 298 | def validate_guidance_grammar( |
| 299 | sampling_params: SamplingParams, tokenizer: llguidance.LLTokenizer | None = None |
| 300 | ) -> None: |
| 301 | # if structured output is not enabled, there is nothing to validate |
| 302 | if sampling_params.structured_outputs is None: |
| 303 | return |
| 304 | tp, grm = get_structured_output_key(sampling_params.structured_outputs) |
| 305 | try: |
| 306 | guidance_grm = serialize_guidance_grammar(tp, grm) |
| 307 | except (ValueError, KeyError, TypeError) as e: |
| 308 | raise VLLMValidationError(f"Invalid grammar specification: {e}") from e |
| 309 | err = llguidance.LLMatcher.validate_grammar(guidance_grm, tokenizer) |
| 310 | if err: |
| 311 | raise VLLMValidationError(f"Grammar error: {err}") |