| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright 2026 The HuggingFace Team. All rights reserved. |
| 5 | # ruff: noqa |
| 6 | # fmt: off |
| 7 | # |
| 8 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 9 | # you may not use this file except in compliance with the License. |
| 10 | # You may obtain a copy of the License at |
| 11 | # |
| 12 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 13 | # |
| 14 | # Unless required by applicable law or agreed to in writing, software |
| 15 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 16 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 17 | # See the License for the specific language governing permissions and |
| 18 | # limitations under the License. |
| 19 | """Template loading and validation for response_template dicts.""" |
| 20 | |
| 21 | from __future__ import annotations |
| 22 | |
| 23 | import copy |
| 24 | import logging |
| 25 | from dataclasses import dataclass |
| 26 | from typing import Any |
| 27 | |
| 28 | import regex as re |
| 29 | |
| 30 | from .content_parsers import CONTENT_PARSERS, validate_transform_strings |
| 31 | |
| 32 | |
| 33 | logger = logging.getLogger(__name__) |
| 34 | |
| 35 | |
| 36 | @dataclass |
| 37 | class ResponseTemplateField: |
| 38 | name: str |
| 39 | open_re: Any |
| 40 | open_literals: list[str] | None |
| 41 | open_literal_can_extend: bool |
| 42 | close_re: Any |
| 43 | close_literals: list[str] | None |
| 44 | close_literal_can_extend: bool |
| 45 | content: str |
| 46 | content_args: dict |
| 47 | repeats: bool |
| 48 | join: str | None |
| 49 | optional: bool |
| 50 | transform: Any |
| 51 | transform_each: bool |
| 52 | |
| 53 | |
| 54 | @dataclass |
| 55 | class ResponseTemplate: |
| 56 | defaults: dict |
| 57 | fields: dict[str, ResponseTemplateField] |
| 58 | start_anchor_re: Any |
| 59 | implicit: str | None = None |
| 60 | start_anchor_literals: list[str] | None = None |
| 61 | |
| 62 | def truncate_past_last_anchor(self, text: str) -> str: |
| 63 | """When parsing responses, we can't just parse the tokens generated |
| 64 | by the model. This is because chat templates often include early parts of the message such as <start_of_turn> |
| 65 | or <think> tokens in the prefill. Therefore, we take the whole chat as input, but truncate to the start |
| 66 | of the most recent assistant message, which can include tokens from the template as well as the output.""" |
| 67 | last_end: int | None = None |
| 68 | for m in self.start_anchor_re.finditer(text): |
| 69 | last_end = m.end() |
| 70 | if last_end is None: |
| 71 | kind = "start_anchor" if self.start_anchor_literals is not None else "start_anchor_pattern" |
| 72 | logger.info( |
| 73 | "response_template defines %s but the anchor was not found " |
| 74 | "in the prefix; the parser will process the entire prefix instead.", |
| 75 | kind, |
| 76 | ) |
| 77 | return text |
| 78 | return text[last_end:] |
| 79 | |
| 80 | |
| 81 | def _compile_anchor(scope: str, field: dict, literal_key: str, pattern_key: str) -> tuple[Any, list[str] | None, bool]: |
| 82 | """Compile a region's start/end anchor. An anchor can be: |
| 83 | |
| 84 | - a regex (`pattern_key`), |
| 85 | - a single literal string (`literal_key`), or |
| 86 | - a list of literal strings (`literal_key`), in which case any of these match |
| 87 | |
| 88 | All forms compile to a regex for matching, but we keep the literal strings around too: when present, |
| 89 | the parser can emit regions faster, since it knows exactly how many trailing bytes might be a |
| 90 | partial-match prefix (and that a finished literal match can never be extended by future bytes, except |
| 91 | when a literal is a strict prefix of another in the same list). |
| 92 | |
| 93 | Returns `(compiled_re, literals, can_extend)`: |
| 94 | |
| 95 | - `compiled_re`: the compiled pattern. |
| 96 | - `literals`: the literal strings, or `None` for a regex anchor. |
| 97 | - `can_extend`: `True` iff one literal is a strict prefix of another, in which case a match at the |
| 98 | buffer edge could still be lengthened by future input. |
| 99 | """ |
| 100 | if literal_key in field and pattern_key in field: |
| 101 | raise ValueError(f"{scope}: cannot specify both '{literal_key}' and '{pattern_key}'") |
| 102 | if literal_key in field: |
| 103 | raw = field[literal_key] |
| 104 | if isinstance(raw, str): |
| 105 | literals = [raw] |
| 106 | elif isinstance(raw, list): |
| 107 | if not raw: |
| 108 | raise ValueError(f"{scope}: '{literal_key}' list must contain at least one literal") |
| 109 | if not all(isinstance(s, str) for s in raw): |
| 110 | raise ValueError(f"{scope}: '{literal_key}' list must contain only strings") |
| 111 | literals = list(dict.fromkeys(raw)) # dedupe, preserve first-seen order |
| 112 | else: |
| 113 | raise ValueError(f"{scope}: '{literal_key}' must be a string or list of strings, got {type(raw).__name__}") |
| 114 | if any(s == "" for s in literals): |
| 115 | raise ValueError(f"{scope}: '{literal_key}' literals cannot be empty strings") |
| 116 | # Sort longest-first so alternation prefers the longer alternative when both could match. |
| 117 | ordered = sorted(literals, key=len, reverse=True) |
| 118 | can_extend = any(a != b and a.startswith(b) for a in literals for b in literals) |
| 119 | pattern = "|".join(re.escape(s) for s in ordered) |
| 120 | return re.compile(pattern, re.DOTALL), literals, can_extend |
| 121 | if pattern_key in field: |
| 122 | try: |
| 123 | return re.compile(field[pattern_key], re.DOTALL), None, False |
| 124 | except re.error as e: |
| 125 | raise ValueError(f"{scope}: invalid {pattern_key} regex: {e}") from e |
| 126 | return None, None, False |
| 127 | |
| 128 | |
| 129 | def _validate_template_shape(spec: dict) -> None: |
| 130 | """Validate the top-level response_template structure: it must be a dict of the expected version with |
| 131 | only known keys, a dict of `defaults`, and a non-empty dict of `fields`. Per-field specs are checked |
| 132 | separately in `_build_field`.""" |
| 133 | if not isinstance(spec, dict): |
| 134 | raise ValueError(f"response_template must be a dict, got {type(spec).__name__}") |
| 135 | version = spec.get("version", 1) |
| 136 | if version != 1: |
| 137 | raise ValueError(f"Unsupported response_template version: {version}") |
| 138 | if unknown_template_keys := set(spec) - {"version", "defaults", "fields", "start_anchor", "start_anchor_pattern"}: |
| 139 | raise ValueError(f"Unknown keys in response_template: {sorted(unknown_template_keys)}") |
| 140 | if not isinstance(spec.get("defaults", {}), dict): |
| 141 | raise ValueError("response_template.defaults must be a dict") |
| 142 | fields_raw = spec.get("fields", {}) |
| 143 | if not isinstance(fields_raw, dict) or not fields_raw: |
| 144 | raise ValueError("response_template.fields must be a non-empty dict") |
| 145 | |
| 146 | |
| 147 | def _build_field(name: str, field: dict) -> ResponseTemplateField: |
| 148 | """Validate a single field spec and compile it into a `ResponseTemplateField`.""" |
| 149 | |
| 150 | _ALLOWED_FIELD_KEYS = { |
| 151 | "open", |
| 152 | "open_pattern", |
| 153 | "close", |
| 154 | "close_pattern", |
| 155 | "content", |
| 156 | "content_args", |
| 157 | "repeats", |
| 158 | "join", |
| 159 | "optional", |
| 160 | "transform", |
| 161 | "transform_each", |
| 162 | } |
| 163 | scope = f"Field '{name}'" |
| 164 | if not isinstance(field, dict): |
| 165 | raise ValueError(f"{scope} must be a dict") |
| 166 | if unknown_field_keys := set(field) - _ALLOWED_FIELD_KEYS: |
| 167 | raise ValueError(f"{scope}: unknown keys {sorted(unknown_field_keys)}") |
| 168 | content = field.get("content", "text") |
| 169 | if content not in CONTENT_PARSERS: |
| 170 | raise ValueError(f"{scope}: unknown content parser '{content}'. Available: {sorted(CONTENT_PARSERS)}") |
| 171 | open_re, open_literals, open_literal_can_extend = _compile_anchor(scope, field, "open", "open_pattern") |
| 172 | close_re, close_literals, close_literal_can_extend = _compile_anchor(scope, field, "close", "close_pattern") |
| 173 | join = field.get("join") |
| 174 | if join is not None and not isinstance(join, str): |
| 175 | raise ValueError(f"{scope}: 'join' must be a string, got {type(join).__name__}") |
| 176 | if join is not None and not field.get("repeats", False): |
| 177 | raise ValueError(f"{scope}: 'join' requires 'repeats': true") |
| 178 | transform = field.get("transform") |
| 179 | transform_each = field.get("transform_each", False) |
| 180 | if not isinstance(transform_each, bool): |
| 181 | raise ValueError(f"{scope}: transform_each must be a bool, got {type(transform_each).__name__}") |
| 182 | if transform_each and transform is None: |
| 183 | raise ValueError(f"{scope}: transform_each is set but no transform was provided") |
| 184 | if transform is not None: |
| 185 | validate_transform_strings(scope, transform) |
| 186 | else: |
| 187 | # Named captures only reach the output through a transform, so flag any that would be silently dropped. |
| 188 | captured_names = set() |
| 189 | if open_re is not None: |
| 190 | captured_names |= set(open_re.groupindex) |
| 191 | if close_re is not None: |
| 192 | captured_names |= set(close_re.groupindex) |
| 193 | if captured_names: |
| 194 | raise ValueError( |
| 195 | f"{scope}: open_pattern/close_pattern declares named group(s) " |
| 196 | f"{sorted(captured_names)}, but the field has no 'transform'. Named captures " |
| 197 | f"are only surfaced through a 'transform' template (where they appear " |
| 198 | f"alongside 'content'). Either add a 'transform' that uses the captures, or " |
| 199 | f"remove the named groups from the pattern." |
| 200 | ) |
| 201 | return ResponseTemplateField( |
| 202 | name=name, |
| 203 | open_re=open_re, |
| 204 | open_literals=open_literals, |
| 205 | open_literal_can_extend=open_literal_can_extend, |
| 206 | close_re=close_re, |
| 207 | close_literals=close_literals, |
| 208 | close_literal_can_extend=close_literal_can_extend, |
| 209 | content=content, |
| 210 | content_args=field.get("content_args", {}), |
| 211 | repeats=field.get("repeats", False), |
| 212 | join=join, |
| 213 | optional=field.get("optional", True), |
| 214 | transform=transform, |
| 215 | transform_each=transform_each, |
| 216 | ) |
| 217 | |
| 218 | |
| 219 | def load_response_template(spec: dict | ResponseTemplate) -> ResponseTemplate: |
| 220 | if isinstance(spec, ResponseTemplate): |
| 221 | return spec |
| 222 | _validate_template_shape(spec) |
| 223 | |
| 224 | fields = {name: _build_field(name, raw) for name, raw in spec["fields"].items()} |
| 225 | # A field without an open anchor is the implicit-open / leftover sink; at most one is allowed. |
| 226 | implicit_fields = [name for name, field in fields.items() if field.open_re is None] |
| 227 | if len(implicit_fields) > 1: |
| 228 | raise ValueError( |
| 229 | "At most one field may omit 'open'/'open_pattern' (that field becomes the " |
| 230 | f"implicit-open / leftover sink). Found: {', '.join(implicit_fields)}" |
| 231 | ) |
| 232 | |
| 233 | start_anchor_re, start_anchor_literals, _ = _compile_anchor( |
| 234 | "response_template", spec, "start_anchor", "start_anchor_pattern" |
| 235 | ) |
| 236 | if start_anchor_re is None: |
| 237 | raise ValueError("response_template must define 'start_anchor' or 'start_anchor_pattern'.") |
| 238 | |
| 239 | return ResponseTemplate( |
| 240 | defaults=copy.deepcopy(spec.get("defaults", {})), |
| 241 | fields=fields, |
| 242 | implicit=implicit_fields[0] if implicit_fields else None, |
| 243 | start_anchor_re=start_anchor_re, |
| 244 | start_anchor_literals=start_anchor_literals, |
| 245 | ) |
| 246 | |
| 247 | |
| 248 | __all__ = ["ResponseTemplate", "ResponseTemplateField", "load_response_template"] |
| 249 | # fmt: on |