| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | import contextlib |
| 5 | import copy |
| 6 | import json |
| 7 | import logging |
| 8 | from abc import ABC, abstractmethod |
| 9 | from contextlib import AsyncExitStack |
| 10 | from dataclasses import replace |
| 11 | from typing import TYPE_CHECKING, Any, Final, Union |
| 12 | |
| 13 | from openai.types.responses import ResponseFunctionToolCall, ResponseOutputItem |
| 14 | from openai.types.responses.response_function_tool_call_output_item import ( |
| 15 | ResponseFunctionToolCallOutputItem, |
| 16 | ) |
| 17 | from openai.types.responses.response_output_item import McpCall |
| 18 | from openai.types.responses.response_output_message import ResponseOutputMessage |
| 19 | from openai.types.responses.response_output_text import ResponseOutputText |
| 20 | from openai.types.responses.tool import Mcp |
| 21 | from openai_harmony import Author, Message, Role, TextContent |
| 22 | |
| 23 | from vllm import envs |
| 24 | from vllm.entrypoints.chat_utils import ( |
| 25 | ChatTemplateContentFormatOption, |
| 26 | ) |
| 27 | from vllm.entrypoints.mcp.tool import Tool |
| 28 | from vllm.entrypoints.mcp.tool_server import ToolServer |
| 29 | from vllm.entrypoints.openai.parser.harmony_utils import render_for_completion |
| 30 | from vllm.entrypoints.openai.responses.protocol import ( |
| 31 | ResponseInputOutputItem, |
| 32 | ResponseRawMessageAndToken, |
| 33 | ResponsesRequest, |
| 34 | ) |
| 35 | from vllm.entrypoints.openai.responses.utils import ( |
| 36 | build_response_output_items, |
| 37 | construct_tool_dicts, |
| 38 | ) |
| 39 | from vllm.entrypoints.serve.utils.tool_calls_utils import ( |
| 40 | maybe_filter_parallel_tool_calls, |
| 41 | ) |
| 42 | from vllm.outputs import RequestOutput |
| 43 | from vllm.parser.abstract_parser import Parser |
| 44 | from vllm.tokenizers import TokenizerLike |
| 45 | from vllm.utils import random_uuid |
| 46 | |
| 47 | if TYPE_CHECKING: |
| 48 | from mcp.client import ClientSession |
| 49 | |
| 50 | from vllm.v1.metrics.stats import RequestStateStats |
| 51 | |
| 52 | logger = logging.getLogger(__name__) |
| 53 | |
| 54 | # This is currently needed as the tool type doesn't 1:1 match the |
| 55 | # tool namespace, which is what is used to look up the |
| 56 | # connection to the tool server |
| 57 | _TOOL_NAME_TO_TYPE_MAP = { |
| 58 | "browser": "web_search_preview", |
| 59 | "python": "code_interpreter", |
| 60 | "container": "container", |
| 61 | } |
| 62 | MCP_PREFIX = "mcp_" |
| 63 | |
| 64 | |
| 65 | def _map_tool_name_to_tool_type(tool_name: str) -> str: |
| 66 | if tool_name not in _TOOL_NAME_TO_TYPE_MAP: |
| 67 | available_tools = ", ".join(_TOOL_NAME_TO_TYPE_MAP.keys()) |
| 68 | raise ValueError( |
| 69 | f"Built-in tool name '{tool_name}' not defined in mapping. " |
| 70 | f"Available tools: {available_tools}" |
| 71 | ) |
| 72 | return _TOOL_NAME_TO_TYPE_MAP[tool_name] |
| 73 | |
| 74 | |
| 75 | class TurnMetrics: |
| 76 | """Tracks token and toolcall details for a single conversation turn.""" |
| 77 | |
| 78 | def __init__( |
| 79 | self, |
| 80 | input_tokens: int = 0, |
| 81 | output_tokens: int = 0, |
| 82 | cached_input_tokens: int = 0, |
| 83 | tool_output_tokens: int = 0, |
| 84 | ) -> None: |
| 85 | self.input_tokens = input_tokens |
| 86 | self.output_tokens = output_tokens |
| 87 | self.cached_input_tokens = cached_input_tokens |
| 88 | self.tool_output_tokens = tool_output_tokens |
| 89 | |
| 90 | def reset(self) -> None: |
| 91 | """Reset counters for a new turn.""" |
| 92 | self.input_tokens = 0 |
| 93 | self.output_tokens = 0 |
| 94 | self.cached_input_tokens = 0 |
| 95 | self.tool_output_tokens = 0 |
| 96 | |
| 97 | def copy(self) -> "TurnMetrics": |
| 98 | """Create a copy of this turn's token counts.""" |
| 99 | return TurnMetrics( |
| 100 | self.input_tokens, |
| 101 | self.output_tokens, |
| 102 | self.cached_input_tokens, |
| 103 | self.tool_output_tokens, |
| 104 | ) |
| 105 | |
| 106 | |
| 107 | class ConversationContext(ABC): |
| 108 | response_parser: Parser | None = None |
| 109 | request_metrics: "RequestStateStats | None" = None |
| 110 | # Built-in tools can trigger additional model generations. In that case, |
| 111 | # the stored engine timestamps cover only one turn, while token usage is |
| 112 | # accumulated across all turns. |
| 113 | request_metrics_cover_all_generation_turns: bool = True |
| 114 | |
| 115 | @abstractmethod |
| 116 | def append_output(self, output: RequestOutput) -> None: |
| 117 | pass |
| 118 | |
| 119 | @abstractmethod |
| 120 | def append_tool_output(self, output) -> None: |
| 121 | pass |
| 122 | |
| 123 | @abstractmethod |
| 124 | async def call_tool(self) -> list[Message]: |
| 125 | pass |
| 126 | |
| 127 | @abstractmethod |
| 128 | def need_builtin_tool_call(self) -> bool: |
| 129 | pass |
| 130 | |
| 131 | @abstractmethod |
| 132 | def render_for_completion(self) -> list[int]: |
| 133 | pass |
| 134 | |
| 135 | @abstractmethod |
| 136 | async def init_tool_sessions( |
| 137 | self, |
| 138 | tool_server: ToolServer | None, |
| 139 | exit_stack: AsyncExitStack, |
| 140 | request_id: str, |
| 141 | mcp_tools: dict[str, Mcp], |
| 142 | ) -> None: |
| 143 | pass |
| 144 | |
| 145 | @abstractmethod |
| 146 | async def cleanup_session(self) -> None: |
| 147 | raise NotImplementedError("Should not be called.") |
| 148 | |
| 149 | |
| 150 | def _create_json_parse_error_messages( |
| 151 | last_msg: Message, e: json.JSONDecodeError |
| 152 | ) -> list[Message]: |
| 153 | """Creates an error message when json parse failed.""" |
| 154 | error_msg = ( |
| 155 | f"Error parsing tool arguments as JSON: {str(e)}. " |
| 156 | "Please ensure the tool call arguments are valid JSON and try again." |
| 157 | ) |
| 158 | content = TextContent(text=error_msg) |
| 159 | author = Author(role=Role.TOOL, name=last_msg.recipient) |
| 160 | return [ |
| 161 | Message( |
| 162 | author=author, |
| 163 | content=[content], |
| 164 | recipient=Role.ASSISTANT, |
| 165 | channel=last_msg.channel, |
| 166 | ) |
| 167 | ] |
| 168 | |
| 169 | |
| 170 | class SimpleContext(ConversationContext): |
| 171 | """This is a context that cannot handle MCP tool calls.""" |
| 172 | |
| 173 | def __init__( |
| 174 | self, |
| 175 | *, |
| 176 | response_parser: Parser | None = None, |
| 177 | parser_cls: type[Parser] | None = None, |
| 178 | tokenizer: TokenizerLike | None = None, |
| 179 | request: ResponsesRequest | None = None, |
| 180 | chat_template_kwargs: dict[str, Any] | None = None, |
| 181 | ): |
| 182 | self.last_output = None |
| 183 | self.response_parser = response_parser or ( |
| 184 | parser_cls( |
| 185 | tokenizer, |
| 186 | request.tools, |
| 187 | chat_template_kwargs=chat_template_kwargs, |
| 188 | ) |
| 189 | if parser_cls is not None and tokenizer is not None and request is not None |
| 190 | else None |
| 191 | ) |
| 192 | |
| 193 | # Accumulated final output for streaming mode |
| 194 | self._accumulated_text: str = "" |
| 195 | self._accumulated_token_ids: list[int] = [] |
| 196 | self._accumulated_logprobs: list = [] |
| 197 | # Items already sent as output_item.done events; the final |
| 198 | # streaming response reuses them instead of reparsing. |
| 199 | self.streamed_output_items: list[ResponseOutputItem] | None = None |
| 200 | |
| 201 | self.num_prompt_tokens = 0 |
| 202 | self.num_output_tokens = 0 |
| 203 | self.num_cached_tokens = 0 |
| 204 | self.num_cache_creation_tokens = 0 |
| 205 | # todo num_reasoning_tokens is not implemented yet. |
| 206 | self.num_reasoning_tokens = 0 |
| 207 | # not implemented yet for SimpleContext |
| 208 | self.all_turn_metrics: list[TurnMetrics] = [] |
| 209 | |
| 210 | self.input_messages: list[ResponseRawMessageAndToken] = [] |
| 211 | self.kv_transfer_params: dict[str, Any] | None = None |
| 212 | self.ec_transfer_params: dict[str, Any] | None = None |
| 213 | |
| 214 | def append_output(self, output) -> None: |
| 215 | self.last_output = output |
| 216 | if not isinstance(output, RequestOutput): |
| 217 | raise ValueError("SimpleContext only supports RequestOutput.") |
| 218 | self.num_prompt_tokens = len(output.prompt_token_ids or []) |
| 219 | self.num_cached_tokens = output.num_cached_tokens or 0 |
| 220 | if output.num_cache_creation_tokens is not None: |
| 221 | self.num_cache_creation_tokens = output.num_cache_creation_tokens |
| 222 | self.num_output_tokens += len(output.outputs[0].token_ids or []) |
| 223 | if output.kv_transfer_params is not None: |
| 224 | self.kv_transfer_params = output.kv_transfer_params |
| 225 | if output.ec_transfer_params is not None: |
| 226 | self.ec_transfer_params = output.ec_transfer_params |
| 227 | |
| 228 | # Accumulate text, token_ids, and logprobs for streaming mode |
| 229 | delta_output = output.outputs[0] |
| 230 | self._accumulated_text += delta_output.text |
| 231 | self._accumulated_token_ids.extend(delta_output.token_ids) |
| 232 | if delta_output.logprobs is not None: |
| 233 | self._accumulated_logprobs.extend(delta_output.logprobs) |
| 234 | |
| 235 | if len(self.input_messages) == 0: |
| 236 | output_prompt = output.prompt or "" |
| 237 | output_prompt_token_ids = output.prompt_token_ids or [] |
| 238 | self.input_messages.append( |
| 239 | ResponseRawMessageAndToken( |
| 240 | message=output_prompt, |
| 241 | tokens=output_prompt_token_ids, |
| 242 | ) |
| 243 | ) |
| 244 | |
| 245 | @property |
| 246 | def output_messages(self) -> list[ResponseRawMessageAndToken]: |
| 247 | """Return consolidated output as a single message. |
| 248 | |
| 249 | In streaming mode, text and tokens are accumulated across many deltas. |
| 250 | This property returns them as a single entry rather than one per delta. |
| 251 | """ |
| 252 | if not self._accumulated_text and not self._accumulated_token_ids: |
| 253 | return [] |
| 254 | return [ |
| 255 | ResponseRawMessageAndToken( |
| 256 | message=self._accumulated_text, |
| 257 | tokens=list(self._accumulated_token_ids), |
| 258 | ) |
| 259 | ] |
| 260 | |
| 261 | @property |
| 262 | def final_output(self) -> RequestOutput | None: |
| 263 | """Return the final output, with complete text/token_ids/logprobs.""" |
| 264 | if self.last_output is not None and self.last_output.outputs: |
| 265 | assert isinstance(self.last_output, RequestOutput) |
| 266 | final_output = copy.copy(self.last_output) |
| 267 | # copy inner item to avoid modify last_output |
| 268 | final_output.outputs = [replace(item) for item in self.last_output.outputs] |
| 269 | final_output.outputs[0].text = self._accumulated_text |
| 270 | final_output.outputs[0].token_ids = tuple(self._accumulated_token_ids) |
| 271 | if self._accumulated_logprobs: |
| 272 | final_output.outputs[0].logprobs = self._accumulated_logprobs |
| 273 | return final_output |
| 274 | return self.last_output |
| 275 | |
| 276 | def append_tool_output(self, output) -> None: |
| 277 | raise NotImplementedError("Should not be called.") |
| 278 | |
| 279 | def need_builtin_tool_call(self) -> bool: |
| 280 | return False |
| 281 | |
| 282 | async def call_tool(self) -> list[Message]: |
| 283 | raise NotImplementedError("Should not be called.") |
| 284 | |
| 285 | def render_for_completion(self) -> list[int]: |
| 286 | raise NotImplementedError("Should not be called.") |
| 287 | |
| 288 | async def init_tool_sessions( |
| 289 | self, |
| 290 | tool_server: ToolServer | None, |
| 291 | exit_stack: AsyncExitStack, |
| 292 | request_id: str, |
| 293 | mcp_tools: dict[str, Mcp], |
| 294 | ) -> None: |
| 295 | pass |
| 296 | |
| 297 | async def cleanup_session(self) -> None: |
| 298 | raise NotImplementedError("Should not be called.") |
| 299 | |
| 300 | |
| 301 | class ParsableContext(ConversationContext): |
| 302 | def __init__( |
| 303 | self, |
| 304 | *, |
| 305 | response_messages: list[ResponseInputOutputItem], |
| 306 | tokenizer: TokenizerLike, |
| 307 | parser_cls: type[Parser] | None, |
| 308 | request: ResponsesRequest, |
| 309 | available_tools: list[str] | None, |
| 310 | chat_template: str | None, |
| 311 | chat_template_content_format: ChatTemplateContentFormatOption, |
| 312 | response_parser: Parser | None = None, |
| 313 | enable_auto_tools: bool = False, |
| 314 | ): |
| 315 | self.num_prompt_tokens = 0 |
| 316 | self.num_output_tokens = 0 |
| 317 | self.num_cached_tokens = 0 |
| 318 | self.num_reasoning_tokens = 0 |
| 319 | # not implemented yet for ParsableContext |
| 320 | self.all_turn_metrics: list[TurnMetrics] = [] |
| 321 | |
| 322 | self.response_messages: list[ResponseInputOutputItem] = response_messages |
| 323 | self.num_init_messages = len(response_messages) |
| 324 | self.finish_reason: str | None = None |
| 325 | self.enable_auto_tools = enable_auto_tools |
| 326 | |
| 327 | self.response_parser = response_parser or ( |
| 328 | parser_cls(tokenizer, request.tools) if parser_cls is not None else None |
| 329 | ) |
| 330 | self.parser_cls = parser_cls |
| 331 | self.request = request |
| 332 | |
| 333 | self.available_tools = available_tools or [] |
| 334 | self._tool_sessions: dict[str, ClientSession | Tool] = {} |
| 335 | self.called_tools: set[str] = set() |
| 336 | |
| 337 | self.tool_dicts = construct_tool_dicts(request.tools, request.tool_choice) |
| 338 | self.chat_template = chat_template |
| 339 | self.chat_template_content_format: Final = chat_template_content_format |
| 340 | |
| 341 | self.input_messages: list[ResponseRawMessageAndToken] = [] |
| 342 | self.output_messages: list[ResponseRawMessageAndToken] = [] |
| 343 | self._accumulated_token_ids: list[int] = [] |
| 344 | self.kv_transfer_params: dict[str, Any] | None = None |
| 345 | self.ec_transfer_params: dict[str, Any] | None = None |
| 346 | |
| 347 | def append_output(self, output: RequestOutput) -> None: |
| 348 | self.num_prompt_tokens = len(output.prompt_token_ids or []) |
| 349 | self.num_cached_tokens = output.num_cached_tokens or 0 |
| 350 | self.num_output_tokens += len(output.outputs[0].token_ids or []) |
| 351 | if output.kv_transfer_params is not None: |
| 352 | self.kv_transfer_params = output.kv_transfer_params |
| 353 | |
| 354 | if output.ec_transfer_params is not None: |
| 355 | self.ec_transfer_params = output.ec_transfer_params |
| 356 | |
| 357 | completion = output.outputs[0] |
| 358 | self.finish_reason = completion.finish_reason |
| 359 | |
| 360 | if self.response_parser is not None: |
| 361 | reasoning, content, tool_calls = self.response_parser.parse( |
| 362 | completion.text, |
| 363 | self.request, |
| 364 | enable_auto_tools=self.enable_auto_tools, |
| 365 | model_output_token_ids=completion.token_ids, |
| 366 | ) |
| 367 | # Each round is a separate generation, so count it on its own ids. |
| 368 | self.num_reasoning_tokens += self.response_parser.count_reasoning_tokens( |
| 369 | completion.token_ids or [] |
| 370 | ) |
| 371 | if not self.request.include_reasoning: |
| 372 | reasoning = None |
| 373 | self.response_messages.extend( |
| 374 | build_response_output_items( |
| 375 | reasoning=reasoning, |
| 376 | content=content, |
| 377 | tool_calls=maybe_filter_parallel_tool_calls( |
| 378 | tool_calls or [], self.request |
| 379 | ), |
| 380 | tools=self.request.tools, |
| 381 | ) |
| 382 | ) |
| 383 | elif completion.text: |
| 384 | self.response_messages.append( |
| 385 | ResponseOutputMessage( |
| 386 | type="message", |
| 387 | id=f"msg_{random_uuid()}", |
| 388 | status="completed", |
| 389 | role="assistant", |
| 390 | content=[ |
| 391 | ResponseOutputText( |
| 392 | annotations=[], |
| 393 | type="output_text", |
| 394 | text=completion.text, |
| 395 | logprobs=None, |
| 396 | ) |
| 397 | ], |
| 398 | ) |
| 399 | ) |
| 400 | |
| 401 | self._accumulated_token_ids.extend(completion.token_ids or []) |
| 402 | |
| 403 | if self.request.enable_response_messages: |
| 404 | output_prompt = output.prompt or "" |
| 405 | output_prompt_token_ids = output.prompt_token_ids or [] |
| 406 | if len(self.input_messages) == 0: |
| 407 | self.input_messages.append( |
| 408 | ResponseRawMessageAndToken( |
| 409 | message=output_prompt, |
| 410 | tokens=output_prompt_token_ids, |
| 411 | ) |
| 412 | ) |
| 413 | else: |
| 414 | self.output_messages.append( |
| 415 | ResponseRawMessageAndToken( |
| 416 | message=output_prompt, |
| 417 | tokens=output_prompt_token_ids, |
| 418 | ) |
| 419 | ) |
| 420 | self.output_messages.append( |
| 421 | ResponseRawMessageAndToken( |
| 422 | message=completion.text, |
| 423 | tokens=completion.token_ids, |
| 424 | ) |
| 425 | ) |
| 426 | |
| 427 | def append_tool_output(self, output: list[ResponseInputOutputItem]) -> None: |
| 428 | self.response_messages.extend(output) |
| 429 | |
| 430 | def need_builtin_tool_call(self) -> bool: |
| 431 | """Return true if the last message is a builtin tool call |
| 432 | that the request has enabled.""" |
| 433 | last_message = self.response_messages[-1] |
| 434 | if last_message.type != "function_call": |
| 435 | return False |
| 436 | if last_message.name in ("code_interpreter", "python"): |
| 437 | return "python" in self.available_tools |
| 438 | if last_message.name == "web_search_preview": |
| 439 | return "browser" in self.available_tools |
| 440 | if last_message.name.startswith("container"): |
| 441 | return "container" in self.available_tools |
| 442 | return False |
| 443 | |
| 444 | async def call_python_tool( |
| 445 | self, |
| 446 | tool_session: Union["ClientSession", Tool], |
| 447 | last_msg: ResponseFunctionToolCall, |
| 448 | ) -> list[ResponseInputOutputItem]: |
| 449 | self.called_tools.add("python") |
| 450 | if isinstance(tool_session, Tool): |
| 451 | return await tool_session.get_result_parsable_context(self) |
| 452 | args = json.loads(last_msg.arguments) |
| 453 | param = { |
| 454 | "code": args["code"], |
| 455 | } |
| 456 | result = await tool_session.call_tool("python", param) |
| 457 | result_str = result.content[0].text |
| 458 | |
| 459 | message = ResponseFunctionToolCallOutputItem( |
| 460 | id=f"mcpo_{random_uuid()}", |
| 461 | type="function_call_output", |
| 462 | call_id=last_msg.call_id, |
| 463 | output=result_str, |
| 464 | status="completed", |
| 465 | ) |
| 466 | |
| 467 | return [message] |
| 468 | |
| 469 | async def call_search_tool( |
| 470 | self, |
| 471 | tool_session: Union["ClientSession", Tool], |
| 472 | last_msg: ResponseFunctionToolCall, |
| 473 | ) -> list[ResponseInputOutputItem]: |
| 474 | self.called_tools.add("browser") |
| 475 | if isinstance(tool_session, Tool): |
| 476 | return await tool_session.get_result_parsable_context(self) |
| 477 | if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: |
| 478 | try: |
| 479 | args = json.loads(last_msg.arguments) |
| 480 | except json.JSONDecodeError as e: |
| 481 | return _create_json_parse_error_messages(last_msg, e) |
| 482 | else: |
| 483 | args = json.loads(last_msg.arguments) |
| 484 | result = await tool_session.call_tool("search", args) |
| 485 | result_str = result.content[0].text |
| 486 | |
| 487 | message = ResponseFunctionToolCallOutputItem( |
| 488 | id=f"fco_{random_uuid()}", |
| 489 | type="function_call_output", |
| 490 | call_id=last_msg.call_id, |
| 491 | output=result_str, |
| 492 | status="completed", |
| 493 | ) |
| 494 | |
| 495 | return [message] |
| 496 | |
| 497 | async def call_container_tool( |
| 498 | self, |
| 499 | tool_session: Union["ClientSession", Tool], |
| 500 | last_msg: ResponseFunctionToolCall, |
| 501 | ) -> list[ResponseInputOutputItem]: |
| 502 | """Call container tool. Expect this to be run in a stateful docker |
| 503 | with command line terminal. |
| 504 | The official container tool would at least |
| 505 | expect the following format: |
| 506 | - for tool name: exec |
| 507 | - args: |
| 508 | { |
| 509 | "cmd":List[str] "command to execute", |
| 510 | "workdir":optional[str] "current working directory", |
| 511 | "env":optional[object/dict] "environment variables", |
| 512 | "session_name":optional[str] "session name", |
| 513 | "timeout":optional[int] "timeout in seconds", |
| 514 | "user":optional[str] "user name", |
| 515 | } |
| 516 | """ |
| 517 | self.called_tools.add("container") |
| 518 | if isinstance(tool_session, Tool): |
| 519 | return await tool_session.get_result_parsable_context(self) |
| 520 | # tool_name = last_msg.recipient.split(".")[1].split(" ")[0] |
| 521 | if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: |
| 522 | try: |
| 523 | args = json.loads(last_msg.arguments) |
| 524 | except json.JSONDecodeError as e: |
| 525 | return _create_json_parse_error_messages(last_msg, e) |
| 526 | else: |
| 527 | args = json.loads(last_msg.arguments) |
| 528 | result = await tool_session.call_tool("exec", args) |
| 529 | result_str = result.content[0].text |
| 530 | |
| 531 | message = ResponseFunctionToolCallOutputItem( |
| 532 | id=f"fco_{random_uuid()}", |
| 533 | type="function_call_output", |
| 534 | call_id=last_msg.call_id, |
| 535 | output=result_str, |
| 536 | status="completed", |
| 537 | ) |
| 538 | |
| 539 | return [message] |
| 540 | |
| 541 | async def call_tool(self) -> list[ResponseInputOutputItem]: |
| 542 | if not self.response_messages: |
| 543 | return [] |
| 544 | last_msg = self.response_messages[-1] |
| 545 | # change this to a mcp_ function call |
| 546 | last_msg.id = f"{MCP_PREFIX}{random_uuid()}" |
| 547 | self.response_messages[-1] = last_msg |
| 548 | if last_msg.name == "code_interpreter": |
| 549 | return await self.call_python_tool(self._tool_sessions["python"], last_msg) |
| 550 | elif last_msg.name == "web_search_preview": |
| 551 | return await self.call_search_tool(self._tool_sessions["browser"], last_msg) |
| 552 | elif last_msg.name.startswith("container"): |
| 553 | return await self.call_container_tool( |
| 554 | self._tool_sessions["container"], last_msg |
| 555 | ) |
| 556 | return [] |
| 557 | |
| 558 | def make_response_output_items(self) -> list[ResponseOutputItem]: |
| 559 | response_messages = self.response_messages[self.num_init_messages :] |
| 560 | output_messages: list[ResponseOutputItem] = [] |
| 561 | for message in response_messages: |
| 562 | if not isinstance(message, ResponseFunctionToolCallOutputItem): |
| 563 | output_messages.append(message) |
| 564 | else: |
| 565 | if len(output_messages) == 0: |
| 566 | raise ValueError( |
| 567 | "Cannot have a FunctionToolCallOutput before FunctionToolCall." |
| 568 | ) |
| 569 | if isinstance(output_messages[-1], ResponseFunctionToolCall): |
| 570 | output_messages[-1] = McpCall( |
| 571 | id=f"{MCP_PREFIX}{random_uuid()}", |
| 572 | arguments=output_messages[-1].arguments, |
| 573 | name=output_messages[-1].name, |
| 574 | server_label=output_messages[-1].name, |
| 575 | type="mcp_call", |
| 576 | status="completed", |
| 577 | output=message.output, |
| 578 | ) |
| 579 | return output_messages |
| 580 | |
| 581 | def render_for_completion(self): |
| 582 | raise NotImplementedError("Should not be called.") |
| 583 | |
| 584 | async def init_tool_sessions( |
| 585 | self, |
| 586 | tool_server: ToolServer | None, |
| 587 | exit_stack: AsyncExitStack, |
| 588 | request_id: str, |
| 589 | mcp_tools: dict[str, Mcp], |
| 590 | ): |
| 591 | if tool_server: |
| 592 | initialized_session = False |
| 593 | for tool_name in self.available_tools: |
| 594 | if tool_name in self._tool_sessions: |
| 595 | continue |
| 596 | |
| 597 | tool_type = _map_tool_name_to_tool_type(tool_name) |
| 598 | headers = ( |
| 599 | mcp_tools[tool_type].headers if tool_type in mcp_tools else None |
| 600 | ) |
| 601 | tool_session = await exit_stack.enter_async_context( |
| 602 | tool_server.new_session(tool_name, request_id, headers) |
| 603 | ) |
| 604 | self._tool_sessions[tool_name] = tool_session |
| 605 | initialized_session = True |
| 606 | if initialized_session: |
| 607 | exit_stack.push_async_exit(self.cleanup_session) |
| 608 | |
| 609 | async def cleanup_session(self, *args, **kwargs) -> None: |
| 610 | """Can be used as coro to used in __aexit__.""" |
| 611 | |
| 612 | async def cleanup_tool_session(tool_session): |
| 613 | if not isinstance(tool_session, Tool): |
| 614 | logger.info( |
| 615 | "Cleaning up tool session for %s", tool_session._client_info |
| 616 | ) |
| 617 | with contextlib.suppress(Exception): |
| 618 | await tool_session.call_tool("cleanup_session", {}) |
| 619 | |
| 620 | await asyncio.gather( |
| 621 | *( |
| 622 | cleanup_tool_session(self._tool_sessions[tool]) |
| 623 | for tool in self.called_tools |
| 624 | ) |
| 625 | ) |
| 626 | |
| 627 | |
| 628 | class HarmonyContext(ConversationContext): |
| 629 | def __init__( |
| 630 | self, |
| 631 | messages: list, |
| 632 | available_tools: list[str], |
| 633 | function_tool_names: frozenset[str], |
| 634 | response_parser: Parser | None = None, |
| 635 | request: ResponsesRequest | None = None, |
| 636 | ): |
| 637 | from vllm.parser.harmony import HarmonyParser, Segment |
| 638 | |
| 639 | assert isinstance(response_parser, HarmonyParser) |
| 640 | |
| 641 | self._messages = messages |
| 642 | self.request = request |
| 643 | self.response_parser: HarmonyParser = response_parser |
| 644 | self.finish_reason: str | None = None |
| 645 | self.available_tools = available_tools |
| 646 | self.function_tool_names = function_tool_names |
| 647 | self._tool_sessions: dict[str, ClientSession | Tool] = {} |
| 648 | self.called_tools: set[str] = set() |
| 649 | |
| 650 | self.num_init_messages = len(messages) |
| 651 | self.num_prompt_tokens = 0 |
| 652 | self.num_output_tokens = 0 |
| 653 | self.num_cached_tokens = 0 |
| 654 | self.num_reasoning_tokens = 0 |
| 655 | self.num_tool_output_tokens = 0 |
| 656 | |
| 657 | self.last_append_segments: list[Segment] = [] |
| 658 | self.last_append_flush_status: bool = False |
| 659 | |
| 660 | # Turn tracking - replaces multiple individual tracking variables |
| 661 | self.current_turn_metrics = TurnMetrics() |
| 662 | # Track metrics for all turns |
| 663 | self.all_turn_metrics: list[TurnMetrics] = [] |
| 664 | self.is_first_turn = True |
| 665 | self.first_tok_of_message = True |
| 666 | self.kv_transfer_params: dict[str, Any] | None = None |
| 667 | self.ec_transfer_params: dict[str, Any] | None = None |
| 668 | |
| 669 | def append_output(self, output: RequestOutput) -> None: |
| 670 | if self.first_tok_of_message: |
| 671 | self.finish_reason = None |
| 672 | self._update_prefill_token_usage(output) |
| 673 | |
| 674 | output_token_ids = output.outputs[0].token_ids |
| 675 | result = self.response_parser.process_chunk(output_token_ids) |
| 676 | segments = result.segments |
| 677 | self.num_reasoning_tokens += result.reasoning_token_count |
| 678 | |
| 679 | self.first_tok_of_message = output.finished |
| 680 | self._update_decode_token_usage(output) |
| 681 | if output.kv_transfer_params is not None: |
| 682 | self.kv_transfer_params = output.kv_transfer_params |
| 683 | if output.ec_transfer_params is not None: |
| 684 | self.ec_transfer_params = output.ec_transfer_params |
| 685 | |
| 686 | if output.finished: |
| 687 | self.finish_reason = output.outputs[0].finish_reason |
| 688 | flushed_segments = self.response_parser.flush() |
| 689 | if flushed_segments: |
| 690 | segments.extend(flushed_segments) |
| 691 | self.last_append_flush_status = len(flushed_segments) > 0 |
| 692 | self.all_turn_metrics.append(self.current_turn_metrics.copy()) |
| 693 | self.current_turn_metrics.reset() |
| 694 | |
| 695 | self.last_append_segments = segments |
| 696 | self._messages.extend( |
| 697 | segment.completed_message |
| 698 | for segment in segments |
| 699 | if segment.completed_message is not None |
| 700 | ) |
| 701 | |
| 702 | def append_tool_output(self, output: list[Message]) -> None: |
| 703 | output_msgs = output |
| 704 | self._messages.extend(output_msgs) |
| 705 | |
| 706 | def _update_prefill_token_usage(self, output: RequestOutput) -> None: |
| 707 | """Update token usage statistics for the prefill phase of generation. |
| 708 | |
| 709 | The prefill phase processes the input prompt tokens. This method: |
| 710 | 1. Counts the prompt tokens for this turn |
| 711 | 2. Calculates tool output tokens for multi-turn conversations |
| 712 | 3. Updates cached token counts |
| 713 | 4. Tracks state for next turn calculations |
| 714 | |
| 715 | Tool output tokens are calculated as: |
| 716 | current_prompt_tokens - last_turn_prompt_tokens - |
| 717 | last_turn_output_tokens |
| 718 | This represents tokens added between turns (typically tool responses). |
| 719 | |
| 720 | Args: |
| 721 | output: The RequestOutput containing prompt token information |
| 722 | |
| 723 | """ |
| 724 | if output.prompt_token_ids is not None: |
| 725 | this_turn_input_tokens = len(output.prompt_token_ids) |
| 726 | else: |
| 727 | this_turn_input_tokens = 0 |
| 728 | logger.error("RequestOutput appended contains no prompt_token_ids.") |
| 729 | |
| 730 | # Update current turn input tokens |
| 731 | self.current_turn_metrics.input_tokens = this_turn_input_tokens |
| 732 | self.num_prompt_tokens += this_turn_input_tokens |
| 733 | |
| 734 | # Calculate tool tokens (except on first turn) |
| 735 | if self.is_first_turn: |
| 736 | self.is_first_turn = False |
| 737 | else: |
| 738 | previous_turn = self.all_turn_metrics[-1] |
| 739 | # start counting tool after first turn |
| 740 | # tool tokens = this turn prefill - last turn prefill - |
| 741 | # last turn decode |
| 742 | this_turn_tool_tokens = ( |
| 743 | self.current_turn_metrics.input_tokens |
| 744 | - previous_turn.input_tokens |
| 745 | - previous_turn.output_tokens |
| 746 | ) |
| 747 | |
| 748 | # Handle negative tool token counts (shouldn't happen in normal |
| 749 | # cases) |
| 750 | if this_turn_tool_tokens < 0: |
| 751 | logger.error( |
| 752 | "Negative tool output tokens calculated: %d " |
| 753 | "(current_input=%d, previous_input=%d, " |
| 754 | "previous_output=%d). Setting to 0.", |
| 755 | this_turn_tool_tokens, |
| 756 | self.current_turn_metrics.input_tokens, |
| 757 | previous_turn.input_tokens, |
| 758 | previous_turn.output_tokens, |
| 759 | ) |
| 760 | this_turn_tool_tokens = 0 |
| 761 | |
| 762 | self.num_tool_output_tokens += this_turn_tool_tokens |
| 763 | self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens |
| 764 | |
| 765 | # Update cached tokens |
| 766 | num_cached_token = output.num_cached_tokens |
| 767 | if num_cached_token is not None: |
| 768 | self.num_cached_tokens += num_cached_token |
| 769 | self.current_turn_metrics.cached_input_tokens = num_cached_token |
| 770 | |
| 771 | def _update_decode_token_usage(self, output: RequestOutput) -> int: |
| 772 | """Update token usage statistics for the decode phase of generation. |
| 773 | |
| 774 | The decode phase processes the generated output tokens. This method: |
| 775 | 1. Counts output tokens from all completion outputs |
| 776 | 2. Updates the total output token count |
| 777 | 3. Tracks tokens generated in the current turn |
| 778 | |
| 779 | In streaming mode, this is called for each token generated. |
| 780 | In non-streaming mode, this is called once with all output tokens. |
| 781 | |
| 782 | Args: |
| 783 | output: The RequestOutput containing generated token information |
| 784 | |
| 785 | Returns: |
| 786 | int: Number of output tokens processed in this call |
| 787 | |
| 788 | """ |
| 789 | updated_output_token_count = 0 |
| 790 | if output.outputs: |
| 791 | for completion_output in output.outputs: |
| 792 | # only keep last round |
| 793 | updated_output_token_count += len(completion_output.token_ids) |
| 794 | self.num_output_tokens += updated_output_token_count |
| 795 | self.current_turn_metrics.output_tokens += updated_output_token_count |
| 796 | return updated_output_token_count |
| 797 | |
| 798 | @property |
| 799 | def messages(self) -> list: |
| 800 | return self._messages |
| 801 | |
| 802 | def need_builtin_tool_call(self) -> bool: |
| 803 | last_msg = self.messages[-1] |
| 804 | recipient = last_msg.recipient |
| 805 | if recipient is None: |
| 806 | return False |
| 807 | if recipient.startswith("browser."): |
| 808 | return "browser" in self.available_tools |
| 809 | if recipient.startswith("python"): |
| 810 | return "python" in self.available_tools |
| 811 | if recipient.startswith("container."): |
| 812 | return "container" in self.available_tools |
| 813 | return False |
| 814 | |
| 815 | async def call_tool(self) -> list[Message]: |
| 816 | if not self.messages: |
| 817 | return [] |
| 818 | last_msg = self.messages[-1] |
| 819 | recipient = last_msg.recipient |
| 820 | if recipient is not None: |
| 821 | if recipient.startswith("browser."): |
| 822 | return await self.call_search_tool( |
| 823 | self._tool_sessions["browser"], last_msg |
| 824 | ) |
| 825 | elif recipient.startswith("python"): |
| 826 | return await self.call_python_tool( |
| 827 | self._tool_sessions["python"], last_msg |
| 828 | ) |
| 829 | elif recipient.startswith("container."): |
| 830 | return await self.call_container_tool( |
| 831 | self._tool_sessions["container"], last_msg |
| 832 | ) |
| 833 | raise ValueError("No tool call found") |
| 834 | |
| 835 | def render_for_completion(self) -> list[int]: |
| 836 | return render_for_completion(self.messages) |
| 837 | |
| 838 | async def call_search_tool( |
| 839 | self, tool_session: Union["ClientSession", Tool], last_msg: Message |
| 840 | ) -> list[Message]: |
| 841 | self.called_tools.add("browser") |
| 842 | if isinstance(tool_session, Tool): |
| 843 | return await tool_session.get_result(self) |
| 844 | tool_name = last_msg.recipient.split(".")[1] |
| 845 | if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: |
| 846 | try: |
| 847 | args = json.loads(last_msg.content[0].text) |
| 848 | except json.JSONDecodeError as e: |
| 849 | return _create_json_parse_error_messages(last_msg, e) |
| 850 | else: |
| 851 | args = json.loads(last_msg.content[0].text) |
| 852 | result = await tool_session.call_tool(tool_name, args) |
| 853 | result_str = result.content[0].text |
| 854 | content = TextContent(text=result_str) |
| 855 | author = Author(role=Role.TOOL, name=last_msg.recipient) |
| 856 | return [ |
| 857 | Message( |
| 858 | author=author, |
| 859 | content=[content], |
| 860 | recipient=Role.ASSISTANT, |
| 861 | channel=last_msg.channel, |
| 862 | ) |
| 863 | ] |
| 864 | |
| 865 | async def call_python_tool( |
| 866 | self, tool_session: Union["ClientSession", Tool], last_msg: Message |
| 867 | ) -> list[Message]: |
| 868 | self.called_tools.add("python") |
| 869 | if isinstance(tool_session, Tool): |
| 870 | return await tool_session.get_result(self) |
| 871 | param = { |
| 872 | "code": last_msg.content[0].text, |
| 873 | } |
| 874 | result = await tool_session.call_tool("python", param) |
| 875 | result_str = result.content[0].text |
| 876 | |
| 877 | content = TextContent(text=result_str) |
| 878 | author = Author(role=Role.TOOL, name="python") |
| 879 | |
| 880 | return [ |
| 881 | Message( |
| 882 | author=author, |
| 883 | content=[content], |
| 884 | channel=last_msg.channel, |
| 885 | recipient=Role.ASSISTANT, |
| 886 | ) |
| 887 | ] |
| 888 | |
| 889 | async def init_tool_sessions( |
| 890 | self, |
| 891 | tool_server: ToolServer | None, |
| 892 | exit_stack: AsyncExitStack, |
| 893 | request_id: str, |
| 894 | mcp_tools: dict[str, Mcp], |
| 895 | ): |
| 896 | if tool_server: |
| 897 | initialized_session = False |
| 898 | for tool_name in self.available_tools: |
| 899 | if tool_name not in self._tool_sessions: |
| 900 | tool_type = _map_tool_name_to_tool_type(tool_name) |
| 901 | headers = ( |
| 902 | mcp_tools[tool_type].headers if tool_type in mcp_tools else None |
| 903 | ) |
| 904 | tool_session = await exit_stack.enter_async_context( |
| 905 | tool_server.new_session(tool_name, request_id, headers) |
| 906 | ) |
| 907 | self._tool_sessions[tool_name] = tool_session |
| 908 | initialized_session = True |
| 909 | if initialized_session: |
| 910 | exit_stack.push_async_exit(self.cleanup_session) |
| 911 | |
| 912 | async def call_container_tool( |
| 913 | self, tool_session: Union["ClientSession", Tool], last_msg: Message |
| 914 | ) -> list[Message]: |
| 915 | """Call container tool. Expect this to be run in a stateful docker |
| 916 | with command line terminal. |
| 917 | The official container tool would at least |
| 918 | expect the following format: |
| 919 | - for tool name: exec |
| 920 | - args: |
| 921 | { |
| 922 | "cmd":List[str] "command to execute", |
| 923 | "workdir":optional[str] "current working directory", |
| 924 | "env":optional[object/dict] "environment variables", |
| 925 | "session_name":optional[str] "session name", |
| 926 | "timeout":optional[int] "timeout in seconds", |
| 927 | "user":optional[str] "user name", |
| 928 | } |
| 929 | """ |
| 930 | self.called_tools.add("container") |
| 931 | if isinstance(tool_session, Tool): |
| 932 | return await tool_session.get_result(self) |
| 933 | tool_name = last_msg.recipient.split(".")[1].split(" ")[0] |
| 934 | if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY: |
| 935 | try: |
| 936 | args = json.loads(last_msg.content[0].text) |
| 937 | except json.JSONDecodeError as e: |
| 938 | return _create_json_parse_error_messages(last_msg, e) |
| 939 | else: |
| 940 | args = json.loads(last_msg.content[0].text) |
| 941 | result = await tool_session.call_tool(tool_name, args) |
| 942 | result_str = result.content[0].text |
| 943 | content = TextContent(text=result_str) |
| 944 | author = Author(role=Role.TOOL, name=last_msg.recipient) |
| 945 | return [ |
| 946 | Message( |
| 947 | author=author, |
| 948 | content=[content], |
| 949 | recipient=Role.ASSISTANT, |
| 950 | channel=last_msg.channel, |
| 951 | ) |
| 952 | ] |
| 953 | |
| 954 | async def cleanup_session(self, *args, **kwargs) -> None: |
| 955 | """Can be used as coro to used in __aexit__.""" |
| 956 | |
| 957 | async def cleanup_tool_session(tool_session): |
| 958 | if not isinstance(tool_session, Tool): |
| 959 | logger.info( |
| 960 | "Cleaning up tool session for %s", tool_session._client_info |
| 961 | ) |
| 962 | with contextlib.suppress(Exception): |
| 963 | await tool_session.call_tool("cleanup_session", {}) |
| 964 | |
| 965 | await asyncio.gather( |
| 966 | *( |
| 967 | cleanup_tool_session(self._tool_sessions[tool]) |
| 968 | for tool in self.called_tools |
| 969 | ) |
| 970 | ) |