811 lines
31 KiB
Python
811 lines
31 KiB
Python
import json
|
|
import unittest
|
|
|
|
from sglang.srt.entrypoints.openai.protocol import Function, Tool
|
|
from sglang.srt.function_call.kimik2_detector import (
|
|
KimiK2Detector as KimiK2FuncDetector,
|
|
)
|
|
from sglang.srt.function_call.kimik2_detector import (
|
|
_strip_special_tokens,
|
|
)
|
|
from sglang.srt.parser.reasoning_parser import KimiK2Detector as KimiK2ReasoningDetector
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(5, "base-a-test-cpu")
|
|
|
|
|
|
def _make_tool(name, parameters=None):
|
|
"""Helper to create a Tool with less boilerplate."""
|
|
return Tool(
|
|
type="function",
|
|
function=Function(
|
|
name=name,
|
|
description=f"{name} tool",
|
|
parameters=parameters
|
|
or {
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {"type": "string", "description": "File path"},
|
|
},
|
|
"required": ["path"],
|
|
},
|
|
),
|
|
)
|
|
|
|
|
|
def _collect_streaming_tool_calls(detector, chunks, tools):
|
|
"""Run streaming chunks through a detector and collect assembled tool calls."""
|
|
tool_calls = []
|
|
all_normal_text = ""
|
|
for chunk in chunks:
|
|
result = detector.parse_streaming_increment(chunk, tools)
|
|
all_normal_text += result.normal_text
|
|
for tc_chunk in result.calls:
|
|
if tc_chunk.tool_index is not None:
|
|
while len(tool_calls) <= tc_chunk.tool_index:
|
|
tool_calls.append({"name": "", "parameters": ""})
|
|
tc = tool_calls[tc_chunk.tool_index]
|
|
if tc_chunk.name:
|
|
tc["name"] = tc_chunk.name
|
|
if tc_chunk.parameters:
|
|
tc["parameters"] += tc_chunk.parameters
|
|
return tool_calls, all_normal_text
|
|
|
|
|
|
# ============================================================
|
|
# Part 1: KimiK2Detector (function call parsing) tests
|
|
# ============================================================
|
|
|
|
|
|
class TestKimiK2DetectorBasic(unittest.TestCase):
|
|
"""Basic non-streaming parsing tests for KimiK2Detector."""
|
|
|
|
def setUp(self):
|
|
self.tools = [
|
|
_make_tool("ReadFile"),
|
|
_make_tool(
|
|
"get_weather",
|
|
{
|
|
"type": "object",
|
|
"properties": {
|
|
"city": {"type": "string"},
|
|
"unit": {"type": "string"},
|
|
},
|
|
"required": ["city"],
|
|
},
|
|
),
|
|
]
|
|
self.detector = KimiK2FuncDetector()
|
|
|
|
def test_single_tool_call(self):
|
|
"""Parse a single complete tool call."""
|
|
text = (
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/test.py"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = self.detector.detect_and_parse(text, self.tools)
|
|
self.assertEqual(len(result.calls), 1)
|
|
self.assertEqual(result.calls[0].name, "ReadFile")
|
|
self.assertEqual(result.calls[0].parameters, '{"path": "/test.py"}')
|
|
self.assertEqual(result.normal_text, "")
|
|
|
|
def test_multiple_tool_calls(self):
|
|
"""Parse two consecutive tool calls."""
|
|
text = (
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/a.py"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_call_begin|>functions.get_weather:1"
|
|
'<|tool_call_argument_begin|>{"city": "Tokyo"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = self.detector.detect_and_parse(text, self.tools)
|
|
self.assertEqual(len(result.calls), 2)
|
|
self.assertEqual(result.calls[0].name, "ReadFile")
|
|
self.assertEqual(result.calls[1].name, "get_weather")
|
|
self.assertEqual(result.calls[1].parameters, '{"city": "Tokyo"}')
|
|
|
|
def test_normal_text_before_tool_call(self):
|
|
"""Normal text before tool call markers is preserved."""
|
|
text = (
|
|
"Let me check the file."
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/test.py"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = self.detector.detect_and_parse(text, self.tools)
|
|
self.assertEqual(len(result.calls), 1)
|
|
self.assertEqual(result.normal_text, "Let me check the file.")
|
|
|
|
def test_no_tool_call(self):
|
|
"""Text without tool call markers returns as normal text."""
|
|
text = "Just a normal response."
|
|
result = self.detector.detect_and_parse(text, self.tools)
|
|
self.assertEqual(len(result.calls), 0)
|
|
self.assertEqual(result.normal_text, text)
|
|
|
|
def test_has_tool_call(self):
|
|
"""has_tool_call correctly detects the presence of tool call markers."""
|
|
self.assertTrue(
|
|
self.detector.has_tool_call("<|tool_calls_section_begin|>stuff")
|
|
)
|
|
self.assertFalse(self.detector.has_tool_call("no markers here"))
|
|
|
|
|
|
class TestKimiK2DetectorHyphenatedNames(unittest.TestCase):
|
|
"""Test support for hyphenated function names (common in MCP tools)."""
|
|
|
|
def setUp(self):
|
|
self.tools = [
|
|
_make_tool("mcp__portal__search-documents"),
|
|
_make_tool("list-files"),
|
|
]
|
|
self.detector = KimiK2FuncDetector()
|
|
|
|
def test_hyphenated_name_non_streaming(self):
|
|
"""Parse tool call with hyphenated function name."""
|
|
text = (
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.mcp__portal__search-documents:0"
|
|
'<|tool_call_argument_begin|>{"path": "/docs"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = self.detector.detect_and_parse(text, self.tools)
|
|
self.assertEqual(len(result.calls), 1)
|
|
self.assertEqual(result.calls[0].name, "mcp__portal__search-documents")
|
|
|
|
def test_hyphenated_name_streaming(self):
|
|
"""Stream tool call with hyphenated function name."""
|
|
chunks = [
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.list-files:0"
|
|
'<|tool_call_argument_begin|>{"path',
|
|
'": "/home"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
tool_calls, _ = _collect_streaming_tool_calls(self.detector, chunks, self.tools)
|
|
self.assertEqual(len(tool_calls), 1)
|
|
self.assertEqual(tool_calls[0]["name"], "list-files")
|
|
params = json.loads(tool_calls[0]["parameters"])
|
|
self.assertEqual(params["path"], "/home")
|
|
|
|
|
|
class TestKimiK2DetectorStreaming(unittest.TestCase):
|
|
"""Streaming incremental parsing tests for KimiK2Detector."""
|
|
|
|
def setUp(self):
|
|
self.tools = [
|
|
_make_tool("ReadFile"),
|
|
_make_tool(
|
|
"get_weather",
|
|
{
|
|
"type": "object",
|
|
"properties": {"city": {"type": "string"}},
|
|
"required": ["city"],
|
|
},
|
|
),
|
|
]
|
|
|
|
def test_streaming_single_tool_call(self):
|
|
"""Stream a single tool call across multiple chunks."""
|
|
detector = KimiK2FuncDetector()
|
|
chunks = [
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
"<|tool_call_argument_begin|>{",
|
|
'"path": "/test.py"',
|
|
"}",
|
|
"<|tool_call_end|><|tool_calls_section_end|>",
|
|
]
|
|
tool_calls, _ = _collect_streaming_tool_calls(detector, chunks, self.tools)
|
|
self.assertEqual(len(tool_calls), 1)
|
|
self.assertEqual(tool_calls[0]["name"], "ReadFile")
|
|
self.assertEqual(tool_calls[0]["parameters"], '{"path": "/test.py"}')
|
|
|
|
def test_streaming_multiple_tool_calls(self):
|
|
"""Stream two tool calls sequentially."""
|
|
detector = KimiK2FuncDetector()
|
|
chunks = [
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/a.py"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_call_begin|>functions.get_weather:1"
|
|
'<|tool_call_argument_begin|>{"city": "Paris"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
tool_calls, _ = _collect_streaming_tool_calls(detector, chunks, self.tools)
|
|
self.assertEqual(len(tool_calls), 2)
|
|
self.assertEqual(tool_calls[0]["name"], "ReadFile")
|
|
self.assertEqual(tool_calls[1]["name"], "get_weather")
|
|
self.assertEqual(json.loads(tool_calls[1]["parameters"]), {"city": "Paris"})
|
|
|
|
def test_streaming_state_reset_after_completion(self):
|
|
"""Buffer and state reset after tool call completes."""
|
|
detector = KimiK2FuncDetector()
|
|
chunks = [
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/x"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
for chunk in chunks:
|
|
detector.parse_streaming_increment(chunk, self.tools)
|
|
|
|
self.assertEqual(detector._buffer, "")
|
|
self.assertEqual(detector.current_tool_id, 1)
|
|
|
|
|
|
class TestKimiK2DetectorSpecialTokenLeakage(unittest.TestCase):
|
|
"""Verify special tokens are never leaked into normal_text output."""
|
|
|
|
def setUp(self):
|
|
self.tools = [_make_tool("ReadFile")]
|
|
|
|
def test_no_leak_in_non_tool_text(self):
|
|
"""End tokens appearing without start tokens are stripped from output."""
|
|
detector = KimiK2FuncDetector()
|
|
result = detector.parse_streaming_increment(
|
|
"normal text<|tool_calls_section_end|>", self.tools
|
|
)
|
|
self.assertNotIn("<|tool_calls_section_end|>", result.normal_text)
|
|
self.assertIn("normal text", result.normal_text)
|
|
|
|
def test_no_leak_of_argument_begin_token(self):
|
|
"""Argument begin token is stripped when leaked."""
|
|
detector = KimiK2FuncDetector()
|
|
result = detector.parse_streaming_increment(
|
|
"text<|tool_call_argument_begin|>more", self.tools
|
|
)
|
|
self.assertNotIn("<|tool_call_argument_begin|>", result.normal_text)
|
|
|
|
def test_no_leak_on_error_fallback(self):
|
|
"""On parse errors, normal_text fallback has tokens stripped."""
|
|
cleaned = _strip_special_tokens(
|
|
"leaked<|tool_calls_section_begin|>" "<|tool_call_end|>content"
|
|
)
|
|
self.assertEqual(cleaned, "leakedcontent")
|
|
|
|
def test_strip_special_tokens_all_tokens(self):
|
|
"""All 5 known special tokens are stripped."""
|
|
dirty = (
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>"
|
|
"<|tool_call_argument_begin|>"
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
self.assertEqual(_strip_special_tokens(dirty), "")
|
|
|
|
def test_strip_preserves_normal_text(self):
|
|
"""Stripping doesn't affect normal text content."""
|
|
text = "Hello world, this is normal text."
|
|
self.assertEqual(_strip_special_tokens(text), text)
|
|
|
|
|
|
# ============================================================
|
|
# Part 2: KimiK2ReasoningDetector tests
|
|
# ============================================================
|
|
|
|
|
|
class TestKimiK2ReasoningDetectorNonStreaming(unittest.TestCase):
|
|
"""Non-streaming tests for KimiK2ReasoningDetector."""
|
|
|
|
def test_normal_reasoning_with_think_end(self):
|
|
"""Standard case: <think>...</think> followed by tool call markers."""
|
|
det = KimiK2ReasoningDetector()
|
|
text = (
|
|
"<think>I need to check the file.</think>"
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/test.py"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = det.detect_and_parse(text)
|
|
self.assertEqual(result.reasoning_text, "I need to check the file.")
|
|
self.assertIn("<|tool_calls_section_begin|>", result.normal_text)
|
|
|
|
def test_tool_call_inside_think_without_close_tag(self):
|
|
"""
|
|
BUG FIX: Model outputs tool call markers inside <think> without </think>.
|
|
|
|
This is the primary scenario that caused special token leakage.
|
|
The model decides to call a tool while reasoning and directly outputs
|
|
<|tool_calls_section_begin|> without first closing with </think>.
|
|
"""
|
|
det = KimiK2ReasoningDetector()
|
|
text = (
|
|
"<think>Let me read this file..."
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/test.py"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = det.detect_and_parse(text)
|
|
|
|
# Reasoning content must NOT contain tool call tokens
|
|
self.assertNotIn("<|tool_calls_section_begin|>", result.reasoning_text)
|
|
self.assertNotIn("<|tool_call_begin|>", result.reasoning_text)
|
|
self.assertIn("Let me read this file...", result.reasoning_text)
|
|
self.assertNotIn("<think>", result.reasoning_text)
|
|
|
|
# Tool call markers must be in normal_text for downstream parsing
|
|
self.assertIn("<|tool_calls_section_begin|>", result.normal_text)
|
|
self.assertIn("<|tool_call_begin|>", result.normal_text)
|
|
|
|
def test_no_reasoning_just_tool_call(self):
|
|
"""No <think> block, just tool call markers — pass through as normal_text."""
|
|
det = KimiK2ReasoningDetector()
|
|
text = (
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/x"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = det.detect_and_parse(text)
|
|
self.assertEqual(result.reasoning_text, "")
|
|
self.assertIn("<|tool_calls_section_begin|>", result.normal_text)
|
|
|
|
def test_normal_text_without_reasoning(self):
|
|
"""Plain text without reasoning or tool calls."""
|
|
det = KimiK2ReasoningDetector()
|
|
result = det.detect_and_parse("Hello, how can I help?")
|
|
self.assertEqual(result.normal_text, "Hello, how can I help?")
|
|
self.assertEqual(result.reasoning_text, "")
|
|
|
|
|
|
class TestKimiK2ReasoningDetectorStreaming(unittest.TestCase):
|
|
"""Streaming tests for KimiK2ReasoningDetector."""
|
|
|
|
def _run_streaming(self, chunks, **kwargs):
|
|
"""Helper: run chunks through streaming detector, collect reasoning and normal text."""
|
|
det = KimiK2ReasoningDetector(**kwargs)
|
|
all_reasoning = ""
|
|
all_normal = ""
|
|
for chunk in chunks:
|
|
r = det.parse_streaming_increment(chunk)
|
|
all_reasoning += r.reasoning_text
|
|
all_normal += r.normal_text
|
|
return all_reasoning, all_normal
|
|
|
|
def test_streaming_normal_think_then_tool_call(self):
|
|
"""Standard streaming: <think>...</think> then tool call markers."""
|
|
reasoning, normal = self._run_streaming(
|
|
[
|
|
"<think>",
|
|
"Analyzing the request...",
|
|
"</think>",
|
|
"<|tool_calls_section_begin|>",
|
|
"<|tool_call_begin|>functions.ReadFile:0",
|
|
]
|
|
)
|
|
self.assertIn("Analyzing the request...", reasoning)
|
|
self.assertIn("<|tool_calls_section_begin|>", normal)
|
|
self.assertNotIn("<|tool_calls_section_begin|>", reasoning)
|
|
|
|
def test_streaming_tool_call_inside_think(self):
|
|
"""
|
|
BUG FIX (streaming): Tool call markers inside <think> without </think>.
|
|
|
|
This is the streaming equivalent of the primary bug. The model streams
|
|
reasoning content, then directly outputs tool call markers without </think>.
|
|
"""
|
|
reasoning, normal = self._run_streaming(
|
|
[
|
|
"<think>",
|
|
"I need to",
|
|
" read the file.",
|
|
"<|tool_calls_section_begin|>",
|
|
"<|tool_call_begin|>functions.ReadFile:5",
|
|
"<|tool_call_argument_begin|>",
|
|
'{"path": "/Users/user/project/file.ts"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
)
|
|
|
|
# Reasoning is clean
|
|
self.assertIn("I need to read the file.", reasoning)
|
|
self.assertNotIn("<|tool_calls_section_begin|>", reasoning)
|
|
self.assertNotIn("<|tool_call_begin|>", reasoning)
|
|
self.assertNotIn("<think>", reasoning)
|
|
|
|
# Tool call markers are in normal_text
|
|
self.assertIn("<|tool_calls_section_begin|>", normal)
|
|
self.assertIn("functions.ReadFile:5", normal)
|
|
|
|
def test_streaming_tool_call_marker_in_single_chunk(self):
|
|
"""Tool call marker arrives in a single chunk while in reasoning mode."""
|
|
reasoning, normal = self._run_streaming(
|
|
[
|
|
"<think>thinking...",
|
|
'<|tool_calls_section_begin|><|tool_call_begin|>functions.ReadFile:0<|tool_call_argument_begin|>{"path": "/x"}',
|
|
]
|
|
)
|
|
self.assertIn("thinking...", reasoning)
|
|
self.assertIn("<|tool_calls_section_begin|>", normal)
|
|
|
|
def test_streaming_partial_marker_buffering(self):
|
|
"""
|
|
Partial tool call marker at end of chunk is buffered to prevent
|
|
premature streaming of marker characters as reasoning content.
|
|
"""
|
|
det = KimiK2ReasoningDetector(stream_reasoning=True)
|
|
|
|
# First chunk: reasoning + partial marker "<|tool_calls"
|
|
det._in_reasoning = True
|
|
det.stripped_think_start = True
|
|
|
|
r1 = det.parse_streaming_increment("some reasoning")
|
|
self.assertEqual(r1.reasoning_text, "some reasoning")
|
|
|
|
# Chunk that ends with start of marker
|
|
r2 = det.parse_streaming_increment("<|tool")
|
|
# Partial marker should be buffered, not streamed
|
|
self.assertNotIn("<|tool", r2.reasoning_text)
|
|
|
|
# Complete the marker
|
|
r3 = det.parse_streaming_increment("_calls_section_begin|>rest")
|
|
# Now it should force-exit reasoning
|
|
self.assertIn("<|tool_calls_section_begin|>", r3.normal_text)
|
|
|
|
def test_streaming_no_reasoning_mode(self):
|
|
"""Normal text without reasoning passes through as normal_text."""
|
|
reasoning, normal = self._run_streaming(
|
|
[
|
|
"Hello, I can help with that.",
|
|
" What do you need?",
|
|
]
|
|
)
|
|
self.assertEqual(reasoning, "")
|
|
self.assertIn("Hello, I can help with that.", normal)
|
|
self.assertIn(" What do you need?", normal)
|
|
|
|
def test_streaming_force_reasoning(self):
|
|
"""With force_reasoning, content before </think> is reasoning."""
|
|
reasoning, normal = self._run_streaming(
|
|
[
|
|
"I should analyze this...",
|
|
"</think>",
|
|
"Here is the answer.",
|
|
],
|
|
force_reasoning=True,
|
|
)
|
|
self.assertIn("I should analyze this...", reasoning)
|
|
self.assertIn("Here is the answer.", normal)
|
|
|
|
|
|
# ============================================================
|
|
# Part 3: End-to-end integration tests
|
|
# ============================================================
|
|
|
|
|
|
class TestKimiK2EndToEnd(unittest.TestCase):
|
|
"""
|
|
End-to-end tests simulating the full flow:
|
|
reasoning parser -> tool call parser.
|
|
|
|
These test the exact bug scenario from the issue: Kimi-K2.5 outputs
|
|
tool call markers inside <think> blocks, which must be correctly
|
|
split between reasoning and tool call parsers.
|
|
"""
|
|
|
|
def setUp(self):
|
|
self.tools = [
|
|
_make_tool("ReadFile"),
|
|
_make_tool(
|
|
"get_weather",
|
|
{
|
|
"type": "object",
|
|
"properties": {"city": {"type": "string"}},
|
|
"required": ["city"],
|
|
},
|
|
),
|
|
]
|
|
|
|
def test_e2e_streaming_reasoning_to_tool_call(self):
|
|
"""
|
|
Full pipeline: streaming reasoning parser feeds into streaming tool call parser.
|
|
|
|
Simulates the exact path through serving_chat.py:
|
|
1. Model outputs <think>reasoning...<|tool_calls_section_begin|>...
|
|
2. ReasoningParser splits: reasoning_text + normal_text
|
|
3. FunctionCallParser receives normal_text and extracts tool calls
|
|
"""
|
|
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
|
|
tc_det = KimiK2FuncDetector()
|
|
|
|
streaming_chunks = [
|
|
"<think>",
|
|
"I need to read the file",
|
|
" to understand the code.",
|
|
"<|tool_calls_section_begin|>",
|
|
"<|tool_call_begin|>functions.ReadFile:0",
|
|
"<|tool_call_argument_begin|>",
|
|
'{"path": "/Users/user/project/file.ts"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
|
|
all_reasoning = ""
|
|
all_tc_calls = []
|
|
|
|
for chunk in streaming_chunks:
|
|
# Step 1: reasoning parser
|
|
r = reasoning_det.parse_streaming_increment(chunk)
|
|
all_reasoning += r.reasoning_text
|
|
|
|
# Step 2: feed normal_text into tool call parser (like serving_chat.py does)
|
|
if r.normal_text:
|
|
tc_result = tc_det.parse_streaming_increment(r.normal_text, self.tools)
|
|
all_tc_calls.extend(tc_result.calls)
|
|
|
|
# Verify reasoning content
|
|
self.assertIn("I need to read the file to understand the code.", all_reasoning)
|
|
self.assertNotIn("<|", all_reasoning)
|
|
|
|
# Verify tool calls were extracted
|
|
name_calls = [c for c in all_tc_calls if c.name]
|
|
self.assertEqual(len(name_calls), 1)
|
|
self.assertEqual(name_calls[0].name, "ReadFile")
|
|
|
|
param_calls = [c for c in all_tc_calls if c.parameters]
|
|
full_params = "".join(c.parameters for c in param_calls)
|
|
self.assertIn("/Users/user/project/file.ts", full_params)
|
|
|
|
def test_e2e_non_streaming_reasoning_to_tool_call(self):
|
|
"""Non-streaming pipeline: reason parser then tool call parser."""
|
|
reasoning_det = KimiK2ReasoningDetector()
|
|
tc_det = KimiK2FuncDetector()
|
|
|
|
text = (
|
|
"<think>Let me check this file."
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/src/main.py"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
|
|
# Step 1: reasoning parser
|
|
r = reasoning_det.detect_and_parse(text)
|
|
self.assertIn("Let me check this file.", r.reasoning_text)
|
|
self.assertNotIn("<|", r.reasoning_text)
|
|
|
|
# Step 2: tool call parser on normal_text
|
|
tc_result = tc_det.detect_and_parse(r.normal_text, self.tools)
|
|
self.assertEqual(len(tc_result.calls), 1)
|
|
self.assertEqual(tc_result.calls[0].name, "ReadFile")
|
|
self.assertEqual(
|
|
json.loads(tc_result.calls[0].parameters),
|
|
{"path": "/src/main.py"},
|
|
)
|
|
|
|
def test_e2e_normal_think_close_then_tool_call(self):
|
|
"""Standard case with </think> — should also work correctly."""
|
|
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
|
|
tc_det = KimiK2FuncDetector()
|
|
|
|
chunks = [
|
|
"<think>",
|
|
"Thinking about it...",
|
|
"</think>",
|
|
"<|tool_calls_section_begin|>",
|
|
"<|tool_call_begin|>functions.get_weather:0",
|
|
'<|tool_call_argument_begin|>{"city": "London"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
|
|
all_reasoning = ""
|
|
all_tc_calls = []
|
|
|
|
for chunk in chunks:
|
|
r = reasoning_det.parse_streaming_increment(chunk)
|
|
all_reasoning += r.reasoning_text
|
|
if r.normal_text:
|
|
tc_result = tc_det.parse_streaming_increment(r.normal_text, self.tools)
|
|
all_tc_calls.extend(tc_result.calls)
|
|
|
|
self.assertIn("Thinking about it...", all_reasoning)
|
|
name_calls = [c for c in all_tc_calls if c.name]
|
|
self.assertEqual(len(name_calls), 1)
|
|
self.assertEqual(name_calls[0].name, "get_weather")
|
|
|
|
def test_e2e_multiple_tool_calls_without_think_close(self):
|
|
"""Multiple tool calls inside <think> without </think>."""
|
|
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
|
|
tc_det = KimiK2FuncDetector()
|
|
|
|
chunks = [
|
|
"<think>",
|
|
"Let me check both files.",
|
|
"<|tool_calls_section_begin|>",
|
|
"<|tool_call_begin|>functions.ReadFile:0"
|
|
'<|tool_call_argument_begin|>{"path": "/a.py"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_call_begin|>functions.ReadFile:1"
|
|
'<|tool_call_argument_begin|>{"path": "/b.py"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
|
|
all_reasoning = ""
|
|
all_tc_calls = []
|
|
|
|
for chunk in chunks:
|
|
r = reasoning_det.parse_streaming_increment(chunk)
|
|
all_reasoning += r.reasoning_text
|
|
if r.normal_text:
|
|
tc_result = tc_det.parse_streaming_increment(r.normal_text, self.tools)
|
|
all_tc_calls.extend(tc_result.calls)
|
|
|
|
self.assertIn("Let me check both files.", all_reasoning)
|
|
self.assertNotIn("<|", all_reasoning)
|
|
|
|
name_calls = [c for c in all_tc_calls if c.name]
|
|
self.assertEqual(len(name_calls), 2)
|
|
self.assertEqual(name_calls[0].name, "ReadFile")
|
|
self.assertEqual(name_calls[1].name, "ReadFile")
|
|
|
|
|
|
# ============================================================
|
|
# Part 3: Bare-counter tool call ID parsing
|
|
# ============================================================
|
|
|
|
|
|
class TestKimiK2BareCounterParsing(unittest.TestCase):
|
|
"""Tests for bare numeric tool_call_id format (e.g., '3' instead of 'functions.ReadFile:0')."""
|
|
|
|
def setUp(self):
|
|
self.detector = KimiK2FuncDetector()
|
|
self.tools = [
|
|
_make_tool("ReadFile"),
|
|
_make_tool(
|
|
"get_weather",
|
|
{
|
|
"type": "object",
|
|
"properties": {
|
|
"city": {"type": "string"},
|
|
"unit": {"type": "string"},
|
|
},
|
|
"required": ["city"],
|
|
},
|
|
),
|
|
]
|
|
|
|
# --- _parse_tool_call_id ---
|
|
|
|
def test_standard_format_with_functions_prefix(self):
|
|
name, idx = self.detector._parse_tool_call_id(
|
|
"functions.ReadFile:0", self.tools
|
|
)
|
|
self.assertEqual(name, "ReadFile")
|
|
self.assertEqual(idx, 0)
|
|
|
|
def test_standard_format_without_functions_prefix(self):
|
|
name, idx = self.detector._parse_tool_call_id("ReadFile:1", self.tools)
|
|
self.assertEqual(name, "ReadFile")
|
|
self.assertEqual(idx, 1)
|
|
|
|
def test_bare_counter_single_tool(self):
|
|
single_tool = [_make_tool("search")]
|
|
name, idx = self.detector._parse_tool_call_id(
|
|
"3", single_tool, '{"query": "test"}'
|
|
)
|
|
self.assertEqual(name, "search")
|
|
self.assertEqual(idx, 3)
|
|
|
|
def test_bare_counter_infers_by_args(self):
|
|
name, idx = self.detector._parse_tool_call_id(
|
|
"0", self.tools, '{"city": "Tokyo"}'
|
|
)
|
|
self.assertEqual(name, "get_weather")
|
|
self.assertEqual(idx, 0)
|
|
|
|
def test_bare_counter_no_tools_returns_none(self):
|
|
name, idx = self.detector._parse_tool_call_id("5", [], '{"x": 1}')
|
|
self.assertIsNone(name)
|
|
self.assertEqual(idx, 5)
|
|
|
|
def test_bare_counter_no_args_multiple_tools_returns_none(self):
|
|
name, idx = self.detector._parse_tool_call_id("2", self.tools, None)
|
|
self.assertIsNone(name)
|
|
self.assertEqual(idx, 2)
|
|
|
|
def test_unexpected_format_returns_none(self):
|
|
name, idx = self.detector._parse_tool_call_id("some_garbage", self.tools)
|
|
self.assertIsNone(name)
|
|
self.assertEqual(idx, 0)
|
|
|
|
# --- _infer_tool_name ---
|
|
|
|
def test_infer_no_tools(self):
|
|
self.assertIsNone(self.detector._infer_tool_name([], '{"x": 1}'))
|
|
|
|
def test_infer_single_tool(self):
|
|
result = self.detector._infer_tool_name([_make_tool("only_one")], '{"x": 1}')
|
|
self.assertEqual(result, "only_one")
|
|
|
|
def test_infer_by_argument_overlap(self):
|
|
result = self.detector._infer_tool_name(
|
|
self.tools, '{"city": "Paris", "unit": "celsius"}'
|
|
)
|
|
self.assertEqual(result, "get_weather")
|
|
|
|
def test_infer_malformed_json_returns_none(self):
|
|
result = self.detector._infer_tool_name(self.tools, '{"city": "Par')
|
|
self.assertIsNone(result)
|
|
|
|
def test_infer_empty_args_returns_none(self):
|
|
self.assertIsNone(self.detector._infer_tool_name(self.tools, None))
|
|
self.assertIsNone(self.detector._infer_tool_name(self.tools, ""))
|
|
|
|
def test_infer_no_matching_props_returns_none(self):
|
|
tools_no_props = [
|
|
_make_tool("a", {"type": "object"}),
|
|
_make_tool("b", {"type": "object"}),
|
|
]
|
|
result = self.detector._infer_tool_name(tools_no_props, '{"x": 1}')
|
|
self.assertIsNone(result)
|
|
|
|
# --- detect_and_parse with bare counter (end-to-end) ---
|
|
|
|
def test_detect_and_parse_bare_counter(self):
|
|
text = (
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>0"
|
|
'<|tool_call_argument_begin|>{"city": "Tokyo"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = self.detector.detect_and_parse(text, self.tools)
|
|
self.assertEqual(len(result.calls), 1)
|
|
self.assertEqual(result.calls[0].name, "get_weather")
|
|
self.assertEqual(result.calls[0].parameters, '{"city": "Tokyo"}')
|
|
|
|
def test_detect_and_parse_bare_counter_skips_unknown(self):
|
|
text = (
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>0"
|
|
'<|tool_call_argument_begin|>{"unknown_key": "value"}'
|
|
"<|tool_call_end|>"
|
|
"<|tool_calls_section_end|>"
|
|
)
|
|
result = self.detector.detect_and_parse(text, self.tools)
|
|
# No tool props match, _infer_tool_name returns None, call is skipped
|
|
self.assertEqual(len(result.calls), 0)
|
|
|
|
def test_streaming_bare_counter_single_tool(self):
|
|
detector = KimiK2FuncDetector()
|
|
single_tool = [_make_tool("search")]
|
|
chunks = [
|
|
"<|tool_calls_section_begin|>"
|
|
"<|tool_call_begin|>0"
|
|
'<|tool_call_argument_begin|>{"path',
|
|
'": "/test"}',
|
|
"<|tool_call_end|>",
|
|
"<|tool_calls_section_end|>",
|
|
]
|
|
tool_calls, _ = _collect_streaming_tool_calls(detector, chunks, single_tool)
|
|
self.assertEqual(len(tool_calls), 1)
|
|
self.assertEqual(tool_calls[0]["name"], "search")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|