wehub-resource-sync
329 lines
10 KiB
Python
329 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
||
import json
|
||
from collections.abc import Sequence
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
|
||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||
ChatCompletionRequest,
|
||
ChatCompletionToolsParam,
|
||
)
|
||
from vllm.tool_parsers.rust_tool_parser import RustToolParser
|
||
|
||
# The PyO3 extension is an optional build artifact; skip when absent.
|
||
_rust_tool_parser = pytest.importorskip("vllm._rust_tool_parser")
|
||
|
||
MOCK_TOKENIZER = MagicMock()
|
||
MOCK_TOKENIZER.get_vocab.return_value = {}
|
||
|
||
TC_START = "<｜DSML｜tool_calls>"
|
||
TC_END = "</｜DSML｜tool_calls>"
|
||
INV_START = '<｜DSML｜invoke name="'
|
||
INV_END = "</｜DSML｜invoke>"
|
||
PARAM_START = '<｜DSML｜parameter name="'
|
||
PARAM_END = "</｜DSML｜parameter>"
|
||
|
||
|
||
class DeepSeekV4RustToolParser(RustToolParser):
|
||
rust_parser_name = "DeepSeekV4ToolParser"
|
||
tool_call_start_token = TC_START
|
||
|
||
|
||
class KimiK2RustToolParser(RustToolParser):
|
||
rust_parser_name = "KimiK2ToolParser"
|
||
tool_call_start_token = "<|tool_calls_section_begin|>"
|
||
|
||
|
||
def sample_tools() -> list[ChatCompletionToolsParam]:
|
||
return [
|
||
ChatCompletionToolsParam(
|
||
type="function",
|
||
function={
|
||
"name": "get_weather",
|
||
"description": "Get weather for a location",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"location": {"type": "string"},
|
||
"date": {"type": "string"},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
ChatCompletionToolsParam(
|
||
type="function",
|
||
function={
|
||
"name": "add",
|
||
"description": "Add two integers",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"x": {"type": "integer"},
|
||
"y": {"type": "integer"},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
]
|
||
|
||
|
||
EXPECTED_CALLS = [
|
||
("get_weather", {"location": "SF", "date": "2024-01-16"}),
|
||
("add", {"x": 3, "y": 5}),
|
||
]
|
||
|
||
|
||
def build_invoke(
|
||
function_name: str,
|
||
params: Sequence[tuple[str, str, bool]],
|
||
) -> str:
|
||
param_text = "\n".join(
|
||
f'{PARAM_START}{name}" string="{str(is_string).lower()}">{value}{PARAM_END}'
|
||
for name, value, is_string in params
|
||
)
|
||
return f'{INV_START}{function_name}">\n{param_text}\n{INV_END}\n'
|
||
|
||
|
||
def build_tool_call() -> str:
|
||
weather = build_invoke(
|
||
"get_weather",
|
||
[
|
||
("location", "SF", True),
|
||
("date", "2024-01-16", True),
|
||
],
|
||
)
|
||
add = build_invoke(
|
||
"add",
|
||
[
|
||
("x", "3", False),
|
||
("y", "5", False),
|
||
],
|
||
)
|
||
return f"{TC_START}\n{weather}{add}{TC_END}"
|
||
|
||
|
||
def parse_streaming(
|
||
parser: DeepSeekV4RustToolParser,
|
||
text: str,
|
||
chunk_size: int,
|
||
) -> list:
|
||
deltas = []
|
||
previous_text = ""
|
||
for start in range(0, len(text), chunk_size):
|
||
delta_text = text[start : start + chunk_size]
|
||
current_text = previous_text + delta_text
|
||
delta = parser.extract_tool_calls_streaming(
|
||
previous_text=previous_text,
|
||
current_text=current_text,
|
||
delta_text=delta_text,
|
||
previous_token_ids=[],
|
||
current_token_ids=[],
|
||
delta_token_ids=[1],
|
||
request=MagicMock(),
|
||
)
|
||
previous_text = current_text
|
||
if delta is not None:
|
||
deltas.append(delta)
|
||
|
||
delta = parser.extract_tool_calls_streaming(
|
||
previous_text=previous_text,
|
||
current_text=previous_text,
|
||
delta_text="",
|
||
previous_token_ids=[],
|
||
current_token_ids=[],
|
||
delta_token_ids=[2],
|
||
request=MagicMock(),
|
||
)
|
||
if delta is not None:
|
||
deltas.append(delta)
|
||
|
||
return deltas
|
||
|
||
|
||
def collect_streamed_arguments(deltas: Sequence, tool_index: int = 0) -> str:
|
||
return "".join(
|
||
tool_call.function.arguments
|
||
for delta in deltas
|
||
for tool_call in delta.tool_calls or []
|
||
if (
|
||
tool_call.index == tool_index
|
||
and tool_call.function is not None
|
||
and tool_call.function.arguments is not None
|
||
)
|
||
)
|
||
|
||
|
||
def test_rust_tool_parser_extension_typed_api() -> None:
|
||
tools = [
|
||
_rust_tool_parser.Tool(
|
||
tool.function.name,
|
||
tool.function.description,
|
||
tool.function.parameters,
|
||
None,
|
||
)
|
||
for tool in sample_tools()
|
||
]
|
||
parser = _rust_tool_parser.ToolParser("DeepSeekV4ToolParser", tools)
|
||
output = _rust_tool_parser.ToolParserOutput()
|
||
|
||
parser.parse_into(build_tool_call(), output)
|
||
output.append(parser.finish())
|
||
output = output.coalesce()
|
||
|
||
assert parser.preserve_special_tokens()
|
||
assert output.normal_text == ""
|
||
assert len(output.calls) == 2
|
||
for call, (name, arguments) in zip(output.calls, EXPECTED_CALLS):
|
||
assert call.name == name
|
||
assert json.loads(call.arguments) == arguments
|
||
|
||
|
||
def test_rust_tool_parser_adapter_extracts_complete_output() -> None:
|
||
tools = sample_tools()
|
||
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools)
|
||
|
||
result = parser.extract_tool_calls(
|
||
"Let me create it. " + build_tool_call(),
|
||
ChatCompletionRequest(messages=[], model="m", tools=tools),
|
||
)
|
||
|
||
assert result.tools_called
|
||
assert result.content == "Let me create it. "
|
||
assert len(result.tool_calls) == 2
|
||
for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS):
|
||
assert tool_call.function.name == name
|
||
assert json.loads(tool_call.function.arguments) == arguments
|
||
|
||
|
||
def test_rust_tool_parser_adapter_streaming_handles_multiple_calls() -> None:
|
||
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())
|
||
|
||
deltas = parse_streaming(parser, build_tool_call(), chunk_size=5)
|
||
|
||
names = [
|
||
tool_call.function.name
|
||
for delta in deltas
|
||
for tool_call in delta.tool_calls or []
|
||
if tool_call.function is not None and tool_call.function.name is not None
|
||
]
|
||
assert names == [name for name, _ in EXPECTED_CALLS]
|
||
for index, (_, arguments) in enumerate(EXPECTED_CALLS):
|
||
assert json.loads(collect_streamed_arguments(deltas, index)) == arguments
|
||
|
||
|
||
def test_rust_tool_parser_adapter_ignores_midstream_empty_delta() -> None:
|
||
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())
|
||
text = build_tool_call()
|
||
split_at = len(TC_START) + 8
|
||
deltas = []
|
||
previous_text = ""
|
||
|
||
for delta_text in (text[:split_at], "", text[split_at:], ""):
|
||
current_text = previous_text + delta_text
|
||
delta = parser.extract_tool_calls_streaming(
|
||
previous_text=previous_text,
|
||
current_text=current_text,
|
||
delta_text=delta_text,
|
||
previous_token_ids=[],
|
||
current_token_ids=[],
|
||
delta_token_ids=[1],
|
||
request=MagicMock(),
|
||
)
|
||
previous_text = current_text
|
||
if delta is not None:
|
||
deltas.append(delta)
|
||
|
||
names = [
|
||
tool_call.function.name
|
||
for delta in deltas
|
||
for tool_call in delta.tool_calls or []
|
||
if tool_call.function is not None and tool_call.function.name is not None
|
||
]
|
||
assert names == [name for name, _ in EXPECTED_CALLS]
|
||
for index, (_, arguments) in enumerate(EXPECTED_CALLS):
|
||
assert json.loads(collect_streamed_arguments(deltas, index)) == arguments
|
||
|
||
|
||
KIMI_EXPECTED_IDS = ["functions.get_weather:0", "functions.add:1"]
|
||
|
||
|
||
def build_kimi_tool_call() -> str:
|
||
return (
|
||
"<|tool_calls_section_begin|>"
|
||
"<|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>"
|
||
'{"location": "SF", "date": "2024-01-16"}<|tool_call_end|>'
|
||
"<|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>"
|
||
'{"x": 3, "y": 5}<|tool_call_end|>'
|
||
"<|tool_calls_section_end|>"
|
||
)
|
||
|
||
|
||
def test_rust_tool_parser_adapter_complete_prefers_model_tool_call_ids() -> None:
|
||
tools = sample_tools()
|
||
parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=tools)
|
||
|
||
result = parser.extract_tool_calls(
|
||
"Let me check. " + build_kimi_tool_call(),
|
||
ChatCompletionRequest(messages=[], model="m", tools=tools),
|
||
)
|
||
|
||
assert result.tools_called
|
||
assert [tool_call.id for tool_call in result.tool_calls] == KIMI_EXPECTED_IDS
|
||
for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS):
|
||
assert tool_call.function.name == name
|
||
assert json.loads(tool_call.function.arguments) == arguments
|
||
|
||
|
||
def test_rust_tool_parser_adapter_streaming_prefers_model_tool_call_ids() -> None:
|
||
parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=sample_tools())
|
||
|
||
deltas = parse_streaming(parser, build_kimi_tool_call(), chunk_size=5)
|
||
|
||
ids = [
|
||
tool_call.id
|
||
for delta in deltas
|
||
for tool_call in delta.tool_calls or []
|
||
if tool_call.id is not None
|
||
]
|
||
assert ids == KIMI_EXPECTED_IDS
|
||
for index, (_, arguments) in enumerate(EXPECTED_CALLS):
|
||
assert json.loads(collect_streamed_arguments(deltas, index)) == arguments
|
||
|
||
|
||
def test_rust_tool_parser_adapter_streaming_generates_ids_as_fallback() -> None:
|
||
# DeepSeekV4 never emits model tool call IDs, so the bridge mints them.
|
||
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools())
|
||
|
||
deltas = parse_streaming(parser, build_tool_call(), chunk_size=5)
|
||
|
||
ids = [
|
||
tool_call.id
|
||
for delta in deltas
|
||
for tool_call in delta.tool_calls or []
|
||
if tool_call.function is not None and tool_call.function.name is not None
|
||
]
|
||
assert len(ids) == len(EXPECTED_CALLS)
|
||
assert all(ids)
|
||
assert len(set(ids)) == len(ids)
|
||
|
||
|
||
def test_rust_tool_parser_adapter_adjust_request_is_opaque() -> None:
|
||
tools = sample_tools()
|
||
parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools)
|
||
request = ChatCompletionRequest(
|
||
messages=[],
|
||
model="m",
|
||
tools=tools,
|
||
tool_choice="required",
|
||
skip_special_tokens=True,
|
||
)
|
||
|
||
adjusted = parser.adjust_request(request)
|
||
|
||
assert adjusted is request
|
||
assert adjusted.skip_special_tokens is False
|
||
assert adjusted.structured_outputs is None
|