| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright 2025-present the Outlines developers |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import ast |
| 7 | import importlib |
| 8 | import json |
| 9 | import sys |
| 10 | from dataclasses import dataclass, field |
| 11 | from typing import TYPE_CHECKING |
| 12 | |
| 13 | import torch |
| 14 | from regex import escape as regex_escape |
| 15 | |
| 16 | from vllm.exceptions import VLLMValidationError |
| 17 | from vllm.sampling_params import SamplingParams |
| 18 | from vllm.utils.import_utils import LazyLoader |
| 19 | from vllm.utils.torch_utils import PIN_MEMORY |
| 20 | from vllm.v1.structured_output.backend_types import ( |
| 21 | StructuredOutputBackend, |
| 22 | StructuredOutputGrammar, |
| 23 | StructuredOutputOptions, |
| 24 | ) |
| 25 | from vllm.v1.structured_output.utils import ( |
| 26 | OutlinesVocabulary, |
| 27 | compile_regex_with_timeout, |
| 28 | get_outlines_cache, |
| 29 | get_outlines_vocabulary, |
| 30 | ) |
| 31 | |
| 32 | if TYPE_CHECKING: |
| 33 | import outlines_core as oc |
| 34 | import outlines_core.json_schema as json_schema |
| 35 | else: |
| 36 | oc = LazyLoader("oc", globals(), "outlines_core") |
| 37 | json_schema = LazyLoader("json_schema", globals(), "outlines_core.json_schema") |
| 38 | |
| 39 | # Python 3.11+ sre_parse and sre_constants |
| 40 | # are deprecated, so we must import them from re |
| 41 | if sys.version_info >= (3, 11): |
| 42 | # Hack to get around pre-commit regex module rule |
| 43 | # because going through re is the only way to get sre_parse |
| 44 | # and sre_constants in Python 3.11+ |
| 45 | _re = importlib.import_module("re") |
| 46 | sre_parse = _re._parser |
| 47 | sre_constants = _re._constants |
| 48 | else: |
| 49 | import sre_constants |
| 50 | import sre_parse |
| 51 | |
| 52 | |
| 53 | @dataclass |
| 54 | class OutlinesBackend(StructuredOutputBackend): |
| 55 | def __post_init__(self): |
| 56 | self.vocabulary = get_outlines_vocabulary(self.tokenizer) |
| 57 | self.cache = get_outlines_cache() |
| 58 | |
| 59 | def _compile_index( |
| 60 | self, regex_string: str, vocabulary: OutlinesVocabulary |
| 61 | ) -> oc.Index: |
| 62 | cache_key = f"{vocabulary._hash}_{regex_string}" |
| 63 | if cache_key in self.cache: |
| 64 | return self.cache[cache_key] |
| 65 | |
| 66 | index = compile_regex_with_timeout( |
| 67 | lambda pat: oc.Index(pat, vocabulary.inner), |
| 68 | regex_string, |
| 69 | ) |
| 70 | self.cache[cache_key] = index |
| 71 | |
| 72 | return index |
| 73 | |
| 74 | def compile_grammar( |
| 75 | self, |
| 76 | request_type: StructuredOutputOptions, |
| 77 | grammar_spec: str, |
| 78 | stop_token_ids: set[int] | None = None, |
| 79 | ) -> StructuredOutputGrammar: |
| 80 | if request_type == StructuredOutputOptions.JSON: |
| 81 | regex = json_schema.build_regex_from_schema(grammar_spec) |
| 82 | elif request_type == StructuredOutputOptions.REGEX: |
| 83 | regex = grammar_spec |
| 84 | elif request_type == StructuredOutputOptions.CHOICE: |
| 85 | choices = ast.literal_eval(grammar_spec) |
| 86 | choices = [regex_escape(c) for c in choices] |
| 87 | regex = "(" + "|".join(choices) + ")" |
| 88 | else: |
| 89 | raise ValueError( |
| 90 | f"Invalid request type for Outlines backend ({request_type!s})" |
| 91 | ) |
| 92 | index = self._compile_index(regex, self.vocabulary) |
| 93 | max_rollback_tokens = ( |
| 94 | self.vllm_config.speculative_config.num_speculative_tokens |
| 95 | if self.vllm_config.speculative_config is not None |
| 96 | else 0 |
| 97 | ) |
| 98 | return OutlinesGrammar( |
| 99 | vocab_size=self.vocab_size, |
| 100 | eos_token_id=self.vocabulary.inner.get_eos_token_id(), |
| 101 | guide=oc.Guide(index, max_rollback=max_rollback_tokens), |
| 102 | ) |
| 103 | |
| 104 | def allocate_token_bitmask(self, max_num_seqs: int) -> torch.Tensor: |
| 105 | return torch.full( |
| 106 | (max_num_seqs, (self.vocab_size + 31) // 32), |
| 107 | -1, |
| 108 | dtype=torch.int32, |
| 109 | pin_memory=PIN_MEMORY, |
| 110 | ) |
| 111 | |
| 112 | def destroy(self): |
| 113 | pass |
| 114 | |
| 115 | |
| 116 | @dataclass |
| 117 | class OutlinesGrammar(StructuredOutputGrammar): |
| 118 | vocab_size: int |
| 119 | eos_token_id: int |
| 120 | guide: oc.Guide = field(hash=False) |
| 121 | num_processed_tokens: int = field( |
| 122 | default_factory=lambda: 0, repr=False, hash=False, init=False |
| 123 | ) |
| 124 | |
| 125 | # outlines_core signals done on DFA accept and never advances on EOS |
| 126 | # (its mask allows only EOS in a final state, but advance(EOS) fails), |
| 127 | # so accepting EOS in a final state is tracked here instead. |
| 128 | _is_terminated: bool = field(default=False, init=False, repr=False, hash=False) |
| 129 | |
| 130 | def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: |
| 131 | """Accepts a list of tokens and advances the FSM. |
| 132 | |
| 133 | Returns True if all grammar-constrained tokens were accepted. |
| 134 | Tokens after termination (EOS) are ignored. Returns False if the FSM |
| 135 | failed to advance. |
| 136 | """ |
| 137 | if self._is_terminated: |
| 138 | return True |
| 139 | eos_index = next( |
| 140 | (i for i, t in enumerate(tokens) if t == self.eos_token_id), None |
| 141 | ) |
| 142 | if eos_index is not None: |
| 143 | tokens = tokens[:eos_index] |
| 144 | # Advance can fail when the next state reached after advancing with |
| 145 | # the current tokens is a dead state. This is because Guide.accepts_tokens() |
| 146 | # only checks whether the current tokens can be accepted, |
| 147 | # whereas guide.advance() additionally checks the next state |
| 148 | # after all tokens are accepted. |
| 149 | # We need to be aware that the FSM must be prepared without dead states. |
| 150 | if tokens and not self.guide.accepts_tokens(tokens): |
| 151 | return False |
| 152 | for t in tokens: |
| 153 | self.guide.advance(t) |
| 154 | if eos_index is not None: |
| 155 | if not self.guide.is_finished(): |
| 156 | if tokens: |
| 157 | self.guide.rollback_state(len(tokens)) |
| 158 | return False |
| 159 | self._is_terminated = True |
| 160 | self.num_processed_tokens += len(tokens) + self._is_terminated |
| 161 | return True |
| 162 | |
| 163 | def rollback(self, num_tokens: int) -> None: |
| 164 | if num_tokens <= 0: |
| 165 | return |
| 166 | self.num_processed_tokens -= num_tokens |
| 167 | if self._is_terminated: |
| 168 | self._is_terminated = False |
| 169 | num_tokens -= 1 |
| 170 | if num_tokens: |
| 171 | self.guide.rollback_state(num_tokens) |
| 172 | |
| 173 | def validate_tokens(self, tokens: list[int]) -> list[int]: |
| 174 | if self._is_terminated: |
| 175 | return [] |
| 176 | accepted: list[int] = [] |
| 177 | for tok in tokens: |
| 178 | if tok == self.eos_token_id: |
| 179 | if self._is_finished_after(accepted): |
| 180 | accepted.append(tok) |
| 181 | break |
| 182 | accepted.append(tok) |
| 183 | if not self.guide.accepts_tokens(accepted): |
| 184 | accepted.pop() |
| 185 | break |
| 186 | return accepted |
| 187 | |
| 188 | def _is_finished_after(self, tokens: list[int]) -> bool: |
| 189 | if not tokens: |
| 190 | return self.guide.is_finished() |
| 191 | for t in tokens: |
| 192 | self.guide.advance(t) |
| 193 | finished = self.guide.is_finished() |
| 194 | self.guide.rollback_state(len(tokens)) |
| 195 | return finished |
| 196 | |
| 197 | def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None: |
| 198 | mask = bitmask[idx] |
| 199 | self.guide.write_mask_into(mask.data_ptr(), mask.numel(), mask.element_size()) |
| 200 | |
| 201 | def is_terminated(self) -> bool: |
| 202 | return self._is_terminated |
| 203 | |
| 204 | def reset(self): |
| 205 | self.num_processed_tokens = 0 |
| 206 | self._is_terminated = False |
| 207 | self.guide.reset() |
| 208 | |
| 209 | |
| 210 | def validate_structured_output_request_outlines(params: SamplingParams): |
| 211 | if params.structured_outputs is None: |
| 212 | return |
| 213 | |
| 214 | so_params = params.structured_outputs |
| 215 | |
| 216 | if so_params.regex: |
| 217 | validate_regex_is_buildable(so_params.regex) |
| 218 | elif so_params.json: |
| 219 | if isinstance(so_params.json, str): |
| 220 | try: |
| 221 | # make sure schema is valid json |
| 222 | json.loads(so_params.json) |
| 223 | schema = so_params.json |
| 224 | except json.JSONDecodeError as e: |
| 225 | raise VLLMValidationError("Invalid JSON grammar specification.") from e |
| 226 | else: |
| 227 | try: |
| 228 | schema = json.dumps(so_params.json) |
| 229 | except Exception as e: |
| 230 | raise VLLMValidationError( |
| 231 | f"Error serializing structured outputs jsonschema: {e}" |
| 232 | ) from e |
| 233 | try: |
| 234 | pattern = json_schema.build_regex_from_schema(schema) |
| 235 | except Exception as e: |
| 236 | raise VLLMValidationError( |
| 237 | f"Failed to transform json schema into a regex: {e}" |
| 238 | ) from e |
| 239 | validate_regex_is_buildable(pattern) |
| 240 | elif so_params.choice: |
| 241 | choices = [regex_escape(str(choice)) for choice in so_params.choice] |
| 242 | regex = "(" + "|".join(choices) + ")" |
| 243 | validate_regex_is_buildable(regex) |
| 244 | elif so_params.json_object: |
| 245 | raise VLLMValidationError( |
| 246 | "Outlines structured outputs backend " |
| 247 | "does not support json_object specifications" |
| 248 | ) |
| 249 | elif so_params.grammar: |
| 250 | raise VLLMValidationError( |
| 251 | "Outlines structured outputs backend " |
| 252 | "does not support grammar specifications" |
| 253 | ) |
| 254 | |
| 255 | |
| 256 | def _prefix_needs_context(parsed) -> bool: |
| 257 | """Return True if there's a look-around/anchor before any consumer.""" |
| 258 | |
| 259 | def subpattern_consumes(parsed) -> bool: |
| 260 | """Return True if subpattern can consume at least one character.""" |
| 261 | tokens = parsed.data if hasattr(parsed, "data") else parsed |
| 262 | for ttype, tval in tokens: |
| 263 | # literal, character class, or dot always consumes |
| 264 | if ttype in (sre_parse.LITERAL, sre_parse.IN, sre_parse.ANY): |
| 265 | return True |
| 266 | # quantified subpattern: check inner pattern |
| 267 | elif ttype == sre_parse.MAX_REPEAT: |
| 268 | _, mx, sub = tval |
| 269 | if mx != 0 and subpattern_consumes(sub): |
| 270 | return True |
| 271 | # alternation: if any branch consumes, the whole does |
| 272 | elif ttype == sre_parse.BRANCH: |
| 273 | _, branches = tval |
| 274 | if any(subpattern_consumes(br) for br in branches): |
| 275 | return True |
| 276 | # grouped subpattern: recurse into its contents |
| 277 | elif ttype == sre_parse.SUBPATTERN and subpattern_consumes(tval[3]): |
| 278 | return True |
| 279 | # No consumers, return False |
| 280 | return False |
| 281 | |
| 282 | tokens = parsed.data if hasattr(parsed, "data") else parsed |
| 283 | for ttype, tval in tokens: |
| 284 | # Direct anchors or look-around |
| 285 | if ttype == sre_parse.AT or ttype in ( |
| 286 | sre_constants.ASSERT, |
| 287 | sre_constants.ASSERT_NOT, |
| 288 | ): |
| 289 | return True |
| 290 | |
| 291 | # Nested subpattern: check |
| 292 | if ttype == sre_parse.SUBPATTERN: |
| 293 | # tval: (group, add_flags, del_flags, subpattern) |
| 294 | if _prefix_needs_context(tval[3]): |
| 295 | return True |
| 296 | if subpattern_consumes(tval[3]): |
| 297 | return False |
| 298 | |
| 299 | # if any branch has a prefix anchor => True, |
| 300 | # else if at least one branch consumes => prefix ends => False |
| 301 | elif ttype == sre_parse.BRANCH: |
| 302 | saw_consumer = False |
| 303 | for br in tval[1]: |
| 304 | if _prefix_needs_context(br): |
| 305 | return True |
| 306 | if subpattern_consumes(br): |
| 307 | saw_consumer = True |
| 308 | if saw_consumer: |
| 309 | return False |
| 310 | |
| 311 | # Immediate consumer tokens |
| 312 | elif ttype in (sre_parse.LITERAL, sre_parse.IN, sre_parse.ANY): |
| 313 | return False |
| 314 | |
| 315 | # if subpattern has anchor => True, if it can consume => stop |
| 316 | elif ttype == sre_parse.MAX_REPEAT: |
| 317 | if _prefix_needs_context(tval[2]): |
| 318 | return True |
| 319 | if subpattern_consumes(tval[2]): |
| 320 | return False |
| 321 | |
| 322 | return False |
| 323 | |
| 324 | |
| 325 | def _check_unsupported(parsed) -> None: |
| 326 | """Check for regex features unsupported by regex-automata.""" |
| 327 | tokens = parsed.data if hasattr(parsed, "data") else parsed |
| 328 | for ttype, tval in tokens: |
| 329 | # backreference |
| 330 | if ttype in (sre_parse.GROUPREF, sre_parse.GROUPREF_EXISTS): |
| 331 | raise ValueError("Backreferences are unsupported.") |
| 332 | |
| 333 | # look-around assertion |
| 334 | elif ttype in (sre_constants.ASSERT, sre_constants.ASSERT_NOT): |
| 335 | raise ValueError("Look-Around assertion are unsupported.") |
| 336 | |
| 337 | # unicode word boundaries |
| 338 | elif ttype == sre_parse.AT: |
| 339 | if tval in (sre_constants.AT_BOUNDARY, sre_constants.AT_NON_BOUNDARY): |
| 340 | raise ValueError("Unicode word boundaries are unsupported.") |
| 341 | |
| 342 | elif ttype == sre_parse.BRANCH: |
| 343 | # tval is (None, branches) |
| 344 | for branch in tval[1]: |
| 345 | _check_unsupported(branch) |
| 346 | |
| 347 | # tval is (min, max, subpattern) |
| 348 | elif ttype == sre_parse.MAX_REPEAT: |
| 349 | _check_unsupported(tval[2]) |
| 350 | |
| 351 | |
| 352 | def validate_regex_is_buildable(pattern: str) -> None: |
| 353 | """Validates that the input regex is not using unsupported features |
| 354 | of the `regex-automata` crate (outlines_core regex engine) and has a |
| 355 | universal start state. |
| 356 | definition of universal start state used can be found at: |
| 357 | https://docs.rs/regex-automata/latest/regex_automata/dfa/trait.Automaton.html#method.universal_start_state |
| 358 | """ |
| 359 | try: |
| 360 | parsed = sre_parse.parse(pattern) |
| 361 | |
| 362 | except sre_constants.error as e: |
| 363 | raise VLLMValidationError(f"Error parsing regex: {e}") from e |
| 364 | |
| 365 | try: |
| 366 | _check_unsupported(parsed) |
| 367 | except ValueError as e: |
| 368 | raise VLLMValidationError( |
| 369 | f"Regex uses unsupported feature for structured outputs: {e}. " |
| 370 | "Only basic matching constructs are supported—lookarounds, " |
| 371 | "backreferences, and unicode boundaries are not." |
| 372 | ) from e |
| 373 | |
| 374 | if _prefix_needs_context(parsed): |
| 375 | raise VLLMValidationError( |
| 376 | "Regex does not have a anchored universal start state" |
| 377 | "This means that the Regex uses anchors (^) or look-arounds " |
| 378 | "in a way which requires context before any token is matched." |
| 379 | "structured outputs needs regexes that can match without needing " |
| 380 | "that context. Try rewriting the pattern without using these " |
| 381 | f"constructs. Pattern:\n{pattern}" |
| 382 | ) |