583 lines
22 KiB
Python
583 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
# ruff: noqa: E501
|
|
|
|
import json
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from tests.tool_parsers.utils import (
|
|
run_tool_extraction,
|
|
run_tool_extraction_streaming,
|
|
)
|
|
from vllm.entrypoints.openai.chat_completion.protocol import (
|
|
ChatCompletionRequest,
|
|
)
|
|
from vllm.tokenizers import get_tokenizer
|
|
from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser
|
|
|
|
MODEL = "moonshotai/Kimi-K2-Instruct"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def kimi_k2_tokenizer():
|
|
return get_tokenizer(tokenizer_name=MODEL, trust_remote_code=True)
|
|
|
|
|
|
@pytest.fixture
|
|
def parser(kimi_k2_tokenizer):
|
|
return KimiK2ToolParser(kimi_k2_tokenizer)
|
|
|
|
|
|
SECTION_BEGIN = "<|tool_calls_section_begin|>"
|
|
SECTION_END = "<|tool_calls_section_end|>"
|
|
TOOL_BEGIN = "<|tool_call_begin|>"
|
|
TOOL_END = "<|tool_call_end|>"
|
|
ARG_BEGIN = "<|tool_call_argument_begin|>"
|
|
|
|
|
|
def _tool(tool_id: str, args: str) -> str:
|
|
return f"{TOOL_BEGIN}{tool_id} {ARG_BEGIN}{args}{TOOL_END}"
|
|
|
|
|
|
def _wrap(*tool_strs: str) -> str:
|
|
return SECTION_BEGIN + "".join(tool_strs) + SECTION_END
|
|
|
|
|
|
class TestExtractToolCalls:
|
|
def test_no_tools(self, parser):
|
|
content, tool_calls = run_tool_extraction(
|
|
parser, "This is a test", streaming=False
|
|
)
|
|
assert content == "This is a test"
|
|
assert tool_calls == []
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_output, expected_names, expected_args_list, expected_content",
|
|
[
|
|
pytest.param(
|
|
"I'll check. "
|
|
+ _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')),
|
|
["get_weather"],
|
|
[{"city": "Beijing"}],
|
|
"I'll check. ",
|
|
id="single_tool_call",
|
|
),
|
|
pytest.param(
|
|
"Compare weather. "
|
|
+ _wrap(
|
|
_tool("functions.get_weather:0", '{"city": "Beijing"}'),
|
|
_tool("functions.get_weather:1", '{"city": "Shanghai"}'),
|
|
),
|
|
["get_weather", "get_weather"],
|
|
[{"city": "Beijing"}, {"city": "Shanghai"}],
|
|
"Compare weather. ",
|
|
id="parallel_tool_calls",
|
|
),
|
|
pytest.param(
|
|
"Multiple tasks. "
|
|
+ _wrap(
|
|
_tool("functions.get_weather:0", '{"city": "New York"}'),
|
|
_tool("functions.get_news:1", '{"topic": "technology"}'),
|
|
_tool(
|
|
"functions.send_email:2",
|
|
'{"to": "user@example.com", "subject": "Daily Update"}',
|
|
),
|
|
),
|
|
["get_weather", "get_news", "send_email"],
|
|
[
|
|
{"city": "New York"},
|
|
{"topic": "technology"},
|
|
{"to": "user@example.com", "subject": "Daily Update"},
|
|
],
|
|
"Multiple tasks. ",
|
|
id="three_tool_calls",
|
|
),
|
|
pytest.param(
|
|
"Process HTML. "
|
|
+ _wrap(
|
|
_tool("functions.process_html:0", '{"html": "<div>content</div>"}')
|
|
),
|
|
["process_html"],
|
|
[{"html": "<div>content</div>"}],
|
|
"Process HTML. ",
|
|
id="angle_brackets_in_json",
|
|
),
|
|
pytest.param(
|
|
"Formatted. "
|
|
+ _wrap(
|
|
_tool(
|
|
"functions.process_data:0",
|
|
'{\n "name": "test",\n "value": 123\n}',
|
|
)
|
|
),
|
|
["process_data"],
|
|
[{"name": "test", "value": 123}],
|
|
"Formatted. ",
|
|
id="multiline_json",
|
|
),
|
|
pytest.param(
|
|
"No prefix. " + _wrap(_tool("get_weather:0", '{"city": "Tokyo"}')),
|
|
["get_weather"],
|
|
[{"city": "Tokyo"}],
|
|
"No prefix. ",
|
|
id="no_functions_prefix",
|
|
),
|
|
pytest.param(
|
|
"Empty args. " + _wrap(_tool("functions.test:0", "{}")),
|
|
["test"],
|
|
[{}],
|
|
"Empty args. ",
|
|
id="empty_arguments",
|
|
),
|
|
],
|
|
)
|
|
def test_extract_tool_calls(
|
|
self, parser, model_output, expected_names, expected_args_list, expected_content
|
|
):
|
|
content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)
|
|
assert content == expected_content
|
|
assert len(tool_calls) == len(expected_names)
|
|
for tc, name, expected_args in zip(
|
|
tool_calls, expected_names, expected_args_list
|
|
):
|
|
assert tc.type == "function"
|
|
assert tc.function.name == name
|
|
assert json.loads(tc.function.arguments) == expected_args
|
|
# id format: "something:digit"
|
|
assert tc.id.split(":")[-1].isdigit()
|
|
|
|
def test_invalid_json_still_extracted(self, parser):
|
|
"""Tool calls with invalid JSON are still returned (arguments as-is)."""
|
|
model_output = (
|
|
"Help. "
|
|
+ SECTION_BEGIN
|
|
+ _tool("functions.bad:0", '{"city": "Beijing"')
|
|
+ _tool("functions.good:1", '{"city": "Shanghai"}')
|
|
+ SECTION_END
|
|
)
|
|
content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)
|
|
assert len(tool_calls) == 2
|
|
assert tool_calls[0].function.name == "bad"
|
|
assert tool_calls[1].function.name == "good"
|
|
|
|
def test_invalid_funcall_id_skipped(self, parser):
|
|
"""Tool calls with malformed id (no colon+digit) are skipped."""
|
|
model_output = (
|
|
"Help. "
|
|
+ SECTION_BEGIN
|
|
+ _tool("functions.invalid.0", '{"city": "Beijing"}')
|
|
+ _tool("functions.valid:1", '{"city": "Shanghai"}')
|
|
+ SECTION_END
|
|
)
|
|
content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)
|
|
assert len(tool_calls) == 1
|
|
assert tool_calls[0].function.name == "valid"
|
|
|
|
def test_native_id_extracted(self, parser):
|
|
"""Regression: parser extracts native ID onto ToolCall (PR #32768)."""
|
|
model_output = "Checking weather. " + _wrap(
|
|
_tool("functions.get_weather:0", '{"city": "Tokyo"}')
|
|
)
|
|
content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)
|
|
assert len(tool_calls) == 1
|
|
assert tool_calls[0].id == "functions.get_weather:0"
|
|
assert tool_calls[0].function.name == "get_weather"
|
|
assert json.loads(tool_calls[0].function.arguments) == {"city": "Tokyo"}
|
|
|
|
def test_multi_turn_native_id_continuity(self, kimi_k2_tokenizer):
|
|
"""Regression: native IDs from turn 1 preserved across turns (PR #32768)."""
|
|
turn1_parser = KimiK2ToolParser(kimi_k2_tokenizer)
|
|
turn1_output = "Let me check. " + _wrap(
|
|
_tool("functions.get_weather:0", '{"city": "Beijing"}')
|
|
)
|
|
_, turn1_tools = run_tool_extraction(
|
|
turn1_parser, turn1_output, streaming=False
|
|
)
|
|
assert len(turn1_tools) == 1
|
|
assert turn1_tools[0].id == "functions.get_weather:0"
|
|
|
|
# Fresh parser for turn 2
|
|
turn2_parser = KimiK2ToolParser(kimi_k2_tokenizer)
|
|
turn2_output = "Now let me get news. " + _wrap(
|
|
_tool("functions.get_news:0", '{"topic": "weather in Beijing"}')
|
|
)
|
|
_, turn2_tools = run_tool_extraction(
|
|
turn2_parser, turn2_output, streaming=False
|
|
)
|
|
assert len(turn2_tools) == 1
|
|
assert turn2_tools[0].id == "functions.get_news:0"
|
|
|
|
|
|
def _split_tool_output_to_deltas(
|
|
content: str, tool_strs: list[tuple[str, str]]
|
|
) -> list[str]:
|
|
"""Build a list of string deltas with special tokens as separate chunks.
|
|
|
|
Args:
|
|
content: text before tool section
|
|
tool_strs: list of (tool_id, args_json)
|
|
"""
|
|
deltas = [content, SECTION_BEGIN]
|
|
for tool_id, args_json in tool_strs:
|
|
deltas.extend(
|
|
[
|
|
TOOL_BEGIN,
|
|
f"{tool_id} ",
|
|
ARG_BEGIN,
|
|
f"{args_json} ",
|
|
TOOL_END,
|
|
]
|
|
)
|
|
deltas.append(SECTION_END)
|
|
return deltas
|
|
|
|
|
|
class TestStreamingHappyPath:
|
|
def test_single_tool_call(self, parser):
|
|
"""Verify DeltaToolCall output: name, id, arguments for one tool."""
|
|
deltas = _split_tool_output_to_deltas(
|
|
"I'll help. ",
|
|
[("functions.get_weather:0", '{"city": "Beijing"}')],
|
|
)
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert len(rec.tool_calls) == 1
|
|
tc = rec.tool_calls[0]
|
|
assert tc.function.name == "get_weather"
|
|
assert tc.id == "functions.get_weather:0"
|
|
assert json.loads(tc.function.arguments) == {"city": "Beijing"}
|
|
|
|
def test_multiple_tool_calls(self, parser):
|
|
"""Two tool calls emitted with correct indices, names, arguments."""
|
|
deltas = _split_tool_output_to_deltas(
|
|
"Compare weather. ",
|
|
[
|
|
("functions.get_weather:0", '{"city": "Tokyo"}'),
|
|
("functions.get_weather:1", '{"city": "NYC"}'),
|
|
],
|
|
)
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert len(rec.tool_calls) == 2
|
|
assert rec.tool_calls[0].function.name == "get_weather"
|
|
assert rec.tool_calls[0].id == "functions.get_weather:0"
|
|
assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"}
|
|
|
|
assert rec.tool_calls[1].function.name == "get_weather"
|
|
assert rec.tool_calls[1].id == "functions.get_weather:1"
|
|
assert json.loads(rec.tool_calls[1].function.arguments) == {"city": "NYC"}
|
|
|
|
def test_content_before_tools(self, parser):
|
|
"""Content before section is streamed; markers/args don't leak."""
|
|
deltas = _split_tool_output_to_deltas(
|
|
"I'll check the weather. ",
|
|
[("functions.get_weather:0", '{"city": "Tokyo"}')],
|
|
)
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert "check the weather" in rec.other_content
|
|
# No markers or tool content leaked
|
|
for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]:
|
|
assert marker not in rec.other_content
|
|
assert "get_weather" not in rec.other_content
|
|
assert "Tokyo" not in rec.other_content
|
|
|
|
def test_no_tool_calls(self, parser):
|
|
"""Plain text streaming returns content only."""
|
|
deltas = ["This is just ", "regular text ", "without tools."]
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert rec.other_content == "This is just regular text without tools."
|
|
assert rec.tool_calls == []
|
|
|
|
def test_incremental_arguments(self, parser):
|
|
"""Arguments split across small chunks accumulate correctly."""
|
|
deltas = [
|
|
"Help. ",
|
|
SECTION_BEGIN,
|
|
TOOL_BEGIN,
|
|
"functions.get_weather:0 ",
|
|
ARG_BEGIN,
|
|
'{"ci',
|
|
'ty": "Be',
|
|
'ijing"}',
|
|
" ",
|
|
TOOL_END,
|
|
SECTION_END,
|
|
]
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert len(rec.tool_calls) == 1
|
|
assert rec.tool_calls[0].function.name == "get_weather"
|
|
assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"}
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_output",
|
|
[
|
|
pytest.param(
|
|
"Single. "
|
|
+ _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')),
|
|
id="single_tool",
|
|
),
|
|
pytest.param(
|
|
"Multi. "
|
|
+ _wrap(
|
|
_tool("functions.get_weather:0", '{"city": "Tokyo"}'),
|
|
_tool("functions.get_news:1", '{"topic": "tech"}'),
|
|
),
|
|
id="parallel_tools",
|
|
),
|
|
pytest.param(
|
|
"No prefix id. " + _wrap(_tool("get_weather:0", '{"city": "NYC"}')),
|
|
id="no_functions_prefix",
|
|
),
|
|
],
|
|
)
|
|
def test_streaming_matches_nonstreaming(self, parser, model_output):
|
|
"""Streaming reconstruction matches non-streaming extraction."""
|
|
content_non, tools_non = run_tool_extraction(
|
|
parser, model_output, streaming=False
|
|
)
|
|
content_stream, tools_stream = run_tool_extraction(
|
|
parser, model_output, streaming=True
|
|
)
|
|
|
|
assert len(tools_non) == len(tools_stream)
|
|
for tc_non, tc_stream in zip(tools_non, tools_stream):
|
|
assert tc_non.function.name == tc_stream.function.name
|
|
assert json.loads(tc_non.function.arguments) == json.loads(
|
|
tc_stream.function.arguments
|
|
)
|
|
|
|
|
|
class TestStreamingEdgeCases:
|
|
def test_marker_suppression(self, parser):
|
|
"""No special-token markers appear in reconstructed content."""
|
|
deltas = _split_tool_output_to_deltas(
|
|
"I'll check. ",
|
|
[("functions.get_weather:0", '{"city": "Tokyo"}')],
|
|
)
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
forbidden = [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]
|
|
for marker in forbidden:
|
|
assert marker not in rec.other_content, (
|
|
f"Marker leaked: {marker!r} in {rec.other_content!r}"
|
|
)
|
|
|
|
def test_noise_between_markers_suppressed(self, parser):
|
|
"""Text between section_begin and tool_call_begin doesn't leak."""
|
|
deltas = [
|
|
"Reasoning. ",
|
|
SECTION_BEGIN,
|
|
" spurious noise ",
|
|
TOOL_BEGIN,
|
|
"functions.test:0 ",
|
|
ARG_BEGIN,
|
|
'{"k": "v"} ',
|
|
TOOL_END,
|
|
SECTION_END,
|
|
]
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert "spurious" not in rec.other_content
|
|
assert "noise" not in rec.other_content
|
|
|
|
def test_empty_tool_section(self, parser):
|
|
"""Empty section (begin immediately followed by end) doesn't crash."""
|
|
deltas = ["Reasoning. ", SECTION_BEGIN, SECTION_END]
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert rec.tool_calls == []
|
|
|
|
def test_three_different_tools(self, parser):
|
|
"""Three tool calls with different functions stream correctly."""
|
|
deltas = _split_tool_output_to_deltas(
|
|
"Multiple tasks. ",
|
|
[
|
|
("functions.get_weather:0", '{"city": "NYC"}'),
|
|
("functions.get_news:1", '{"topic": "tech"}'),
|
|
("functions.send_email:2", '{"to": "a@b.com"}'),
|
|
],
|
|
)
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
assert len(rec.tool_calls) == 3
|
|
names = [tc.function.name for tc in rec.tool_calls]
|
|
assert names == ["get_weather", "get_news", "send_email"]
|
|
ids = [tc.id for tc in rec.tool_calls]
|
|
assert len(set(ids)) == 3 # unique ids
|
|
|
|
def test_truncated_tool_call_no_end_marker(self, parser):
|
|
"""Stream ending mid-tool-call (max_tokens) doesn't crash."""
|
|
deltas = [
|
|
"I'll check. ",
|
|
SECTION_BEGIN,
|
|
TOOL_BEGIN,
|
|
"functions.get_weather:0 ",
|
|
ARG_BEGIN,
|
|
'{"city": "Bei',
|
|
# Stream ends here — no TOOL_END, no SECTION_END
|
|
]
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
# Should not crash; tool name and partial args extracted
|
|
assert len(rec.tool_calls) == 1
|
|
assert rec.tool_calls[0].function.name == "get_weather"
|
|
assert rec.tool_calls[0].id == "functions.get_weather:0"
|
|
assert rec.tool_calls[0].function.arguments == '{"city": "Bei'
|
|
# No markers leaked into content
|
|
for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]:
|
|
assert marker not in rec.other_content
|
|
|
|
def test_content_after_tool_section(self, parser):
|
|
"""Trailing text after section_end doesn't crash or leak markers."""
|
|
deltas = [
|
|
"Before. ",
|
|
SECTION_BEGIN,
|
|
TOOL_BEGIN,
|
|
"functions.get_weather:0 ",
|
|
ARG_BEGIN,
|
|
'{"city": "Tokyo"} ',
|
|
TOOL_END,
|
|
SECTION_END,
|
|
" After tools.",
|
|
]
|
|
rec = run_tool_extraction_streaming(parser, deltas)
|
|
|
|
# Tool call extracted correctly
|
|
assert len(rec.tool_calls) == 1
|
|
assert rec.tool_calls[0].function.name == "get_weather"
|
|
assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"}
|
|
# Trailing content after tool section is dropped
|
|
assert "After tools." not in rec.other_content
|
|
# No markers leaked into content
|
|
for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]:
|
|
assert marker not in rec.other_content
|
|
|
|
|
|
class TestAdjustRequest:
|
|
def test_sets_skip_special_tokens_false(self, parser):
|
|
request = MagicMock(spec=ChatCompletionRequest)
|
|
request.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
request.tool_choice = "auto"
|
|
request.skip_special_tokens = True
|
|
|
|
result = parser.adjust_request(request)
|
|
assert result.skip_special_tokens is False
|
|
|
|
def test_no_change_when_tool_choice_none(self, parser):
|
|
request = MagicMock(spec=ChatCompletionRequest)
|
|
request.tools = [{"type": "function", "function": {"name": "test"}}]
|
|
request.tool_choice = "none"
|
|
request.skip_special_tokens = True
|
|
|
|
result = parser.adjust_request(request)
|
|
assert result.skip_special_tokens is True
|
|
|
|
def test_no_change_when_no_tools(self, parser):
|
|
request = MagicMock(spec=ChatCompletionRequest)
|
|
request.tools = None
|
|
request.tool_choice = "auto"
|
|
request.skip_special_tokens = True
|
|
|
|
result = parser.adjust_request(request)
|
|
assert result.skip_special_tokens is True
|
|
|
|
|
|
def _chunk_tokenized_deltas(tokenizer, text: str, stream_interval: int) -> list[str]:
|
|
"""Encode text, group tokens into chunks of stream_interval, decode each."""
|
|
token_ids = tokenizer.encode(text, add_special_tokens=False)
|
|
deltas = []
|
|
prev = ""
|
|
for i in range(0, len(token_ids), stream_interval):
|
|
decoded = tokenizer.decode(
|
|
token_ids[: i + stream_interval], skip_special_tokens=False
|
|
)
|
|
deltas.append(decoded[len(prev) :])
|
|
prev = decoded
|
|
return deltas
|
|
|
|
|
|
class TestStreamingIntervals:
|
|
"""Test streaming at various token-chunk sizes to catch boundary bugs."""
|
|
|
|
@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])
|
|
def test_single_tool_call_at_interval(self, kimi_k2_tokenizer, stream_interval):
|
|
text = "Help. " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}'))
|
|
deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)
|
|
parser = KimiK2ToolParser(kimi_k2_tokenizer)
|
|
rec = run_tool_extraction_streaming(
|
|
parser, deltas, assert_one_tool_per_delta=False
|
|
)
|
|
|
|
assert len(rec.tool_calls) == 1
|
|
assert rec.tool_calls[0].function.name == "get_weather"
|
|
assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"}
|
|
|
|
@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])
|
|
def test_content_then_tool_call_at_interval(
|
|
self, kimi_k2_tokenizer, stream_interval
|
|
):
|
|
text = "Sure, let me check. " + _wrap(
|
|
_tool("functions.get_weather:0", '{"city": "Tokyo"}')
|
|
)
|
|
deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)
|
|
parser = KimiK2ToolParser(kimi_k2_tokenizer)
|
|
rec = run_tool_extraction_streaming(
|
|
parser, deltas, assert_one_tool_per_delta=False
|
|
)
|
|
|
|
assert "let me check" in rec.other_content
|
|
assert "get_weather" not in rec.other_content
|
|
assert len(rec.tool_calls) == 1
|
|
assert rec.tool_calls[0].function.name == "get_weather"
|
|
assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"}
|
|
|
|
@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])
|
|
def test_multiple_tool_calls_at_interval(self, kimi_k2_tokenizer, stream_interval):
|
|
text = "Compare. " + _wrap(
|
|
_tool("functions.search:0", '{"q": "cats"}'),
|
|
_tool("functions.search:1", '{"q": "dogs"}'),
|
|
)
|
|
deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)
|
|
parser = KimiK2ToolParser(kimi_k2_tokenizer)
|
|
rec = run_tool_extraction_streaming(
|
|
parser, deltas, assert_one_tool_per_delta=False
|
|
)
|
|
|
|
assert len(rec.tool_calls) == 2
|
|
assert rec.tool_calls[0].function.name == "search"
|
|
assert json.loads(rec.tool_calls[0].function.arguments) == {"q": "cats"}
|
|
assert rec.tool_calls[1].function.name == "search"
|
|
assert json.loads(rec.tool_calls[1].function.arguments) == {"q": "dogs"}
|
|
|
|
@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])
|
|
def test_plain_text_at_interval(self, kimi_k2_tokenizer, stream_interval):
|
|
text = "This is plain text with no tool calling involved."
|
|
deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)
|
|
parser = KimiK2ToolParser(kimi_k2_tokenizer)
|
|
rec = run_tool_extraction_streaming(
|
|
parser, deltas, assert_one_tool_per_delta=False
|
|
)
|
|
|
|
assert rec.other_content == text
|
|
assert rec.tool_calls == []
|
|
|
|
def test_content_and_tool_call_in_single_chunk(self, kimi_k2_tokenizer):
|
|
"""Content + complete tool call in one chunk must both be emitted."""
|
|
text = "Hi! " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}'))
|
|
deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval=9999)
|
|
parser = KimiK2ToolParser(kimi_k2_tokenizer)
|
|
rec = run_tool_extraction_streaming(
|
|
parser, deltas, assert_one_tool_per_delta=False
|
|
)
|
|
|
|
assert "Hi!" in rec.other_content
|
|
assert "get_weather" not in rec.other_content
|
|
assert len(rec.tool_calls) == 1
|
|
assert rec.tool_calls[0].function.name == "get_weather"
|
|
assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"}
|