73a5831127
Signed-off-by: Yoni Gozlan <yonigozlan@users.noreply.github.com> Co-authored-by: Yoni Gozlan <yonigozlan@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
250 lines
11 KiB
Python
250 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
# Copyright 2026 The HuggingFace Team. All rights reserved.
|
|
# ruff: noqa
|
|
# fmt: off
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Template loading and validation for response_template dicts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import regex as re
|
|
|
|
from .content_parsers import CONTENT_PARSERS, validate_transform_strings
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ResponseTemplateField:
|
|
name: str
|
|
open_re: Any
|
|
open_literals: list[str] | None
|
|
open_literal_can_extend: bool
|
|
close_re: Any
|
|
close_literals: list[str] | None
|
|
close_literal_can_extend: bool
|
|
content: str
|
|
content_args: dict
|
|
repeats: bool
|
|
join: str | None
|
|
optional: bool
|
|
transform: Any
|
|
transform_each: bool
|
|
|
|
|
|
@dataclass
|
|
class ResponseTemplate:
|
|
defaults: dict
|
|
fields: dict[str, ResponseTemplateField]
|
|
start_anchor_re: Any
|
|
implicit: str | None = None
|
|
start_anchor_literals: list[str] | None = None
|
|
|
|
def truncate_past_last_anchor(self, text: str) -> str:
|
|
"""When parsing responses, we can't just parse the tokens generated
|
|
by the model. This is because chat templates often include early parts of the message such as <start_of_turn>
|
|
or <think> tokens in the prefill. Therefore, we take the whole chat as input, but truncate to the start
|
|
of the most recent assistant message, which can include tokens from the template as well as the output."""
|
|
last_end: int | None = None
|
|
for m in self.start_anchor_re.finditer(text):
|
|
last_end = m.end()
|
|
if last_end is None:
|
|
kind = "start_anchor" if self.start_anchor_literals is not None else "start_anchor_pattern"
|
|
logger.info(
|
|
"response_template defines %s but the anchor was not found "
|
|
"in the prefix; the parser will process the entire prefix instead.",
|
|
kind,
|
|
)
|
|
return text
|
|
return text[last_end:]
|
|
|
|
|
|
def _compile_anchor(scope: str, field: dict, literal_key: str, pattern_key: str) -> tuple[Any, list[str] | None, bool]:
|
|
"""Compile a region's start/end anchor. An anchor can be:
|
|
|
|
- a regex (`pattern_key`),
|
|
- a single literal string (`literal_key`), or
|
|
- a list of literal strings (`literal_key`), in which case any of these match
|
|
|
|
All forms compile to a regex for matching, but we keep the literal strings around too: when present,
|
|
the parser can emit regions faster, since it knows exactly how many trailing bytes might be a
|
|
partial-match prefix (and that a finished literal match can never be extended by future bytes, except
|
|
when a literal is a strict prefix of another in the same list).
|
|
|
|
Returns `(compiled_re, literals, can_extend)`:
|
|
|
|
- `compiled_re`: the compiled pattern.
|
|
- `literals`: the literal strings, or `None` for a regex anchor.
|
|
- `can_extend`: `True` iff one literal is a strict prefix of another, in which case a match at the
|
|
buffer edge could still be lengthened by future input.
|
|
"""
|
|
if literal_key in field and pattern_key in field:
|
|
raise ValueError(f"{scope}: cannot specify both '{literal_key}' and '{pattern_key}'")
|
|
if literal_key in field:
|
|
raw = field[literal_key]
|
|
if isinstance(raw, str):
|
|
literals = [raw]
|
|
elif isinstance(raw, list):
|
|
if not raw:
|
|
raise ValueError(f"{scope}: '{literal_key}' list must contain at least one literal")
|
|
if not all(isinstance(s, str) for s in raw):
|
|
raise ValueError(f"{scope}: '{literal_key}' list must contain only strings")
|
|
literals = list(dict.fromkeys(raw)) # dedupe, preserve first-seen order
|
|
else:
|
|
raise ValueError(f"{scope}: '{literal_key}' must be a string or list of strings, got {type(raw).__name__}")
|
|
if any(s == "" for s in literals):
|
|
raise ValueError(f"{scope}: '{literal_key}' literals cannot be empty strings")
|
|
# Sort longest-first so alternation prefers the longer alternative when both could match.
|
|
ordered = sorted(literals, key=len, reverse=True)
|
|
can_extend = any(a != b and a.startswith(b) for a in literals for b in literals)
|
|
pattern = "|".join(re.escape(s) for s in ordered)
|
|
return re.compile(pattern, re.DOTALL), literals, can_extend
|
|
if pattern_key in field:
|
|
try:
|
|
return re.compile(field[pattern_key], re.DOTALL), None, False
|
|
except re.error as e:
|
|
raise ValueError(f"{scope}: invalid {pattern_key} regex: {e}") from e
|
|
return None, None, False
|
|
|
|
|
|
def _validate_template_shape(spec: dict) -> None:
|
|
"""Validate the top-level response_template structure: it must be a dict of the expected version with
|
|
only known keys, a dict of `defaults`, and a non-empty dict of `fields`. Per-field specs are checked
|
|
separately in `_build_field`."""
|
|
if not isinstance(spec, dict):
|
|
raise ValueError(f"response_template must be a dict, got {type(spec).__name__}")
|
|
version = spec.get("version", 1)
|
|
if version != 1:
|
|
raise ValueError(f"Unsupported response_template version: {version}")
|
|
if unknown_template_keys := set(spec) - {"version", "defaults", "fields", "start_anchor", "start_anchor_pattern"}:
|
|
raise ValueError(f"Unknown keys in response_template: {sorted(unknown_template_keys)}")
|
|
if not isinstance(spec.get("defaults", {}), dict):
|
|
raise ValueError("response_template.defaults must be a dict")
|
|
fields_raw = spec.get("fields", {})
|
|
if not isinstance(fields_raw, dict) or not fields_raw:
|
|
raise ValueError("response_template.fields must be a non-empty dict")
|
|
|
|
|
|
def _build_field(name: str, field: dict) -> ResponseTemplateField:
|
|
"""Validate a single field spec and compile it into a `ResponseTemplateField`."""
|
|
|
|
_ALLOWED_FIELD_KEYS = {
|
|
"open",
|
|
"open_pattern",
|
|
"close",
|
|
"close_pattern",
|
|
"content",
|
|
"content_args",
|
|
"repeats",
|
|
"join",
|
|
"optional",
|
|
"transform",
|
|
"transform_each",
|
|
}
|
|
scope = f"Field '{name}'"
|
|
if not isinstance(field, dict):
|
|
raise ValueError(f"{scope} must be a dict")
|
|
if unknown_field_keys := set(field) - _ALLOWED_FIELD_KEYS:
|
|
raise ValueError(f"{scope}: unknown keys {sorted(unknown_field_keys)}")
|
|
content = field.get("content", "text")
|
|
if content not in CONTENT_PARSERS:
|
|
raise ValueError(f"{scope}: unknown content parser '{content}'. Available: {sorted(CONTENT_PARSERS)}")
|
|
open_re, open_literals, open_literal_can_extend = _compile_anchor(scope, field, "open", "open_pattern")
|
|
close_re, close_literals, close_literal_can_extend = _compile_anchor(scope, field, "close", "close_pattern")
|
|
join = field.get("join")
|
|
if join is not None and not isinstance(join, str):
|
|
raise ValueError(f"{scope}: 'join' must be a string, got {type(join).__name__}")
|
|
if join is not None and not field.get("repeats", False):
|
|
raise ValueError(f"{scope}: 'join' requires 'repeats': true")
|
|
transform = field.get("transform")
|
|
transform_each = field.get("transform_each", False)
|
|
if not isinstance(transform_each, bool):
|
|
raise ValueError(f"{scope}: transform_each must be a bool, got {type(transform_each).__name__}")
|
|
if transform_each and transform is None:
|
|
raise ValueError(f"{scope}: transform_each is set but no transform was provided")
|
|
if transform is not None:
|
|
validate_transform_strings(scope, transform)
|
|
else:
|
|
# Named captures only reach the output through a transform, so flag any that would be silently dropped.
|
|
captured_names = set()
|
|
if open_re is not None:
|
|
captured_names |= set(open_re.groupindex)
|
|
if close_re is not None:
|
|
captured_names |= set(close_re.groupindex)
|
|
if captured_names:
|
|
raise ValueError(
|
|
f"{scope}: open_pattern/close_pattern declares named group(s) "
|
|
f"{sorted(captured_names)}, but the field has no 'transform'. Named captures "
|
|
f"are only surfaced through a 'transform' template (where they appear "
|
|
f"alongside 'content'). Either add a 'transform' that uses the captures, or "
|
|
f"remove the named groups from the pattern."
|
|
)
|
|
return ResponseTemplateField(
|
|
name=name,
|
|
open_re=open_re,
|
|
open_literals=open_literals,
|
|
open_literal_can_extend=open_literal_can_extend,
|
|
close_re=close_re,
|
|
close_literals=close_literals,
|
|
close_literal_can_extend=close_literal_can_extend,
|
|
content=content,
|
|
content_args=field.get("content_args", {}),
|
|
repeats=field.get("repeats", False),
|
|
join=join,
|
|
optional=field.get("optional", True),
|
|
transform=transform,
|
|
transform_each=transform_each,
|
|
)
|
|
|
|
|
|
def load_response_template(spec: dict | ResponseTemplate) -> ResponseTemplate:
|
|
if isinstance(spec, ResponseTemplate):
|
|
return spec
|
|
_validate_template_shape(spec)
|
|
|
|
fields = {name: _build_field(name, raw) for name, raw in spec["fields"].items()}
|
|
# A field without an open anchor is the implicit-open / leftover sink; at most one is allowed.
|
|
implicit_fields = [name for name, field in fields.items() if field.open_re is None]
|
|
if len(implicit_fields) > 1:
|
|
raise ValueError(
|
|
"At most one field may omit 'open'/'open_pattern' (that field becomes the "
|
|
f"implicit-open / leftover sink). Found: {', '.join(implicit_fields)}"
|
|
)
|
|
|
|
start_anchor_re, start_anchor_literals, _ = _compile_anchor(
|
|
"response_template", spec, "start_anchor", "start_anchor_pattern"
|
|
)
|
|
if start_anchor_re is None:
|
|
raise ValueError("response_template must define 'start_anchor' or 'start_anchor_pattern'.")
|
|
|
|
return ResponseTemplate(
|
|
defaults=copy.deepcopy(spec.get("defaults", {})),
|
|
fields=fields,
|
|
implicit=implicit_fields[0] if implicit_fields else None,
|
|
start_anchor_re=start_anchor_re,
|
|
start_anchor_literals=start_anchor_literals,
|
|
)
|
|
|
|
|
|
__all__ = ["ResponseTemplate", "ResponseTemplateField", "load_response_template"]
|
|
# fmt: on
|