| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Unit tests for DeepSeekV4EngineToolParser.""" |
| 5 | |
| 6 | import json |
| 7 | from unittest.mock import MagicMock |
| 8 | |
| 9 | import pytest |
| 10 | from xgrammar import StructuralTag |
| 11 | |
| 12 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 13 | ChatCompletionNamedFunction, |
| 14 | ChatCompletionNamedToolChoiceParam, |
| 15 | ChatCompletionRequest, |
| 16 | ChatCompletionToolsParam, |
| 17 | FunctionDefinition, |
| 18 | ) |
| 19 | from vllm.tool_parsers import ToolParserManager |
| 20 | from vllm.tool_parsers.deepseekv4_engine_tool_parser import ( |
| 21 | DeepSeekV4EngineToolParser, |
| 22 | ) |
| 23 | |
| 24 | pytestmark = pytest.mark.skip_global_cleanup |
| 25 | |
| 26 | MOCK_TOKENIZER = MagicMock() |
| 27 | MOCK_TOKENIZER.get_vocab.return_value = {} |
| 28 | |
| 29 | TC_START = "<｜DSML｜tool_calls>" |
| 30 | TC_END = "</｜DSML｜tool_calls>" |
| 31 | INV_START = '<｜DSML｜invoke name="' |
| 32 | INV_END = "</｜DSML｜invoke>" |
| 33 | PARAM_START = '<｜DSML｜parameter name="' |
| 34 | PARAM_END = "</｜DSML｜parameter>" |
| 35 | |
| 36 | |
| 37 | @pytest.fixture |
| 38 | def sample_tools() -> list[ChatCompletionToolsParam]: |
| 39 | return [ |
| 40 | ChatCompletionToolsParam( |
| 41 | type="function", |
| 42 | function={ |
| 43 | "name": "get_current_weather", |
| 44 | "description": "Get the current weather", |
| 45 | "parameters": { |
| 46 | "type": "object", |
| 47 | "properties": { |
| 48 | "city": {"type": "string", "description": "The city name"}, |
| 49 | "state": {"type": "string", "description": "The state code"}, |
| 50 | "unit": {"type": "string", "enum": ["fahrenheit", "celsius"]}, |
| 51 | }, |
| 52 | "required": ["city", "state"], |
| 53 | }, |
| 54 | }, |
| 55 | ), |
| 56 | ChatCompletionToolsParam( |
| 57 | type="function", |
| 58 | function={ |
| 59 | "name": "calculate_area", |
| 60 | "description": "Calculate area of a shape", |
| 61 | "parameters": { |
| 62 | "type": "object", |
| 63 | "properties": { |
| 64 | "shape": {"type": "string"}, |
| 65 | "dimensions": {"type": "object"}, |
| 66 | "precision": {"type": "integer"}, |
| 67 | }, |
| 68 | }, |
| 69 | }, |
| 70 | ), |
| 71 | ] |
| 72 | |
| 73 | |
| 74 | def make_parser(tools=None) -> DeepSeekV4EngineToolParser: |
| 75 | return DeepSeekV4EngineToolParser(MOCK_TOKENIZER, tools=tools) |
| 76 | |
| 77 | |
| 78 | def make_request(tools=None) -> MagicMock: |
| 79 | req = MagicMock() |
| 80 | req.tools = tools |
| 81 | return req |
| 82 | |
| 83 | |
| 84 | def build_tool_call(func_name: str, params: dict[str, str]) -> str: |
| 85 | param_strs = "".join( |
| 86 | f'{PARAM_START}{k}" string="true">{v}{PARAM_END}\n' for k, v in params.items() |
| 87 | ) |
| 88 | return f'{TC_START}\n{INV_START}{func_name}">\n{param_strs}{INV_END}\n{TC_END}' |
| 89 | |
| 90 | |
| 91 | def stream(parser: DeepSeekV4EngineToolParser, full_text: str, chunk_size: int = 7): |
| 92 | deltas = [] |
| 93 | previous_text = "" |
| 94 | for start in range(0, len(full_text), chunk_size): |
| 95 | delta_text = full_text[start : start + chunk_size] |
| 96 | current_text = previous_text + delta_text |
| 97 | delta = parser.extract_tool_calls_streaming( |
| 98 | previous_text=previous_text, |
| 99 | current_text=current_text, |
| 100 | delta_text=delta_text, |
| 101 | previous_token_ids=[], |
| 102 | current_token_ids=[], |
| 103 | delta_token_ids=[1], |
| 104 | request=make_request(), |
| 105 | ) |
| 106 | previous_text = current_text |
| 107 | if delta is not None: |
| 108 | deltas.append(delta) |
| 109 | return deltas |
| 110 | |
| 111 | |
| 112 | def reconstruct_args(deltas, tool_index: int = 0) -> str: |
| 113 | fragments = [] |
| 114 | for delta in deltas: |
| 115 | if delta.tool_calls: |
| 116 | for tool_call in delta.tool_calls: |
| 117 | if ( |
| 118 | tool_call.index == tool_index |
| 119 | and tool_call.function |
| 120 | and tool_call.function.arguments |
| 121 | ): |
| 122 | fragments.append(tool_call.function.arguments) |
| 123 | return "".join(fragments) |
| 124 | |
| 125 | |
| 126 | def test_registered(): |
| 127 | assert ( |
| 128 | ToolParserManager.get_tool_parser("deepseek_v4") is DeepSeekV4EngineToolParser |
| 129 | ) |
| 130 | |
| 131 | |
| 132 | def test_extract_tool_calls(): |
| 133 | parser = make_parser() |
| 134 | model_output = "Let me check. " + build_tool_call( |
| 135 | "get_weather", {"location": "Beijing", "unit": "celsius"} |
| 136 | ) |
| 137 | |
| 138 | result = parser.extract_tool_calls(model_output, make_request()) |
| 139 | |
| 140 | assert result.tools_called |
| 141 | assert result.content == "Let me check. " |
| 142 | assert len(result.tool_calls) == 1 |
| 143 | tool_call = result.tool_calls[0] |
| 144 | assert tool_call.function.name == "get_weather" |
| 145 | assert json.loads(tool_call.function.arguments) == { |
| 146 | "location": "Beijing", |
| 147 | "unit": "celsius", |
| 148 | } |
| 149 | |
| 150 | |
| 151 | def test_function_calls_wrapper_is_not_recognized(): |
| 152 | # The V3.2 wrapper is not a V4 terminal, so it passes through as content, |
| 153 | # but the invoke inside it is still parsed (tool calls anchor on the invoke). |
| 154 | parser = make_parser() |
| 155 | model_output = build_tool_call("search", {"query": "vllm"}).replace( |
| 156 | "tool_calls", "function_calls" |
| 157 | ) |
| 158 | |
| 159 | result = parser.extract_tool_calls(model_output, make_request()) |
| 160 | |
| 161 | assert result.tools_called |
| 162 | assert result.tool_calls[0].function.name == "search" |
| 163 | assert result.content == "<｜DSML｜function_calls>\n" |
| 164 | |
| 165 | |
| 166 | def test_missing_tool_calls_wrapper_is_recovered(): |
| 167 | # Regression for #48931: at long context the model omits the |
| 168 | # <｜DSML｜tool_calls> START token but still emits a complete invoke. |
| 169 | parser = make_parser() |
| 170 | model_output = build_tool_call("search", {"query": "vllm"}).replace( |
| 171 | TC_START + "\n", "" |
| 172 | ) |
| 173 | assert TC_START not in model_output |
| 174 | |
| 175 | result = parser.extract_tool_calls(model_output, make_request()) |
| 176 | |
| 177 | assert result.tools_called |
| 178 | assert result.tool_calls[0].function.name == "search" |
| 179 | assert json.loads(result.tool_calls[0].function.arguments) == {"query": "vllm"} |
| 180 | assert result.content is None |
| 181 | |
| 182 | deltas = stream(make_parser(), model_output, chunk_size=3) |
| 183 | names = [ |
| 184 | tc.function.name |
| 185 | for d in deltas |
| 186 | if d.tool_calls |
| 187 | for tc in d.tool_calls |
| 188 | if tc.function.name |
| 189 | ] |
| 190 | assert names == ["search"] |
| 191 | assert json.loads(reconstruct_args(deltas)) == {"query": "vllm"} |
| 192 | assert not any(d.content and "DSML" in d.content for d in deltas) |
| 193 | |
| 194 | |
| 195 | def test_streaming_extracts_complete_invokes(): |
| 196 | parser = make_parser() |
| 197 | full_text = build_tool_call("search", {"query": "deepseek v4"}) |
| 198 | |
| 199 | deltas = stream(parser, full_text, chunk_size=5) |
| 200 | |
| 201 | names = [ |
| 202 | tool_call.function.name |
| 203 | for delta in deltas |
| 204 | if delta.tool_calls |
| 205 | for tool_call in delta.tool_calls |
| 206 | if tool_call.function.name |
| 207 | ] |
| 208 | assert names == ["search"] |
| 209 | assert json.loads(reconstruct_args(deltas)) == {"query": "deepseek v4"} |
| 210 | |
| 211 | |
| 212 | def test_streaming_emits_incremental_argument_chunks(): |
| 213 | tool = ChatCompletionToolsParam( |
| 214 | function=FunctionDefinition( |
| 215 | name="plan_trip", |
| 216 | parameters={ |
| 217 | "type": "object", |
| 218 | "properties": { |
| 219 | "days": {"type": "integer"}, |
| 220 | "flexible": {"type": "boolean"}, |
| 221 | "cities": {"type": "array", "items": {"type": "string"}}, |
| 222 | "notes": {"type": "string"}, |
| 223 | }, |
| 224 | }, |
| 225 | ), |
| 226 | ) |
| 227 | parser = make_parser(tools=[tool]) |
| 228 | full_text = ( |
| 229 | f"{TC_START}\n" |
| 230 | f'{INV_START}plan_trip">\n' |
| 231 | f'{PARAM_START}days" string="false">3{PARAM_END}\n' |
| 232 | f'{PARAM_START}flexible" string="false">false{PARAM_END}\n' |
| 233 | f'{PARAM_START}cities" string="false">' |
| 234 | f'["Beijing","Shanghai","Tokyo","New York"]{PARAM_END}\n' |
| 235 | f'{PARAM_START}notes" string="true">靠窗座位{PARAM_END}\n' |
| 236 | f"{INV_END}\n" |
| 237 | f"{TC_END}" |
| 238 | ) |
| 239 | |
| 240 | deltas = stream(parser, full_text, chunk_size=4) |
| 241 | arg_chunks = [ |
| 242 | tool_call.function.arguments |
| 243 | for delta in deltas |
| 244 | for tool_call in delta.tool_calls or [] |
| 245 | if tool_call.function and tool_call.function.arguments is not None |
| 246 | ] |
| 247 | |
| 248 | assert len([chunk for chunk in arg_chunks if chunk]) > 2 |
| 249 | assert json.loads("".join(arg_chunks)) == { |
| 250 | "days": 3, |
| 251 | "flexible": False, |
| 252 | "cities": ["Beijing", "Shanghai", "Tokyo", "New York"], |
| 253 | "notes": "靠窗座位", |
| 254 | } |
| 255 | |
| 256 | |
| 257 | def _with_strict( |
| 258 | tools: list[ChatCompletionToolsParam], |
| 259 | ) -> list[ChatCompletionToolsParam]: |
| 260 | return [ |
| 261 | ChatCompletionToolsParam( |
| 262 | type=t.type, |
| 263 | function=FunctionDefinition( |
| 264 | name=t.function.name, |
| 265 | description=t.function.description, |
| 266 | parameters=t.function.parameters, |
| 267 | strict=True, |
| 268 | ), |
| 269 | ) |
| 270 | for t in tools |
| 271 | ] |
| 272 | |
| 273 | |
| 274 | def test_get_vllm_registry_structural_tag_returns_structural_tag( |
| 275 | sample_tools: list[ChatCompletionToolsParam], |
| 276 | ) -> None: |
| 277 | parser = make_parser() |
| 278 | strict_tools = _with_strict(sample_tools) |
| 279 | req = ChatCompletionRequest( |
| 280 | messages=[], |
| 281 | model="m", |
| 282 | tools=strict_tools, |
| 283 | tool_choice="auto", |
| 284 | ) |
| 285 | tag = parser.get_structural_tag(req) |
| 286 | assert isinstance(tag, StructuralTag) |
| 287 | |
| 288 | req = ChatCompletionRequest( |
| 289 | messages=[], |
| 290 | model="m", |
| 291 | tools=sample_tools, |
| 292 | tool_choice="required", |
| 293 | ) |
| 294 | tag = parser.get_structural_tag(req) |
| 295 | assert isinstance(tag, StructuralTag) |
| 296 | |
| 297 | if sample_tools: |
| 298 | tool = sample_tools[0] |
| 299 | req = ChatCompletionRequest( |
| 300 | messages=[], |
| 301 | model="m", |
| 302 | tools=sample_tools, |
| 303 | ) |
| 304 | req.tool_choice = ChatCompletionNamedToolChoiceParam( |
| 305 | function=ChatCompletionNamedFunction(name=tool.function.name) |
| 306 | ) |
| 307 | tag = parser.get_structural_tag(req) |
| 308 | assert isinstance(tag, StructuralTag) |
| 309 | |
| 310 | |
| 311 | def test_extract_tool_calls_arguments_wrapper(): |
| 312 | mock_tokenizer = MagicMock() |
| 313 | mock_tokenizer.get_vocab.return_value = {} |
| 314 | |
| 315 | tool = ChatCompletionToolsParam( |
| 316 | type="function", |
| 317 | function={ |
| 318 | "name": "get_weather", |
| 319 | "parameters": { |
| 320 | "type": "object", |
| 321 | "properties": {"location": {"type": "string"}}, |
| 322 | }, |
| 323 | }, |
| 324 | ) |
| 325 | |
| 326 | parser = DeepSeekV4EngineToolParser(mock_tokenizer, tools=[tool]) |
| 327 | request = MagicMock() |
| 328 | request.tools = [tool] |
| 329 | |
| 330 | model_output = ( |
| 331 | f"{TC_START}" |
| 332 | f'{INV_START}get_weather">' |
| 333 | f'{PARAM_START}arguments" string="false">{{"location":"Beijing"}}{PARAM_END}' |
| 334 | f"{INV_END}" |
| 335 | f"{TC_END}" |
| 336 | ) |
| 337 | |
| 338 | result = parser.extract_tool_calls(model_output, request) |
| 339 | assert result.tools_called |
| 340 | args = json.loads(result.tool_calls[0].function.arguments) |
| 341 | assert args == {"location": "Beijing"} |
| 342 | |
| 343 | |
| 344 | _ANGLE_BRACKET_TOOL = ChatCompletionToolsParam( |
| 345 | function=FunctionDefinition( |
| 346 | name="run_command", |
| 347 | parameters={ |
| 348 | "type": "object", |
| 349 | "properties": { |
| 350 | "command": {"type": "string"}, |
| 351 | }, |
| 352 | }, |
| 353 | ), |
| 354 | ) |
| 355 | |
| 356 | |
| 357 | @pytest.mark.parametrize( |
| 358 | "tools", |
| 359 | [[_ANGLE_BRACKET_TOOL], None], |
| 360 | ids=["with_tools", "without_tools"], |
| 361 | ) |
| 362 | def test_no_dsml_closing_tag_leak_in_streamed_args(tools): |
| 363 | """Streaming must not leak </｜DSML｜parameter> into argument values. |
| 364 | |
| 365 | When a parameter value contains '>' (e.g. shell redirects like |
| 366 | '2>&1'), certain chunk boundaries cause the incremental lexer to |
| 367 | emit the closing delimiter text as part of the content token. The |
| 368 | partial regex then captures it as part of the value, violating the |
| 369 | prefix invariant and corrupting the streamed JSON. |
| 370 | """ |
| 371 | full_text = build_tool_call("run_command", {"command": "git --version 2>&1"}) |
| 372 | expected = {"command": "git --version 2>&1"} |
| 373 | |
| 374 | for chunk_size in range(1, len(full_text) + 1): |
| 375 | parser = make_parser(tools=tools) |
| 376 | deltas = stream(parser, full_text, chunk_size=chunk_size) |
| 377 | args_str = reconstruct_args(deltas) |
| 378 | assert args_str, f"No args emitted at chunk_size={chunk_size}" |
| 379 | assert "DSML" not in args_str, ( |
| 380 | f"DSML marker leaked into args at chunk_size={chunk_size}: {args_str!r}" |
| 381 | ) |
| 382 | parsed = json.loads(args_str) |
| 383 | assert parsed == expected, ( |
| 384 | f"Args mismatch at chunk_size={chunk_size}: " |
| 385 | f"got {parsed!r}, expected {expected!r}" |
| 386 | ) |
| 387 | |
| 388 | |
| 389 | def test_non_streaming_extract_with_angle_brackets(): |
| 390 | """Non-streaming extraction must correctly handle '>' in values.""" |
| 391 | parser = make_parser() |
| 392 | full_text = build_tool_call("run_command", {"command": "git --version 2>&1"}) |
| 393 | result = parser.extract_tool_calls(full_text, make_request()) |
| 394 | |
| 395 | assert result.tools_called |
| 396 | assert len(result.tool_calls) == 1 |
| 397 | args = json.loads(result.tool_calls[0].function.arguments) |
| 398 | assert args == {"command": "git --version 2>&1"} |
| 399 | assert "DSML" not in result.tool_calls[0].function.arguments |
| 400 | |
| 401 | |
| 402 | def test_composed_schema_converts_object_and_array_params(): |
| 403 | tool = ChatCompletionToolsParam( |
| 404 | type="function", |
| 405 | function={ |
| 406 | "name": "set_timer", |
| 407 | "parameters": { |
| 408 | "type": "object", |
| 409 | "properties": { |
| 410 | "wait": { |
| 411 | "anyOf": [ |
| 412 | {"type": "object"}, |
| 413 | {"type": "null"}, |
| 414 | ], |
| 415 | }, |
| 416 | "patches": { |
| 417 | "allOf": [ |
| 418 | {"type": "array", "items": {"type": "object"}}, |
| 419 | ], |
| 420 | }, |
| 421 | }, |
| 422 | }, |
| 423 | }, |
| 424 | ) |
| 425 | parser = make_parser(tools=[tool]) |
| 426 | request = make_request(tools=[tool]) |
| 427 | model_output = ( |
| 428 | f"{TC_START}\n" |
| 429 | f'{INV_START}set_timer">\n' |
| 430 | f'{PARAM_START}wait" string="false">' |
| 431 | f'{{"type":"for","minutes":2880}}' |
| 432 | f"{PARAM_END}\n" |
| 433 | f'{PARAM_START}patches" string="false">' |
| 434 | f'[{{"op":"replace","path":"/schedule","value":"quiet"}}]' |
| 435 | f"{PARAM_END}\n" |
| 436 | f"{INV_END}\n" |
| 437 | f"{TC_END}" |
| 438 | ) |
| 439 | |
| 440 | result = parser.extract_tool_calls(model_output, request) |
| 441 | |
| 442 | assert result.tools_called |
| 443 | args = json.loads(result.tool_calls[0].function.arguments) |
| 444 | assert args == { |
| 445 | "wait": {"type": "for", "minutes": 2880}, |
| 446 | "patches": [{"op": "replace", "path": "/schedule", "value": "quiet"}], |
| 447 | } |