minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

kimik2_detector fix the normal text detection before tool call. (#25071)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Rohit Harkhani <rharkhani@gmail.com>
This commit is contained in:
Rohit Harkhani
2026-06-24 15:53:22 +08:00
committed by GitHub
co-authored by gemini-code-assist[bot] Rohit Harkhani
parent b6a8000473
commit 5f76736427
2 changed files with 771 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/function_call/kimik2_detector.py
+184 -98
View File
@@ -15,7 +15,6 @@ from sglang.srt.function_call.core_types import (
ToolCallItem,
_GetInfoFunc,
)
from sglang.srt.function_call.utils import _is_complete_json
logger = logging.getLogger(__name__)
@@ -177,9 +176,15 @@ class KimiK2Detector(BaseFormatDetector):
logger.debug("function_call_tuples: %s", function_call_tuples)
tool_calls = []
# ``tool_index`` is the per-response 0-based position of the call
# (OpenAI spec); enumerate parsed calls locally and ignore the
# model's ``:N`` suffix, which is a conversation-level counter.
# ``serving_chat._process_tool_call_id()`` later offsets these by
# ``history_tool_calls_cnt`` for multi-turn responses.
local_tool_index = 0
for match in function_call_tuples:
function_id, function_args = match
function_name, function_idx = self._parse_tool_call_id(
function_name, _ = self._parse_tool_call_id(
function_id, tools, function_args
)
if function_name is None:
@@ -189,11 +194,12 @@ class KimiK2Detector(BaseFormatDetector):
tool_calls.append(
ToolCallItem(
tool_index=function_idx,
tool_index=local_tool_index,
name=function_name,
parameters=function_args,
)
)
local_tool_index += 1
content = text[: text.find(self.bot_token)]
return StreamingParseResult(normal_text=content, calls=tool_calls)
@@ -205,127 +211,207 @@ class KimiK2Detector(BaseFormatDetector):
def parse_streaming_increment(
self, new_text: str, tools: List[Tool]
) -> StreamingParseResult:
"""
Streaming incremental parsing tool calls for KimiK2 format.
"""
"""Streaming incremental parsing tool calls for KimiK2 format."""
self._buffer += new_text
current_text = self._buffer
# Check if we have a tool call (either the start token or individual tool call)
has_tool_call = (
self.bot_token in current_text or self.tool_call_start_token in current_text
)
if not has_tool_call:
self._buffer = ""
normal_text = _strip_special_tokens(new_text)
return StreamingParseResult(normal_text=normal_text)
# Fast path: no tool call in flight and no markers yet -- emit as
# normal text, holding back any trailing partial start token.
if (
self._current_stream_function_name is None
and self.bot_token not in self._buffer
and self.tool_call_start_token not in self._buffer
):
emit, hold = self._split_pending_start(self._buffer)
self._buffer = hold
return StreamingParseResult(normal_text=_strip_special_tokens(emit))
if not hasattr(self, "_tool_indices"):
self._tool_indices = self._get_tool_indices(tools)
normal_text_parts: list[str] = []
calls: list[ToolCallItem] = []
try:
match = self.stream_tool_call_portion_regex.search(current_text)
if match:
function_id = match.group("tool_call_id")
function_args = match.group("function_arguments")
while True:
buffer = self._buffer
# Reuse cached name for current tool call to avoid repeated
# json.loads on partial JSON in _infer_tool_name.
if self._current_stream_function_name is not None:
function_name = self._current_stream_function_name
else:
function_name, _ = self._parse_tool_call_id(
function_id, tools, function_args
# Locate next <|tool_call_begin|>, draining any prefix as text.
begin_idx = self._locate_tool_call_start(buffer, normal_text_parts)
if begin_idx is None:
break
buffer = self._buffer
# If another <|tool_call_begin|> appears before the header
# closes with <|tool_call_argument_begin|>, the section is
# malformed -- discard and restart at the orphan.
arg_begin_idx = buffer.find(self.tool_call_argument_begin_token)
next_begin = buffer.find(
self.tool_call_start_token, len(self.tool_call_start_token)
)
if next_begin != -1 and (
arg_begin_idx == -1 or next_begin < arg_begin_idx
):
logger.warning(
"Kimi-K2 tool_call_begin without preceding tool_call_end; "
"discarding incomplete section."
)
if function_name is None:
return StreamingParseResult(normal_text="", calls=calls)
self._buffer = buffer[next_begin:]
self._reset_inflight_call_state()
continue
# Initialize state if this is the first tool call
if self.current_tool_id == -1:
self.current_tool_id = 0
self.prev_tool_call_arr = []
self.streamed_args_for_tool = [""]
if arg_begin_idx == -1:
# Header not fully arrived yet.
break
# Ensure we have enough entries in our tracking arrays
while len(self.prev_tool_call_arr) <= self.current_tool_id:
self.prev_tool_call_arr.append({})
while len(self.streamed_args_for_tool) <= self.current_tool_id:
self.streamed_args_for_tool.append("")
id_start = len(self.tool_call_start_token)
function_id = buffer[id_start:arg_begin_idx].strip()
args_start = arg_begin_idx + len(self.tool_call_argument_begin_token)
end_idx = buffer.find(self.tool_call_end_token)
if not self.current_tool_name_sent:
# Resolve function name (cached across chunks within a section).
name_just_resolved = False
if self._current_stream_function_name is None:
args_for_inference = (
buffer[args_start:end_idx]
if end_idx != -1
else buffer[args_start:]
)
resolved = self._resolve_function_name(
function_id, tools, args_for_inference
)
if resolved is None:
if end_idx == -1:
# Wait for the end marker before deciding.
break
logger.warning(
"Kimi-K2 unrecognized tool_call_id %r; skipping section.",
function_id,
)
self._buffer = buffer[end_idx + len(self.tool_call_end_token) :]
self._reset_inflight_call_state()
continue
name = resolved
self._current_stream_function_name = name
name_just_resolved = True
# ``tool_index`` is the per-response 0-based position
# (OpenAI streaming spec); ignore the model's ``:N`` suffix
# which is a conversation-level counter.
if self.current_tool_id == -1:
self.current_tool_id = 0
self.prev_tool_call_arr = []
self.streamed_args_for_tool = [""]
while len(self.prev_tool_call_arr) <= self.current_tool_id:
self.prev_tool_call_arr.append({})
while len(self.streamed_args_for_tool) <= self.current_tool_id:
self.streamed_args_for_tool.append("")
self.prev_tool_call_arr[self.current_tool_id] = {
"name": name,
"arguments": {},
}
self.current_tool_name_sent = True
# Stream newly-arrived args, combining the first event with
# the freshly-resolved name.
if end_idx != -1:
args_full = buffer[args_start:end_idx]
else:
args_full = buffer[args_start:]
argument_diff = args_full[len(self._last_arguments) :]
if argument_diff or name_just_resolved:
calls.append(
ToolCallItem(
tool_index=self.current_tool_id,
name=function_name,
parameters="",
name=(
self._current_stream_function_name
if name_just_resolved
else None
),
parameters=argument_diff,
)
)
self.current_tool_name_sent = True
self._current_stream_function_name = function_name
self.prev_tool_call_arr[self.current_tool_id] = {
"name": function_name,
"arguments": {},
}
else:
argument_diff = (
function_args[len(self._last_arguments) :]
if function_args.startswith(self._last_arguments)
else function_args
)
parsed_args_diff = argument_diff.split(self.tool_call_end_token, 1)[
0
]
if parsed_args_diff:
calls.append(
ToolCallItem(
tool_index=self.current_tool_id,
name=None,
parameters=parsed_args_diff,
)
)
self._last_arguments += parsed_args_diff
if argument_diff:
self._last_arguments += argument_diff
self.streamed_args_for_tool[
self.current_tool_id
] += parsed_args_diff
] += argument_diff
parsed_args = function_args.split(self.tool_call_end_token, 1)[0]
if _is_complete_json(parsed_args):
try:
parsed_args = json.loads(parsed_args)
self.prev_tool_call_arr[self.current_tool_id][
"arguments"
] = parsed_args
except json.JSONDecodeError:
pass
if end_idx == -1:
# Args still streaming.
break
# Find the end of the current tool call and remove only that part from buffer
tool_call_end_pattern = (
r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>"
)
end_match = re.search(
tool_call_end_pattern, current_text, re.DOTALL
)
if end_match:
self._buffer = current_text[end_match.end() :]
else:
self._buffer = ""
# Section finalized -- advance buffer and prepare next call.
self._buffer = buffer[end_idx + len(self.tool_call_end_token) :]
self.current_tool_id += 1
self._reset_inflight_call_state()
result = StreamingParseResult(normal_text="", calls=calls)
self.current_tool_id += 1
self._last_arguments = ""
self.current_tool_name_sent = False
self._current_stream_function_name = None
return result
return StreamingParseResult(normal_text="", calls=calls)
return StreamingParseResult(
normal_text="".join(normal_text_parts), calls=calls
)
except Exception as e:
logger.error("Error in parse_streaming_increment: %s", e, exc_info=True)
return StreamingParseResult(normal_text=_strip_special_tokens(current_text))
# Drop the buffer to avoid leaking raw special tokens.
self._buffer = ""
self._reset_inflight_call_state()
return StreamingParseResult(
normal_text="".join(normal_text_parts), calls=calls
)
def _reset_inflight_call_state(self) -> None:
"""Reset per-section streaming state after finalize/discard."""
self._last_arguments = ""
self.current_tool_name_sent = False
self._current_stream_function_name = None
def _locate_tool_call_start(
self, buffer: str, normal_text_parts: list
) -> int | None:
"""Find the next <|tool_call_begin|>; drain any prefix as normal text.
Returns 0 on success, or ``None`` when no start token is present yet.
"""
begin_idx = buffer.find(self.tool_call_start_token)
if begin_idx == -1:
emit, hold = self._split_pending_start(buffer)
if emit:
normal_text_parts.append(_strip_special_tokens(emit))
self._buffer = hold
return None
if begin_idx > 0:
normal_text_parts.append(_strip_special_tokens(buffer[:begin_idx]))
self._buffer = buffer[begin_idx:]
return 0
def _split_pending_start(self, text: str) -> tuple[str, str]:
"""Hold back a trailing fragment that could be the start of
<|tool_calls_section_begin|> or <|tool_call_begin|>. Everything
before it is safe to emit as normal text.
"""
candidates = (self.bot_token, self.tool_call_start_token)
max_tail = max(len(t) for t in candidates) - 1
for n in range(min(len(text), max_tail), 1, -1):
tail = text[-n:]
if any(t.startswith(tail) for t in candidates):
return text[:-n], tail
return text, ""
def _resolve_function_name(
self, function_id: str, tools: List[Tool], function_args: str
) -> Optional[str]:
"""Map a Kimi-K2 tool_call_id to a tool name, or ``None`` if unknown."""
if not function_id:
return self._infer_tool_name(tools, function_args)
m = self.tool_call_id_regex.match(function_id)
if m:
return m.group("name")
if self.tool_call_id_counter_regex.match(function_id):
return self._infer_tool_name(tools, function_args)
return None
def structure_info(self) -> _GetInfoFunc:
"""Return function that creates StructureInfo for guided generation."""
test/registered/function_call/test_kimik2_detector.py
+587
View File
@@ -111,6 +111,32 @@ class TestKimiK2DetectorBasic(unittest.TestCase):
self.assertEqual(result.calls[1].name, "get_weather")
self.assertEqual(result.calls[1].parameters, '{"city": "Tokyo"}')
def test_non_streaming_tool_index_is_local(self):
"""tool_index is the per-response 0-based position, not the model's :N suffix.
The model may emit conversation-level ``:N`` counters (e.g. ``:5``, ``:6``)
in a multi-turn conversation. The non-streaming parser must enumerate
parsed calls locally (0, 1, ...) so that
``serving_chat._process_tool_call_id()`` can offset them by
``history_tool_calls_cnt`` without double-counting.
"""
text = (
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.ReadFile:5"
'<|tool_call_argument_begin|>{"path": "/a.py"}'
"<|tool_call_end|>"
"<|tool_call_begin|>functions.get_weather:6"
'<|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].tool_index, 0)
self.assertEqual(result.calls[0].name, "ReadFile")
self.assertEqual(result.calls[1].tool_index, 1)
self.assertEqual(result.calls[1].name, "get_weather")
def test_normal_text_before_tool_call(self):
"""Normal text before tool call markers is preserved."""
text = (
@@ -183,6 +209,15 @@ class TestKimiK2DetectorHyphenatedNames(unittest.TestCase):
class TestKimiK2DetectorStreaming(unittest.TestCase):
"""Streaming incremental parsing tests for KimiK2Detector."""
def test_streaming_trailing_literal_left_angle_is_not_dropped(self):
"""A final literal '<' must remain in normal_text instead of being buffered away."""
detector = KimiK2FuncDetector()
result = detector.parse_streaming_increment("normal text <", [])
self.assertEqual(result.normal_text, "normal text <")
self.assertEqual(detector._buffer, "")
def setUp(self):
self.tools = [
_make_tool("ReadFile"),
@@ -627,6 +662,409 @@ class TestKimiK2EndToEnd(unittest.TestCase):
self.assertEqual(len(name_calls), 1)
self.assertEqual(name_calls[0].name, "get_weather")
def test_e2e_normal_think_close_then_content_overlap_tool_call(self):
"""Trailing content between ``</think>`` and the tool-call markers
(e.g. ``"This is a content:"``) must be surfaced as ``normal_text``
by the tool-call parser and not stripped — this is the exact bug the
PR fixes.
"""
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
tc_det = KimiK2FuncDetector()
chunks = [
"<think>",
"Thinking about it...",
"</think>This is a ",
"content:<|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_content = ""
toolcall_chunks = []
for chunk in chunks:
r = reasoning_det.parse_streaming_increment(chunk)
all_reasoning += r.reasoning_text
if r.normal_text:
toolcall_chunks.append(r.normal_text)
tool_calls, all_content = _collect_streaming_tool_calls(
tc_det, toolcall_chunks, self.tools
)
self.assertEqual("Thinking about it...", all_reasoning)
self.assertEqual("This is a content:", all_content)
self.assertEqual(len(tool_calls), 1)
first_call = tool_calls.pop()
self.assertEqual(first_call["name"], "get_weather")
self.assertEqual(first_call["parameters"], '{"city": "London"}')
def test_e2e_normal_think_close_then_content_overlap_tool_call_multi_token(self):
"""Speculative decoding: a single chunk may contain normal text followed
by tool-call markers and even the tool_call_begin/id. The normal-text
prefix must still be emitted (not stripped) by the tool-call parser.
"""
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
tc_det = KimiK2FuncDetector()
chunks = [
"<think>",
"Thinking about it...",
"</think>This is a ",
"content:<|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_content = ""
toolcall_chunks = []
for chunk in chunks:
r = reasoning_det.parse_streaming_increment(chunk)
all_reasoning += r.reasoning_text
if r.normal_text:
toolcall_chunks.append(r.normal_text)
tool_calls, all_content = _collect_streaming_tool_calls(
tc_det, toolcall_chunks, self.tools
)
self.assertEqual("Thinking about it...", all_reasoning)
self.assertEqual("This is a content:", all_content)
self.assertEqual(len(tool_calls), 1)
first_call = tool_calls.pop()
self.assertEqual(first_call["name"], "get_weather")
self.assertEqual(first_call["parameters"], '{"city": "London"}')
def test_e2e_normal_think_close_then_content_overlap_tool_call_multi_token_multi_calls(
self,
):
"""Speculative decoding: a single chunk may contain normal text followed
by tool-call markers and even the tool_call_begin/id. Additionally, this
single chunk packs two complete tool-call sections back-to-back, which
exercises the ``while True:`` drain loop introduced by this PR — both
calls must be emitted from one ``parse_streaming_increment`` invocation.
"""
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
tc_det = KimiK2FuncDetector()
chunks = [
"<think>",
"Thinking about it...",
"</think>This is a ",
'content:<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "London"}<|tool_call_end|><|tool_calls_section_end|><|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:1<|tool_call_argument_begin|>'
'{"city": "Delhi"}<|tool_call_end|><|tool_calls_section_end|>',
]
all_reasoning = ""
all_content = ""
toolcall_chunks = []
for chunk in chunks:
r = reasoning_det.parse_streaming_increment(chunk)
all_reasoning += r.reasoning_text
if r.normal_text:
toolcall_chunks.append(r.normal_text)
tool_calls, all_content = _collect_streaming_tool_calls(
tc_det, toolcall_chunks, self.tools
)
self.assertEqual("Thinking about it...", all_reasoning)
self.assertEqual("This is a content:", all_content)
self.assertEqual(len(tool_calls), 2)
first_call = tool_calls.pop(0)
self.assertEqual(first_call["name"], "get_weather")
self.assertEqual(first_call["parameters"], '{"city": "London"}')
second_call = tool_calls.pop(0)
self.assertEqual(second_call["name"], "get_weather")
self.assertEqual(second_call["parameters"], '{"city": "Delhi"}')
def test_e2e_chunk_split_invariance(self):
"""The detector must produce identical results across a few realistic
chunking variants. Special tokens (e.g. ``<|tool_calls_section_begin|>``)
are atomic and never split, so cuts only fall on token boundaries or
inside JSON args.
"""
prefix = "<think>Thinking about it...</think>This is a content:"
call1 = (
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:0"
'<|tool_call_argument_begin|>{"city": "London"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
)
call2 = (
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:1"
'<|tool_call_argument_begin|>{"city": "Delhi"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
)
expected_reasoning = "Thinking about it..."
expected_content = "This is a content:"
expected_calls = [
{"name": "get_weather", "parameters": '{"city": "London"}'},
{"name": "get_weather", "parameters": '{"city": "Delhi"}'},
]
variants = {
# One complete tool call per chunk.
"first_complete_then_second_complete": [prefix + call1, call2],
# Both tool calls arrive in a single chunk.
"both_in_one_chunk": [prefix + call1 + call2],
# First call complete + second call partial (cut inside JSON args),
# then the rest of the second call.
"first_complete_second_partial_then_rest": [
prefix
+ call1
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:1"
+ '<|tool_call_argument_begin|>{"city": "De',
'lhi"}<|tool_call_end|><|tool_calls_section_end|>',
],
# First call partial (cut inside JSON args), then rest of first +
# full second call.
"first_partial_then_first_complete_second_complete": [
prefix
+ "<|tool_calls_section_begin|>"
+ "<|tool_call_begin|>functions.get_weather:0"
+ '<|tool_call_argument_begin|>{"city": "Lon',
'don"}<|tool_call_end|><|tool_calls_section_end|>' + call2,
],
}
for name, chunks in variants.items():
with self.subTest(variant=name):
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
tc_det = KimiK2FuncDetector()
all_reasoning = ""
toolcall_chunks = []
for chunk in chunks:
r = reasoning_det.parse_streaming_increment(chunk)
all_reasoning += r.reasoning_text
if r.normal_text:
toolcall_chunks.append(r.normal_text)
tool_calls, all_content = _collect_streaming_tool_calls(
tc_det, toolcall_chunks, self.tools
)
self.assertEqual(all_reasoning, expected_reasoning)
self.assertEqual(all_content, expected_content)
self.assertEqual(tool_calls, expected_calls)
def test_e2e_normal_text_between_two_tool_calls(self):
"""Normal text appearing BETWEEN two tool-call sections must be
surfaced as ``normal_text``.
"""
tc_det = KimiK2FuncDetector()
chunks = [
"Prefix text:"
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:0"
'<|tool_call_argument_begin|>{"city": "London"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
" Now calling next: "
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:1"
"<|tool_call_argument_begin|>"
'{"city":'
' "Delhi"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
]
tool_calls, all_content = _collect_streaming_tool_calls(
tc_det, chunks, self.tools
)
self.assertIn("Prefix text:", all_content)
self.assertIn("Now calling next:", all_content)
self.assertEqual(len(tool_calls), 2)
self.assertEqual(tool_calls[0]["name"], "get_weather")
self.assertEqual(tool_calls[0]["parameters"], '{"city": "London"}')
self.assertEqual(tool_calls[1]["name"], "get_weather")
self.assertEqual(tool_calls[1]["parameters"], '{"city": "Delhi"}')
def test_e2e_unparsable_tool_id_does_not_wedge_stream(self):
"""A tool_call header with an unparsable ID must not wedge the
streaming parser.
"""
tc_det = KimiK2FuncDetector()
# ``weird@id`` matches the broad ``[^\\s<|]+`` capture in
# ``stream_tool_call_portion_regex`` but fails both the standard
# ``name:idx`` form and the bare-counter form, so
# ``_parse_tool_call_id`` returns ``(None, 0)``.
chunks = [
"normal text before",
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>weird@id"
'<|tool_call_argument_begin|>{"city"'
': "London"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>",
# A valid follow-up call: must still be parsed.
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:0"
'<|tool_call_argument_begin|>{"city": "Delhi"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>",
]
tool_calls, all_content = _collect_streaming_tool_calls(
tc_det, chunks, self.tools
)
# The bad call may be skipped/logged, but the follow-up MUST be
# emitted. The stream must not be wedged on the bad header.
valid = [c for c in tool_calls if c.get("name") == "get_weather"]
self.assertEqual(len(valid), 1)
self.assertEqual(valid[0]["parameters"], '{"city": "Delhi"}')
self.assertIn("normal text before", all_content)
def test_e2e_malformed_json_args_passes_through_to_client(self):
"""A tool call with malformed JSON args (e.g. unclosed brace,
spelling mistake) is the **client's** problem to
validate/repair — the parser's job is to locate boundaries and
hand back the raw argument string. This mirrors the
``detect_and_parse`` (non-streaming) contract.
Required behavior:
1. The malformed call is emitted unchanged (raw bytes
preserved, name + index intact) so the client can decide
how to handle it (reject, repair, replay-prompt, etc).
2. The stream is not wedged — the trailing valid call must
still parse.
3. ``current_tool_id`` advances normally — the trailing valid
call sits at index 2, not 1 or 3.
Asserts the SAME outcome under three chunk layouts:
* **single-chunk / MTP path** — all three sections in one
forward step (mimics speculative / multi-token-prediction).
* **per-call split path** — one section per chunk.
* **bad section split mid-payload** — the malformed section
itself is fragmented across three chunks (header + partial
args; more args; end-token + trailing valid call). Verifies
the atomic-section buffer correctly defers emission until
``<|tool_call_end|>`` arrives.
"""
good_args_0 = '{"city": "London"}'
# JSON keyword (the model misspelled ``false``).
bad_args = '{"city": "Bad", "valid": fasle'
good_args_1 = '{"city": "Delhi"}'
good_section_0 = (
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:0"
f"<|tool_call_argument_begin|>{good_args_0}"
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
)
bad_section = (
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:1"
f"<|tool_call_argument_begin|>{bad_args}"
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
)
good_section_1 = (
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:2"
f"<|tool_call_argument_begin|>{good_args_1}"
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
)
layouts = {
"mtp_single_chunk": [good_section_0 + bad_section + good_section_1],
"per_call_chunks": [good_section_0, bad_section, good_section_1],
"bad_section_split_mid_payload": [
good_section_0,
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:1"
f'<|tool_call_argument_begin|>{{"city":',
' "Bad", "valid": fasle',
"<|tool_call_end|>" "<|tool_calls_section_end|>" + good_section_1,
],
}
for layout_name, chunks in layouts.items():
with self.subTest(layout=layout_name):
tc_det = KimiK2FuncDetector()
tool_calls, _ = _collect_streaming_tool_calls(
tc_det, chunks, self.tools
)
# Three contiguous tool calls (good, bad-passthrough, good).
self.assertEqual(
len(tool_calls),
3,
f"[{layout_name}] expected 3 calls, got: {tool_calls!r}",
)
self.assertEqual(tool_calls[0]["name"], "get_weather")
self.assertEqual(tool_calls[1]["name"], "get_weather")
self.assertEqual(tool_calls[2]["name"], "get_weather")
self.assertEqual(tool_calls[0]["parameters"], good_args_0)
# Bad payload preserved byte-for-byte for the client.
self.assertEqual(tool_calls[1]["parameters"], bad_args)
self.assertEqual(tool_calls[2]["parameters"], good_args_1)
def test_e2e_exception_mid_drain_preserves_accumulated_calls(self):
"""An exception raised mid-drain must not discard tool calls already
finalized in the same ``parse_streaming_increment`` invocation.
"""
import unittest.mock as mock
tc_det = KimiK2FuncDetector()
chunks = [
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:0"
'<|tool_call_argument_begin|>{"city": "London"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:1"
'<|tool_call_argument_begin|>{"city": "Delhi"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
]
# Force an exception on the SECOND resolution step by making
# ``_resolve_function_name`` raise the second time it is invoked.
call_count = {"n": 0}
real_resolve = tc_det._resolve_function_name
def flaky_resolve(function_id, tools, function_args=None):
call_count["n"] += 1
if call_count["n"] >= 2:
raise RuntimeError("forced mid-drain failure")
return real_resolve(function_id, tools, function_args)
with mock.patch.object(
tc_det, "_resolve_function_name", side_effect=flaky_resolve
):
tool_calls, all_content = _collect_streaming_tool_calls(
tc_det, chunks, self.tools
)
# First call must survive the mid-drain exception.
named = [c for c in tool_calls if c.get("name")]
self.assertGreaterEqual(len(named), 1)
self.assertEqual(named[0]["name"], "get_weather")
self.assertEqual(named[0]["parameters"], '{"city": "London"}')
# Already-finalized call payload must NOT leak into normal_text.
self.assertNotIn('{"city": "London"}', all_content)
self.assertNotIn("<|tool_call_begin|>", all_content)
def test_e2e_multiple_tool_calls_without_think_close(self):
"""Multiple tool calls inside <think> without </think>."""
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
@@ -664,6 +1102,155 @@ class TestKimiK2EndToEnd(unittest.TestCase):
self.assertEqual(name_calls[1].name, "ReadFile")
# ============================================================
# Part 2b: OpenAI streaming-spec compliance for ``tool_index``
# ============================================================
class TestKimiK2DetectorOpenAIIndexCompliance(unittest.TestCase):
"""The detector must emit ``tool_index`` as a dense, 0-based position
within the *current response* (per the OpenAI streaming spec), regardless
of the value of the model's conversation-level ``:N`` counter in the
tool_call header. The serving layer is responsible for adding any
history offset back when synthesizing the public ``id`` field.
These tests pin the detector contract so multi-turn conversations
(where the model continues an auto-incrementing counter across turns)
can never produce sparse / non-zero-based ``index`` values in the
streamed delta.
"""
def setUp(self):
self.tools = [
_make_tool(
"get_weather",
{
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
),
]
def _run(self, chunks):
det = KimiK2FuncDetector()
events = []
for chunk in chunks:
r = det.parse_streaming_increment(chunk, self.tools)
events.extend(r.calls)
return events
def test_single_call_with_nonzero_model_counter_starts_at_index_0(self):
"""Model continues a conversation-level counter (``:5``) across turns.
The detector must still emit ``tool_index=0`` for the first call in
this response — the model's ``:N`` suffix MUST NOT leak into ``index``.
"""
chunks = [
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:5"
'<|tool_call_argument_begin|>{"city":',
' "Paris"}<|tool_call_end|><|tool_calls_section_end|>',
]
events = self._run(chunks)
self.assertGreaterEqual(len(events), 1)
# Every emitted delta (name event + arg deltas) must use index 0.
for ev in events:
self.assertEqual(
ev.tool_index,
0,
f"tool_index must be 0-based per response, got {ev.tool_index}",
)
first = events[0]
self.assertEqual(first.name, "get_weather")
def test_multi_call_response_uses_dense_0_based_indices(self):
"""Two calls in one response, model emits ``:7`` then ``:8``. The
detector must emit ``tool_index=0`` then ``tool_index=1`` (dense,
0-based), independent of the model's counter.
"""
chunks = [
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:7"
'<|tool_call_argument_begin|>{"city": "London"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:8"
'<|tool_call_argument_begin|>{"city": "Berlin"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
]
events = self._run(chunks)
name_events = [e for e in events if e.name]
self.assertEqual(len(name_events), 2)
self.assertEqual(name_events[0].tool_index, 0)
self.assertEqual(name_events[1].tool_index, 1)
# Every delta for call N must carry tool_index == N.
for ev in events:
self.assertIn(ev.tool_index, (0, 1))
def test_continuation_chunks_keep_same_tool_index(self):
"""Per OpenAI spec, all argument-delta chunks for a given call
must share the same ``index``. Split the args across multiple
chunks and verify ``tool_index`` stays at 0 throughout.
"""
chunks = [
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>functions.get_weather:10"
'<|tool_call_argument_begin|>{"city":',
' "Pa',
'ris"}',
"<|tool_call_end|><|tool_calls_section_end|>",
]
events = self._run(chunks)
self.assertGreater(len(events), 1, "expected name event + arg deltas")
for ev in events:
self.assertEqual(ev.tool_index, 0)
# Reassembled args must round-trip.
joined = "".join(ev.parameters or "" for ev in events)
self.assertEqual(joined, '{"city": "Paris"}')
def test_three_calls_with_nonzero_model_counter_indices_are_dense(self):
"""Multi-turn worst case: model continues at ``:10`` and emits three
calls in this response. Indices must be 0, 1, 2 — not 10, 11, 12.
"""
sections = []
for offset, city in enumerate(("Paris", "Berlin", "Madrid")):
sections.append(
"<|tool_calls_section_begin|>"
f"<|tool_call_begin|>functions.get_weather:{10 + offset}"
f'<|tool_call_argument_begin|>{{"city": "{city}"}}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
)
events = self._run(["".join(sections)])
name_events = [e for e in events if e.name]
self.assertEqual([e.tool_index for e in name_events], [0, 1, 2])
def test_bare_counter_id_also_starts_at_index_0(self):
"""Same invariant for the bare-counter ID form (model omits function
name and emits just a numeric counter, e.g. ``:5``).
"""
chunks = [
"<|tool_calls_section_begin|>"
"<|tool_call_begin|>5"
'<|tool_call_argument_begin|>{"city": "Paris"}'
"<|tool_call_end|>"
"<|tool_calls_section_end|>"
]
events = self._run(chunks)
self.assertGreaterEqual(len(events), 1)
for ev in events:
self.assertEqual(ev.tool_index, 0)
# ============================================================
# Part 3: Bare-counter tool call ID parsing
# ============================================================