304 lines
12 KiB
Python
304 lines
12 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
|||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|||
|
|
|
|||
|
|
import contextlib
|
|||
|
|
import json
|
|||
|
|
from collections.abc import Sequence
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import regex as re
|
|||
|
|
|
|||
|
|
from vllm.entrypoints.openai.protocol import (
|
|||
|
|
ChatCompletionRequest,
|
|||
|
|
DeltaFunctionCall,
|
|||
|
|
DeltaMessage,
|
|||
|
|
DeltaToolCall,
|
|||
|
|
ExtractedToolCallInformation,
|
|||
|
|
FunctionCall,
|
|||
|
|
ToolCall,
|
|||
|
|
)
|
|||
|
|
from vllm.logger import init_logger
|
|||
|
|
from vllm.tokenizers import TokenizerLike
|
|||
|
|
from vllm.tool_parsers.abstract_tool_parser import (
|
|||
|
|
ToolParser,
|
|||
|
|
)
|
|||
|
|
from vllm.utils import random_uuid
|
|||
|
|
|
|||
|
|
logger = init_logger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Step3ToolParser(ToolParser):
|
|||
|
|
"""
|
|||
|
|
Tool parser for a model that uses a specific XML-like format for tool calls.
|
|||
|
|
This version uses a robust, stateful, cursor-based streaming parser and
|
|||
|
|
consolidates tool arguments into a single message.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
TOOL_CALLS_BEGIN = "<｜tool_calls_begin｜>"
|
|||
|
|
TOOL_CALLS_END = "<｜tool_calls_end｜>"
|
|||
|
|
TOOL_CALL_BEGIN = "<｜tool_call_begin｜>"
|
|||
|
|
TOOL_CALL_END = "<｜tool_call_end｜>"
|
|||
|
|
TOOL_SEP = "<｜tool_sep｜>"
|
|||
|
|
SPECIAL_TOKENS = [TOOL_CALLS_BEGIN, TOOL_CALLS_END, TOOL_CALL_BEGIN, TOOL_CALL_END]
|
|||
|
|
|
|||
|
|
def __init__(self, tokenizer: TokenizerLike):
|
|||
|
|
super().__init__(tokenizer)
|
|||
|
|
self.position = 0
|
|||
|
|
# Explicit state flags for robust streaming
|
|||
|
|
self.tool_block_started = False
|
|||
|
|
self.tool_block_finished = False
|
|||
|
|
|
|||
|
|
def adjust_request(self, request: ChatCompletionRequest) -> ChatCompletionRequest:
|
|||
|
|
request = super().adjust_request(request)
|
|||
|
|
if request.tools and request.tool_choice != "none":
|
|||
|
|
request.skip_special_tokens = False
|
|||
|
|
return request
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _parse_steptml_invoke(
|
|||
|
|
action_text: str,
|
|||
|
|
) -> tuple[str | None, dict[str, str] | None]:
|
|||
|
|
func_name_match = re.search(r'<steptml:invoke name="([^"]+)">', action_text)
|
|||
|
|
if not func_name_match:
|
|||
|
|
return None, None
|
|||
|
|
func_name = func_name_match.group(1)
|
|||
|
|
|
|||
|
|
params: dict[str, str] = {}
|
|||
|
|
param_matches = re.findall(
|
|||
|
|
r'<steptml:parameter name="([^"]+)">([^<]*)</steptml:parameter>',
|
|||
|
|
action_text,
|
|||
|
|
)
|
|||
|
|
for name, value in param_matches:
|
|||
|
|
params[name] = value.strip()
|
|||
|
|
return func_name, params
|
|||
|
|
|
|||
|
|
def _cast_arguments(
|
|||
|
|
self,
|
|||
|
|
func_name: str,
|
|||
|
|
params: dict[str, Any],
|
|||
|
|
request: ChatCompletionRequest,
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
for tool in request.tools or []:
|
|||
|
|
if tool.function.name == func_name:
|
|||
|
|
schema = tool.function.parameters or {}
|
|||
|
|
properties = schema.get("properties", {})
|
|||
|
|
for key, value in params.items():
|
|||
|
|
if not isinstance(value, str):
|
|||
|
|
continue
|
|||
|
|
prop = properties.get(key, {})
|
|||
|
|
typ = prop.get("type")
|
|||
|
|
if typ == "string":
|
|||
|
|
params[key] = value.strip()
|
|||
|
|
elif typ == "integer":
|
|||
|
|
with contextlib.suppress(ValueError):
|
|||
|
|
params[key] = int(value)
|
|||
|
|
elif typ == "number":
|
|||
|
|
with contextlib.suppress(ValueError):
|
|||
|
|
params[key] = float(value)
|
|||
|
|
elif typ == "boolean":
|
|||
|
|
lower_val = value.lower()
|
|||
|
|
params[key] = (
|
|||
|
|
lower_val == "true"
|
|||
|
|
if lower_val in ("true", "false")
|
|||
|
|
else value
|
|||
|
|
)
|
|||
|
|
elif typ == "null":
|
|||
|
|
params[key] = None if value.lower() == "null" else value
|
|||
|
|
break
|
|||
|
|
return params
|
|||
|
|
|
|||
|
|
def extract_tool_calls_streaming(
|
|||
|
|
self,
|
|||
|
|
previous_text: str,
|
|||
|
|
current_text: str,
|
|||
|
|
delta_text: str,
|
|||
|
|
previous_token_ids: Sequence[int],
|
|||
|
|
current_token_ids: Sequence[int],
|
|||
|
|
delta_token_ids: Sequence[int],
|
|||
|
|
request: ChatCompletionRequest,
|
|||
|
|
) -> DeltaMessage | None:
|
|||
|
|
# The main loop processes the stream from the last known position.
|
|||
|
|
while True:
|
|||
|
|
if self.position >= len(current_text):
|
|||
|
|
return None # We've processed the entire stream.
|
|||
|
|
|
|||
|
|
unprocessed_text = current_text[self.position :]
|
|||
|
|
|
|||
|
|
# STATE: After all tools are done, all subsequent text is content.
|
|||
|
|
if self.tool_block_finished:
|
|||
|
|
self.position = len(current_text)
|
|||
|
|
return DeltaMessage(content=unprocessed_text)
|
|||
|
|
|
|||
|
|
# STATE: Before the tool block has started.
|
|||
|
|
if not self.tool_block_started:
|
|||
|
|
if unprocessed_text.startswith(self.TOOL_CALLS_BEGIN):
|
|||
|
|
self.position += len(self.TOOL_CALLS_BEGIN)
|
|||
|
|
self.tool_block_started = True
|
|||
|
|
continue # Token consumed, re-loop.
|
|||
|
|
|
|||
|
|
start_pos = unprocessed_text.find(self.TOOL_CALLS_BEGIN)
|
|||
|
|
if start_pos == -1:
|
|||
|
|
if (
|
|||
|
|
self.TOOL_CALLS_BEGIN.startswith(unprocessed_text.strip())
|
|||
|
|
and unprocessed_text
|
|||
|
|
):
|
|||
|
|
return None # It's a prefix, wait.
|
|||
|
|
self.position = len(current_text)
|
|||
|
|
return DeltaMessage(content=unprocessed_text)
|
|||
|
|
else:
|
|||
|
|
content = unprocessed_text[:start_pos]
|
|||
|
|
self.position += len(content)
|
|||
|
|
return DeltaMessage(content=content)
|
|||
|
|
|
|||
|
|
# STATE: Inside the main tool block.
|
|||
|
|
offset = len(unprocessed_text) - len(unprocessed_text.lstrip())
|
|||
|
|
unprocessed_text = unprocessed_text.lstrip()
|
|||
|
|
self.position += offset
|
|||
|
|
|
|||
|
|
if unprocessed_text.startswith(self.TOOL_CALLS_END):
|
|||
|
|
self.position += len(self.TOOL_CALLS_END)
|
|||
|
|
self.tool_block_finished = True
|
|||
|
|
self.current_tool_id = -1
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# Check if we are between tool calls.
|
|||
|
|
tool_finished = self.current_tool_id != -1 and self.prev_tool_call_arr[
|
|||
|
|
self.current_tool_id
|
|||
|
|
].get("finished")
|
|||
|
|
if self.current_tool_id == -1 or tool_finished:
|
|||
|
|
if unprocessed_text.startswith(self.TOOL_CALL_BEGIN):
|
|||
|
|
self.position += len(self.TOOL_CALL_BEGIN)
|
|||
|
|
if self.current_tool_id == -1:
|
|||
|
|
self.current_tool_id = 0
|
|||
|
|
else:
|
|||
|
|
self.current_tool_id += 1
|
|||
|
|
self.current_tool_name_sent = False
|
|||
|
|
while len(self.prev_tool_call_arr) <= self.current_tool_id:
|
|||
|
|
self.prev_tool_call_arr.append({})
|
|||
|
|
self.prev_tool_call_arr[self.current_tool_id]["finished"] = False
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
if self.TOOL_CALL_BEGIN.startswith(unprocessed_text):
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# STATE: Parsing an active tool call.
|
|||
|
|
if self.current_tool_id != -1 and not self.prev_tool_call_arr[
|
|||
|
|
self.current_tool_id
|
|||
|
|
].get("finished", False):
|
|||
|
|
end_tool_pos = unprocessed_text.find(self.TOOL_CALL_END)
|
|||
|
|
if end_tool_pos == -1:
|
|||
|
|
tool_body = unprocessed_text
|
|||
|
|
else:
|
|||
|
|
tool_body = unprocessed_text[:end_tool_pos]
|
|||
|
|
|
|||
|
|
if end_tool_pos == -1 and self.TOOL_CALL_END.startswith(tool_body):
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
function_name, arguments = self._parse_steptml_invoke(tool_body)
|
|||
|
|
if not function_name:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
tool_call_arr = {"name": function_name, "parameters": arguments or {}}
|
|||
|
|
|
|||
|
|
# Send the function name as soon as it's parsed.
|
|||
|
|
if not self.current_tool_name_sent:
|
|||
|
|
self.current_tool_name_sent = True
|
|||
|
|
self.prev_tool_call_arr[self.current_tool_id].update(tool_call_arr)
|
|||
|
|
return DeltaMessage(
|
|||
|
|
tool_calls=[
|
|||
|
|
DeltaToolCall(
|
|||
|
|
index=self.current_tool_id,
|
|||
|
|
type="function",
|
|||
|
|
id=f"chatcmpl-tool-{random_uuid()}",
|
|||
|
|
function=DeltaFunctionCall(name=function_name),
|
|||
|
|
)
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Update our internal state with the latest parsed arguments.
|
|||
|
|
self.prev_tool_call_arr[self.current_tool_id].update( # noqa: E501
|
|||
|
|
tool_call_arr
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Only send arguments when the tool call is complete.
|
|||
|
|
if end_tool_pos != -1:
|
|||
|
|
self.position += end_tool_pos + len(self.TOOL_CALL_END)
|
|||
|
|
self.prev_tool_call_arr[self.current_tool_id]["finished"] = True
|
|||
|
|
|
|||
|
|
final_args = self._cast_arguments(
|
|||
|
|
function_name,
|
|||
|
|
tool_call_arr.get("parameters", {}), # type: ignore
|
|||
|
|
request,
|
|||
|
|
)
|
|||
|
|
if final_args:
|
|||
|
|
final_args_json = json.dumps(final_args, ensure_ascii=False)
|
|||
|
|
return DeltaMessage(
|
|||
|
|
tool_calls=[
|
|||
|
|
DeltaToolCall(
|
|||
|
|
index=self.current_tool_id,
|
|||
|
|
function=DeltaFunctionCall(
|
|||
|
|
arguments=final_args_json
|
|||
|
|
),
|
|||
|
|
)
|
|||
|
|
]
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# If tool is not finished, return None to wait for more tokens.
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def extract_tool_calls(
|
|||
|
|
self,
|
|||
|
|
model_output: str,
|
|||
|
|
request: ChatCompletionRequest,
|
|||
|
|
) -> ExtractedToolCallInformation:
|
|||
|
|
if self.TOOL_CALLS_BEGIN not in model_output:
|
|||
|
|
return ExtractedToolCallInformation(
|
|||
|
|
tools_called=False, tool_calls=[], content=model_output
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
pre_text, rest = model_output.split(self.TOOL_CALLS_BEGIN, 1)
|
|||
|
|
if self.TOOL_CALLS_END not in rest:
|
|||
|
|
return ExtractedToolCallInformation(
|
|||
|
|
tools_called=False, tool_calls=[], content=model_output
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
tool_block, post_text = rest.split(self.TOOL_CALLS_END, 1)
|
|||
|
|
content = (pre_text + post_text).strip()
|
|||
|
|
|
|||
|
|
tool_calls: list[ToolCall] = []
|
|||
|
|
call_parts = tool_block.split(self.TOOL_CALL_BEGIN)
|
|||
|
|
|
|||
|
|
for part in call_parts:
|
|||
|
|
if not part or self.TOOL_CALL_END not in part:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
call_content = part.split(self.TOOL_CALL_END, 1)[0]
|
|||
|
|
if self.TOOL_SEP not in call_content:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
type_part, invoke_part = call_content.split(self.TOOL_SEP, 1)
|
|||
|
|
if type_part.strip() != "function":
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
function_name, params_dict = self._parse_steptml_invoke(invoke_part)
|
|||
|
|
|
|||
|
|
if function_name and params_dict is not None:
|
|||
|
|
params_dict = self._cast_arguments(function_name, params_dict, request)
|
|||
|
|
params_str = json.dumps(params_dict, ensure_ascii=False)
|
|||
|
|
tool_calls.append(
|
|||
|
|
ToolCall(
|
|||
|
|
function=FunctionCall(name=function_name, arguments=params_str)
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
if tool_calls:
|
|||
|
|
return ExtractedToolCallInformation(
|
|||
|
|
tools_called=True,
|
|||
|
|
tool_calls=tool_calls,
|
|||
|
|
content=content if content else None,
|
|||
|
|
)
|
|||
|
|
return ExtractedToolCallInformation(
|
|||
|
|
tools_called=False, tool_calls=[], content=model_output
|
|||
|
|
)
|