wehub-resource-sync
940 lines
35 KiB
Python
940 lines
35 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import asyncio
|
|
import contextlib
|
|
import copy
|
|
import json
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from contextlib import AsyncExitStack
|
|
from dataclasses import replace
|
|
from typing import TYPE_CHECKING, Any, Final, Union
|
|
|
|
from openai.types.responses import ResponseFunctionToolCall, ResponseOutputItem
|
|
from openai.types.responses.response_function_tool_call_output_item import (
|
|
ResponseFunctionToolCallOutputItem,
|
|
)
|
|
from openai.types.responses.response_output_item import McpCall
|
|
from openai.types.responses.response_output_message import ResponseOutputMessage
|
|
from openai.types.responses.response_output_text import ResponseOutputText
|
|
from openai.types.responses.tool import Mcp
|
|
from openai_harmony import Author, Message, Role, TextContent
|
|
|
|
from vllm import envs
|
|
from vllm.entrypoints.chat_utils import (
|
|
ChatTemplateContentFormatOption,
|
|
)
|
|
from vllm.entrypoints.mcp.tool import Tool
|
|
from vllm.entrypoints.mcp.tool_server import ToolServer
|
|
from vllm.entrypoints.openai.engine.protocol import (
|
|
FunctionCall,
|
|
)
|
|
from vllm.entrypoints.openai.parser.harmony_utils import render_for_completion
|
|
from vllm.entrypoints.openai.responses.protocol import (
|
|
ResponseInputOutputItem,
|
|
ResponseRawMessageAndToken,
|
|
ResponsesRequest,
|
|
)
|
|
from vllm.entrypoints.openai.responses.utils import (
|
|
build_response_output_items,
|
|
construct_tool_dicts,
|
|
)
|
|
from vllm.entrypoints.serve.utils.constants import MCP_PREFIX
|
|
from vllm.outputs import RequestOutput
|
|
from vllm.parser.abstract_parser import Parser
|
|
from vllm.tokenizers import TokenizerLike
|
|
from vllm.utils import random_uuid
|
|
|
|
if TYPE_CHECKING:
|
|
from mcp.client import ClientSession
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# This is currently needed as the tool type doesn't 1:1 match the
|
|
# tool namespace, which is what is used to look up the
|
|
# connection to the tool server
|
|
_TOOL_NAME_TO_TYPE_MAP = {
|
|
"browser": "web_search_preview",
|
|
"python": "code_interpreter",
|
|
"container": "container",
|
|
}
|
|
|
|
|
|
def _map_tool_name_to_tool_type(tool_name: str) -> str:
|
|
if tool_name not in _TOOL_NAME_TO_TYPE_MAP:
|
|
available_tools = ", ".join(_TOOL_NAME_TO_TYPE_MAP.keys())
|
|
raise ValueError(
|
|
f"Built-in tool name '{tool_name}' not defined in mapping. "
|
|
f"Available tools: {available_tools}"
|
|
)
|
|
return _TOOL_NAME_TO_TYPE_MAP[tool_name]
|
|
|
|
|
|
class TurnMetrics:
|
|
"""Tracks token and toolcall details for a single conversation turn."""
|
|
|
|
def __init__(
|
|
self,
|
|
input_tokens: int = 0,
|
|
output_tokens: int = 0,
|
|
cached_input_tokens: int = 0,
|
|
tool_output_tokens: int = 0,
|
|
) -> None:
|
|
self.input_tokens = input_tokens
|
|
self.output_tokens = output_tokens
|
|
self.cached_input_tokens = cached_input_tokens
|
|
self.tool_output_tokens = tool_output_tokens
|
|
|
|
def reset(self) -> None:
|
|
"""Reset counters for a new turn."""
|
|
self.input_tokens = 0
|
|
self.output_tokens = 0
|
|
self.cached_input_tokens = 0
|
|
self.tool_output_tokens = 0
|
|
|
|
def copy(self) -> "TurnMetrics":
|
|
"""Create a copy of this turn's token counts."""
|
|
return TurnMetrics(
|
|
self.input_tokens,
|
|
self.output_tokens,
|
|
self.cached_input_tokens,
|
|
self.tool_output_tokens,
|
|
)
|
|
|
|
|
|
class ConversationContext(ABC):
|
|
response_parser: Parser | None = None
|
|
|
|
@abstractmethod
|
|
def append_output(self, output: RequestOutput) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def append_tool_output(self, output) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def call_tool(self) -> list[Message]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def need_builtin_tool_call(self) -> bool:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def render_for_completion(self) -> list[int]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def init_tool_sessions(
|
|
self,
|
|
tool_server: ToolServer | None,
|
|
exit_stack: AsyncExitStack,
|
|
request_id: str,
|
|
mcp_tools: dict[str, Mcp],
|
|
) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def cleanup_session(self) -> None:
|
|
raise NotImplementedError("Should not be called.")
|
|
|
|
|
|
def _create_json_parse_error_messages(
|
|
last_msg: Message, e: json.JSONDecodeError
|
|
) -> list[Message]:
|
|
"""
|
|
Creates an error message when json parse failed.
|
|
"""
|
|
error_msg = (
|
|
f"Error parsing tool arguments as JSON: {str(e)}. "
|
|
"Please ensure the tool call arguments are valid JSON and try again."
|
|
)
|
|
content = TextContent(text=error_msg)
|
|
author = Author(role=Role.TOOL, name=last_msg.recipient)
|
|
return [
|
|
Message(
|
|
author=author,
|
|
content=[content],
|
|
recipient=Role.ASSISTANT,
|
|
channel=last_msg.channel,
|
|
)
|
|
]
|
|
|
|
|
|
class SimpleContext(ConversationContext):
|
|
"""This is a context that cannot handle MCP tool calls"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
response_parser: Parser | None = None,
|
|
parser_cls: type[Parser] | None = None,
|
|
tokenizer: TokenizerLike | None = None,
|
|
request: ResponsesRequest | None = None,
|
|
chat_template_kwargs: dict[str, Any] | None = None,
|
|
):
|
|
self.last_output = None
|
|
self.response_parser = response_parser or (
|
|
parser_cls(
|
|
tokenizer,
|
|
request.tools,
|
|
chat_template_kwargs=chat_template_kwargs,
|
|
)
|
|
if parser_cls is not None and tokenizer is not None and request is not None
|
|
else None
|
|
)
|
|
|
|
# Accumulated final output for streaming mode
|
|
self._accumulated_text: str = ""
|
|
self._accumulated_token_ids: list[int] = []
|
|
self._accumulated_logprobs: list = []
|
|
|
|
self.num_prompt_tokens = 0
|
|
self.num_output_tokens = 0
|
|
self.num_cached_tokens = 0
|
|
# todo num_reasoning_tokens is not implemented yet.
|
|
self.num_reasoning_tokens = 0
|
|
# not implemented yet for SimpleContext
|
|
self.all_turn_metrics: list[TurnMetrics] = []
|
|
|
|
self.input_messages: list[ResponseRawMessageAndToken] = []
|
|
self.kv_transfer_params: dict[str, Any] | None = None
|
|
self.ec_transfer_params: dict[str, Any] | None = None
|
|
|
|
def append_output(self, output) -> None:
|
|
self.last_output = output
|
|
if not isinstance(output, RequestOutput):
|
|
raise ValueError("SimpleContext only supports RequestOutput.")
|
|
self.num_prompt_tokens = len(output.prompt_token_ids or [])
|
|
self.num_cached_tokens = output.num_cached_tokens or 0
|
|
self.num_output_tokens += len(output.outputs[0].token_ids or [])
|
|
if output.kv_transfer_params is not None:
|
|
self.kv_transfer_params = output.kv_transfer_params
|
|
if output.ec_transfer_params is not None:
|
|
self.ec_transfer_params = output.ec_transfer_params
|
|
|
|
# Accumulate text, token_ids, and logprobs for streaming mode
|
|
delta_output = output.outputs[0]
|
|
self._accumulated_text += delta_output.text
|
|
self._accumulated_token_ids.extend(delta_output.token_ids)
|
|
if delta_output.logprobs is not None:
|
|
self._accumulated_logprobs.extend(delta_output.logprobs)
|
|
|
|
if len(self.input_messages) == 0:
|
|
output_prompt = output.prompt or ""
|
|
output_prompt_token_ids = output.prompt_token_ids or []
|
|
self.input_messages.append(
|
|
ResponseRawMessageAndToken(
|
|
message=output_prompt,
|
|
tokens=output_prompt_token_ids,
|
|
)
|
|
)
|
|
|
|
@property
|
|
def output_messages(self) -> list[ResponseRawMessageAndToken]:
|
|
"""Return consolidated output as a single message.
|
|
|
|
In streaming mode, text and tokens are accumulated across many deltas.
|
|
This property returns them as a single entry rather than one per delta.
|
|
"""
|
|
if not self._accumulated_text and not self._accumulated_token_ids:
|
|
return []
|
|
return [
|
|
ResponseRawMessageAndToken(
|
|
message=self._accumulated_text,
|
|
tokens=list(self._accumulated_token_ids),
|
|
)
|
|
]
|
|
|
|
@property
|
|
def final_output(self) -> RequestOutput | None:
|
|
"""Return the final output, with complete text/token_ids/logprobs."""
|
|
if self.last_output is not None and self.last_output.outputs:
|
|
assert isinstance(self.last_output, RequestOutput)
|
|
final_output = copy.copy(self.last_output)
|
|
# copy inner item to avoid modify last_output
|
|
final_output.outputs = [replace(item) for item in self.last_output.outputs]
|
|
final_output.outputs[0].text = self._accumulated_text
|
|
final_output.outputs[0].token_ids = tuple(self._accumulated_token_ids)
|
|
if self._accumulated_logprobs:
|
|
final_output.outputs[0].logprobs = self._accumulated_logprobs
|
|
return final_output
|
|
return self.last_output
|
|
|
|
def append_tool_output(self, output) -> None:
|
|
raise NotImplementedError("Should not be called.")
|
|
|
|
def need_builtin_tool_call(self) -> bool:
|
|
return False
|
|
|
|
async def call_tool(self) -> list[Message]:
|
|
raise NotImplementedError("Should not be called.")
|
|
|
|
def render_for_completion(self) -> list[int]:
|
|
raise NotImplementedError("Should not be called.")
|
|
|
|
async def init_tool_sessions(
|
|
self,
|
|
tool_server: ToolServer | None,
|
|
exit_stack: AsyncExitStack,
|
|
request_id: str,
|
|
mcp_tools: dict[str, Mcp],
|
|
) -> None:
|
|
pass
|
|
|
|
async def cleanup_session(self) -> None:
|
|
raise NotImplementedError("Should not be called.")
|
|
|
|
|
|
class ParsableContext(ConversationContext):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
response_messages: list[ResponseInputOutputItem],
|
|
tokenizer: TokenizerLike,
|
|
parser_cls: type[Parser] | None,
|
|
request: ResponsesRequest,
|
|
available_tools: list[str] | None,
|
|
chat_template: str | None,
|
|
chat_template_content_format: ChatTemplateContentFormatOption,
|
|
response_parser: Parser | None = None,
|
|
enable_auto_tools: bool = False,
|
|
):
|
|
self.num_prompt_tokens = 0
|
|
self.num_output_tokens = 0
|
|
self.num_cached_tokens = 0
|
|
self.num_reasoning_tokens = 0
|
|
# not implemented yet for ParsableContext
|
|
self.all_turn_metrics: list[TurnMetrics] = []
|
|
|
|
self.response_messages: list[ResponseInputOutputItem] = response_messages
|
|
self.num_init_messages = len(response_messages)
|
|
self.finish_reason: str | None = None
|
|
self.enable_auto_tools = enable_auto_tools
|
|
|
|
self.response_parser = response_parser or (
|
|
parser_cls(tokenizer, request.tools) if parser_cls is not None else None
|
|
)
|
|
self.parser_cls = parser_cls
|
|
self.request = request
|
|
|
|
self.available_tools = available_tools or []
|
|
self._tool_sessions: dict[str, ClientSession | Tool] = {}
|
|
self.called_tools: set[str] = set()
|
|
|
|
self.tool_dicts = construct_tool_dicts(request.tools, request.tool_choice)
|
|
self.chat_template = chat_template
|
|
self.chat_template_content_format: Final = chat_template_content_format
|
|
|
|
self.input_messages: list[ResponseRawMessageAndToken] = []
|
|
self.output_messages: list[ResponseRawMessageAndToken] = []
|
|
self._accumulated_token_ids: list[int] = []
|
|
self.kv_transfer_params: dict[str, Any] | None = None
|
|
self.ec_transfer_params: dict[str, Any] | None = None
|
|
|
|
def append_output(self, output: RequestOutput) -> None:
|
|
self.num_prompt_tokens = len(output.prompt_token_ids or [])
|
|
self.num_cached_tokens = output.num_cached_tokens or 0
|
|
self.num_output_tokens += len(output.outputs[0].token_ids or [])
|
|
if output.kv_transfer_params is not None:
|
|
self.kv_transfer_params = output.kv_transfer_params
|
|
|
|
if output.ec_transfer_params is not None:
|
|
self.ec_transfer_params = output.ec_transfer_params
|
|
|
|
completion = output.outputs[0]
|
|
self.finish_reason = completion.finish_reason
|
|
|
|
if self.response_parser is not None:
|
|
reasoning, content, tool_calls = self.response_parser.parse(
|
|
completion.text,
|
|
self.request,
|
|
enable_auto_tools=self.enable_auto_tools,
|
|
model_output_token_ids=completion.token_ids,
|
|
)
|
|
if not self.request.include_reasoning:
|
|
reasoning = None
|
|
self.response_messages.extend(
|
|
build_response_output_items(
|
|
reasoning=reasoning,
|
|
content=content,
|
|
tool_calls=tool_calls,
|
|
tools=self.request.tools,
|
|
)
|
|
)
|
|
elif completion.text:
|
|
self.response_messages.append(
|
|
ResponseOutputMessage(
|
|
type="message",
|
|
id=f"msg_{random_uuid()}",
|
|
status="completed",
|
|
role="assistant",
|
|
content=[
|
|
ResponseOutputText(
|
|
annotations=[],
|
|
type="output_text",
|
|
text=completion.text,
|
|
logprobs=None,
|
|
)
|
|
],
|
|
)
|
|
)
|
|
|
|
self._accumulated_token_ids.extend(completion.token_ids or [])
|
|
|
|
if self.request.enable_response_messages:
|
|
output_prompt = output.prompt or ""
|
|
output_prompt_token_ids = output.prompt_token_ids or []
|
|
if len(self.input_messages) == 0:
|
|
self.input_messages.append(
|
|
ResponseRawMessageAndToken(
|
|
message=output_prompt,
|
|
tokens=output_prompt_token_ids,
|
|
)
|
|
)
|
|
else:
|
|
self.output_messages.append(
|
|
ResponseRawMessageAndToken(
|
|
message=output_prompt,
|
|
tokens=output_prompt_token_ids,
|
|
)
|
|
)
|
|
self.output_messages.append(
|
|
ResponseRawMessageAndToken(
|
|
message=completion.text,
|
|
tokens=completion.token_ids,
|
|
)
|
|
)
|
|
|
|
def append_tool_output(self, output: list[ResponseInputOutputItem]) -> None:
|
|
self.response_messages.extend(output)
|
|
|
|
def need_builtin_tool_call(self) -> bool:
|
|
"""Return true if the last message is a builtin tool call
|
|
that the request has enabled."""
|
|
last_message = self.response_messages[-1]
|
|
if last_message.type != "function_call":
|
|
return False
|
|
if last_message.name in ("code_interpreter", "python"):
|
|
return "python" in self.available_tools
|
|
if last_message.name == "web_search_preview":
|
|
return "browser" in self.available_tools
|
|
if last_message.name.startswith("container"):
|
|
return "container" in self.available_tools
|
|
return False
|
|
|
|
async def call_python_tool(
|
|
self, tool_session: Union["ClientSession", Tool], last_msg: FunctionCall
|
|
) -> list[ResponseInputOutputItem]:
|
|
self.called_tools.add("python")
|
|
if isinstance(tool_session, Tool):
|
|
return await tool_session.get_result_parsable_context(self)
|
|
args = json.loads(last_msg.arguments)
|
|
param = {
|
|
"code": args["code"],
|
|
}
|
|
result = await tool_session.call_tool("python", param)
|
|
result_str = result.content[0].text
|
|
|
|
message = ResponseFunctionToolCallOutputItem(
|
|
id=f"mcpo_{random_uuid()}",
|
|
type="function_call_output",
|
|
call_id=f"call_{random_uuid()}",
|
|
output=result_str,
|
|
status="completed",
|
|
)
|
|
|
|
return [message]
|
|
|
|
async def call_search_tool(
|
|
self, tool_session: Union["ClientSession", Tool], last_msg: FunctionCall
|
|
) -> list[ResponseInputOutputItem]:
|
|
self.called_tools.add("browser")
|
|
if isinstance(tool_session, Tool):
|
|
return await tool_session.get_result_parsable_context(self)
|
|
if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
|
|
try:
|
|
args = json.loads(last_msg.arguments)
|
|
except json.JSONDecodeError as e:
|
|
return _create_json_parse_error_messages(last_msg, e)
|
|
else:
|
|
args = json.loads(last_msg.arguments)
|
|
result = await tool_session.call_tool("search", args)
|
|
result_str = result.content[0].text
|
|
|
|
message = ResponseFunctionToolCallOutputItem(
|
|
id=f"fco_{random_uuid()}",
|
|
type="function_call_output",
|
|
call_id=f"call_{random_uuid()}",
|
|
output=result_str,
|
|
status="completed",
|
|
)
|
|
|
|
return [message]
|
|
|
|
async def call_container_tool(
|
|
self, tool_session: Union["ClientSession", Tool], last_msg: Message
|
|
) -> list[Message]:
|
|
"""
|
|
Call container tool. Expect this to be run in a stateful docker
|
|
with command line terminal.
|
|
The official container tool would at least
|
|
expect the following format:
|
|
- for tool name: exec
|
|
- args:
|
|
{
|
|
"cmd":List[str] "command to execute",
|
|
"workdir":optional[str] "current working directory",
|
|
"env":optional[object/dict] "environment variables",
|
|
"session_name":optional[str] "session name",
|
|
"timeout":optional[int] "timeout in seconds",
|
|
"user":optional[str] "user name",
|
|
}
|
|
"""
|
|
self.called_tools.add("container")
|
|
if isinstance(tool_session, Tool):
|
|
return await tool_session.get_result_parsable_context(self)
|
|
# tool_name = last_msg.recipient.split(".")[1].split(" ")[0]
|
|
if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
|
|
try:
|
|
args = json.loads(last_msg.arguments)
|
|
except json.JSONDecodeError as e:
|
|
return _create_json_parse_error_messages(last_msg, e)
|
|
else:
|
|
args = json.loads(last_msg.arguments)
|
|
result = await tool_session.call_tool("exec", args)
|
|
result_str = result.content[0].text
|
|
|
|
message = ResponseFunctionToolCallOutputItem(
|
|
id=f"fco_{random_uuid()}",
|
|
type="function_call_output",
|
|
call_id=f"call_{random_uuid()}",
|
|
output=result_str,
|
|
status="completed",
|
|
)
|
|
|
|
return [message]
|
|
|
|
async def call_tool(self) -> list[ResponseInputOutputItem]:
|
|
if not self.response_messages:
|
|
return []
|
|
last_msg = self.response_messages[-1]
|
|
# change this to a mcp_ function call
|
|
last_msg.id = f"{MCP_PREFIX}{random_uuid()}"
|
|
self.response_messages[-1] = last_msg
|
|
if last_msg.name == "code_interpreter":
|
|
return await self.call_python_tool(self._tool_sessions["python"], last_msg)
|
|
elif last_msg.name == "web_search_preview":
|
|
return await self.call_search_tool(self._tool_sessions["browser"], last_msg)
|
|
elif last_msg.name.startswith("container"):
|
|
return await self.call_container_tool(
|
|
self._tool_sessions["container"], last_msg
|
|
)
|
|
return []
|
|
|
|
def make_response_output_items(self) -> list[ResponseOutputItem]:
|
|
response_messages = self.response_messages[self.num_init_messages :]
|
|
output_messages: list[ResponseOutputItem] = []
|
|
for message in response_messages:
|
|
if not isinstance(message, ResponseFunctionToolCallOutputItem):
|
|
output_messages.append(message)
|
|
else:
|
|
if len(output_messages) == 0:
|
|
raise ValueError(
|
|
"Cannot have a FunctionToolCallOutput before FunctionToolCall."
|
|
)
|
|
if isinstance(output_messages[-1], ResponseFunctionToolCall):
|
|
output_messages[-1] = McpCall(
|
|
id=f"{MCP_PREFIX}{random_uuid()}",
|
|
arguments=output_messages[-1].arguments,
|
|
name=output_messages[-1].name,
|
|
server_label=output_messages[-1].name,
|
|
type="mcp_call",
|
|
status="completed",
|
|
output=message.output,
|
|
)
|
|
return output_messages
|
|
|
|
def render_for_completion(self):
|
|
raise NotImplementedError("Should not be called.")
|
|
|
|
async def init_tool_sessions(
|
|
self,
|
|
tool_server: ToolServer | None,
|
|
exit_stack: AsyncExitStack,
|
|
request_id: str,
|
|
mcp_tools: dict[str, Mcp],
|
|
):
|
|
if tool_server:
|
|
for tool_name in self.available_tools:
|
|
if tool_name in self._tool_sessions:
|
|
continue
|
|
|
|
tool_type = _map_tool_name_to_tool_type(tool_name)
|
|
headers = (
|
|
mcp_tools[tool_type].headers if tool_type in mcp_tools else None
|
|
)
|
|
tool_session = await exit_stack.enter_async_context(
|
|
tool_server.new_session(tool_name, request_id, headers)
|
|
)
|
|
self._tool_sessions[tool_name] = tool_session
|
|
exit_stack.push_async_exit(self.cleanup_session)
|
|
|
|
async def cleanup_session(self, *args, **kwargs) -> None:
|
|
"""Can be used as coro to used in __aexit__"""
|
|
|
|
async def cleanup_tool_session(tool_session):
|
|
if not isinstance(tool_session, Tool):
|
|
logger.info(
|
|
"Cleaning up tool session for %s", tool_session._client_info
|
|
)
|
|
with contextlib.suppress(Exception):
|
|
await tool_session.call_tool("cleanup_session", {})
|
|
|
|
await asyncio.gather(
|
|
*(
|
|
cleanup_tool_session(self._tool_sessions[tool])
|
|
for tool in self.called_tools
|
|
)
|
|
)
|
|
|
|
|
|
class HarmonyContext(ConversationContext):
|
|
def __init__(
|
|
self,
|
|
messages: list,
|
|
available_tools: list[str],
|
|
function_tool_names: frozenset[str],
|
|
response_parser: Parser | None = None,
|
|
):
|
|
from vllm.parser.harmony import HarmonyParser, Segment
|
|
|
|
assert isinstance(response_parser, HarmonyParser)
|
|
|
|
self._messages = messages
|
|
self.response_parser: HarmonyParser = response_parser
|
|
self.finish_reason: str | None = None
|
|
self.available_tools = available_tools
|
|
self.function_tool_names = function_tool_names
|
|
self._tool_sessions: dict[str, ClientSession | Tool] = {}
|
|
self.called_tools: set[str] = set()
|
|
|
|
self.num_init_messages = len(messages)
|
|
self.num_prompt_tokens = 0
|
|
self.num_output_tokens = 0
|
|
self.num_cached_tokens = 0
|
|
self.num_reasoning_tokens = 0
|
|
self.num_tool_output_tokens = 0
|
|
|
|
self.last_append_segments: list[Segment] = []
|
|
self.last_append_flush_status: bool = False
|
|
|
|
# Turn tracking - replaces multiple individual tracking variables
|
|
self.current_turn_metrics = TurnMetrics()
|
|
# Track metrics for all turns
|
|
self.all_turn_metrics: list[TurnMetrics] = []
|
|
self.is_first_turn = True
|
|
self.first_tok_of_message = True
|
|
self.kv_transfer_params: dict[str, Any] | None = None
|
|
self.ec_transfer_params: dict[str, Any] | None = None
|
|
|
|
def append_output(self, output: RequestOutput) -> None:
|
|
if self.first_tok_of_message:
|
|
self.finish_reason = None
|
|
self._update_prefill_token_usage(output)
|
|
|
|
output_token_ids = output.outputs[0].token_ids
|
|
result = self.response_parser.process_chunk(output_token_ids)
|
|
segments = result.segments
|
|
self.num_reasoning_tokens += result.reasoning_token_count
|
|
|
|
self.first_tok_of_message = output.finished
|
|
self._update_decode_token_usage(output)
|
|
if output.kv_transfer_params is not None:
|
|
self.kv_transfer_params = output.kv_transfer_params
|
|
if output.ec_transfer_params is not None:
|
|
self.ec_transfer_params = output.ec_transfer_params
|
|
|
|
if output.finished:
|
|
self.finish_reason = output.outputs[0].finish_reason
|
|
flushed_segments = self.response_parser.flush()
|
|
if flushed_segments:
|
|
segments.extend(flushed_segments)
|
|
self.last_append_flush_status = len(flushed_segments) > 0
|
|
self.all_turn_metrics.append(self.current_turn_metrics.copy())
|
|
self.current_turn_metrics.reset()
|
|
|
|
self.last_append_segments = segments
|
|
self._messages.extend(
|
|
segment.completed_message
|
|
for segment in segments
|
|
if segment.completed_message is not None
|
|
)
|
|
|
|
def append_tool_output(self, output: list[Message]) -> None:
|
|
output_msgs = output
|
|
self._messages.extend(output_msgs)
|
|
|
|
def _update_prefill_token_usage(self, output: RequestOutput) -> None:
|
|
"""Update token usage statistics for the prefill phase of generation.
|
|
|
|
The prefill phase processes the input prompt tokens. This method:
|
|
1. Counts the prompt tokens for this turn
|
|
2. Calculates tool output tokens for multi-turn conversations
|
|
3. Updates cached token counts
|
|
4. Tracks state for next turn calculations
|
|
|
|
Tool output tokens are calculated as:
|
|
current_prompt_tokens - last_turn_prompt_tokens -
|
|
last_turn_output_tokens
|
|
This represents tokens added between turns (typically tool responses).
|
|
|
|
Args:
|
|
output: The RequestOutput containing prompt token information
|
|
"""
|
|
if output.prompt_token_ids is not None:
|
|
this_turn_input_tokens = len(output.prompt_token_ids)
|
|
else:
|
|
this_turn_input_tokens = 0
|
|
logger.error("RequestOutput appended contains no prompt_token_ids.")
|
|
|
|
# Update current turn input tokens
|
|
self.current_turn_metrics.input_tokens = this_turn_input_tokens
|
|
self.num_prompt_tokens += this_turn_input_tokens
|
|
|
|
# Calculate tool tokens (except on first turn)
|
|
if self.is_first_turn:
|
|
self.is_first_turn = False
|
|
else:
|
|
previous_turn = self.all_turn_metrics[-1]
|
|
# start counting tool after first turn
|
|
# tool tokens = this turn prefill - last turn prefill -
|
|
# last turn decode
|
|
this_turn_tool_tokens = (
|
|
self.current_turn_metrics.input_tokens
|
|
- previous_turn.input_tokens
|
|
- previous_turn.output_tokens
|
|
)
|
|
|
|
# Handle negative tool token counts (shouldn't happen in normal
|
|
# cases)
|
|
if this_turn_tool_tokens < 0:
|
|
logger.error(
|
|
"Negative tool output tokens calculated: %d "
|
|
"(current_input=%d, previous_input=%d, "
|
|
"previous_output=%d). Setting to 0.",
|
|
this_turn_tool_tokens,
|
|
self.current_turn_metrics.input_tokens,
|
|
previous_turn.input_tokens,
|
|
previous_turn.output_tokens,
|
|
)
|
|
this_turn_tool_tokens = 0
|
|
|
|
self.num_tool_output_tokens += this_turn_tool_tokens
|
|
self.current_turn_metrics.tool_output_tokens = this_turn_tool_tokens
|
|
|
|
# Update cached tokens
|
|
num_cached_token = output.num_cached_tokens
|
|
if num_cached_token is not None:
|
|
self.num_cached_tokens += num_cached_token
|
|
self.current_turn_metrics.cached_input_tokens = num_cached_token
|
|
|
|
def _update_decode_token_usage(self, output: RequestOutput) -> int:
|
|
"""Update token usage statistics for the decode phase of generation.
|
|
|
|
The decode phase processes the generated output tokens. This method:
|
|
1. Counts output tokens from all completion outputs
|
|
2. Updates the total output token count
|
|
3. Tracks tokens generated in the current turn
|
|
|
|
In streaming mode, this is called for each token generated.
|
|
In non-streaming mode, this is called once with all output tokens.
|
|
|
|
Args:
|
|
output: The RequestOutput containing generated token information
|
|
|
|
Returns:
|
|
int: Number of output tokens processed in this call
|
|
"""
|
|
updated_output_token_count = 0
|
|
if output.outputs:
|
|
for completion_output in output.outputs:
|
|
# only keep last round
|
|
updated_output_token_count += len(completion_output.token_ids)
|
|
self.num_output_tokens += updated_output_token_count
|
|
self.current_turn_metrics.output_tokens += updated_output_token_count
|
|
return updated_output_token_count
|
|
|
|
@property
|
|
def messages(self) -> list:
|
|
return self._messages
|
|
|
|
def need_builtin_tool_call(self) -> bool:
|
|
last_msg = self.messages[-1]
|
|
recipient = last_msg.recipient
|
|
if recipient is None:
|
|
return False
|
|
if recipient.startswith("browser."):
|
|
return "browser" in self.available_tools
|
|
if recipient.startswith("python"):
|
|
return "python" in self.available_tools
|
|
if recipient.startswith("container."):
|
|
return "container" in self.available_tools
|
|
return False
|
|
|
|
async def call_tool(self) -> list[Message]:
|
|
if not self.messages:
|
|
return []
|
|
last_msg = self.messages[-1]
|
|
recipient = last_msg.recipient
|
|
if recipient is not None:
|
|
if recipient.startswith("browser."):
|
|
return await self.call_search_tool(
|
|
self._tool_sessions["browser"], last_msg
|
|
)
|
|
elif recipient.startswith("python"):
|
|
return await self.call_python_tool(
|
|
self._tool_sessions["python"], last_msg
|
|
)
|
|
elif recipient.startswith("container."):
|
|
return await self.call_container_tool(
|
|
self._tool_sessions["container"], last_msg
|
|
)
|
|
raise ValueError("No tool call found")
|
|
|
|
def render_for_completion(self) -> list[int]:
|
|
return render_for_completion(self.messages)
|
|
|
|
async def call_search_tool(
|
|
self, tool_session: Union["ClientSession", Tool], last_msg: Message
|
|
) -> list[Message]:
|
|
self.called_tools.add("browser")
|
|
if isinstance(tool_session, Tool):
|
|
return await tool_session.get_result(self)
|
|
tool_name = last_msg.recipient.split(".")[1]
|
|
if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
|
|
try:
|
|
args = json.loads(last_msg.content[0].text)
|
|
except json.JSONDecodeError as e:
|
|
return _create_json_parse_error_messages(last_msg, e)
|
|
else:
|
|
args = json.loads(last_msg.content[0].text)
|
|
result = await tool_session.call_tool(tool_name, args)
|
|
result_str = result.content[0].text
|
|
content = TextContent(text=result_str)
|
|
author = Author(role=Role.TOOL, name=last_msg.recipient)
|
|
return [
|
|
Message(
|
|
author=author,
|
|
content=[content],
|
|
recipient=Role.ASSISTANT,
|
|
channel=last_msg.channel,
|
|
)
|
|
]
|
|
|
|
async def call_python_tool(
|
|
self, tool_session: Union["ClientSession", Tool], last_msg: Message
|
|
) -> list[Message]:
|
|
self.called_tools.add("python")
|
|
if isinstance(tool_session, Tool):
|
|
return await tool_session.get_result(self)
|
|
param = {
|
|
"code": last_msg.content[0].text,
|
|
}
|
|
result = await tool_session.call_tool("python", param)
|
|
result_str = result.content[0].text
|
|
|
|
content = TextContent(text=result_str)
|
|
author = Author(role=Role.TOOL, name="python")
|
|
|
|
return [
|
|
Message(
|
|
author=author,
|
|
content=[content],
|
|
channel=last_msg.channel,
|
|
recipient=Role.ASSISTANT,
|
|
)
|
|
]
|
|
|
|
async def init_tool_sessions(
|
|
self,
|
|
tool_server: ToolServer | None,
|
|
exit_stack: AsyncExitStack,
|
|
request_id: str,
|
|
mcp_tools: dict[str, Mcp],
|
|
):
|
|
if tool_server:
|
|
for tool_name in self.available_tools:
|
|
if tool_name not in self._tool_sessions:
|
|
tool_type = _map_tool_name_to_tool_type(tool_name)
|
|
headers = (
|
|
mcp_tools[tool_type].headers if tool_type in mcp_tools else None
|
|
)
|
|
tool_session = await exit_stack.enter_async_context(
|
|
tool_server.new_session(tool_name, request_id, headers)
|
|
)
|
|
self._tool_sessions[tool_name] = tool_session
|
|
exit_stack.push_async_exit(self.cleanup_session)
|
|
|
|
async def call_container_tool(
|
|
self, tool_session: Union["ClientSession", Tool], last_msg: Message
|
|
) -> list[Message]:
|
|
"""
|
|
Call container tool. Expect this to be run in a stateful docker
|
|
with command line terminal.
|
|
The official container tool would at least
|
|
expect the following format:
|
|
- for tool name: exec
|
|
- args:
|
|
{
|
|
"cmd":List[str] "command to execute",
|
|
"workdir":optional[str] "current working directory",
|
|
"env":optional[object/dict] "environment variables",
|
|
"session_name":optional[str] "session name",
|
|
"timeout":optional[int] "timeout in seconds",
|
|
"user":optional[str] "user name",
|
|
}
|
|
"""
|
|
self.called_tools.add("container")
|
|
if isinstance(tool_session, Tool):
|
|
return await tool_session.get_result(self)
|
|
tool_name = last_msg.recipient.split(".")[1].split(" ")[0]
|
|
if envs.VLLM_TOOL_JSON_ERROR_AUTOMATIC_RETRY:
|
|
try:
|
|
args = json.loads(last_msg.content[0].text)
|
|
except json.JSONDecodeError as e:
|
|
return _create_json_parse_error_messages(last_msg, e)
|
|
else:
|
|
args = json.loads(last_msg.content[0].text)
|
|
result = await tool_session.call_tool(tool_name, args)
|
|
result_str = result.content[0].text
|
|
content = TextContent(text=result_str)
|
|
author = Author(role=Role.TOOL, name=last_msg.recipient)
|
|
return [
|
|
Message(
|
|
author=author,
|
|
content=[content],
|
|
recipient=Role.ASSISTANT,
|
|
channel=last_msg.channel,
|
|
)
|
|
]
|
|
|
|
async def cleanup_session(self, *args, **kwargs) -> None:
|
|
"""Can be used as coro to used in __aexit__"""
|
|
|
|
async def cleanup_tool_session(tool_session):
|
|
if not isinstance(tool_session, Tool):
|
|
logger.info(
|
|
"Cleaning up tool session for %s", tool_session._client_info
|
|
)
|
|
with contextlib.suppress(Exception):
|
|
await tool_session.call_tool("cleanup_session", {})
|
|
|
|
await asyncio.gather(
|
|
*(
|
|
cleanup_tool_session(self._tool_sessions[tool])
|
|
for tool in self.called_tools
|
|
)
|
|
)
|