wehub-resource-sync
416 lines
13 KiB
Python
416 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
||
"""Unit tests for DeepSeekV4EngineToolParser."""
|
||
|
||
import json
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
from xgrammar import StructuralTag
|
||
|
||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||
ChatCompletionNamedFunction,
|
||
ChatCompletionNamedToolChoiceParam,
|
||
ChatCompletionRequest,
|
||
ChatCompletionToolsParam,
|
||
FunctionDefinition,
|
||
)
|
||
from vllm.tool_parsers import ToolParserManager
|
||
from vllm.tool_parsers.deepseekv4_engine_tool_parser import (
|
||
DeepSeekV4EngineToolParser,
|
||
)
|
||
|
||
pytestmark = pytest.mark.skip_global_cleanup
|
||
|
||
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>"
|
||
|
||
|
||
@pytest.fixture
|
||
def sample_tools() -> list[ChatCompletionToolsParam]:
|
||
return [
|
||
ChatCompletionToolsParam(
|
||
type="function",
|
||
function={
|
||
"name": "get_current_weather",
|
||
"description": "Get the current weather",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"city": {"type": "string", "description": "The city name"},
|
||
"state": {"type": "string", "description": "The state code"},
|
||
"unit": {"type": "string", "enum": ["fahrenheit", "celsius"]},
|
||
},
|
||
"required": ["city", "state"],
|
||
},
|
||
},
|
||
),
|
||
ChatCompletionToolsParam(
|
||
type="function",
|
||
function={
|
||
"name": "calculate_area",
|
||
"description": "Calculate area of a shape",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"shape": {"type": "string"},
|
||
"dimensions": {"type": "object"},
|
||
"precision": {"type": "integer"},
|
||
},
|
||
},
|
||
},
|
||
),
|
||
]
|
||
|
||
|
||
def make_parser(tools=None) -> DeepSeekV4EngineToolParser:
|
||
return DeepSeekV4EngineToolParser(MOCK_TOKENIZER, tools=tools)
|
||
|
||
|
||
def make_request(tools=None) -> MagicMock:
|
||
req = MagicMock()
|
||
req.tools = tools
|
||
return req
|
||
|
||
|
||
def build_tool_call(func_name: str, params: dict[str, str]) -> str:
|
||
param_strs = "".join(
|
||
f'{PARAM_START}{k}" string="true">{v}{PARAM_END}\n' for k, v in params.items()
|
||
)
|
||
return f'{TC_START}\n{INV_START}{func_name}">\n{param_strs}{INV_END}\n{TC_END}'
|
||
|
||
|
||
def stream(parser: DeepSeekV4EngineToolParser, full_text: str, chunk_size: int = 7):
|
||
deltas = []
|
||
previous_text = ""
|
||
for start in range(0, len(full_text), chunk_size):
|
||
delta_text = full_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=make_request(),
|
||
)
|
||
previous_text = current_text
|
||
if delta is not None:
|
||
deltas.append(delta)
|
||
return deltas
|
||
|
||
|
||
def reconstruct_args(deltas, tool_index: int = 0) -> str:
|
||
fragments = []
|
||
for delta in deltas:
|
||
if delta.tool_calls:
|
||
for tool_call in delta.tool_calls:
|
||
if (
|
||
tool_call.index == tool_index
|
||
and tool_call.function
|
||
and tool_call.function.arguments
|
||
):
|
||
fragments.append(tool_call.function.arguments)
|
||
return "".join(fragments)
|
||
|
||
|
||
def test_registered():
|
||
assert (
|
||
ToolParserManager.get_tool_parser("deepseek_v4") is DeepSeekV4EngineToolParser
|
||
)
|
||
|
||
|
||
def test_extract_tool_calls():
|
||
parser = make_parser()
|
||
model_output = "Let me check. " + build_tool_call(
|
||
"get_weather", {"location": "Beijing", "unit": "celsius"}
|
||
)
|
||
|
||
result = parser.extract_tool_calls(model_output, make_request())
|
||
|
||
assert result.tools_called
|
||
assert result.content == "Let me check. "
|
||
assert len(result.tool_calls) == 1
|
||
tool_call = result.tool_calls[0]
|
||
assert tool_call.function.name == "get_weather"
|
||
assert json.loads(tool_call.function.arguments) == {
|
||
"location": "Beijing",
|
||
"unit": "celsius",
|
||
}
|
||
|
||
|
||
def test_function_calls_block_is_not_accepted():
|
||
parser = make_parser()
|
||
model_output = build_tool_call("search", {"query": "vllm"}).replace(
|
||
"tool_calls", "function_calls"
|
||
)
|
||
|
||
result = parser.extract_tool_calls(model_output, make_request())
|
||
|
||
assert not result.tools_called
|
||
assert result.content == model_output
|
||
|
||
|
||
def test_streaming_extracts_complete_invokes():
|
||
parser = make_parser()
|
||
full_text = build_tool_call("search", {"query": "deepseek v4"})
|
||
|
||
deltas = stream(parser, full_text, chunk_size=5)
|
||
|
||
names = [
|
||
tool_call.function.name
|
||
for delta in deltas
|
||
if delta.tool_calls
|
||
for tool_call in delta.tool_calls
|
||
if tool_call.function.name
|
||
]
|
||
assert names == ["search"]
|
||
assert json.loads(reconstruct_args(deltas)) == {"query": "deepseek v4"}
|
||
|
||
|
||
def test_streaming_emits_incremental_argument_chunks():
|
||
tool = ChatCompletionToolsParam(
|
||
function=FunctionDefinition(
|
||
name="plan_trip",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"days": {"type": "integer"},
|
||
"flexible": {"type": "boolean"},
|
||
"cities": {"type": "array", "items": {"type": "string"}},
|
||
"notes": {"type": "string"},
|
||
},
|
||
},
|
||
),
|
||
)
|
||
parser = make_parser(tools=[tool])
|
||
full_text = (
|
||
f"{TC_START}\n"
|
||
f'{INV_START}plan_trip">\n'
|
||
f'{PARAM_START}days" string="false">3{PARAM_END}\n'
|
||
f'{PARAM_START}flexible" string="false">false{PARAM_END}\n'
|
||
f'{PARAM_START}cities" string="false">'
|
||
f'["Beijing","Shanghai","Tokyo","New York"]{PARAM_END}\n'
|
||
f'{PARAM_START}notes" string="true">靠窗座位{PARAM_END}\n'
|
||
f"{INV_END}\n"
|
||
f"{TC_END}"
|
||
)
|
||
|
||
deltas = stream(parser, full_text, chunk_size=4)
|
||
arg_chunks = [
|
||
tool_call.function.arguments
|
||
for delta in deltas
|
||
for tool_call in delta.tool_calls or []
|
||
if tool_call.function and tool_call.function.arguments is not None
|
||
]
|
||
|
||
assert len([chunk for chunk in arg_chunks if chunk]) > 2
|
||
assert json.loads("".join(arg_chunks)) == {
|
||
"days": 3,
|
||
"flexible": False,
|
||
"cities": ["Beijing", "Shanghai", "Tokyo", "New York"],
|
||
"notes": "靠窗座位",
|
||
}
|
||
|
||
|
||
def _with_strict(
|
||
tools: list[ChatCompletionToolsParam],
|
||
) -> list[ChatCompletionToolsParam]:
|
||
return [
|
||
ChatCompletionToolsParam(
|
||
type=t.type,
|
||
function=FunctionDefinition(
|
||
name=t.function.name,
|
||
description=t.function.description,
|
||
parameters=t.function.parameters,
|
||
strict=True,
|
||
),
|
||
)
|
||
for t in tools
|
||
]
|
||
|
||
|
||
def test_get_vllm_registry_structural_tag_returns_structural_tag(
|
||
sample_tools: list[ChatCompletionToolsParam],
|
||
) -> None:
|
||
parser = make_parser()
|
||
strict_tools = _with_strict(sample_tools)
|
||
req = ChatCompletionRequest(
|
||
messages=[],
|
||
model="m",
|
||
tools=strict_tools,
|
||
tool_choice="auto",
|
||
)
|
||
tag = parser.get_structural_tag(req)
|
||
assert isinstance(tag, StructuralTag)
|
||
|
||
req = ChatCompletionRequest(
|
||
messages=[],
|
||
model="m",
|
||
tools=sample_tools,
|
||
tool_choice="required",
|
||
)
|
||
tag = parser.get_structural_tag(req)
|
||
assert isinstance(tag, StructuralTag)
|
||
|
||
if sample_tools:
|
||
tool = sample_tools[0]
|
||
req = ChatCompletionRequest(
|
||
messages=[],
|
||
model="m",
|
||
tools=sample_tools,
|
||
)
|
||
req.tool_choice = ChatCompletionNamedToolChoiceParam(
|
||
function=ChatCompletionNamedFunction(name=tool.function.name)
|
||
)
|
||
tag = parser.get_structural_tag(req)
|
||
assert isinstance(tag, StructuralTag)
|
||
|
||
|
||
def test_extract_tool_calls_arguments_wrapper():
|
||
mock_tokenizer = MagicMock()
|
||
mock_tokenizer.get_vocab.return_value = {}
|
||
|
||
tool = ChatCompletionToolsParam(
|
||
type="function",
|
||
function={
|
||
"name": "get_weather",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {"location": {"type": "string"}},
|
||
},
|
||
},
|
||
)
|
||
|
||
parser = DeepSeekV4EngineToolParser(mock_tokenizer, tools=[tool])
|
||
request = MagicMock()
|
||
request.tools = [tool]
|
||
|
||
model_output = (
|
||
f"{TC_START}"
|
||
f'{INV_START}get_weather">'
|
||
f'{PARAM_START}arguments" string="false">{{"location":"Beijing"}}{PARAM_END}'
|
||
f"{INV_END}"
|
||
f"{TC_END}"
|
||
)
|
||
|
||
result = parser.extract_tool_calls(model_output, request)
|
||
assert result.tools_called
|
||
args = json.loads(result.tool_calls[0].function.arguments)
|
||
assert args == {"location": "Beijing"}
|
||
|
||
|
||
_ANGLE_BRACKET_TOOL = ChatCompletionToolsParam(
|
||
function=FunctionDefinition(
|
||
name="run_command",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"command": {"type": "string"},
|
||
},
|
||
},
|
||
),
|
||
)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"tools",
|
||
[[_ANGLE_BRACKET_TOOL], None],
|
||
ids=["with_tools", "without_tools"],
|
||
)
|
||
def test_no_dsml_closing_tag_leak_in_streamed_args(tools):
|
||
"""Streaming must not leak </｜DSML｜parameter> into argument values.
|
||
|
||
When a parameter value contains '>' (e.g. shell redirects like
|
||
'2>&1'), certain chunk boundaries cause the incremental lexer to
|
||
emit the closing delimiter text as part of the content token. The
|
||
partial regex then captures it as part of the value, violating the
|
||
prefix invariant and corrupting the streamed JSON.
|
||
"""
|
||
full_text = build_tool_call("run_command", {"command": "git --version 2>&1"})
|
||
expected = {"command": "git --version 2>&1"}
|
||
|
||
for chunk_size in range(1, len(full_text) + 1):
|
||
parser = make_parser(tools=tools)
|
||
deltas = stream(parser, full_text, chunk_size=chunk_size)
|
||
args_str = reconstruct_args(deltas)
|
||
assert args_str, f"No args emitted at chunk_size={chunk_size}"
|
||
assert "DSML" not in args_str, (
|
||
f"DSML marker leaked into args at chunk_size={chunk_size}: {args_str!r}"
|
||
)
|
||
parsed = json.loads(args_str)
|
||
assert parsed == expected, (
|
||
f"Args mismatch at chunk_size={chunk_size}: "
|
||
f"got {parsed!r}, expected {expected!r}"
|
||
)
|
||
|
||
|
||
def test_non_streaming_extract_with_angle_brackets():
|
||
"""Non-streaming extraction must correctly handle '>' in values."""
|
||
parser = make_parser()
|
||
full_text = build_tool_call("run_command", {"command": "git --version 2>&1"})
|
||
result = parser.extract_tool_calls(full_text, make_request())
|
||
|
||
assert result.tools_called
|
||
assert len(result.tool_calls) == 1
|
||
args = json.loads(result.tool_calls[0].function.arguments)
|
||
assert args == {"command": "git --version 2>&1"}
|
||
assert "DSML" not in result.tool_calls[0].function.arguments
|
||
|
||
|
||
def test_composed_schema_converts_object_and_array_params():
|
||
tool = ChatCompletionToolsParam(
|
||
type="function",
|
||
function={
|
||
"name": "set_timer",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"wait": {
|
||
"anyOf": [
|
||
{"type": "object"},
|
||
{"type": "null"},
|
||
],
|
||
},
|
||
"patches": {
|
||
"allOf": [
|
||
{"type": "array", "items": {"type": "object"}},
|
||
],
|
||
},
|
||
},
|
||
},
|
||
},
|
||
)
|
||
parser = make_parser(tools=[tool])
|
||
request = make_request(tools=[tool])
|
||
model_output = (
|
||
f"{TC_START}\n"
|
||
f'{INV_START}set_timer">\n'
|
||
f'{PARAM_START}wait" string="false">'
|
||
f'{{"type":"for","minutes":2880}}'
|
||
f"{PARAM_END}\n"
|
||
f'{PARAM_START}patches" string="false">'
|
||
f'[{{"op":"replace","path":"/schedule","value":"quiet"}}]'
|
||
f"{PARAM_END}\n"
|
||
f"{INV_END}\n"
|
||
f"{TC_END}"
|
||
)
|
||
|
||
result = parser.extract_tool_calls(model_output, request)
|
||
|
||
assert result.tools_called
|
||
args = json.loads(result.tool_calls[0].function.arguments)
|
||
assert args == {
|
||
"wait": {"type": "for", "minutes": 2880},
|
||
"patches": [{"op": "replace", "path": "/schedule", "value": "quiet"}],
|
||
}
|