| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | from collections.abc import Sequence |
| 6 | from unittest.mock import MagicMock |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 11 | ChatCompletionRequest, |
| 12 | ChatCompletionToolsParam, |
| 13 | ) |
| 14 | from vllm.tool_parsers.rust_tool_parser import RustToolParser |
| 15 | |
| 16 | # The PyO3 extension is an optional build artifact; skip when absent. |
| 17 | _rust_tool_parser = pytest.importorskip("vllm._rust_tool_parser") |
| 18 | |
| 19 | MOCK_TOKENIZER = MagicMock() |
| 20 | MOCK_TOKENIZER.get_vocab.return_value = {} |
| 21 | |
| 22 | TC_START = "<｜DSML｜tool_calls>" |
| 23 | TC_END = "</｜DSML｜tool_calls>" |
| 24 | INV_START = '<｜DSML｜invoke name="' |
| 25 | INV_END = "</｜DSML｜invoke>" |
| 26 | PARAM_START = '<｜DSML｜parameter name="' |
| 27 | PARAM_END = "</｜DSML｜parameter>" |
| 28 | |
| 29 | |
| 30 | class DeepSeekV4RustToolParser(RustToolParser): |
| 31 | rust_parser_name = "DeepSeekV4ToolParser" |
| 32 | tool_call_start_token = TC_START |
| 33 | |
| 34 | |
| 35 | class KimiK2RustToolParser(RustToolParser): |
| 36 | rust_parser_name = "KimiK2ToolParser" |
| 37 | tool_call_start_token = "<|tool_calls_section_begin|>" |
| 38 | |
| 39 | |
| 40 | def sample_tools() -> list[ChatCompletionToolsParam]: |
| 41 | return [ |
| 42 | ChatCompletionToolsParam( |
| 43 | type="function", |
| 44 | function={ |
| 45 | "name": "get_weather", |
| 46 | "description": "Get weather for a location", |
| 47 | "parameters": { |
| 48 | "type": "object", |
| 49 | "properties": { |
| 50 | "location": {"type": "string"}, |
| 51 | "date": {"type": "string"}, |
| 52 | }, |
| 53 | }, |
| 54 | }, |
| 55 | ), |
| 56 | ChatCompletionToolsParam( |
| 57 | type="function", |
| 58 | function={ |
| 59 | "name": "add", |
| 60 | "description": "Add two integers", |
| 61 | "parameters": { |
| 62 | "type": "object", |
| 63 | "properties": { |
| 64 | "x": {"type": "integer"}, |
| 65 | "y": {"type": "integer"}, |
| 66 | }, |
| 67 | }, |
| 68 | }, |
| 69 | ), |
| 70 | ] |
| 71 | |
| 72 | |
| 73 | EXPECTED_CALLS = [ |
| 74 | ("get_weather", {"location": "SF", "date": "2024-01-16"}), |
| 75 | ("add", {"x": 3, "y": 5}), |
| 76 | ] |
| 77 | |
| 78 | |
| 79 | def build_invoke( |
| 80 | function_name: str, |
| 81 | params: Sequence[tuple[str, str, bool]], |
| 82 | ) -> str: |
| 83 | param_text = "\n".join( |
| 84 | f'{PARAM_START}{name}" string="{str(is_string).lower()}">{value}{PARAM_END}' |
| 85 | for name, value, is_string in params |
| 86 | ) |
| 87 | return f'{INV_START}{function_name}">\n{param_text}\n{INV_END}\n' |
| 88 | |
| 89 | |
| 90 | def build_tool_call() -> str: |
| 91 | weather = build_invoke( |
| 92 | "get_weather", |
| 93 | [ |
| 94 | ("location", "SF", True), |
| 95 | ("date", "2024-01-16", True), |
| 96 | ], |
| 97 | ) |
| 98 | add = build_invoke( |
| 99 | "add", |
| 100 | [ |
| 101 | ("x", "3", False), |
| 102 | ("y", "5", False), |
| 103 | ], |
| 104 | ) |
| 105 | return f"{TC_START}\n{weather}{add}{TC_END}" |
| 106 | |
| 107 | |
| 108 | def parse_streaming( |
| 109 | parser: DeepSeekV4RustToolParser, |
| 110 | text: str, |
| 111 | chunk_size: int, |
| 112 | ) -> list: |
| 113 | deltas = [] |
| 114 | previous_text = "" |
| 115 | for start in range(0, len(text), chunk_size): |
| 116 | delta_text = text[start : start + chunk_size] |
| 117 | current_text = previous_text + delta_text |
| 118 | delta = parser.extract_tool_calls_streaming( |
| 119 | previous_text=previous_text, |
| 120 | current_text=current_text, |
| 121 | delta_text=delta_text, |
| 122 | previous_token_ids=[], |
| 123 | current_token_ids=[], |
| 124 | delta_token_ids=[1], |
| 125 | request=MagicMock(), |
| 126 | ) |
| 127 | previous_text = current_text |
| 128 | if delta is not None: |
| 129 | deltas.append(delta) |
| 130 | |
| 131 | delta = parser.extract_tool_calls_streaming( |
| 132 | previous_text=previous_text, |
| 133 | current_text=previous_text, |
| 134 | delta_text="", |
| 135 | previous_token_ids=[], |
| 136 | current_token_ids=[], |
| 137 | delta_token_ids=[2], |
| 138 | request=MagicMock(), |
| 139 | ) |
| 140 | if delta is not None: |
| 141 | deltas.append(delta) |
| 142 | |
| 143 | return deltas |
| 144 | |
| 145 | |
| 146 | def collect_streamed_arguments(deltas: Sequence, tool_index: int = 0) -> str: |
| 147 | return "".join( |
| 148 | tool_call.function.arguments |
| 149 | for delta in deltas |
| 150 | for tool_call in delta.tool_calls or [] |
| 151 | if ( |
| 152 | tool_call.index == tool_index |
| 153 | and tool_call.function is not None |
| 154 | and tool_call.function.arguments is not None |
| 155 | ) |
| 156 | ) |
| 157 | |
| 158 | |
| 159 | def test_rust_tool_parser_extension_typed_api() -> None: |
| 160 | tools = [ |
| 161 | _rust_tool_parser.Tool( |
| 162 | tool.function.name, |
| 163 | tool.function.description, |
| 164 | tool.function.parameters, |
| 165 | None, |
| 166 | ) |
| 167 | for tool in sample_tools() |
| 168 | ] |
| 169 | parser = _rust_tool_parser.ToolParser("DeepSeekV4ToolParser", tools) |
| 170 | output = _rust_tool_parser.ToolParserOutput() |
| 171 | |
| 172 | parser.parse_into(build_tool_call(), output) |
| 173 | output.append(parser.finish()) |
| 174 | output = output.coalesce() |
| 175 | |
| 176 | assert parser.preserve_special_tokens() |
| 177 | assert output.normal_text == "" |
| 178 | assert len(output.calls) == 2 |
| 179 | for call, (name, arguments) in zip(output.calls, EXPECTED_CALLS): |
| 180 | assert call.name == name |
| 181 | assert json.loads(call.arguments) == arguments |
| 182 | |
| 183 | |
| 184 | def test_rust_tool_parser_adapter_extracts_complete_output() -> None: |
| 185 | tools = sample_tools() |
| 186 | parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools) |
| 187 | |
| 188 | result = parser.extract_tool_calls( |
| 189 | "Let me create it. " + build_tool_call(), |
| 190 | ChatCompletionRequest(messages=[], model="m", tools=tools), |
| 191 | ) |
| 192 | |
| 193 | assert result.tools_called |
| 194 | assert result.content == "Let me create it. " |
| 195 | assert len(result.tool_calls) == 2 |
| 196 | for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS): |
| 197 | assert tool_call.function.name == name |
| 198 | assert json.loads(tool_call.function.arguments) == arguments |
| 199 | |
| 200 | |
| 201 | def test_rust_tool_parser_adapter_streaming_handles_multiple_calls() -> None: |
| 202 | parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) |
| 203 | |
| 204 | deltas = parse_streaming(parser, build_tool_call(), chunk_size=5) |
| 205 | |
| 206 | names = [ |
| 207 | tool_call.function.name |
| 208 | for delta in deltas |
| 209 | for tool_call in delta.tool_calls or [] |
| 210 | if tool_call.function is not None and tool_call.function.name is not None |
| 211 | ] |
| 212 | assert names == [name for name, _ in EXPECTED_CALLS] |
| 213 | for index, (_, arguments) in enumerate(EXPECTED_CALLS): |
| 214 | assert json.loads(collect_streamed_arguments(deltas, index)) == arguments |
| 215 | |
| 216 | |
| 217 | def test_rust_tool_parser_adapter_ignores_midstream_empty_delta() -> None: |
| 218 | parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) |
| 219 | text = build_tool_call() |
| 220 | split_at = len(TC_START) + 8 |
| 221 | deltas = [] |
| 222 | previous_text = "" |
| 223 | |
| 224 | for delta_text in (text[:split_at], "", text[split_at:], ""): |
| 225 | current_text = previous_text + delta_text |
| 226 | delta = parser.extract_tool_calls_streaming( |
| 227 | previous_text=previous_text, |
| 228 | current_text=current_text, |
| 229 | delta_text=delta_text, |
| 230 | previous_token_ids=[], |
| 231 | current_token_ids=[], |
| 232 | delta_token_ids=[1], |
| 233 | request=MagicMock(), |
| 234 | ) |
| 235 | previous_text = current_text |
| 236 | if delta is not None: |
| 237 | deltas.append(delta) |
| 238 | |
| 239 | names = [ |
| 240 | tool_call.function.name |
| 241 | for delta in deltas |
| 242 | for tool_call in delta.tool_calls or [] |
| 243 | if tool_call.function is not None and tool_call.function.name is not None |
| 244 | ] |
| 245 | assert names == [name for name, _ in EXPECTED_CALLS] |
| 246 | for index, (_, arguments) in enumerate(EXPECTED_CALLS): |
| 247 | assert json.loads(collect_streamed_arguments(deltas, index)) == arguments |
| 248 | |
| 249 | |
| 250 | KIMI_EXPECTED_IDS = ["functions.get_weather:0", "functions.add:1"] |
| 251 | |
| 252 | |
| 253 | def build_kimi_tool_call() -> str: |
| 254 | return ( |
| 255 | "<|tool_calls_section_begin|>" |
| 256 | "<|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>" |
| 257 | '{"location": "SF", "date": "2024-01-16"}<|tool_call_end|>' |
| 258 | "<|tool_call_begin|>functions.add:1<|tool_call_argument_begin|>" |
| 259 | '{"x": 3, "y": 5}<|tool_call_end|>' |
| 260 | "<|tool_calls_section_end|>" |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | def test_rust_tool_parser_adapter_complete_prefers_model_tool_call_ids() -> None: |
| 265 | tools = sample_tools() |
| 266 | parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=tools) |
| 267 | |
| 268 | result = parser.extract_tool_calls( |
| 269 | "Let me check. " + build_kimi_tool_call(), |
| 270 | ChatCompletionRequest(messages=[], model="m", tools=tools), |
| 271 | ) |
| 272 | |
| 273 | assert result.tools_called |
| 274 | assert [tool_call.id for tool_call in result.tool_calls] == KIMI_EXPECTED_IDS |
| 275 | for tool_call, (name, arguments) in zip(result.tool_calls, EXPECTED_CALLS): |
| 276 | assert tool_call.function.name == name |
| 277 | assert json.loads(tool_call.function.arguments) == arguments |
| 278 | |
| 279 | |
| 280 | def test_rust_tool_parser_adapter_streaming_prefers_model_tool_call_ids() -> None: |
| 281 | parser = KimiK2RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) |
| 282 | |
| 283 | deltas = parse_streaming(parser, build_kimi_tool_call(), chunk_size=5) |
| 284 | |
| 285 | ids = [ |
| 286 | tool_call.id |
| 287 | for delta in deltas |
| 288 | for tool_call in delta.tool_calls or [] |
| 289 | if tool_call.id is not None |
| 290 | ] |
| 291 | assert ids == KIMI_EXPECTED_IDS |
| 292 | for index, (_, arguments) in enumerate(EXPECTED_CALLS): |
| 293 | assert json.loads(collect_streamed_arguments(deltas, index)) == arguments |
| 294 | |
| 295 | |
| 296 | def test_rust_tool_parser_adapter_streaming_generates_ids_as_fallback() -> None: |
| 297 | # DeepSeekV4 never emits model tool call IDs, so the bridge mints them. |
| 298 | parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=sample_tools()) |
| 299 | |
| 300 | deltas = parse_streaming(parser, build_tool_call(), chunk_size=5) |
| 301 | |
| 302 | ids = [ |
| 303 | tool_call.id |
| 304 | for delta in deltas |
| 305 | for tool_call in delta.tool_calls or [] |
| 306 | if tool_call.function is not None and tool_call.function.name is not None |
| 307 | ] |
| 308 | assert len(ids) == len(EXPECTED_CALLS) |
| 309 | assert all(ids) |
| 310 | assert len(set(ids)) == len(ids) |
| 311 | |
| 312 | |
| 313 | def test_rust_tool_parser_adapter_adjust_request_is_opaque() -> None: |
| 314 | tools = sample_tools() |
| 315 | parser = DeepSeekV4RustToolParser(MOCK_TOKENIZER, tools=tools) |
| 316 | request = ChatCompletionRequest( |
| 317 | messages=[], |
| 318 | model="m", |
| 319 | tools=tools, |
| 320 | tool_choice="required", |
| 321 | skip_special_tokens=True, |
| 322 | ) |
| 323 | |
| 324 | adjusted = parser.adjust_request(request) |
| 325 | |
| 326 | assert adjusted is request |
| 327 | assert adjusted.skip_special_tokens is False |
| 328 | assert adjusted.structured_outputs is None |