From 5f767364279ccc483aceaee1c5fce9c0fbbdf7f3 Mon Sep 17 00:00:00 2001 From: Rohit Harkhani Date: Wed, 24 Jun 2026 13:23:22 +0530 Subject: [PATCH] kimik2_detector fix the normal text detection before tool call. (#25071) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Rohit Harkhani --- .../srt/function_call/kimik2_detector.py | 282 ++++++--- .../function_call/test_kimik2_detector.py | 587 ++++++++++++++++++ 2 files changed, 771 insertions(+), 98 deletions(-) diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py index 2345817de..11909fa6f 100644 --- a/python/sglang/srt/function_call/kimik2_detector.py +++ b/python/sglang/srt/function_call/kimik2_detector.py @@ -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.""" diff --git a/test/registered/function_call/test_kimik2_detector.py b/test/registered/function_call/test_kimik2_detector.py index ee1cccf16..736ed1624 100644 --- a/test/registered/function_call/test_kimik2_detector.py +++ b/test/registered/function_call/test_kimik2_detector.py @@ -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 ```` 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 = [ + "", + "Thinking about it...", + "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 = [ + "", + "Thinking about it...", + "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 = [ + "", + "Thinking about it...", + "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 = "Thinking about it...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 without .""" 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 # ============================================================