| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import contextlib |
| 5 | import json |
| 6 | from collections.abc import Sequence |
| 7 | from typing import Any |
| 8 | |
| 9 | import regex as re |
| 10 | |
| 11 | from vllm.entrypoints.generate.base.protocol import ( |
| 12 | DeltaFunctionCall, |
| 13 | DeltaMessage, |
| 14 | DeltaToolCall, |
| 15 | ExtractedToolCallInformation, |
| 16 | FunctionCall, |
| 17 | ToolCall, |
| 18 | ) |
| 19 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 20 | ChatCompletionRequest, |
| 21 | ) |
| 22 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 23 | from vllm.logger import init_logger |
| 24 | from vllm.tokenizers import TokenizerLike |
| 25 | from vllm.tool_parsers.abstract_tool_parser import ( |
| 26 | Tool, |
| 27 | ToolParser, |
| 28 | ) |
| 29 | from vllm.utils import random_uuid |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | |
| 34 | class Step3ToolParser(ToolParser): |
| 35 | """Tool parser for a model that uses a specific XML-like format for tool calls. |
| 36 | This version uses a robust, stateful, cursor-based streaming parser and |
| 37 | consolidates tool arguments into a single message. |
| 38 | """ |
| 39 | |
| 40 | TOOL_CALLS_BEGIN = "<｜tool_calls_begin｜>" |
| 41 | TOOL_CALLS_END = "<｜tool_calls_end｜>" |
| 42 | TOOL_CALL_BEGIN = "<｜tool_call_begin｜>" |
| 43 | TOOL_CALL_END = "<｜tool_call_end｜>" |
| 44 | TOOL_SEP = "<｜tool_sep｜>" |
| 45 | SPECIAL_TOKENS = [TOOL_CALLS_BEGIN, TOOL_CALLS_END, TOOL_CALL_BEGIN, TOOL_CALL_END] |
| 46 | |
| 47 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 48 | super().__init__(tokenizer, tools) |
| 49 | self.position = 0 |
| 50 | # Explicit state flags for robust streaming |
| 51 | self.tool_block_started = False |
| 52 | self.tool_block_finished = False |
| 53 | |
| 54 | def adjust_request( |
| 55 | self, request: ChatCompletionRequest | ResponsesRequest |
| 56 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 57 | request = super().adjust_request(request) |
| 58 | if request.tools and request.tool_choice != "none": |
| 59 | request.skip_special_tokens = False |
| 60 | return request |
| 61 | |
| 62 | @staticmethod |
| 63 | def _parse_steptml_invoke( |
| 64 | action_text: str, |
| 65 | ) -> tuple[str | None, dict[str, str] | None]: |
| 66 | func_name_match = re.search(r'<steptml:invoke name="([^"]+)">', action_text) |
| 67 | if not func_name_match: |
| 68 | return None, None |
| 69 | func_name = func_name_match.group(1) |
| 70 | |
| 71 | params: dict[str, str] = {} |
| 72 | param_matches = re.findall( |
| 73 | r'<steptml:parameter name="([^"]+)">([^<]*)</steptml:parameter>', |
| 74 | action_text, |
| 75 | ) |
| 76 | for name, value in param_matches: |
| 77 | params[name] = value.strip() |
| 78 | return func_name, params |
| 79 | |
| 80 | def _cast_arguments( |
| 81 | self, |
| 82 | func_name: str, |
| 83 | params: dict[str, Any], |
| 84 | ) -> dict[str, Any]: |
| 85 | for tool in self.tools or []: |
| 86 | if tool.function.name == func_name: |
| 87 | schema = tool.function.parameters or {} |
| 88 | properties = schema.get("properties", {}) |
| 89 | for key, value in params.items(): |
| 90 | if not isinstance(value, str): |
| 91 | continue |
| 92 | prop = properties.get(key, {}) |
| 93 | typ = prop.get("type") |
| 94 | if typ == "string": |
| 95 | params[key] = value.strip() |
| 96 | elif typ == "integer": |
| 97 | with contextlib.suppress(ValueError): |
| 98 | params[key] = int(value) |
| 99 | elif typ == "number": |
| 100 | with contextlib.suppress(ValueError): |
| 101 | params[key] = float(value) |
| 102 | elif typ == "boolean": |
| 103 | lower_val = value.lower() |
| 104 | params[key] = ( |
| 105 | lower_val == "true" |
| 106 | if lower_val in ("true", "false") |
| 107 | else value |
| 108 | ) |
| 109 | elif typ == "null": |
| 110 | params[key] = None if value.lower() == "null" else value |
| 111 | break |
| 112 | return params |
| 113 | |
| 114 | def extract_tool_calls_streaming( |
| 115 | self, |
| 116 | previous_text: str, |
| 117 | current_text: str, |
| 118 | delta_text: str, |
| 119 | previous_token_ids: Sequence[int], |
| 120 | current_token_ids: Sequence[int], |
| 121 | delta_token_ids: Sequence[int], |
| 122 | request: ChatCompletionRequest, |
| 123 | ) -> DeltaMessage | None: |
| 124 | # The main loop processes the stream from the last known position. |
| 125 | while True: |
| 126 | if self.position >= len(current_text): |
| 127 | return None # We've processed the entire stream. |
| 128 | |
| 129 | unprocessed_text = current_text[self.position :] |
| 130 | |
| 131 | # STATE: After all tools are done, all subsequent text is content. |
| 132 | if self.tool_block_finished: |
| 133 | self.position = len(current_text) |
| 134 | return DeltaMessage(content=unprocessed_text) |
| 135 | |
| 136 | # STATE: Before the tool block has started. |
| 137 | if not self.tool_block_started: |
| 138 | if unprocessed_text.startswith(self.TOOL_CALLS_BEGIN): |
| 139 | self.position += len(self.TOOL_CALLS_BEGIN) |
| 140 | self.tool_block_started = True |
| 141 | continue # Token consumed, re-loop. |
| 142 | |
| 143 | start_pos = unprocessed_text.find(self.TOOL_CALLS_BEGIN) |
| 144 | if start_pos == -1: |
| 145 | if ( |
| 146 | self.TOOL_CALLS_BEGIN.startswith(unprocessed_text.strip()) |
| 147 | and unprocessed_text |
| 148 | ): |
| 149 | return None # It's a prefix, wait. |
| 150 | self.position = len(current_text) |
| 151 | return DeltaMessage(content=unprocessed_text) |
| 152 | else: |
| 153 | content = unprocessed_text[:start_pos] |
| 154 | self.position += len(content) |
| 155 | return DeltaMessage(content=content) |
| 156 | |
| 157 | # STATE: Inside the main tool block. |
| 158 | offset = len(unprocessed_text) - len(unprocessed_text.lstrip()) |
| 159 | unprocessed_text = unprocessed_text.lstrip() |
| 160 | self.position += offset |
| 161 | |
| 162 | if unprocessed_text.startswith(self.TOOL_CALLS_END): |
| 163 | self.position += len(self.TOOL_CALLS_END) |
| 164 | self.tool_block_finished = True |
| 165 | self.current_tool_id = -1 |
| 166 | continue |
| 167 | |
| 168 | # Check if we are between tool calls. |
| 169 | tool_finished = self.current_tool_id != -1 and self.prev_tool_call_arr[ |
| 170 | self.current_tool_id |
| 171 | ].get("finished") |
| 172 | if self.current_tool_id == -1 or tool_finished: |
| 173 | if unprocessed_text.startswith(self.TOOL_CALL_BEGIN): |
| 174 | self.position += len(self.TOOL_CALL_BEGIN) |
| 175 | if self.current_tool_id == -1: |
| 176 | self.current_tool_id = 0 |
| 177 | else: |
| 178 | self.current_tool_id += 1 |
| 179 | self.current_tool_name_sent = False |
| 180 | while len(self.prev_tool_call_arr) <= self.current_tool_id: |
| 181 | self.prev_tool_call_arr.append({}) |
| 182 | self.prev_tool_call_arr[self.current_tool_id]["finished"] = False |
| 183 | continue |
| 184 | |
| 185 | if self.TOOL_CALL_BEGIN.startswith(unprocessed_text): |
| 186 | return None |
| 187 | |
| 188 | # STATE: Parsing an active tool call. |
| 189 | if self.current_tool_id != -1 and not self.prev_tool_call_arr[ |
| 190 | self.current_tool_id |
| 191 | ].get("finished", False): |
| 192 | end_tool_pos = unprocessed_text.find(self.TOOL_CALL_END) |
| 193 | if end_tool_pos == -1: |
| 194 | tool_body = unprocessed_text |
| 195 | else: |
| 196 | tool_body = unprocessed_text[:end_tool_pos] |
| 197 | |
| 198 | if end_tool_pos == -1 and self.TOOL_CALL_END.startswith(tool_body): |
| 199 | return None |
| 200 | |
| 201 | function_name, arguments = self._parse_steptml_invoke(tool_body) |
| 202 | if not function_name: |
| 203 | return None |
| 204 | |
| 205 | tool_call_arr = {"name": function_name, "parameters": arguments or {}} |
| 206 | |
| 207 | # Send the function name as soon as it's parsed. |
| 208 | if not self.current_tool_name_sent: |
| 209 | self.current_tool_name_sent = True |
| 210 | self.prev_tool_call_arr[self.current_tool_id].update(tool_call_arr) |
| 211 | return DeltaMessage( |
| 212 | tool_calls=[ |
| 213 | DeltaToolCall( |
| 214 | index=self.current_tool_id, |
| 215 | type="function", |
| 216 | id=f"chatcmpl-tool-{random_uuid()}", |
| 217 | function=DeltaFunctionCall(name=function_name), |
| 218 | ) |
| 219 | ] |
| 220 | ) |
| 221 | |
| 222 | # Update our internal state with the latest parsed arguments. |
| 223 | self.prev_tool_call_arr[self.current_tool_id].update( # noqa: E501 |
| 224 | tool_call_arr |
| 225 | ) |
| 226 | |
| 227 | # Only send arguments when the tool call is complete. |
| 228 | if end_tool_pos != -1: |
| 229 | self.position += end_tool_pos + len(self.TOOL_CALL_END) |
| 230 | self.prev_tool_call_arr[self.current_tool_id]["finished"] = True |
| 231 | |
| 232 | final_args = self._cast_arguments( |
| 233 | function_name, |
| 234 | tool_call_arr.get("parameters", {}), # type: ignore |
| 235 | ) |
| 236 | if final_args: |
| 237 | final_args_json = json.dumps(final_args, ensure_ascii=False) |
| 238 | return DeltaMessage( |
| 239 | tool_calls=[ |
| 240 | DeltaToolCall( |
| 241 | index=self.current_tool_id, |
| 242 | function=DeltaFunctionCall( |
| 243 | arguments=final_args_json |
| 244 | ), |
| 245 | ) |
| 246 | ] |
| 247 | ) |
| 248 | |
| 249 | # If tool is not finished, return None to wait for more tokens. |
| 250 | return None |
| 251 | |
| 252 | return None |
| 253 | |
| 254 | def extract_tool_calls( |
| 255 | self, |
| 256 | model_output: str, |
| 257 | request: ChatCompletionRequest, |
| 258 | ) -> ExtractedToolCallInformation: |
| 259 | if self.TOOL_CALLS_BEGIN not in model_output: |
| 260 | return ExtractedToolCallInformation( |
| 261 | tools_called=False, tool_calls=[], content=model_output |
| 262 | ) |
| 263 | |
| 264 | pre_text, rest = model_output.split(self.TOOL_CALLS_BEGIN, 1) |
| 265 | if self.TOOL_CALLS_END not in rest: |
| 266 | return ExtractedToolCallInformation( |
| 267 | tools_called=False, tool_calls=[], content=model_output |
| 268 | ) |
| 269 | |
| 270 | tool_block, post_text = rest.split(self.TOOL_CALLS_END, 1) |
| 271 | content = (pre_text + post_text).strip() |
| 272 | |
| 273 | tool_calls: list[ToolCall] = [] |
| 274 | call_parts = tool_block.split(self.TOOL_CALL_BEGIN) |
| 275 | |
| 276 | for part in call_parts: |
| 277 | if not part or self.TOOL_CALL_END not in part: |
| 278 | continue |
| 279 | |
| 280 | call_content = part.split(self.TOOL_CALL_END, 1)[0] |
| 281 | if self.TOOL_SEP not in call_content: |
| 282 | continue |
| 283 | |
| 284 | type_part, invoke_part = call_content.split(self.TOOL_SEP, 1) |
| 285 | if type_part.strip() != "function": |
| 286 | continue |
| 287 | |
| 288 | function_name, params_dict = self._parse_steptml_invoke(invoke_part) |
| 289 | |
| 290 | if function_name and params_dict is not None: |
| 291 | params_dict = self._cast_arguments(function_name, params_dict) |
| 292 | params_str = json.dumps(params_dict, ensure_ascii=False) |
| 293 | tool_calls.append( |
| 294 | ToolCall( |
| 295 | function=FunctionCall(name=function_name, arguments=params_str) |
| 296 | ) |
| 297 | ) |
| 298 | if tool_calls: |
| 299 | return ExtractedToolCallInformation( |
| 300 | tools_called=True, |
| 301 | tool_calls=tool_calls, |
| 302 | content=content if content else None, |
| 303 | ) |
| 304 | return ExtractedToolCallInformation( |
| 305 | tools_called=False, tool_calls=[], content=model_output |
| 306 | ) |