| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # ruff: noqa: E501 |
| 4 | |
| 5 | import json |
| 6 | from unittest.mock import MagicMock |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from tests.tool_parsers.utils import ( |
| 11 | run_tool_extraction, |
| 12 | run_tool_extraction_streaming, |
| 13 | ) |
| 14 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 15 | ChatCompletionRequest, |
| 16 | ) |
| 17 | from vllm.tokenizers import get_tokenizer |
| 18 | from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser |
| 19 | |
| 20 | MODEL = "moonshotai/Kimi-K2-Instruct" |
| 21 | |
| 22 | |
| 23 | @pytest.fixture(scope="module") |
| 24 | def kimi_k2_tokenizer(): |
| 25 | return get_tokenizer(tokenizer_name=MODEL, trust_remote_code=True) |
| 26 | |
| 27 | |
| 28 | @pytest.fixture |
| 29 | def parser(kimi_k2_tokenizer): |
| 30 | return KimiK2ToolParser(kimi_k2_tokenizer) |
| 31 | |
| 32 | |
| 33 | SECTION_BEGIN = "<|tool_calls_section_begin|>" |
| 34 | SECTION_END = "<|tool_calls_section_end|>" |
| 35 | TOOL_BEGIN = "<|tool_call_begin|>" |
| 36 | TOOL_END = "<|tool_call_end|>" |
| 37 | ARG_BEGIN = "<|tool_call_argument_begin|>" |
| 38 | |
| 39 | |
| 40 | def _tool(tool_id: str, args: str) -> str: |
| 41 | return f"{TOOL_BEGIN}{tool_id} {ARG_BEGIN}{args}{TOOL_END}" |
| 42 | |
| 43 | |
| 44 | def _wrap(*tool_strs: str) -> str: |
| 45 | return SECTION_BEGIN + "".join(tool_strs) + SECTION_END |
| 46 | |
| 47 | |
| 48 | class TestExtractToolCalls: |
| 49 | def test_no_tools(self, parser): |
| 50 | content, tool_calls = run_tool_extraction( |
| 51 | parser, "This is a test", streaming=False |
| 52 | ) |
| 53 | assert content == "This is a test" |
| 54 | assert tool_calls == [] |
| 55 | |
| 56 | @pytest.mark.parametrize( |
| 57 | "model_output, expected_names, expected_args_list, expected_content", |
| 58 | [ |
| 59 | pytest.param( |
| 60 | "I'll check. " |
| 61 | + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')), |
| 62 | ["get_weather"], |
| 63 | [{"city": "Beijing"}], |
| 64 | "I'll check. ", |
| 65 | id="single_tool_call", |
| 66 | ), |
| 67 | pytest.param( |
| 68 | "Compare weather. " |
| 69 | + _wrap( |
| 70 | _tool("functions.get_weather:0", '{"city": "Beijing"}'), |
| 71 | _tool("functions.get_weather:1", '{"city": "Shanghai"}'), |
| 72 | ), |
| 73 | ["get_weather", "get_weather"], |
| 74 | [{"city": "Beijing"}, {"city": "Shanghai"}], |
| 75 | "Compare weather. ", |
| 76 | id="parallel_tool_calls", |
| 77 | ), |
| 78 | pytest.param( |
| 79 | "Multiple tasks. " |
| 80 | + _wrap( |
| 81 | _tool("functions.get_weather:0", '{"city": "New York"}'), |
| 82 | _tool("functions.get_news:1", '{"topic": "technology"}'), |
| 83 | _tool( |
| 84 | "functions.send_email:2", |
| 85 | '{"to": "user@example.com", "subject": "Daily Update"}', |
| 86 | ), |
| 87 | ), |
| 88 | ["get_weather", "get_news", "send_email"], |
| 89 | [ |
| 90 | {"city": "New York"}, |
| 91 | {"topic": "technology"}, |
| 92 | {"to": "user@example.com", "subject": "Daily Update"}, |
| 93 | ], |
| 94 | "Multiple tasks. ", |
| 95 | id="three_tool_calls", |
| 96 | ), |
| 97 | pytest.param( |
| 98 | "Process HTML. " |
| 99 | + _wrap( |
| 100 | _tool("functions.process_html:0", '{"html": "<div>content</div>"}') |
| 101 | ), |
| 102 | ["process_html"], |
| 103 | [{"html": "<div>content</div>"}], |
| 104 | "Process HTML. ", |
| 105 | id="angle_brackets_in_json", |
| 106 | ), |
| 107 | pytest.param( |
| 108 | "Formatted. " |
| 109 | + _wrap( |
| 110 | _tool( |
| 111 | "functions.process_data:0", |
| 112 | '{\n "name": "test",\n "value": 123\n}', |
| 113 | ) |
| 114 | ), |
| 115 | ["process_data"], |
| 116 | [{"name": "test", "value": 123}], |
| 117 | "Formatted. ", |
| 118 | id="multiline_json", |
| 119 | ), |
| 120 | pytest.param( |
| 121 | "No prefix. " + _wrap(_tool("get_weather:0", '{"city": "Tokyo"}')), |
| 122 | ["get_weather"], |
| 123 | [{"city": "Tokyo"}], |
| 124 | "No prefix. ", |
| 125 | id="no_functions_prefix", |
| 126 | ), |
| 127 | pytest.param( |
| 128 | "Empty args. " + _wrap(_tool("functions.test:0", "{}")), |
| 129 | ["test"], |
| 130 | [{}], |
| 131 | "Empty args. ", |
| 132 | id="empty_arguments", |
| 133 | ), |
| 134 | ], |
| 135 | ) |
| 136 | def test_extract_tool_calls( |
| 137 | self, parser, model_output, expected_names, expected_args_list, expected_content |
| 138 | ): |
| 139 | content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) |
| 140 | assert content == expected_content |
| 141 | assert len(tool_calls) == len(expected_names) |
| 142 | for tc, name, expected_args in zip( |
| 143 | tool_calls, expected_names, expected_args_list |
| 144 | ): |
| 145 | assert tc.type == "function" |
| 146 | assert tc.function.name == name |
| 147 | assert json.loads(tc.function.arguments) == expected_args |
| 148 | # id format: "something:digit" |
| 149 | assert tc.id.split(":")[-1].isdigit() |
| 150 | |
| 151 | def test_invalid_json_still_extracted(self, parser): |
| 152 | """Tool calls with invalid JSON are still returned (arguments as-is).""" |
| 153 | model_output = ( |
| 154 | "Help. " |
| 155 | + SECTION_BEGIN |
| 156 | + _tool("functions.bad:0", '{"city": "Beijing"') |
| 157 | + _tool("functions.good:1", '{"city": "Shanghai"}') |
| 158 | + SECTION_END |
| 159 | ) |
| 160 | content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) |
| 161 | assert len(tool_calls) == 2 |
| 162 | assert tool_calls[0].function.name == "bad" |
| 163 | assert tool_calls[1].function.name == "good" |
| 164 | |
| 165 | def test_invalid_funcall_id_skipped(self, parser): |
| 166 | """Tool calls with malformed id (no colon+digit) are skipped.""" |
| 167 | model_output = ( |
| 168 | "Help. " |
| 169 | + SECTION_BEGIN |
| 170 | + _tool("functions.invalid.0", '{"city": "Beijing"}') |
| 171 | + _tool("functions.valid:1", '{"city": "Shanghai"}') |
| 172 | + SECTION_END |
| 173 | ) |
| 174 | content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) |
| 175 | assert len(tool_calls) == 1 |
| 176 | assert tool_calls[0].function.name == "valid" |
| 177 | |
| 178 | def test_native_id_extracted(self, parser): |
| 179 | """Regression: parser extracts native ID onto ToolCall (PR #32768).""" |
| 180 | model_output = "Checking weather. " + _wrap( |
| 181 | _tool("functions.get_weather:0", '{"city": "Tokyo"}') |
| 182 | ) |
| 183 | content, tool_calls = run_tool_extraction(parser, model_output, streaming=False) |
| 184 | assert len(tool_calls) == 1 |
| 185 | assert tool_calls[0].id == "functions.get_weather:0" |
| 186 | assert tool_calls[0].function.name == "get_weather" |
| 187 | assert json.loads(tool_calls[0].function.arguments) == {"city": "Tokyo"} |
| 188 | |
| 189 | def test_multi_turn_native_id_continuity(self, kimi_k2_tokenizer): |
| 190 | """Regression: native IDs from turn 1 preserved across turns (PR #32768).""" |
| 191 | turn1_parser = KimiK2ToolParser(kimi_k2_tokenizer) |
| 192 | turn1_output = "Let me check. " + _wrap( |
| 193 | _tool("functions.get_weather:0", '{"city": "Beijing"}') |
| 194 | ) |
| 195 | _, turn1_tools = run_tool_extraction( |
| 196 | turn1_parser, turn1_output, streaming=False |
| 197 | ) |
| 198 | assert len(turn1_tools) == 1 |
| 199 | assert turn1_tools[0].id == "functions.get_weather:0" |
| 200 | |
| 201 | # Fresh parser for turn 2 |
| 202 | turn2_parser = KimiK2ToolParser(kimi_k2_tokenizer) |
| 203 | turn2_output = "Now let me get news. " + _wrap( |
| 204 | _tool("functions.get_news:0", '{"topic": "weather in Beijing"}') |
| 205 | ) |
| 206 | _, turn2_tools = run_tool_extraction( |
| 207 | turn2_parser, turn2_output, streaming=False |
| 208 | ) |
| 209 | assert len(turn2_tools) == 1 |
| 210 | assert turn2_tools[0].id == "functions.get_news:0" |
| 211 | |
| 212 | |
| 213 | def _split_tool_output_to_deltas( |
| 214 | content: str, tool_strs: list[tuple[str, str]] |
| 215 | ) -> list[str]: |
| 216 | """Build a list of string deltas with special tokens as separate chunks. |
| 217 | |
| 218 | Args: |
| 219 | content: text before tool section |
| 220 | tool_strs: list of (tool_id, args_json) |
| 221 | |
| 222 | """ |
| 223 | deltas = [content, SECTION_BEGIN] |
| 224 | for tool_id, args_json in tool_strs: |
| 225 | deltas.extend( |
| 226 | [ |
| 227 | TOOL_BEGIN, |
| 228 | f"{tool_id} ", |
| 229 | ARG_BEGIN, |
| 230 | f"{args_json} ", |
| 231 | TOOL_END, |
| 232 | ] |
| 233 | ) |
| 234 | deltas.append(SECTION_END) |
| 235 | return deltas |
| 236 | |
| 237 | |
| 238 | class TestStreamingHappyPath: |
| 239 | def test_single_tool_call(self, parser): |
| 240 | """Verify DeltaToolCall output: name, id, arguments for one tool.""" |
| 241 | deltas = _split_tool_output_to_deltas( |
| 242 | "I'll help. ", |
| 243 | [("functions.get_weather:0", '{"city": "Beijing"}')], |
| 244 | ) |
| 245 | rec = run_tool_extraction_streaming(parser, deltas) |
| 246 | |
| 247 | assert len(rec.tool_calls) == 1 |
| 248 | tc = rec.tool_calls[0] |
| 249 | assert tc.function.name == "get_weather" |
| 250 | assert tc.id == "functions.get_weather:0" |
| 251 | assert json.loads(tc.function.arguments) == {"city": "Beijing"} |
| 252 | |
| 253 | def test_multiple_tool_calls(self, parser): |
| 254 | """Two tool calls emitted with correct indices, names, arguments.""" |
| 255 | deltas = _split_tool_output_to_deltas( |
| 256 | "Compare weather. ", |
| 257 | [ |
| 258 | ("functions.get_weather:0", '{"city": "Tokyo"}'), |
| 259 | ("functions.get_weather:1", '{"city": "NYC"}'), |
| 260 | ], |
| 261 | ) |
| 262 | rec = run_tool_extraction_streaming(parser, deltas) |
| 263 | |
| 264 | assert len(rec.tool_calls) == 2 |
| 265 | assert rec.tool_calls[0].function.name == "get_weather" |
| 266 | assert rec.tool_calls[0].id == "functions.get_weather:0" |
| 267 | assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"} |
| 268 | |
| 269 | assert rec.tool_calls[1].function.name == "get_weather" |
| 270 | assert rec.tool_calls[1].id == "functions.get_weather:1" |
| 271 | assert json.loads(rec.tool_calls[1].function.arguments) == {"city": "NYC"} |
| 272 | |
| 273 | def test_content_before_tools(self, parser): |
| 274 | """Content before section is streamed; markers/args don't leak.""" |
| 275 | deltas = _split_tool_output_to_deltas( |
| 276 | "I'll check the weather. ", |
| 277 | [("functions.get_weather:0", '{"city": "Tokyo"}')], |
| 278 | ) |
| 279 | rec = run_tool_extraction_streaming(parser, deltas) |
| 280 | |
| 281 | assert "check the weather" in rec.other_content |
| 282 | # No markers or tool content leaked |
| 283 | for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]: |
| 284 | assert marker not in rec.other_content |
| 285 | assert "get_weather" not in rec.other_content |
| 286 | assert "Tokyo" not in rec.other_content |
| 287 | |
| 288 | def test_no_tool_calls(self, parser): |
| 289 | """Plain text streaming returns content only.""" |
| 290 | deltas = ["This is just ", "regular text ", "without tools."] |
| 291 | rec = run_tool_extraction_streaming(parser, deltas) |
| 292 | |
| 293 | assert rec.other_content == "This is just regular text without tools." |
| 294 | assert rec.tool_calls == [] |
| 295 | |
| 296 | def test_incremental_arguments(self, parser): |
| 297 | """Arguments split across small chunks accumulate correctly.""" |
| 298 | deltas = [ |
| 299 | "Help. ", |
| 300 | SECTION_BEGIN, |
| 301 | TOOL_BEGIN, |
| 302 | "functions.get_weather:0 ", |
| 303 | ARG_BEGIN, |
| 304 | '{"ci', |
| 305 | 'ty": "Be', |
| 306 | 'ijing"}', |
| 307 | " ", |
| 308 | TOOL_END, |
| 309 | SECTION_END, |
| 310 | ] |
| 311 | rec = run_tool_extraction_streaming(parser, deltas) |
| 312 | |
| 313 | assert len(rec.tool_calls) == 1 |
| 314 | assert rec.tool_calls[0].function.name == "get_weather" |
| 315 | assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"} |
| 316 | |
| 317 | @pytest.mark.parametrize( |
| 318 | "model_output", |
| 319 | [ |
| 320 | pytest.param( |
| 321 | "Single. " |
| 322 | + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')), |
| 323 | id="single_tool", |
| 324 | ), |
| 325 | pytest.param( |
| 326 | "Multi. " |
| 327 | + _wrap( |
| 328 | _tool("functions.get_weather:0", '{"city": "Tokyo"}'), |
| 329 | _tool("functions.get_news:1", '{"topic": "tech"}'), |
| 330 | ), |
| 331 | id="parallel_tools", |
| 332 | ), |
| 333 | pytest.param( |
| 334 | "No prefix id. " + _wrap(_tool("get_weather:0", '{"city": "NYC"}')), |
| 335 | id="no_functions_prefix", |
| 336 | ), |
| 337 | ], |
| 338 | ) |
| 339 | def test_streaming_matches_nonstreaming(self, parser, model_output): |
| 340 | """Streaming reconstruction matches non-streaming extraction.""" |
| 341 | content_non, tools_non = run_tool_extraction( |
| 342 | parser, model_output, streaming=False |
| 343 | ) |
| 344 | content_stream, tools_stream = run_tool_extraction( |
| 345 | parser, model_output, streaming=True |
| 346 | ) |
| 347 | |
| 348 | assert len(tools_non) == len(tools_stream) |
| 349 | for tc_non, tc_stream in zip(tools_non, tools_stream): |
| 350 | assert tc_non.function.name == tc_stream.function.name |
| 351 | assert json.loads(tc_non.function.arguments) == json.loads( |
| 352 | tc_stream.function.arguments |
| 353 | ) |
| 354 | |
| 355 | |
| 356 | class TestStreamingEdgeCases: |
| 357 | def test_marker_suppression(self, parser): |
| 358 | """No special-token markers appear in reconstructed content.""" |
| 359 | deltas = _split_tool_output_to_deltas( |
| 360 | "I'll check. ", |
| 361 | [("functions.get_weather:0", '{"city": "Tokyo"}')], |
| 362 | ) |
| 363 | rec = run_tool_extraction_streaming(parser, deltas) |
| 364 | |
| 365 | forbidden = [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN] |
| 366 | for marker in forbidden: |
| 367 | assert marker not in rec.other_content, ( |
| 368 | f"Marker leaked: {marker!r} in {rec.other_content!r}" |
| 369 | ) |
| 370 | |
| 371 | def test_noise_between_markers_suppressed(self, parser): |
| 372 | """Text between section_begin and tool_call_begin doesn't leak.""" |
| 373 | deltas = [ |
| 374 | "Reasoning. ", |
| 375 | SECTION_BEGIN, |
| 376 | " spurious noise ", |
| 377 | TOOL_BEGIN, |
| 378 | "functions.test:0 ", |
| 379 | ARG_BEGIN, |
| 380 | '{"k": "v"} ', |
| 381 | TOOL_END, |
| 382 | SECTION_END, |
| 383 | ] |
| 384 | rec = run_tool_extraction_streaming(parser, deltas) |
| 385 | |
| 386 | assert "spurious" not in rec.other_content |
| 387 | assert "noise" not in rec.other_content |
| 388 | |
| 389 | def test_empty_tool_section(self, parser): |
| 390 | """Empty section (begin immediately followed by end) doesn't crash.""" |
| 391 | deltas = ["Reasoning. ", SECTION_BEGIN, SECTION_END] |
| 392 | rec = run_tool_extraction_streaming(parser, deltas) |
| 393 | |
| 394 | assert rec.tool_calls == [] |
| 395 | |
| 396 | def test_three_different_tools(self, parser): |
| 397 | """Three tool calls with different functions stream correctly.""" |
| 398 | deltas = _split_tool_output_to_deltas( |
| 399 | "Multiple tasks. ", |
| 400 | [ |
| 401 | ("functions.get_weather:0", '{"city": "NYC"}'), |
| 402 | ("functions.get_news:1", '{"topic": "tech"}'), |
| 403 | ("functions.send_email:2", '{"to": "a@b.com"}'), |
| 404 | ], |
| 405 | ) |
| 406 | rec = run_tool_extraction_streaming(parser, deltas) |
| 407 | |
| 408 | assert len(rec.tool_calls) == 3 |
| 409 | names = [tc.function.name for tc in rec.tool_calls] |
| 410 | assert names == ["get_weather", "get_news", "send_email"] |
| 411 | ids = [tc.id for tc in rec.tool_calls] |
| 412 | assert len(set(ids)) == 3 # unique ids |
| 413 | |
| 414 | def test_truncated_tool_call_no_end_marker(self, parser): |
| 415 | """Stream ending mid-tool-call (max_tokens) doesn't crash.""" |
| 416 | deltas = [ |
| 417 | "I'll check. ", |
| 418 | SECTION_BEGIN, |
| 419 | TOOL_BEGIN, |
| 420 | "functions.get_weather:0 ", |
| 421 | ARG_BEGIN, |
| 422 | '{"city": "Bei', |
| 423 | # Stream ends here — no TOOL_END, no SECTION_END |
| 424 | ] |
| 425 | rec = run_tool_extraction_streaming(parser, deltas) |
| 426 | |
| 427 | # Should not crash; tool name and partial args extracted |
| 428 | assert len(rec.tool_calls) == 1 |
| 429 | assert rec.tool_calls[0].function.name == "get_weather" |
| 430 | assert rec.tool_calls[0].id == "functions.get_weather:0" |
| 431 | assert rec.tool_calls[0].function.arguments == '{"city": "Bei' |
| 432 | # No markers leaked into content |
| 433 | for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]: |
| 434 | assert marker not in rec.other_content |
| 435 | |
| 436 | def test_content_after_tool_section(self, parser): |
| 437 | """Trailing text after section_end doesn't crash or leak markers.""" |
| 438 | deltas = [ |
| 439 | "Before. ", |
| 440 | SECTION_BEGIN, |
| 441 | TOOL_BEGIN, |
| 442 | "functions.get_weather:0 ", |
| 443 | ARG_BEGIN, |
| 444 | '{"city": "Tokyo"} ', |
| 445 | TOOL_END, |
| 446 | SECTION_END, |
| 447 | " After tools.", |
| 448 | ] |
| 449 | rec = run_tool_extraction_streaming(parser, deltas) |
| 450 | |
| 451 | # Tool call extracted correctly |
| 452 | assert len(rec.tool_calls) == 1 |
| 453 | assert rec.tool_calls[0].function.name == "get_weather" |
| 454 | assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"} |
| 455 | # Trailing content after tool section is dropped |
| 456 | assert "After tools." not in rec.other_content |
| 457 | # No markers leaked into content |
| 458 | for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]: |
| 459 | assert marker not in rec.other_content |
| 460 | |
| 461 | |
| 462 | class TestAdjustRequest: |
| 463 | def test_sets_skip_special_tokens_false(self, parser): |
| 464 | request = MagicMock(spec=ChatCompletionRequest) |
| 465 | request.tools = [{"type": "function", "function": {"name": "test"}}] |
| 466 | request.tool_choice = "auto" |
| 467 | request.skip_special_tokens = True |
| 468 | |
| 469 | result = parser.adjust_request(request) |
| 470 | assert result.skip_special_tokens is False |
| 471 | |
| 472 | def test_no_change_when_tool_choice_none(self, parser): |
| 473 | request = MagicMock(spec=ChatCompletionRequest) |
| 474 | request.tools = [{"type": "function", "function": {"name": "test"}}] |
| 475 | request.tool_choice = "none" |
| 476 | request.skip_special_tokens = True |
| 477 | |
| 478 | result = parser.adjust_request(request) |
| 479 | assert result.skip_special_tokens is True |
| 480 | |
| 481 | def test_no_change_when_no_tools(self, parser): |
| 482 | request = MagicMock(spec=ChatCompletionRequest) |
| 483 | request.tools = None |
| 484 | request.tool_choice = "auto" |
| 485 | request.skip_special_tokens = True |
| 486 | |
| 487 | result = parser.adjust_request(request) |
| 488 | assert result.skip_special_tokens is True |
| 489 | |
| 490 | |
| 491 | def _chunk_tokenized_deltas(tokenizer, text: str, stream_interval: int) -> list[str]: |
| 492 | """Encode text, group tokens into chunks of stream_interval, decode each.""" |
| 493 | token_ids = tokenizer.encode(text, add_special_tokens=False) |
| 494 | deltas = [] |
| 495 | prev = "" |
| 496 | for i in range(0, len(token_ids), stream_interval): |
| 497 | decoded = tokenizer.decode( |
| 498 | token_ids[: i + stream_interval], skip_special_tokens=False |
| 499 | ) |
| 500 | deltas.append(decoded[len(prev) :]) |
| 501 | prev = decoded |
| 502 | return deltas |
| 503 | |
| 504 | |
| 505 | class TestStreamingIntervals: |
| 506 | """Test streaming at various token-chunk sizes to catch boundary bugs.""" |
| 507 | |
| 508 | @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) |
| 509 | def test_single_tool_call_at_interval(self, kimi_k2_tokenizer, stream_interval): |
| 510 | text = "Help. " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')) |
| 511 | deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) |
| 512 | parser = KimiK2ToolParser(kimi_k2_tokenizer) |
| 513 | rec = run_tool_extraction_streaming( |
| 514 | parser, deltas, assert_one_tool_per_delta=False |
| 515 | ) |
| 516 | |
| 517 | assert len(rec.tool_calls) == 1 |
| 518 | assert rec.tool_calls[0].function.name == "get_weather" |
| 519 | assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"} |
| 520 | |
| 521 | @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) |
| 522 | def test_content_then_tool_call_at_interval( |
| 523 | self, kimi_k2_tokenizer, stream_interval |
| 524 | ): |
| 525 | text = "Sure, let me check. " + _wrap( |
| 526 | _tool("functions.get_weather:0", '{"city": "Tokyo"}') |
| 527 | ) |
| 528 | deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) |
| 529 | parser = KimiK2ToolParser(kimi_k2_tokenizer) |
| 530 | rec = run_tool_extraction_streaming( |
| 531 | parser, deltas, assert_one_tool_per_delta=False |
| 532 | ) |
| 533 | |
| 534 | assert "let me check" in rec.other_content |
| 535 | assert "get_weather" not in rec.other_content |
| 536 | assert len(rec.tool_calls) == 1 |
| 537 | assert rec.tool_calls[0].function.name == "get_weather" |
| 538 | assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"} |
| 539 | |
| 540 | @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) |
| 541 | def test_multiple_tool_calls_at_interval(self, kimi_k2_tokenizer, stream_interval): |
| 542 | text = "Compare. " + _wrap( |
| 543 | _tool("functions.search:0", '{"q": "cats"}'), |
| 544 | _tool("functions.search:1", '{"q": "dogs"}'), |
| 545 | ) |
| 546 | deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) |
| 547 | parser = KimiK2ToolParser(kimi_k2_tokenizer) |
| 548 | rec = run_tool_extraction_streaming( |
| 549 | parser, deltas, assert_one_tool_per_delta=False |
| 550 | ) |
| 551 | |
| 552 | assert len(rec.tool_calls) == 2 |
| 553 | assert rec.tool_calls[0].function.name == "search" |
| 554 | assert json.loads(rec.tool_calls[0].function.arguments) == {"q": "cats"} |
| 555 | assert rec.tool_calls[1].function.name == "search" |
| 556 | assert json.loads(rec.tool_calls[1].function.arguments) == {"q": "dogs"} |
| 557 | |
| 558 | @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) |
| 559 | def test_plain_text_at_interval(self, kimi_k2_tokenizer, stream_interval): |
| 560 | text = "This is plain text with no tool calling involved." |
| 561 | deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval) |
| 562 | parser = KimiK2ToolParser(kimi_k2_tokenizer) |
| 563 | rec = run_tool_extraction_streaming( |
| 564 | parser, deltas, assert_one_tool_per_delta=False |
| 565 | ) |
| 566 | |
| 567 | assert rec.other_content == text |
| 568 | assert rec.tool_calls == [] |
| 569 | |
| 570 | def test_content_and_tool_call_in_single_chunk(self, kimi_k2_tokenizer): |
| 571 | """Content + complete tool call in one chunk must both be emitted.""" |
| 572 | text = "Hi! " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')) |
| 573 | deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval=9999) |
| 574 | parser = KimiK2ToolParser(kimi_k2_tokenizer) |
| 575 | rec = run_tool_extraction_streaming( |
| 576 | parser, deltas, assert_one_tool_per_delta=False |
| 577 | ) |
| 578 | |
| 579 | assert "Hi!" in rec.other_content |
| 580 | assert "get_weather" not in rec.other_content |
| 581 | assert len(rec.tool_calls) == 1 |
| 582 | assert rec.tool_calls[0].function.name == "get_weather" |
| 583 | assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"} |