| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | from collections.abc import Sequence |
| 6 | from typing import Any |
| 7 | |
| 8 | import regex as re |
| 9 | from openai.types.responses import ToolChoiceFunction |
| 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 | ChatCompletionNamedToolChoiceParam, |
| 21 | ChatCompletionRequest, |
| 22 | ) |
| 23 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 24 | from vllm.logger import init_logger |
| 25 | from vllm.tokenizers import TokenizerLike |
| 26 | from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser |
| 27 | from vllm.tool_parsers.utils import ( |
| 28 | coerce_to_schema_type, |
| 29 | extract_types_from_schema, |
| 30 | find_tool_name, |
| 31 | find_tool_properties, |
| 32 | partial_tag_overlap, |
| 33 | ) |
| 34 | |
| 35 | logger = init_logger(__name__) |
| 36 | |
| 37 | ToolParserRequest = ChatCompletionRequest | ResponsesRequest |
| 38 | |
| 39 | |
| 40 | class K2HorizonToolParser(ToolParser): |
| 41 | """Parser for K2 Horizon's IFM tool-call formats.""" |
| 42 | |
| 43 | supports_required_and_named = False |
| 44 | engine_based_streaming = True |
| 45 | |
| 46 | _SUPPORTED_TOOL_FORMATS = frozenset({"json", "xml", "xml_typed"}) |
| 47 | |
| 48 | _TOOL_CALLS_START = "<ifm|tool_calls>" |
| 49 | _TOOL_CALLS_END = "</ifm|tool_calls>" |
| 50 | _TOOL_CALL_START = "<ifm|tool_call>" |
| 51 | _TOOL_CALL_END = "</ifm|tool_call>" |
| 52 | _ARG_KEY_START = "<ifm|arg_key>" |
| 53 | _ARG_KEY_END = "</ifm|arg_key>" |
| 54 | _ARG_TYPE_START = "<ifm|arg_type>" |
| 55 | _ARG_TYPE_END = "</ifm|arg_type>" |
| 56 | _ARG_VALUE_START = "<ifm|arg_value>" |
| 57 | _ARG_VALUE_END = "</ifm|arg_value>" |
| 58 | |
| 59 | _TOOL_CALL_RE = re.compile( |
| 60 | re.escape(_TOOL_CALL_START) + r"(.*?)" + re.escape(_TOOL_CALL_END), |
| 61 | re.DOTALL, |
| 62 | ) |
| 63 | _ARG_RE = re.compile( |
| 64 | re.escape(_ARG_KEY_START) |
| 65 | + r"(.*?)" |
| 66 | + re.escape(_ARG_KEY_END) |
| 67 | + r"\s*(?:" |
| 68 | + re.escape(_ARG_TYPE_START) |
| 69 | + r"(.*?)" |
| 70 | + re.escape(_ARG_TYPE_END) |
| 71 | + r"\s*)?" |
| 72 | + re.escape(_ARG_VALUE_START) |
| 73 | + r"(.*?)" |
| 74 | + re.escape(_ARG_VALUE_END), |
| 75 | re.DOTALL, |
| 76 | ) |
| 77 | _ARG_MARKERS = ( |
| 78 | _ARG_KEY_START, |
| 79 | _ARG_KEY_END, |
| 80 | _ARG_TYPE_START, |
| 81 | _ARG_TYPE_END, |
| 82 | _ARG_VALUE_START, |
| 83 | _ARG_VALUE_END, |
| 84 | ) |
| 85 | |
| 86 | def __init__( |
| 87 | self, |
| 88 | tokenizer: TokenizerLike, |
| 89 | tools: list[Tool] | None = None, |
| 90 | ) -> None: |
| 91 | super().__init__(tokenizer, tools) |
| 92 | self.tool_format = "xml" |
| 93 | self._stream_text = "" |
| 94 | self._stream_cursor = 0 |
| 95 | self._stream_group_start: int | None = None |
| 96 | self._stream_resolved = False |
| 97 | self._pending_content = "" |
| 98 | self._content_seen = False |
| 99 | |
| 100 | @classmethod |
| 101 | def _request_tool_format(cls, request: ToolParserRequest) -> str: |
| 102 | chat_template_kwargs = request.chat_template_kwargs or {} |
| 103 | tool_format = chat_template_kwargs.get("tool_call_format", "xml") |
| 104 | if ( |
| 105 | not isinstance(tool_format, str) |
| 106 | or tool_format not in cls._SUPPORTED_TOOL_FORMATS |
| 107 | ): |
| 108 | supported = ", ".join(sorted(cls._SUPPORTED_TOOL_FORMATS)) |
| 109 | raise ValueError( |
| 110 | f"Unsupported tool_call_format {tool_format!r}. " |
| 111 | f"Supported values: {supported}." |
| 112 | ) |
| 113 | return tool_format |
| 114 | |
| 115 | def adjust_request(self, request: ToolParserRequest) -> ToolParserRequest: |
| 116 | self._request_tool_format(request) |
| 117 | return request |
| 118 | |
| 119 | @staticmethod |
| 120 | def _content_or_none(content: str) -> str | None: |
| 121 | return content if content.strip() else None |
| 122 | |
| 123 | @classmethod |
| 124 | def _no_tool_calls(cls, content: str) -> ExtractedToolCallInformation: |
| 125 | return ExtractedToolCallInformation( |
| 126 | tools_called=False, |
| 127 | tool_calls=[], |
| 128 | content=cls._content_or_none(content), |
| 129 | ) |
| 130 | |
| 131 | @staticmethod |
| 132 | def _make_tool_call(name: str, arguments: dict[str, Any]) -> ToolCall: |
| 133 | return ToolCall( |
| 134 | type="function", |
| 135 | function=FunctionCall( |
| 136 | name=name, |
| 137 | arguments=json.dumps( |
| 138 | arguments, |
| 139 | ensure_ascii=False, |
| 140 | allow_nan=False, |
| 141 | ), |
| 142 | ), |
| 143 | ) |
| 144 | |
| 145 | @staticmethod |
| 146 | def _base_arg_type(arg_type: str | None) -> str | None: |
| 147 | if not arg_type: |
| 148 | return None |
| 149 | return arg_type.split("[", 1)[0].strip().lower() |
| 150 | |
| 151 | @staticmethod |
| 152 | def _coerce_json_value(value: Any, schema: Any) -> Any: |
| 153 | if not isinstance(schema, dict): |
| 154 | return value |
| 155 | raw_value = ( |
| 156 | value |
| 157 | if isinstance(value, str) |
| 158 | else json.dumps(value, ensure_ascii=False, allow_nan=False) |
| 159 | ) |
| 160 | return coerce_to_schema_type(raw_value, extract_types_from_schema(schema)) |
| 161 | |
| 162 | @classmethod |
| 163 | def _coerce_xml_value( |
| 164 | cls, |
| 165 | value: str, |
| 166 | schema: Any, |
| 167 | explicit_type: str | None, |
| 168 | ) -> Any: |
| 169 | explicit_base = cls._base_arg_type(explicit_type) |
| 170 | target_types = extract_types_from_schema(schema) |
| 171 | if isinstance(schema, dict): |
| 172 | if explicit_base in target_types and len(set(target_types)) > 1: |
| 173 | target_types = [explicit_base] |
| 174 | elif explicit_base == "any": |
| 175 | return value |
| 176 | elif explicit_base: |
| 177 | target_types = [explicit_base] |
| 178 | |
| 179 | raw_value = value if "string" in target_types else value.strip() |
| 180 | return coerce_to_schema_type(raw_value, target_types) |
| 181 | |
| 182 | @staticmethod |
| 183 | def _named_tool(request: ToolParserRequest) -> str | None: |
| 184 | tool_choice = request.tool_choice |
| 185 | if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam): |
| 186 | return tool_choice.function.name |
| 187 | if isinstance(tool_choice, ToolChoiceFunction): |
| 188 | return tool_choice.name |
| 189 | return None |
| 190 | |
| 191 | @classmethod |
| 192 | def _validate_tool_name( |
| 193 | cls, |
| 194 | name: Any, |
| 195 | request: ToolParserRequest, |
| 196 | ) -> str: |
| 197 | if ( |
| 198 | not isinstance(name, str) |
| 199 | or not name |
| 200 | or any(char.isspace() for char in name) |
| 201 | ): |
| 202 | raise ValueError("IFM tool call has an invalid function name.") |
| 203 | if request.tools and not find_tool_name(request.tools, name): |
| 204 | raise ValueError(f"Unknown IFM tool name: {name}.") |
| 205 | named_tool = cls._named_tool(request) |
| 206 | if named_tool is not None and name != named_tool: |
| 207 | raise ValueError(f"Unexpected IFM tool name: {name}.") |
| 208 | return name |
| 209 | |
| 210 | @classmethod |
| 211 | def _parse_json_call( |
| 212 | cls, |
| 213 | body: str, |
| 214 | request: ToolParserRequest, |
| 215 | ) -> tuple[str, dict[str, Any]]: |
| 216 | raw_call = json.loads(body.strip()) |
| 217 | if not isinstance(raw_call, dict): |
| 218 | raise ValueError("IFM JSON tool call must be an object.") |
| 219 | |
| 220 | name = cls._validate_tool_name(raw_call.get("name"), request) |
| 221 | arguments = raw_call.get("arguments", {}) |
| 222 | if not isinstance(arguments, dict): |
| 223 | raise ValueError("IFM JSON tool-call arguments must be an object.") |
| 224 | |
| 225 | properties = find_tool_properties(request.tools, name) |
| 226 | arguments = { |
| 227 | arg_name: cls._coerce_json_value(arg_value, properties.get(arg_name)) |
| 228 | for arg_name, arg_value in arguments.items() |
| 229 | } |
| 230 | return name, arguments |
| 231 | |
| 232 | @classmethod |
| 233 | def _parse_xml_call( |
| 234 | cls, |
| 235 | body: str, |
| 236 | request: ToolParserRequest, |
| 237 | tool_format: str, |
| 238 | ) -> tuple[str, dict[str, Any]]: |
| 239 | first_arg = body.find(cls._ARG_KEY_START) |
| 240 | if first_arg == -1: |
| 241 | if any(marker in body for marker in cls._ARG_MARKERS): |
| 242 | raise ValueError("Malformed IFM XML argument tags.") |
| 243 | name = cls._validate_tool_name(body.strip(), request) |
| 244 | return name, {} |
| 245 | |
| 246 | name = cls._validate_tool_name(body[:first_arg].strip(), request) |
| 247 | properties = find_tool_properties(request.tools, name) |
| 248 | arguments: dict[str, Any] = {} |
| 249 | position = first_arg |
| 250 | for match in cls._ARG_RE.finditer(body, first_arg): |
| 251 | if body[position : match.start()].strip(): |
| 252 | raise ValueError("Malformed IFM XML argument tags.") |
| 253 | |
| 254 | arg_name = match.group(1).strip() |
| 255 | if not arg_name: |
| 256 | raise ValueError("IFM XML argument is missing a name.") |
| 257 | if arg_name in arguments: |
| 258 | raise ValueError(f"Duplicate IFM XML argument: {arg_name}.") |
| 259 | |
| 260 | explicit_type = match.group(2) |
| 261 | if explicit_type is not None: |
| 262 | explicit_type = explicit_type.strip() |
| 263 | if tool_format == "xml" and explicit_type is not None: |
| 264 | raise ValueError("Unexpected type tag in IFM XML tool call.") |
| 265 | if tool_format == "xml_typed" and not explicit_type: |
| 266 | raise ValueError("IFM typed XML argument is missing a type.") |
| 267 | arguments[arg_name] = cls._coerce_xml_value( |
| 268 | match.group(3), |
| 269 | properties.get(arg_name), |
| 270 | explicit_type, |
| 271 | ) |
| 272 | position = match.end() |
| 273 | |
| 274 | if position == first_arg or body[position:].strip(): |
| 275 | raise ValueError("Malformed IFM XML argument tags.") |
| 276 | return name, arguments |
| 277 | |
| 278 | @classmethod |
| 279 | def _parse_call( |
| 280 | cls, |
| 281 | body: str, |
| 282 | request: ToolParserRequest, |
| 283 | tool_format: str, |
| 284 | ) -> tuple[str, dict[str, Any]]: |
| 285 | if tool_format == "json": |
| 286 | return cls._parse_json_call(body, request) |
| 287 | return cls._parse_xml_call(body, request, tool_format) |
| 288 | |
| 289 | def extract_tool_calls( |
| 290 | self, |
| 291 | model_output: str, |
| 292 | request: ToolParserRequest, |
| 293 | ) -> ExtractedToolCallInformation: |
| 294 | tool_format = self._request_tool_format(request) |
| 295 | if request.tool_choice == "none": |
| 296 | return self._no_tool_calls(model_output) |
| 297 | |
| 298 | group_start = model_output.find(self._TOOL_CALLS_START) |
| 299 | if group_start == -1: |
| 300 | return self._no_tool_calls(model_output) |
| 301 | |
| 302 | body_start = group_start + len(self._TOOL_CALLS_START) |
| 303 | group_end = model_output.find(self._TOOL_CALLS_END, body_start) |
| 304 | if group_end == -1: |
| 305 | return self._no_tool_calls(model_output) |
| 306 | |
| 307 | group_body = model_output[body_start:group_end] |
| 308 | matches = list(self._TOOL_CALL_RE.finditer(group_body)) |
| 309 | if not matches or self._TOOL_CALL_RE.sub("", group_body).strip(): |
| 310 | logger.warning("Malformed K2 Horizon IFM tool-call group.") |
| 311 | return self._no_tool_calls(model_output) |
| 312 | |
| 313 | try: |
| 314 | tool_calls = [ |
| 315 | self._make_tool_call( |
| 316 | *self._parse_call(match.group(1), request, tool_format) |
| 317 | ) |
| 318 | for match in matches |
| 319 | ] |
| 320 | except (json.JSONDecodeError, TypeError, ValueError): |
| 321 | logger.warning("Failed to parse K2 Horizon IFM tool calls.", exc_info=True) |
| 322 | return self._no_tool_calls(model_output) |
| 323 | |
| 324 | suffix = model_output[group_end + len(self._TOOL_CALLS_END) :] |
| 325 | content = model_output[:group_start] + suffix |
| 326 | return ExtractedToolCallInformation( |
| 327 | tools_called=True, |
| 328 | tool_calls=tool_calls, |
| 329 | content=self._content_or_none(content), |
| 330 | ) |
| 331 | |
| 332 | def _emit_content(self, addition: str, *, final: bool = False) -> str | None: |
| 333 | if addition: |
| 334 | if self._content_seen: |
| 335 | return addition |
| 336 | self._pending_content += addition |
| 337 | if self._pending_content.strip(): |
| 338 | self._content_seen = True |
| 339 | content = self._pending_content |
| 340 | self._pending_content = "" |
| 341 | return content |
| 342 | if final: |
| 343 | self._pending_content = "" |
| 344 | return None |
| 345 | |
| 346 | def _streaming_tool_calls( |
| 347 | self, |
| 348 | tool_calls: list[ToolCall], |
| 349 | ) -> list[DeltaToolCall]: |
| 350 | deltas: list[DeltaToolCall] = [] |
| 351 | for index, tool_call in enumerate(tool_calls): |
| 352 | arguments_json = tool_call.function.arguments |
| 353 | arguments = json.loads(arguments_json) |
| 354 | self.prev_tool_call_arr.append( |
| 355 | {"name": tool_call.function.name, "arguments": arguments} |
| 356 | ) |
| 357 | self.streamed_args_for_tool.append(arguments_json) |
| 358 | deltas.append( |
| 359 | DeltaToolCall( |
| 360 | index=index, |
| 361 | id=tool_call.id, |
| 362 | type="function", |
| 363 | function=DeltaFunctionCall( |
| 364 | name=tool_call.function.name, |
| 365 | arguments=arguments_json, |
| 366 | ), |
| 367 | ) |
| 368 | ) |
| 369 | return deltas |
| 370 | |
| 371 | def extract_tool_calls_streaming( |
| 372 | self, |
| 373 | previous_text: str, |
| 374 | current_text: str, |
| 375 | delta_text: str, |
| 376 | previous_token_ids: Sequence[int], |
| 377 | current_token_ids: Sequence[int], |
| 378 | delta_token_ids: Sequence[int], |
| 379 | request: ToolParserRequest, |
| 380 | ) -> DeltaMessage | None: |
| 381 | del ( |
| 382 | previous_text, |
| 383 | current_text, |
| 384 | previous_token_ids, |
| 385 | current_token_ids, |
| 386 | delta_token_ids, |
| 387 | ) |
| 388 | |
| 389 | self._request_tool_format(request) |
| 390 | if request.tool_choice == "none": |
| 391 | emitted_content = self._emit_content(delta_text) |
| 392 | return ( |
| 393 | DeltaMessage(content=emitted_content) |
| 394 | if emitted_content is not None |
| 395 | else None |
| 396 | ) |
| 397 | |
| 398 | self._stream_text += delta_text |
| 399 | if self._stream_resolved: |
| 400 | emitted_content = self._emit_content(delta_text) |
| 401 | return ( |
| 402 | DeltaMessage(content=emitted_content) |
| 403 | if emitted_content is not None |
| 404 | else None |
| 405 | ) |
| 406 | |
| 407 | content: str | None = None |
| 408 | if self._stream_group_start is None: |
| 409 | group_start = self._stream_text.find( |
| 410 | self._TOOL_CALLS_START, self._stream_cursor |
| 411 | ) |
| 412 | if group_start == -1: |
| 413 | pending = self._stream_text[self._stream_cursor :] |
| 414 | overlap = partial_tag_overlap(pending, self._TOOL_CALLS_START) |
| 415 | sendable_end = len(self._stream_text) - overlap |
| 416 | content = self._emit_content( |
| 417 | self._stream_text[self._stream_cursor : sendable_end] |
| 418 | ) |
| 419 | self._stream_cursor = sendable_end |
| 420 | return DeltaMessage(content=content) if content is not None else None |
| 421 | |
| 422 | self._stream_group_start = group_start |
| 423 | content = self._emit_content( |
| 424 | self._stream_text[self._stream_cursor : group_start] |
| 425 | ) |
| 426 | self._stream_cursor = group_start |
| 427 | |
| 428 | group_start = self._stream_group_start |
| 429 | body_start = group_start + len(self._TOOL_CALLS_START) |
| 430 | group_end = self._stream_text.find(self._TOOL_CALLS_END, body_start) |
| 431 | if group_end == -1: |
| 432 | return DeltaMessage(content=content) if content is not None else None |
| 433 | |
| 434 | group_end += len(self._TOOL_CALLS_END) |
| 435 | group_text = self._stream_text[group_start:group_end] |
| 436 | parsed = self.extract_tool_calls(group_text, request) |
| 437 | self._stream_resolved = True |
| 438 | |
| 439 | if not parsed.tools_called: |
| 440 | failed_content = self._emit_content( |
| 441 | self._stream_text[self._stream_cursor :] |
| 442 | ) |
| 443 | self._stream_cursor = len(self._stream_text) |
| 444 | combined = (content or "") + (failed_content or "") |
| 445 | return DeltaMessage(content=combined) if combined else None |
| 446 | |
| 447 | tool_calls = self._streaming_tool_calls(parsed.tool_calls) |
| 448 | suffix = self._emit_content(self._stream_text[group_end:]) |
| 449 | self._stream_cursor = len(self._stream_text) |
| 450 | combined = (content or "") + (suffix or "") |
| 451 | return DeltaMessage(content=combined or None, tool_calls=tool_calls) |
| 452 | |
| 453 | def finish_streaming(self) -> DeltaMessage | None: |
| 454 | if not self._stream_resolved: |
| 455 | content = self._emit_content( |
| 456 | self._stream_text[self._stream_cursor :], final=True |
| 457 | ) |
| 458 | self._stream_cursor = len(self._stream_text) |
| 459 | self._stream_resolved = True |
| 460 | else: |
| 461 | content = self._emit_content("", final=True) |
| 462 | return DeltaMessage(content=content) if content is not None else None |