From f76730bc883f49aded7896e3b2cb7dda31bbeec6 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 17 Apr 2026 06:05:13 +0000 Subject: [PATCH] glm stuff --- Dockerfile | 9 + glm4_moe_tool_parser.py | 491 +++++++++++++++++++++++++ hf.py | 771 ++++++++++++++++++++++++++++++++++++++++ utils.py | 438 +++++++++++++++++++++++ 4 files changed, 1709 insertions(+) create mode 100644 glm4_moe_tool_parser.py create mode 100644 hf.py create mode 100644 utils.py diff --git a/Dockerfile b/Dockerfile index 46d329b..77673b6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -34,3 +34,12 @@ 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 \ No newline at end of file diff --git a/glm4_moe_tool_parser.py b/glm4_moe_tool_parser.py new file mode 100644 index 0000000..f183a1d --- /dev/null +++ b/glm4_moe_tool_parser.py @@ -0,0 +1,491 @@ +# 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 . The model's chat template instructs: + {name}……… + 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 + ```` 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 = "" + self.tool_call_end_token: str = "" + self.arg_key_start: str = "" + self.arg_key_end: str = "" + self.arg_val_start: str = "" + self.arg_val_end: str = "" + + self.tool_calls_start_token = self.tool_call_start_token + + # ---- FIXED regexes ------------------------------------------------ + # Match the whole … block (unchanged). + self.func_call_regex = re.compile( + r".*?", 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 , and we + # make the body group optional so zero-argument calls don't fail. + self.func_detail_regex = re.compile( + r"\s*" # opening tag + optional whitespace + r"([\w.\-]+)" # group 1: tool/function name (word chars, dots, hyphens) + r"\s*" # optional whitespace / newline + r"((?:.*)?)" # group 2: everything from first onward (may be empty) + r"\s*", # closing tag + re.DOTALL, + ) + + self.func_arg_regex = re.compile( + r"(.*?)\s*(.*?)", 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 ... + # before a partial (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``, ````, or + ````. 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 + 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 \ No newline at end of file diff --git a/hf.py b/hf.py new file mode 100644 index 0000000..5879563 --- /dev/null +++ b/hf.py @@ -0,0 +1,771 @@ +# 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 diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..4711f05 --- /dev/null +++ b/utils.py @@ -0,0 +1,438 @@ +# 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 ``""``. + 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 + )