biondizzle/vllm-with-lmcache
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: biondizzle:glm51-cu130
Branches Tags
biondizzle:master biondizzle:dream-build-glm biondizzle:dream-build biondizzle:rocm-mla-multi-group-kv-cache biondizzle:dream-build-glm-save-decode biondizzle:mla-multi-group-kv-cache biondizzle:glm51-cu130-nightly biondizzle:openai-rocm-nightly biondizzle:deepseekv4-cu130 biondizzle:custom-weights-cu130 biondizzle:glm51-cu130
..
compare: biondizzle:mla-multi-group-kv-cache
Branches Tags
biondizzle:dream-build-glm biondizzle:dream-build biondizzle:rocm-mla-multi-group-kv-cache biondizzle:dream-build-glm-save-decode biondizzle:mla-multi-group-kv-cache biondizzle:glm51-cu130-nightly biondizzle:openai-rocm-nightly biondizzle:deepseekv4-cu130 biondizzle:custom-weights-cu130 biondizzle:master biondizzle:glm51-cu130

2 Commits
glm51-cu13 ... mla-multi-

Author SHA1 Message Date
Matthew Harris
e52c6f0a98 patch to fix MLA multi-group KV cache 2026-05-05 21:10:05 -04:00
biondizzle
82330a31b1 add await to redis write 2026-04-17 07:19:15 +00:00
4 changed files with 6 additions and 1713 deletions
Expand all files Collapse all files

19
Dockerfile
View File

@@ -1,6 +1,8 @@
#FROM vllm/vllm-openai:v0.19.0-cu130
#FROM vllm/vllm-openai:cu130-nightly-x86_64
FROM vllm/vllm-openai:glm51-cu130
FROM vllm/vllm-openai:cu130-nightly-x86_64
# Fix the broken ass nightly build that forgot to include pandas
RUN pip install --no-cache-dir pandas
# Install LMCache for KV cache offloading / sharing across nodes
# Build with system CUDA 13.0 for Blackwell (B200)
@@ -11,9 +13,9 @@ RUN apt-get update && apt-get install -y git \
libcurand-dev-13-0 \
libcufft-dev-13-0 \
libnvjitlink-dev-13-0 && \
git clone https://github.com/biondizzle/LMCache.git /tmp/lmcache && \
git clone https://github.com/Byteflux/LMCache.git /tmp/lmcache && \
cd /tmp/lmcache && \
git checkout feat/redis-ttl && \
git checkout mla-multi-group-kv-cache-with-redis && \
CUDA_HOME=/usr/local/cuda \
TORCH_CUDA_ARCH_LIST="10.0" \
pip install --no-cache-dir --no-build-isolation . && \
@@ -31,12 +33,3 @@ COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_pa
# Copy over minimax parsers with kwargs fixes
COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py
COPY minimax_m2_parser.py /usr/local/lib/python3.12/dist-packages/vllm/parser/minimax_m2_parser.py
# Patch tool parser for GLM regex fix
COPY glm4_moe_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/glm4_moe_tool_parser.py
COPY utils.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/utils.py
# Patch hf renderer to force string content format for GLM models
# This fixes the issue where tool response content is dropped
COPY hf.py /usr/local/lib/python3.12/dist-packages/vllm/renderers/hf.py

491
glm4_moe_tool_parser.py
View File

@@ -1,491 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
GLM-4/5 Tool Call Parser — fixed version.
Fixes applied over the upstream vLLM + sweetapi patch:
1. **func_detail_regex no longer requires a newline** between tool name and
first <arg_key>. The model's chat template instructs:
<tool_call>{name}<arg_key>…</arg_key><arg_value>…</arg_value>…</tool_call>
with NO mandatory newline, but the original regex used ``[^\\n]*\\n`` which
silently failed when the model omitted it.
2. **Zero-argument tool calls no longer crash** (TypeError on NoneType).
3. **extract_tool_calls uses the same robust extraction helpers** as the
streaming path, so both paths parse identically.
4. **_extract_tool_name_from_region** is more tolerant of whitespace /
formatting variants the model may produce.
Drop this file into your vLLM install as a --tool-parser-plugin, or replace
the built-in glm4_moe_tool_parser.py.
"""
import ast
import json
from collections.abc import Sequence
from typing import Any
import regex as re
from vllm.entrypoints.chat_utils import make_tool_call_id
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.logger import init_logger
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import (
Tool,
ToolParser,
)
from vllm.tool_parsers.utils import partial_tag_overlap
logger = init_logger(__name__)
class Glm4MoeModelToolParser(ToolParser):
"""Tool parser for GLM-4/5 models with incremental string streaming.
On every streaming call the parser re-parses ``current_text`` to find
``<tool_call>`` regions, builds the JSON arguments string for each tool
call, and diffs against what was previously sent to emit only new content.
"""
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
super().__init__(tokenizer, tools)
# Stateful streaming fields
self.current_tool_name_sent: bool = False
self.prev_tool_call_arr: list[dict[str, Any]] = []
self.current_tool_id: int = -1
self.streamed_args_for_tool: list[str] = []
self.tool_call_start_token: str = "<tool_call>"
self.tool_call_end_token: str = "</tool_call>"
self.arg_key_start: str = "<arg_key>"
self.arg_key_end: str = "</arg_key>"
self.arg_val_start: str = "<arg_value>"
self.arg_val_end: str = "</arg_value>"
self.tool_calls_start_token = self.tool_call_start_token
# ---- FIXED regexes ------------------------------------------------
# Match the whole <tool_call>…</tool_call> block (unchanged).
self.func_call_regex = re.compile(
r"<tool_call>.*?</tool_call>", re.DOTALL
)
# FIX 1: The original regex required a literal \n between tool name
# and the body. The model often omits it. We now accept any
# whitespace (including none) before the first <arg_key>, and we
# make the body group optional so zero-argument calls don't fail.
self.func_detail_regex = re.compile(
r"<tool_call>\s*" # opening tag + optional whitespace
r"([\w.\-]+)" # group 1: tool/function name (word chars, dots, hyphens)
r"\s*" # optional whitespace / newline
r"((?:<arg_key>.*)?)" # group 2: everything from first <arg_key> onward (may be empty)
r"\s*</tool_call>", # closing tag
re.DOTALL,
)
self.func_arg_regex = re.compile(
r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>", re.DOTALL
)
if not self.model_tokenizer:
raise ValueError(
"The model tokenizer must be passed to the ToolParser "
"constructor during construction."
)
self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
# Pre-compiled pattern for finding the last <arg_key>...</arg_key>
# before a partial <arg_value> (used in _build_args_json_so_far).
self._arg_key_pattern = re.compile(
re.escape(self.arg_key_start) + r"(.*?)" + re.escape(self.arg_key_end),
re.DOTALL,
)
# Streaming state for re-parse-and-diff approach
self._sent_content_idx: int = 0
self._tool_call_ids: list[str] = []
# ------------------------------------------------------------------
# Static helpers
# ------------------------------------------------------------------
@staticmethod
def _deserialize(value: str) -> Any:
try:
return json.loads(value)
except json.JSONDecodeError:
pass
try:
return ast.literal_eval(value)
except (ValueError, SyntaxError):
pass
return value
@staticmethod
def _json_escape_string_content(s: str) -> str:
"""JSON-escape string content (without surrounding quotes)."""
if not s:
return ""
return json.dumps(s, ensure_ascii=False)[1:-1]
@staticmethod
def _is_string_type(
tool_name: str,
arg_name: str,
tools: list[Tool] | None,
) -> bool:
if tools is None:
return False
for tool in tools:
if tool.function.name != tool_name:
continue
if tool.function.parameters is None:
return False
arg_type = (
tool.function.parameters.get("properties", {})
.get(arg_name, {})
.get("type", None)
)
return arg_type == "string"
logger.debug("No tool named '%s'.", tool_name)
return False
@staticmethod
def _tools_enabled(request: ChatCompletionRequest) -> bool:
try:
tools = getattr(request, "tools", None)
tool_choice = getattr(request, "tool_choice", None)
return bool(tools) and tool_choice != "none"
except Exception:
logger.exception("Failed to determine if tools are enabled.")
return False
# ------------------------------------------------------------------
# Request adjustment
# ------------------------------------------------------------------
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
request = super().adjust_request(request)
if request.tools and request.tool_choice != "none":
request.skip_special_tokens = False
return request
# ------------------------------------------------------------------
# Non-streaming extraction
# ------------------------------------------------------------------
def extract_tool_calls(
self,
model_output: str,
request: ChatCompletionRequest,
) -> ExtractedToolCallInformation:
matched_tool_calls = self.func_call_regex.findall(model_output)
logger.debug("model_output: %s", model_output)
try:
tool_calls: list[ToolCall] = []
for match in matched_tool_calls:
tc_detail = self.func_detail_regex.search(match)
if not tc_detail:
logger.warning(
"Failed to parse tool call details from: %s", match
)
continue
tc_name = tc_detail.group(1).strip()
tc_args_raw = tc_detail.group(2) or "" # FIX 2: default to ""
pairs = self.func_arg_regex.findall(tc_args_raw) if tc_args_raw else []
arg_dct: dict[str, Any] = {}
for key, value in pairs:
arg_key = key.strip()
arg_val = value.strip()
if not self._is_string_type(tc_name, arg_key, self.tools):
arg_val = self._deserialize(arg_val)
logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val)
arg_dct[arg_key] = arg_val
tool_calls.append(
ToolCall(
type="function",
function=FunctionCall(
name=tc_name,
arguments=json.dumps(arg_dct, ensure_ascii=False),
),
)
)
except Exception:
logger.exception("Failed to extract tool call spec")
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
if tool_calls:
content: str | None = model_output[
: model_output.find(self.tool_calls_start_token)
]
if not content or not content.strip():
content = None
return ExtractedToolCallInformation(
tools_called=True, tool_calls=tool_calls, content=content
)
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
# ------------------------------------------------------------------
# Streaming helpers
# ------------------------------------------------------------------
def _extract_content(self, current_text: str) -> str | None:
content_segments: list[str] = []
pos = self._sent_content_idx
while pos < len(current_text):
start = current_text.find(self.tool_call_start_token, pos)
if start == -1:
tail = current_text[pos:]
overlap = partial_tag_overlap(tail, self.tool_call_start_token)
sendable = tail[: len(tail) - overlap] if overlap else tail
if sendable:
content_segments.append(sendable)
pos = len(current_text) - overlap
break
if start > pos:
content_segments.append(current_text[pos:start])
end = current_text.find(self.tool_call_end_token, start)
if end != -1:
pos = end + len(self.tool_call_end_token)
else:
pos = start
break
if content_segments:
self._sent_content_idx = pos
return "".join(content_segments)
if pos > self._sent_content_idx:
self._sent_content_idx = pos
return None
def _extract_tool_call_regions(self, text: str) -> list[tuple[str, bool]]:
results: list[tuple[str, bool]] = []
pos = 0
while True:
start = text.find(self.tool_call_start_token, pos)
if start == -1:
break
inner_start = start + len(self.tool_call_start_token)
end = text.find(self.tool_call_end_token, inner_start)
if end != -1:
results.append((text[inner_start:end], True))
pos = end + len(self.tool_call_end_token)
else:
raw = text[inner_start:]
overlap = partial_tag_overlap(raw, self.tool_call_end_token)
if overlap:
raw = raw[:-overlap]
results.append((raw, False))
break
return results
def _extract_tool_name_from_region(self, inner_text: str) -> str | None:
"""Extract the tool name from the beginning of a tool-call region.
The name is everything before the first ``\\n``, ``<arg_key>``, or
``</tool_call>``. We also accept the name being the only content
(for zero-argument calls that are still in-flight).
"""
# Strip leading whitespace — model may emit \n after <tool_call>
stripped = inner_text.lstrip()
if not stripped:
return None
nl = stripped.find("\n")
ak = stripped.find(self.arg_key_start)
candidates = [i for i in [nl, ak] if i != -1]
if not candidates:
# No delimiter yet — if the text looks like a partial name
# (only word chars / dots / hyphens), return None to wait.
# If it's a complete name with no args (zero-arg call, complete),
# it will be handled when is_complete is True.
candidate_name = stripped.strip()
if re.fullmatch(r'[\w.\-]+', candidate_name):
# Could be a complete name or still arriving — return it
# so zero-arg complete calls work; the caller checks is_complete.
return candidate_name
return None
cut = min(candidates)
name = stripped[:cut].strip()
return name if name else None
def _build_args_json_so_far(
self,
tool_name: str,
inner_text: str,
is_complete: bool,
) -> str:
pairs = self.func_arg_regex.findall(inner_text)
parts: list[str] = []
for key, value in pairs:
key = key.strip()
key_json = json.dumps(key, ensure_ascii=False)
if self._is_string_type(tool_name, key, self.tools):
val_json = json.dumps(value, ensure_ascii=False)
else:
val_json = json.dumps(
self._deserialize(value.strip()), ensure_ascii=False
)
parts.append(f"{key_json}: {val_json}")
# Check for a partial (incomplete) arg value
last_val_start = inner_text.rfind(self.arg_val_start)
last_val_end = inner_text.rfind(self.arg_val_end)
has_partial_value = last_val_start != -1 and (
last_val_end == -1 or last_val_end < last_val_start
)
if has_partial_value:
last_key_match = None
for m in self._arg_key_pattern.finditer(inner_text[:last_val_start]):
last_key_match = m
if last_key_match:
partial_key = last_key_match.group(1).strip()
partial_content_start = last_val_start + len(self.arg_val_start)
partial_content = inner_text[partial_content_start:]
overlap = partial_tag_overlap(partial_content, self.arg_val_end)
if overlap:
partial_content = partial_content[:-overlap]
key_json = json.dumps(partial_key, ensure_ascii=False)
if is_complete:
if self._is_string_type(tool_name, partial_key, self.tools):
val_json = json.dumps(partial_content, ensure_ascii=False)
else:
val_json = json.dumps(
self._deserialize(partial_content.strip()),
ensure_ascii=False,
)
parts.append(f"{key_json}: {val_json}")
elif self._is_string_type(tool_name, partial_key, self.tools):
escaped = self._json_escape_string_content(partial_content)
parts.append(f'{key_json}: "{escaped}')
else:
parts.append(f"{key_json}: {partial_content}")
if not parts:
return "{}" if is_complete else ""
joined = "{" + ", ".join(parts)
if is_complete:
joined += "}"
return joined
def _compute_args_diff(self, index: int, args_so_far: str) -> str | None:
if not args_so_far or len(args_so_far) <= len(
self.streamed_args_for_tool[index]
):
return None
diff = args_so_far[len(self.streamed_args_for_tool[index]) :]
self.streamed_args_for_tool[index] = args_so_far
self.prev_tool_call_arr[index]["arguments"] = args_so_far
return diff
def _ensure_tool_state_for(self, index: int) -> None:
while len(self._tool_call_ids) <= index:
self._tool_call_ids.append(
make_tool_call_id(id_type="random", func_name=None, idx=None)
)
while len(self.streamed_args_for_tool) <= index:
self.streamed_args_for_tool.append("")
while len(self.prev_tool_call_arr) <= index:
self.prev_tool_call_arr.append({})
# ------------------------------------------------------------------
# Main streaming entry point
# ------------------------------------------------------------------
def extract_tool_calls_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
request: ChatCompletionRequest,
) -> DeltaMessage | None:
if not self._tools_enabled(request):
return DeltaMessage(content=delta_text) if delta_text else None
content = self._extract_content(current_text)
regions = self._extract_tool_call_regions(current_text)
tool_call_deltas: list[DeltaToolCall] = []
for i, (inner_text, is_complete) in enumerate(regions):
self._ensure_tool_state_for(i)
tool_name = self._extract_tool_name_from_region(inner_text)
if not tool_name:
break
# Emit tool name (once per tool call)
if "name" not in self.prev_tool_call_arr[i]:
self.prev_tool_call_arr[i]["name"] = tool_name
tool_call_deltas.append(
DeltaToolCall(
index=i,
id=self._tool_call_ids[i],
type="function",
function=DeltaFunctionCall(
name=tool_name,
arguments="",
).model_dump(exclude_none=True),
)
)
# Build args JSON so far, diff, emit
args_so_far = self._build_args_json_so_far(
tool_name, inner_text, is_complete
)
diff = self._compute_args_diff(i, args_so_far)
if diff:
tool_call_deltas.append(
DeltaToolCall(
index=i,
function=DeltaFunctionCall(arguments=diff).model_dump(
exclude_none=True
),
)
)
if regions:
self.current_tool_id = len(regions) - 1
if content or tool_call_deltas:
return DeltaMessage(
content=content,
tool_calls=tool_call_deltas,
)
return None

771
hf.py
View File

@@ -1,771 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import inspect
import itertools
from collections import defaultdict, deque
from collections.abc import Set
from functools import lru_cache
from typing import Any, Literal, cast, overload
import jinja2
import jinja2.ext
import jinja2.meta
import jinja2.nodes
import jinja2.parser
import jinja2.sandbox
from vllm.config import ModelConfig, VllmConfig
from vllm.entrypoints.chat_utils import (
ChatCompletionMessageParam,
ChatTemplateContentFormat,
ChatTemplateContentFormatOption,
ChatTemplateResolutionError,
ConversationMessage,
load_chat_template,
parse_chat_messages,
parse_chat_messages_async,
)
from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict
from vllm.logger import init_logger
from vllm.tokenizers.hf import HfTokenizer
from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
from vllm.transformers_utils.processor import cached_get_processor
from vllm.utils.async_utils import make_async
from vllm.utils.func_utils import supports_kw
from .base import BaseRenderer
from .inputs import DictPrompt
from .inputs.preprocess import parse_dec_only_prompt
from .params import ChatParams
logger = init_logger(__name__)
_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()
"""
Used in `_try_get_processor_chat_template` to avoid calling
`cached_get_processor` again if the processor fails to be loaded.
This is needed because `lru_cache` does not cache when an exception happens.
"""
def _try_get_processor_chat_template(
tokenizer: HfTokenizer,
*,
trust_remote_code: bool,
) -> str | None:
cache_key = (tokenizer.name_or_path, trust_remote_code)
if cache_key in _PROCESSOR_CHAT_TEMPLATES:
return _PROCESSOR_CHAT_TEMPLATES[cache_key]
from transformers import (
PreTrainedTokenizer,
PreTrainedTokenizerFast,
ProcessorMixin,
)
try:
processor = cached_get_processor(
tokenizer.name_or_path,
processor_cls=(
PreTrainedTokenizer,
PreTrainedTokenizerFast,
ProcessorMixin,
),
trust_remote_code=trust_remote_code,
)
if (
isinstance(processor, ProcessorMixin)
and hasattr(processor, "chat_template")
and (chat_template := processor.chat_template) is not None
):
_PROCESSOR_CHAT_TEMPLATES[cache_key] = chat_template
return chat_template
except Exception:
logger.debug(
"Failed to load AutoProcessor chat template for %s",
tokenizer.name_or_path,
exc_info=True,
)
_PROCESSOR_CHAT_TEMPLATES[cache_key] = None
return None
def resolve_chat_template(
tokenizer: HfTokenizer,
chat_template: str | None,
tools: list[dict[str, Any]] | None,
*,
model_config: "ModelConfig",
) -> str | None:
# 1st priority: The given chat template
if chat_template is not None:
# Resolve template names (e.g. "tool_use") to actual Jinja content
# so that downstream kwargs detection can parse template variables.
return tokenizer.get_chat_template(chat_template, tools=tools)
# 2nd priority: AutoProcessor chat template, unless tool calling is enabled
if tools is None:
chat_template = _try_get_processor_chat_template(
tokenizer,
trust_remote_code=model_config.trust_remote_code,
)
if chat_template is not None:
return chat_template
# 3rd priority: AutoTokenizer chat template
try:
return tokenizer.get_chat_template(chat_template, tools=tools)
except Exception:
logger.debug(
"Failed to load AutoTokenizer chat template for %s",
tokenizer.name_or_path,
exc_info=True,
)
# 4th priority: Predefined fallbacks
path = get_chat_template_fallback_path(
model_type=model_config.hf_config.model_type,
tokenizer_name_or_path=tokenizer.name_or_path,
)
if path is not None:
logger.info_once(
"Loading chat template fallback for %s as there isn't one "
"defined on HF Hub.",
tokenizer.name_or_path,
)
chat_template = load_chat_template(path)
else:
logger.debug_once(
"There is no chat template fallback for %s", tokenizer.name_or_path
)
return chat_template
def _is_var_access(node: jinja2.nodes.Node, varname: str) -> bool:
if isinstance(node, jinja2.nodes.Name):
return node.ctx == "load" and node.name == varname
return False
def _is_attr_access(node: jinja2.nodes.Node, varname: str, key: str) -> bool:
if isinstance(node, jinja2.nodes.Getitem):
return (
_is_var_access(node.node, varname)
and isinstance(node.arg, jinja2.nodes.Const)
and node.arg.value == key
)
if isinstance(node, jinja2.nodes.Getattr):
return _is_var_access(node.node, varname) and node.attr == key
return False
def _is_var_or_elems_access(
node: jinja2.nodes.Node,
varname: str,
key: str | None = None,
) -> bool:
if isinstance(node, jinja2.nodes.Filter):
return node.node is not None and _is_var_or_elems_access(
node.node, varname, key
)
if isinstance(node, jinja2.nodes.Test):
return _is_var_or_elems_access(node.node, varname, key)
if isinstance(node, jinja2.nodes.Getitem) and isinstance(
node.arg, jinja2.nodes.Slice
):
return _is_var_or_elems_access(node.node, varname, key)
return _is_attr_access(node, varname, key) if key else _is_var_access(node, varname)
def _iter_nodes_assign_var_or_elems(root: jinja2.nodes.Node, varname: str):
# Global variable that is implicitly defined at the root
yield root, varname
# Iterative BFS
related_varnames = deque([varname])
while related_varnames:
related_varname = related_varnames.popleft()
for assign_ast in root.find_all(jinja2.nodes.Assign):
lhs = assign_ast.target
rhs = assign_ast.node
if _is_var_or_elems_access(rhs, related_varname):
assert isinstance(lhs, jinja2.nodes.Name)
yield assign_ast, lhs.name
# Avoid infinite looping for self-assignment
if lhs.name != related_varname:
related_varnames.append(lhs.name)
# NOTE: The proper way to handle this is to build a CFG so that we can handle
# the scope in which each variable is defined, but that is too complicated
def _iter_nodes_assign_messages_item(root: jinja2.nodes.Node):
messages_varnames = [
varname for _, varname in _iter_nodes_assign_var_or_elems(root, "messages")
]
# Search for {%- for message in messages -%} loops
for loop_ast in root.find_all(jinja2.nodes.For):
loop_iter = loop_ast.iter
loop_target = loop_ast.target
for varname in messages_varnames:
if _is_var_or_elems_access(loop_iter, varname):
assert isinstance(loop_target, jinja2.nodes.Name)
yield loop_ast, loop_target.name
break
def _iter_nodes_assign_content_item(root: jinja2.nodes.Node):
message_varnames = [
varname for _, varname in _iter_nodes_assign_messages_item(root)
]
# Search for {%- for content in message['content'] -%} loops
for loop_ast in root.find_all(jinja2.nodes.For):
loop_iter = loop_ast.iter
loop_target = loop_ast.target
for varname in message_varnames:
if _is_var_or_elems_access(loop_iter, varname, "content"):
assert isinstance(loop_target, jinja2.nodes.Name)
yield loop_ast, loop_target.name
break
def _try_extract_ast(chat_template: str) -> jinja2.nodes.Template | None:
import transformers.utils.chat_template_utils as hf_chat_utils
try:
jinja_compiled = hf_chat_utils._compile_jinja_template(chat_template)
return jinja_compiled.environment.parse(chat_template)
except Exception:
logger.exception("Error when compiling Jinja template")
return None
@lru_cache(maxsize=32)
def _detect_content_format(
chat_template: str,
*,
default: ChatTemplateContentFormat,
) -> ChatTemplateContentFormat:
jinja_ast = _try_extract_ast(chat_template)
if jinja_ast is None:
return default
try:
next(_iter_nodes_assign_content_item(jinja_ast))
except StopIteration:
return "string"
except Exception:
logger.exception("Error when parsing AST of Jinja template")
return default
else:
return "openai"
def _is_glm_model(tokenizer: HfTokenizer, model_config: "ModelConfig") -> bool:
"""Check if this is a GLM model that requires string content format.
GLM models (GLM-4, GLM-4.5, GLM-5.x) have a chat template that incorrectly
triggers "openai" content format detection because they iterate over
m.content for tool responses. However, the template expects string content
for tool messages (checking `m.content is string`).
This detection ensures we force "string" format for GLM models.
"""
# Check tokenizer name/path for GLM indicators
name_or_path = tokenizer.name_or_path.lower()
glm_indicators = ["glm-4", "glm-5", "glm4", "glm5", "zai-org/glm"]
if any(ind in name_or_path for ind in glm_indicators):
return True
# Check model type in config
if hasattr(model_config, "hf_config") and hasattr(model_config.hf_config, "model_type"):
model_type = model_config.hf_config.model_type.lower()
if "glm" in model_type:
return True
return False
def _resolve_chat_template_content_format(
chat_template: str | None,
tools: list[dict[str, Any]] | None,
tokenizer: HfTokenizer,
*,
model_config: "ModelConfig",
) -> ChatTemplateContentFormat:
# GLM models require "string" content format for tool responses to work
# The template has `{% for tr in m.content %}` which triggers "openai"
# detection, but then checks `m.content is string` which fails for arrays.
if _is_glm_model(tokenizer, model_config):
logger.debug(
"Forcing 'string' content format for GLM model: %s",
tokenizer.name_or_path,
)
return "string"
resolved_chat_template = resolve_chat_template(
tokenizer,
chat_template=chat_template,
tools=tools,
model_config=model_config,
)
jinja_text = (
resolved_chat_template
if isinstance(resolved_chat_template, str)
else load_chat_template(chat_template, is_literal=True)
)
detected_format = (
"string"
if jinja_text is None
else _detect_content_format(jinja_text, default="string")
)
return detected_format
@lru_cache
def _log_chat_template_content_format(
chat_template: str | None, # For caching purposes
given_format: ChatTemplateContentFormatOption,
detected_format: ChatTemplateContentFormatOption,
):
logger.info(
"Detected the chat template content format to be '%s'. "
"You can set `--chat-template-content-format` to override this.",
detected_format,
)
if given_format != "auto" and given_format != detected_format:
logger.warning(
"You specified `--chat-template-content-format %s` "
"which is different from the detected format '%s'. "
"If our automatic detection is incorrect, please consider "
"opening a GitHub issue so that we can improve it: "
"https://github.com/vllm-project/vllm/issues/new/choose",
given_format,
detected_format,
)
def resolve_chat_template_content_format(
chat_template: str | None,
tools: list[dict[str, Any]] | None,
given_format: ChatTemplateContentFormatOption,
tokenizer: HfTokenizer,
*,
model_config: "ModelConfig",
) -> ChatTemplateContentFormat:
if given_format != "auto":
return given_format
detected_format = _resolve_chat_template_content_format(
chat_template,
tools,
tokenizer,
model_config=model_config,
)
_log_chat_template_content_format(
chat_template,
given_format=given_format,
detected_format=detected_format,
)
return detected_format
# adapted from https://github.com/huggingface/transformers/blob/v4.56.2/src/transformers/utils/chat_template_utils.py#L398-L412
# only preserve the parse function used to resolve chat template kwargs
class AssistantTracker(jinja2.ext.Extension):
tags = {"generation"}
def parse(self, parser: jinja2.parser.Parser) -> jinja2.nodes.Node:
lineno = next(parser.stream).lineno
body = parser.parse_statements(("name:endgeneration",), drop_needle=True)
call = self.call_method("_generation_support")
call_block = jinja2.nodes.CallBlock(call, [], [], body)
return call_block.set_lineno(lineno)
def _resolve_chat_template_kwargs(chat_template: str) -> Set[str]:
env = jinja2.sandbox.ImmutableSandboxedEnvironment(
trim_blocks=True,
lstrip_blocks=True,
extensions=[AssistantTracker, jinja2.ext.loopcontrols],
)
parsed_content = env.parse(chat_template)
template_vars = jinja2.meta.find_undeclared_variables(parsed_content)
return template_vars
_cached_resolve_chat_template_kwargs = lru_cache(_resolve_chat_template_kwargs)
@lru_cache
def _get_hf_base_chat_template_params() -> frozenset[str]:
from transformers import PreTrainedTokenizer
# Get standard parameters from HuggingFace's base tokenizer class.
# This dynamically extracts parameters from PreTrainedTokenizer's
# apply_chat_template method, ensuring compatibility with tokenizers
# that use **kwargs to receive standard parameters.
# Read signature from HF's base class - the single source of truth
base_sig = inspect.signature(PreTrainedTokenizer.apply_chat_template)
# Exclude VAR_KEYWORD (**kwargs) and VAR_POSITIONAL (*args) placeholders
return frozenset(
p.name
for p in base_sig.parameters.values()
if p.kind
not in (inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL)
)
def resolve_chat_template_kwargs(
tokenizer: HfTokenizer,
chat_template: str,
chat_template_kwargs: dict[str, Any],
raise_on_unexpected: bool = True,
) -> dict[str, Any]:
# We exclude chat_template from kwargs here, because
# chat template has been already resolved at this stage
unexpected_vars = {"chat_template", "tokenize"}
if raise_on_unexpected and (
unexpected_in_kwargs := unexpected_vars & chat_template_kwargs.keys()
):
raise ValueError(
"Found unexpected chat template kwargs from request: "
f"{unexpected_in_kwargs}"
)
fn_kw = {
k
for k in chat_template_kwargs
if supports_kw(tokenizer.apply_chat_template, k, allow_var_kwargs=False)
}
template_vars = _cached_resolve_chat_template_kwargs(chat_template)
# Allow standard HF parameters even if tokenizer uses **kwargs to receive them
hf_base_params = _get_hf_base_chat_template_params()
accept_vars = (fn_kw | template_vars | hf_base_params) - unexpected_vars
return {k: v for k, v in chat_template_kwargs.items() if k in accept_vars}
@overload
def safe_apply_chat_template(
model_config: "ModelConfig",
tokenizer: HfTokenizer,
conversation: list[ConversationMessage],
*,
tools: list[dict[str, Any]] | None = ...,
chat_template: str | None = ...,
tokenize: Literal[True] = ...,
**kwargs,
) -> list[int]: ...
@overload
def safe_apply_chat_template(
model_config: "ModelConfig",
tokenizer: HfTokenizer,
conversation: list[ConversationMessage],
*,
tools: list[dict[str, Any]] | None = ...,
chat_template: str | None = ...,
tokenize: Literal[False] = ...,
**kwargs,
) -> str: ...
def safe_apply_chat_template(
model_config: "ModelConfig",
tokenizer: HfTokenizer,
conversation: list[ConversationMessage],
*,
tools: list[dict[str, Any]] | None = None,
chat_template: str | None = None,
tokenize: bool = True,
**kwargs,
) -> str | list[int]:
chat_template = resolve_chat_template(
tokenizer,
chat_template=chat_template,
tools=tools,
model_config=model_config,
)
if chat_template is None:
raise ChatTemplateResolutionError(
"As of transformers v4.44, default chat template is no longer "
"allowed, so you must provide a chat template if the tokenizer "
"does not define one."
)
resolved_kwargs = resolve_chat_template_kwargs(
tokenizer=tokenizer,
chat_template=chat_template,
chat_template_kwargs=kwargs,
)
try:
return tokenizer.apply_chat_template(
conversation=conversation, # type: ignore[arg-type]
tools=tools, # type: ignore[arg-type]
chat_template=chat_template,
tokenize=tokenize,
**resolved_kwargs,
)
# External library exceptions can sometimes occur despite the framework's
# internal exception management capabilities.
except Exception as e:
# Log and report any library-related exceptions for further
# investigation.
logger.exception(
"An error occurred in `transformers` while applying chat template"
)
raise ValueError(str(e)) from e
def rebuild_mm_uuids_from_mm_data(
mm_uuids: MultiModalUUIDDict,
mm_data: MultiModalDataDict,
) -> MultiModalUUIDDict:
"""Rebuild mm_uuids after vision_chunk processing.
When videos are split into chunks, the original UUIDs need to be updated
to reflect the new UUIDs generated for each chunk.
Args:
mm_uuids: Original UUIDs dictionary
mm_data: Processed multimodal data with vision_chunk items
Returns:
Updated UUIDs dictionary with chunk UUIDs
"""
vision_chunks = mm_data.get("vision_chunk")
if vision_chunks is None:
return mm_uuids
assert all(isinstance(item, dict) for item in vision_chunks), (
"Expected all vision_chunk items to be dicts"
)
vision_chunks = cast(list[dict[str, Any]], vision_chunks)
vision_chunk_uuids = [
uuid_val for item in vision_chunks if (uuid_val := item.get("uuid")) is not None
]
if vision_chunk_uuids:
mm_uuids = dict(mm_uuids)
mm_uuids["vision_chunk"] = vision_chunk_uuids
return mm_uuids
def build_video_prompts_from_mm_data(
mm_data: MultiModalDataDict,
) -> list[str]:
"""Build video prompts from vision_chunk data.
Collects prompts from video chunks and groups them by video_idx.
Args:
mm_data: Processed multimodal data with vision_chunk items
Returns:
List of video prompts, one per video.
"""
vision_chunks = mm_data.get("vision_chunk")
if vision_chunks is None:
return []
# Group chunks by video_idx
video_prompts_dict: dict[int, list[str]] = defaultdict(list)
for item in vision_chunks:
# vision_chunk items are always dicts (VisionChunkImage/VisionChunkVideo)
assert isinstance(item, dict)
if item.get("type") == "video_chunk":
video_idx = item.get("video_idx", 0)
prompt = item.get("prompt", "")
video_prompts_dict[video_idx].append(prompt)
# Build prompts in video order
video_prompts = [
"".join(video_prompts_dict[video_idx])
for video_idx in sorted(video_prompts_dict.keys())
]
return video_prompts
def replace_vision_chunk_video_placeholder(
prompt_raw: str | list[int],
mm_data: MultiModalDataDict,
video_placeholder: str | None,
) -> str | list[int]:
# get video placeholder, replace it with runtime video-chunk prompts
if video_placeholder and isinstance(prompt_raw, str):
video_prompts = build_video_prompts_from_mm_data(mm_data)
# replace in order
prompt_raw_parts = prompt_raw.split(video_placeholder)
if len(prompt_raw_parts) == len(video_prompts) + 1:
prompt_raw = "".join(
itertools.chain.from_iterable(zip(prompt_raw_parts, video_prompts))
)
prompt_raw += prompt_raw_parts[-1]
else:
logger.warning(
"Number of video placeholders (%d) does not match "
"number of videos (%d) in the request.",
len(prompt_raw_parts) - 1,
len(video_prompts),
)
return prompt_raw
class HfRenderer(BaseRenderer[HfTokenizer]):
def __init__(
self,
config: VllmConfig,
tokenizer: HfTokenizer | None,
) -> None:
super().__init__(config, tokenizer)
self.use_unified_vision_chunk = getattr(
config.model_config.hf_config, "use_unified_vision_chunk", False
)
self._apply_chat_template_async = make_async(
safe_apply_chat_template, executor=self._executor
)
def render_messages(
self,
messages: list[ChatCompletionMessageParam],
params: ChatParams,
) -> tuple[list[ConversationMessage], DictPrompt]:
model_config = self.model_config
tokenizer = self.get_tokenizer()
conversation, mm_data, mm_uuids = parse_chat_messages(
messages,
model_config,
content_format=resolve_chat_template_content_format(
chat_template=params.chat_template,
tools=params.chat_template_kwargs.get("tools"),
given_format=params.chat_template_content_format,
tokenizer=tokenizer,
model_config=model_config,
),
media_io_kwargs=params.media_io_kwargs,
mm_processor_kwargs=params.mm_processor_kwargs,
)
prompt_raw = safe_apply_chat_template(
model_config,
tokenizer,
conversation,
**params.get_apply_chat_template_kwargs(),
)
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
# model which uses unified vision chunks for both images and videos.
if (
self.use_unified_vision_chunk
and mm_uuids is not None
and mm_data is not None
):
mm_uuids = rebuild_mm_uuids_from_mm_data(mm_uuids, mm_data)
# get video placeholder, replace it with runtime video-chunk prompts
video_placeholder = getattr(
model_config.hf_config, "video_placeholder", None
)
prompt_raw = cast(
list[int],
replace_vision_chunk_video_placeholder(
prompt_raw,
mm_data,
video_placeholder,
),
)
prompt = parse_dec_only_prompt(prompt_raw)
if mm_data is not None:
prompt["multi_modal_data"] = mm_data
if mm_uuids is not None:
prompt["multi_modal_uuids"] = mm_uuids
return conversation, prompt
async def render_messages_async(
self,
messages: list[ChatCompletionMessageParam],
params: ChatParams,
) -> tuple[list[ConversationMessage], DictPrompt]:
model_config = self.model_config
tokenizer = self.get_tokenizer()
conversation, mm_data, mm_uuids = await parse_chat_messages_async(
messages,
model_config,
content_format=resolve_chat_template_content_format(
chat_template=params.chat_template,
tools=params.chat_template_kwargs.get("tools"),
given_format=params.chat_template_content_format,
tokenizer=tokenizer,
model_config=model_config,
),
media_io_kwargs=params.media_io_kwargs,
mm_processor_kwargs=params.mm_processor_kwargs,
)
prompt_raw = await self._apply_chat_template_async(
model_config,
tokenizer,
conversation,
**params.get_apply_chat_template_kwargs(),
)
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
# model which uses unified vision chunks for both images and videos.
if (
self.use_unified_vision_chunk
and mm_uuids is not None
and mm_data is not None
):
# get video placeholder, replace it with runtime video-chunk prompts
video_placeholder = getattr(
model_config.hf_config, "video_placeholder", None
)
prompt_raw = cast(
list[int],
replace_vision_chunk_video_placeholder(
prompt_raw,
mm_data,
video_placeholder,
),
)
prompt = parse_dec_only_prompt(prompt_raw)
if mm_data is not None:
prompt["multi_modal_data"] = mm_data
if mm_uuids is not None:
prompt["multi_modal_uuids"] = mm_uuids
return conversation, prompt

438
utils.py
View File

@@ -1,438 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import ast
import json
from json import JSONDecodeError, JSONDecoder
from typing import Any, TypeAlias
import partial_json_parser
from openai.types.responses import (
FunctionTool,
ToolChoiceFunction,
)
from openai.types.responses.tool import Tool as ResponsesTool
from partial_json_parser.core.options import Allow
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionNamedToolChoiceParam,
ChatCompletionToolsParam,
)
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaToolCall,
FunctionCall,
ToolCall,
)
from vllm.logger import init_logger
Tool: TypeAlias = ChatCompletionToolsParam | ResponsesTool
logger = init_logger(__name__)
def partial_tag_overlap(text: str, tag: str) -> int:
"""Length of the longest prefix of *tag* that matches a suffix of *text*.
E.g. text ending in ``"<tool_"`` returns 6 when tag is ``"<tool_call>"``.
Returns 0 when there is no overlap.
"""
max_check = min(len(tag) - 1, len(text))
for k in range(max_check, 0, -1):
if text.endswith(tag[:k]):
return k
return 0
def find_common_prefix(s1: str, s2: str) -> str:
"""
Finds a common prefix that is shared between two strings, if there is one.
Order of arguments is NOT important.
This function is provided as a UTILITY for extracting information from JSON
generated by partial_json_parser, to help in ensuring that the right tokens
are returned in streaming, so that close-quotes, close-brackets and
close-braces are not returned prematurely.
e.g. find_common_prefix('{"fruit": "ap"}', '{"fruit": "apple"}') ->
'{"fruit": "ap'
"""
prefix = ""
min_length = min(len(s1), len(s2))
for i in range(0, min_length):
if s1[i] == s2[i]:
prefix += s1[i]
else:
break
return prefix
def find_common_suffix(s1: str, s2: str) -> str:
"""
Finds a common suffix shared between two strings, if there is one. Order of
arguments is NOT important.
Stops when the suffix ends OR it hits an alphanumeric character
e.g. find_common_suffix('{"fruit": "ap"}', '{"fruit": "apple"}') -> '"}'
"""
suffix = ""
min_length = min(len(s1), len(s2))
for i in range(1, min_length + 1):
if s1[-i] == s2[-i] and not s1[-i].isalnum():
suffix = s1[-i] + suffix
else:
break
return suffix
def extract_intermediate_diff(curr: str, old: str) -> str:
"""
Given two strings, extract the difference in the middle between two strings
that are known to have a common prefix and/or suffix.
This function is provided as a UTILITY for extracting information from JSON
generated by partial_json_parser, to help in ensuring that the right tokens
are returned in streaming, so that close-quotes, close-brackets and
close-braces are not returned prematurely. The order of arguments IS
important - the new version of the partially-parsed JSON must be the first
argument, and the secnod argument must be from the previous generation.
What it returns, is tokens that should be streamed to the client.
e.g. extract_intermediate_diff('{"fruit": "apple"}', '{"fruit": "ap"}')
-> 'ple'
"""
suffix = find_common_suffix(curr, old)
old = old[::-1].replace(suffix[::-1], "", 1)[::-1]
prefix = find_common_prefix(curr, old)
diff = curr
if len(suffix):
diff = diff[::-1].replace(suffix[::-1], "", 1)[::-1]
if len(prefix):
# replace the prefix only once in case it's mirrored
diff = diff.replace(prefix, "", 1)
return diff
# partial_json_parser doesn't support extra data and
# JSONDecoder.raw_decode doesn't support partial JSON
def partial_json_loads(input_str: str, flags: Allow) -> tuple[Any, int]:
try:
return (partial_json_parser.loads(input_str, flags), len(input_str))
except JSONDecodeError as e:
if "Extra data" in e.msg:
dec = JSONDecoder()
return dec.raw_decode(input_str)
raise
def is_complete_json(input_str: str) -> bool:
try:
json.loads(input_str)
return True
except JSONDecodeError:
return False
def consume_space(i: int, s: str) -> int:
while i < len(s) and s[i].isspace():
i += 1
return i
def _extract_tool_info(
tool: Tool,
) -> tuple[str, dict[str, Any] | None]:
if isinstance(tool, FunctionTool):
return tool.name, tool.parameters
elif isinstance(tool, ChatCompletionToolsParam):
return tool.function.name, tool.function.parameters
else:
raise TypeError(f"Unsupported tool type: {type(tool)}")
def _get_tool_schema_from_tool(tool: Tool) -> dict:
name, params = _extract_tool_info(tool)
params = params if params else {"type": "object", "properties": {}}
return {
"properties": {
"name": {"type": "string", "enum": [name]},
"parameters": params,
},
"required": ["name", "parameters"],
}
def _get_tool_schema_defs(
tools: list[Tool],
) -> dict:
all_defs: dict[str, dict[str, Any]] = {}
for tool in tools:
_, params = _extract_tool_info(tool)
if params is None:
continue
defs = params.pop("$defs", {})
for def_name, def_schema in defs.items():
if def_name in all_defs and all_defs[def_name] != def_schema:
raise ValueError(
f"Tool definition '{def_name}' has multiple schemas, "
"which is not supported."
)
all_defs[def_name] = def_schema
return all_defs
def _get_json_schema_from_tools(
tools: list[Tool],
) -> dict:
json_schema = {
"type": "array",
"minItems": 1,
"items": {
"type": "object",
"anyOf": [_get_tool_schema_from_tool(tool) for tool in tools],
},
}
json_schema_defs = _get_tool_schema_defs(tools)
if json_schema_defs:
json_schema["$defs"] = json_schema_defs
return json_schema
def get_json_schema_from_tools(
tool_choice: str | ToolChoiceFunction | ChatCompletionNamedToolChoiceParam,
tools: list[Tool] | None,
) -> str | dict | None:
# tool_choice: "none"
if tool_choice in ("none", None) or tools is None:
return None
# tool_choice: Forced Function (Responses)
if (not isinstance(tool_choice, str)) and isinstance(
tool_choice, ToolChoiceFunction
):
tool_name = tool_choice.name
tool_map = {tool.name: tool for tool in tools if isinstance(tool, FunctionTool)}
if tool_name not in tool_map:
raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.")
return tool_map[tool_name].parameters
# tool_choice: Forced Function (ChatCompletion)
if (not isinstance(tool_choice, str)) and isinstance(
tool_choice, ChatCompletionNamedToolChoiceParam
):
tool_name = tool_choice.function.name
tool_map = {
tool.function.name: tool
for tool in tools
if isinstance(tool, ChatCompletionToolsParam)
}
if tool_name not in tool_map:
raise ValueError(f"Tool '{tool_name}' has not been passed in `tools`.")
return tool_map[tool_name].function.parameters
# tool_choice: "required"
if tool_choice == "required":
return _get_json_schema_from_tools(tools)
# tool_choice: "auto"
return None
# ---------------------------------------------------------------------------
# Shared utilities for pythonic-style tool call parsers
# (PythonicToolParser, Llama4PythonicToolParser, Olmo3PythonicToolParser)
# ---------------------------------------------------------------------------
class UnexpectedAstError(Exception):
"""Raised when the AST structure does not match the expected
pythonic tool call format."""
pass
_JSON_NAME_LITERALS = {
"null": None,
"true": True,
"false": False,
}
def get_parameter_value(val: ast.expr) -> Any:
"""Extract a Python literal value from an AST expression node.
Handles constants, dicts, lists, and JSON-style name literals
(null, true, false) that some models produce instead of Python
literals (None, True, False).
Raises:
UnexpectedAstError: If the AST node is not a supported literal type.
"""
if isinstance(val, ast.Constant):
return val.value
elif isinstance(val, ast.Dict):
if not all(isinstance(k, ast.Constant) for k in val.keys):
logger.warning(
"Dict argument keys are not all literals: %s",
ast.dump(val),
)
raise UnexpectedAstError("Dict tool call arguments must have literal keys")
return {
k.value: get_parameter_value(v) # type: ignore
for k, v in zip(val.keys, val.values)
}
elif isinstance(val, ast.List):
return [get_parameter_value(v) for v in val.elts]
elif isinstance(val, ast.Name) and val.id in _JSON_NAME_LITERALS:
return _JSON_NAME_LITERALS[val.id]
else:
logger.warning(
"Unsupported AST node type in tool call arguments: %s",
ast.dump(val),
)
raise UnexpectedAstError("Tool call arguments must be literals")
def handle_single_tool(call: ast.Call) -> ToolCall:
"""Convert a single AST function call node into a ToolCall object.
Raises:
UnexpectedAstError: If the call node does not have a simple
function name (e.g. it's an attribute access or subscript).
"""
if not isinstance(call.func, ast.Name):
logger.warning(
"Tool call has non-simple function name: %s",
ast.dump(call.func),
)
raise UnexpectedAstError("Invalid tool call name")
function_name = call.func.id
arguments = {}
for keyword in call.keywords:
arguments[keyword.arg] = get_parameter_value(keyword.value)
return ToolCall(
type="function",
function=FunctionCall(
name=function_name,
arguments=json.dumps(arguments, ensure_ascii=False),
),
)
def make_valid_python(text: str) -> tuple[str, str] | None:
"""Attempt to close all open brackets/quotes to make partial Python valid.
Used during streaming to parse incomplete tool call expressions by
appending the necessary closing characters.
Returns:
A tuple of (completed_text, added_suffix) if the text can be
made valid, or None if the text is too incomplete to complete
meaningfully (e.g. mid-parameter-name or mid-dict-key).
Raises:
UnexpectedAstError: If mismatched brackets or parentheses
are detected.
"""
bracket_stack: list[str] = []
for index, char in enumerate(text):
if char in {"[", "(", "{"}:
bracket_stack.append(char)
elif char == "]":
if not bracket_stack or bracket_stack.pop() != "[":
raise UnexpectedAstError("Mismatched square brackets")
elif char == ")":
if not bracket_stack or bracket_stack.pop() != "(":
raise UnexpectedAstError("Mismatched parentheses")
elif char == "}":
if not bracket_stack or bracket_stack.pop() != "{":
raise UnexpectedAstError("Mismatched curly braces")
elif char in {"'", '"'}:
if bracket_stack and bracket_stack[-1] == char:
if index > 0 and text[index - 1] == "\\":
pass
else:
bracket_stack.pop()
elif bracket_stack and bracket_stack[-1] in {"'", '"'}:
pass
else:
bracket_stack.append(char)
text = text.rstrip()
if text.endswith("=") or text.endswith(":"):
return None
if bracket_stack and bracket_stack[-1] == "{":
trailing_dict_text = text[: text.rfind("{")]
num_keys = trailing_dict_text.count(":")
num_values = trailing_dict_text.count(",")
if num_keys <= num_values:
return None
if bracket_stack and bracket_stack[-1] == "(":
trailing_params_text = text[: text.rfind("(")]
num_full_param_names = trailing_params_text.count("=")
num_full_param_values = trailing_params_text.count(",")
if num_full_param_names <= num_full_param_values:
return None
if text.endswith(","):
text = text[:-1]
if (
bracket_stack
and bracket_stack[-1] == "["
and not text.endswith("[")
and not text.endswith(")")
):
return None
_CLOSING = {"[": "]", "(": ")", "{": "}", "'": "'", '"': '"'}
added_text = ""
for char in reversed(bracket_stack):
added_text += _CLOSING[char]
return text + added_text, added_text
def compute_tool_delta(
previously_sent_args: str,
new_call: ToolCall,
index: int,
withheld_suffix: str,
) -> DeltaToolCall | None:
"""Compute the incremental delta between previously streamed arguments
and the current tool call state.
Returns:
A DeltaToolCall with only the new argument characters, or None
if there is no difference from what was previously sent.
"""
new_call_args = new_call.function.arguments
if withheld_suffix:
if not new_call_args.endswith(withheld_suffix):
msg = (
f"Tool call arguments '{new_call_args}' do not end with "
f"expected withheld suffix '{withheld_suffix}'"
)
logger.error(msg)
raise ValueError(msg)
new_call_args = new_call_args[: -len(withheld_suffix)]
if not previously_sent_args:
return DeltaToolCall(
id=new_call.id,
type="function",
index=index,
function=DeltaFunctionCall(
name=new_call.function.name,
arguments=new_call_args,
),
)
arg_diff = new_call_args[len(previously_sent_args) :]
return (
DeltaToolCall(
id=None,
index=index,
function=DeltaFunctionCall(arguments=arg_diff),
)
if arg_diff
else None
)