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
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
# ============================================================