|
|
import ast
|
|
|
import json
|
|
|
import regex as re
|
|
|
from collections.abc import Sequence
|
|
|
from typing import List, Any
|
|
|
|
|
|
from transformers import PreTrainedTokenizerBase
|
|
|
from vllm.entrypoints.openai.protocol import (
|
|
|
ChatCompletionRequest,
|
|
|
ChatCompletionToolsParam,
|
|
|
DeltaFunctionCall,
|
|
|
DeltaMessage,
|
|
|
DeltaToolCall,
|
|
|
ExtractedToolCallInformation,
|
|
|
FunctionCall,
|
|
|
ToolCall,
|
|
|
)
|
|
|
from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
|
|
|
ToolParser,
|
|
|
ToolParserManager,
|
|
|
)
|
|
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
|
|
|
def _is_string_type(
|
|
|
tool_name: str, arg_name: str, tools: List[ChatCompletionToolsParam] | None
|
|
|
):
|
|
|
if tools is None:
|
|
|
return False
|
|
|
for tool in tools:
|
|
|
if tool.function.name == tool_name:
|
|
|
if tool.function.parameters is None:
|
|
|
return False
|
|
|
arg_type = (
|
|
|
tool.function.parameters.get("properties", {})
|
|
|
.get(arg_name, {})
|
|
|
.get("type", None)
|
|
|
)
|
|
|
return arg_type == "string"
|
|
|
logger.debug("No tool named '%s'.", tool_name)
|
|
|
return False
|
|
|
|
|
|
|
|
|
def _deserialize(value: str) -> Any:
|
|
|
try:
|
|
|
return json.loads(value)
|
|
|
except Exception:
|
|
|
pass
|
|
|
|
|
|
try:
|
|
|
return ast.literal_eval(value)
|
|
|
except Exception:
|
|
|
pass
|
|
|
return value
|
|
|
|
|
|
|
|
|
@ToolParserManager.register_module("telechat3")
|
|
|
class TeleChat3ModelToolParser(ToolParser):
|
|
|
"""
|
|
|
Tool call parser for TeleChat3-36B models.
|
|
|
Used when --enable-auto-tool-choice --tool-call-parser telechat3
|
|
|
"""
|
|
|
|
|
|
def __init__(self, tokenizer: PreTrainedTokenizerBase):
|
|
|
super().__init__(tokenizer)
|
|
|
|
|
|
|
|
|
|
|
|
self.current_tool_id: int = -1
|
|
|
|
|
|
self.tool_start_token = "<tool_call>"
|
|
|
self.tool_end_token = "</tool_call>"
|
|
|
|
|
|
self.func_detail_regex = re.compile(
|
|
|
r"<tool_call>(.*?)(<param_key>.*?)?</tool_call>", re.DOTALL
|
|
|
)
|
|
|
self.func_arg_regex = re.compile(
|
|
|
r"<param_key>(.*?)</param_key>(?:\\n|\s)*<param_value>(.*?)</param_value>",
|
|
|
re.DOTALL,
|
|
|
)
|
|
|
self._buffer = ""
|
|
|
|
|
|
def extract_tool_calls(self, model_output: str, request: ChatCompletionRequest):
|
|
|
|
|
|
matched_tool_calls = self.func_detail_regex.findall(model_output)
|
|
|
logger.debug("model_output: %s", model_output)
|
|
|
|
|
|
tool_calls = []
|
|
|
try:
|
|
|
for match in matched_tool_calls:
|
|
|
tc_name = match[0].strip()
|
|
|
arg_dict = {}
|
|
|
if len(match) > 1:
|
|
|
for key, value in self.func_arg_regex.findall(match[1]):
|
|
|
arg_key = key.strip()
|
|
|
arg_val = value.strip()
|
|
|
if not _is_string_type(tc_name, key, request.tools):
|
|
|
arg_val = _deserialize(arg_val)
|
|
|
logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val)
|
|
|
arg_dict[arg_key] = arg_val
|
|
|
tool_calls.append(
|
|
|
ToolCall(
|
|
|
type="function",
|
|
|
function=FunctionCall(
|
|
|
name=tc_name,
|
|
|
arguments=json.dumps(arg_dict, ensure_ascii=False),
|
|
|
),
|
|
|
)
|
|
|
)
|
|
|
except Exception:
|
|
|
logger.exception("Failed to extract tool call spec")
|
|
|
return ExtractedToolCallInformation(
|
|
|
tools_called=False, tool_calls=[], content=model_output
|
|
|
)
|
|
|
else:
|
|
|
if len(tool_calls) > 0:
|
|
|
content = model_output[: model_output.find(self.tool_start_token)]
|
|
|
return ExtractedToolCallInformation(
|
|
|
tools_called=True, tool_calls=tool_calls, content=content
|
|
|
)
|
|
|
return ExtractedToolCallInformation(
|
|
|
tools_called=False, tool_calls=[], content=model_output
|
|
|
)
|
|
|
|
|
|
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:
|
|
|
self._buffer += delta_text
|
|
|
cur_text = self._buffer
|
|
|
start_idx = cur_text.find(self.tool_start_token)
|
|
|
if start_idx == -1:
|
|
|
self._buffer = ""
|
|
|
return DeltaMessage(content=cur_text)
|
|
|
logger.debug("cur_text = %s", cur_text)
|
|
|
end_idx = cur_text.find(self.tool_end_token)
|
|
|
if end_idx != -1:
|
|
|
extracted_tool_calls = self.extract_tool_calls(
|
|
|
cur_text[: end_idx + len(self.tool_end_token)], request
|
|
|
)
|
|
|
if len(extracted_tool_calls.tool_calls) == 0:
|
|
|
logger.warning("Failed to extract any tool calls.")
|
|
|
return None
|
|
|
self.current_tool_id += 1
|
|
|
tool_call = extracted_tool_calls.tool_calls[0]
|
|
|
delta = DeltaMessage(
|
|
|
content=extracted_tool_calls.content,
|
|
|
tool_calls=[
|
|
|
DeltaToolCall(
|
|
|
index=self.current_tool_id,
|
|
|
id=tool_call.id,
|
|
|
type=tool_call.type,
|
|
|
function=DeltaFunctionCall(
|
|
|
name=tool_call.function.name,
|
|
|
arguments=tool_call.function.arguments,
|
|
|
),
|
|
|
)
|
|
|
],
|
|
|
)
|
|
|
self._buffer = cur_text[end_idx + len(self.tool_end_token) :]
|
|
|
return delta
|
|
|
self._buffer = cur_text[start_idx:]
|
|
|
return DeltaMessage(content=cur_text[:start_idx])
|
|
|
|
|
|
|
|
|
def register_tool_parser(): ...
|
|
|
|
