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 <rharkhani@gmail.com>
This commit is contained in:
co-authored by
gemini-code-assist[bot]
Rohit Harkhani
parent
b6a8000473
commit
5f76736427
@@ -15,7 +15,6 @@ from sglang.srt.function_call.core_types import (
|
|||||||
ToolCallItem,
|
ToolCallItem,
|
||||||
_GetInfoFunc,
|
_GetInfoFunc,
|
||||||
)
|
)
|
||||||
from sglang.srt.function_call.utils import _is_complete_json
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -177,9 +176,15 @@ class KimiK2Detector(BaseFormatDetector):
|
|||||||
logger.debug("function_call_tuples: %s", function_call_tuples)
|
logger.debug("function_call_tuples: %s", function_call_tuples)
|
||||||
|
|
||||||
tool_calls = []
|
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:
|
for match in function_call_tuples:
|
||||||
function_id, function_args = match
|
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
|
function_id, tools, function_args
|
||||||
)
|
)
|
||||||
if function_name is None:
|
if function_name is None:
|
||||||
@@ -189,11 +194,12 @@ class KimiK2Detector(BaseFormatDetector):
|
|||||||
|
|
||||||
tool_calls.append(
|
tool_calls.append(
|
||||||
ToolCallItem(
|
ToolCallItem(
|
||||||
tool_index=function_idx,
|
tool_index=local_tool_index,
|
||||||
name=function_name,
|
name=function_name,
|
||||||
parameters=function_args,
|
parameters=function_args,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
local_tool_index += 1
|
||||||
|
|
||||||
content = text[: text.find(self.bot_token)]
|
content = text[: text.find(self.bot_token)]
|
||||||
return StreamingParseResult(normal_text=content, calls=tool_calls)
|
return StreamingParseResult(normal_text=content, calls=tool_calls)
|
||||||
@@ -205,127 +211,207 @@ class KimiK2Detector(BaseFormatDetector):
|
|||||||
def parse_streaming_increment(
|
def parse_streaming_increment(
|
||||||
self, new_text: str, tools: List[Tool]
|
self, new_text: str, tools: List[Tool]
|
||||||
) -> StreamingParseResult:
|
) -> StreamingParseResult:
|
||||||
"""
|
"""Streaming incremental parsing tool calls for KimiK2 format."""
|
||||||
Streaming incremental parsing tool calls for KimiK2 format.
|
|
||||||
"""
|
|
||||||
self._buffer += new_text
|
self._buffer += new_text
|
||||||
current_text = self._buffer
|
|
||||||
|
|
||||||
# Check if we have a tool call (either the start token or individual tool call)
|
# Fast path: no tool call in flight and no markers yet -- emit as
|
||||||
has_tool_call = (
|
# normal text, holding back any trailing partial start token.
|
||||||
self.bot_token in current_text or self.tool_call_start_token in current_text
|
if (
|
||||||
)
|
self._current_stream_function_name is None
|
||||||
|
and self.bot_token not in self._buffer
|
||||||
if not has_tool_call:
|
and self.tool_call_start_token not in self._buffer
|
||||||
self._buffer = ""
|
):
|
||||||
normal_text = _strip_special_tokens(new_text)
|
emit, hold = self._split_pending_start(self._buffer)
|
||||||
return StreamingParseResult(normal_text=normal_text)
|
self._buffer = hold
|
||||||
|
return StreamingParseResult(normal_text=_strip_special_tokens(emit))
|
||||||
|
|
||||||
if not hasattr(self, "_tool_indices"):
|
if not hasattr(self, "_tool_indices"):
|
||||||
self._tool_indices = self._get_tool_indices(tools)
|
self._tool_indices = self._get_tool_indices(tools)
|
||||||
|
|
||||||
|
normal_text_parts: list[str] = []
|
||||||
calls: list[ToolCallItem] = []
|
calls: list[ToolCallItem] = []
|
||||||
|
|
||||||
try:
|
try:
|
||||||
match = self.stream_tool_call_portion_regex.search(current_text)
|
while True:
|
||||||
if match:
|
buffer = self._buffer
|
||||||
function_id = match.group("tool_call_id")
|
|
||||||
function_args = match.group("function_arguments")
|
|
||||||
|
|
||||||
# Reuse cached name for current tool call to avoid repeated
|
# Locate next <|tool_call_begin|>, draining any prefix as text.
|
||||||
# json.loads on partial JSON in _infer_tool_name.
|
begin_idx = self._locate_tool_call_start(buffer, normal_text_parts)
|
||||||
if self._current_stream_function_name is not None:
|
if begin_idx is None:
|
||||||
function_name = self._current_stream_function_name
|
break
|
||||||
else:
|
buffer = self._buffer
|
||||||
function_name, _ = self._parse_tool_call_id(
|
|
||||||
function_id, tools, function_args
|
# 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 function_name is None:
|
if next_begin != -1 and (
|
||||||
return StreamingParseResult(normal_text="", calls=calls)
|
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."
|
||||||
|
)
|
||||||
|
self._buffer = buffer[next_begin:]
|
||||||
|
self._reset_inflight_call_state()
|
||||||
|
continue
|
||||||
|
|
||||||
# Initialize state if this is the first tool call
|
if arg_begin_idx == -1:
|
||||||
|
# Header not fully arrived yet.
|
||||||
|
break
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
# 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:
|
if self.current_tool_id == -1:
|
||||||
self.current_tool_id = 0
|
self.current_tool_id = 0
|
||||||
self.prev_tool_call_arr = []
|
self.prev_tool_call_arr = []
|
||||||
self.streamed_args_for_tool = [""]
|
self.streamed_args_for_tool = [""]
|
||||||
|
|
||||||
# Ensure we have enough entries in our tracking arrays
|
|
||||||
while len(self.prev_tool_call_arr) <= self.current_tool_id:
|
while len(self.prev_tool_call_arr) <= self.current_tool_id:
|
||||||
self.prev_tool_call_arr.append({})
|
self.prev_tool_call_arr.append({})
|
||||||
while len(self.streamed_args_for_tool) <= self.current_tool_id:
|
while len(self.streamed_args_for_tool) <= self.current_tool_id:
|
||||||
self.streamed_args_for_tool.append("")
|
self.streamed_args_for_tool.append("")
|
||||||
|
|
||||||
if not self.current_tool_name_sent:
|
|
||||||
calls.append(
|
|
||||||
ToolCallItem(
|
|
||||||
tool_index=self.current_tool_id,
|
|
||||||
name=function_name,
|
|
||||||
parameters="",
|
|
||||||
)
|
|
||||||
)
|
|
||||||
self.current_tool_name_sent = True
|
|
||||||
self._current_stream_function_name = function_name
|
|
||||||
self.prev_tool_call_arr[self.current_tool_id] = {
|
self.prev_tool_call_arr[self.current_tool_id] = {
|
||||||
"name": function_name,
|
"name": name,
|
||||||
"arguments": {},
|
"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:
|
else:
|
||||||
argument_diff = (
|
args_full = buffer[args_start:]
|
||||||
function_args[len(self._last_arguments) :]
|
argument_diff = args_full[len(self._last_arguments) :]
|
||||||
if function_args.startswith(self._last_arguments)
|
if argument_diff or name_just_resolved:
|
||||||
else function_args
|
|
||||||
)
|
|
||||||
|
|
||||||
parsed_args_diff = argument_diff.split(self.tool_call_end_token, 1)[
|
|
||||||
0
|
|
||||||
]
|
|
||||||
|
|
||||||
if parsed_args_diff:
|
|
||||||
calls.append(
|
calls.append(
|
||||||
ToolCallItem(
|
ToolCallItem(
|
||||||
tool_index=self.current_tool_id,
|
tool_index=self.current_tool_id,
|
||||||
name=None,
|
name=(
|
||||||
parameters=parsed_args_diff,
|
self._current_stream_function_name
|
||||||
|
if name_just_resolved
|
||||||
|
else None
|
||||||
|
),
|
||||||
|
parameters=argument_diff,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
self._last_arguments += parsed_args_diff
|
if argument_diff:
|
||||||
|
self._last_arguments += argument_diff
|
||||||
self.streamed_args_for_tool[
|
self.streamed_args_for_tool[
|
||||||
self.current_tool_id
|
self.current_tool_id
|
||||||
] += parsed_args_diff
|
] += argument_diff
|
||||||
|
|
||||||
parsed_args = function_args.split(self.tool_call_end_token, 1)[0]
|
if end_idx == -1:
|
||||||
if _is_complete_json(parsed_args):
|
# Args still streaming.
|
||||||
try:
|
break
|
||||||
parsed_args = json.loads(parsed_args)
|
|
||||||
self.prev_tool_call_arr[self.current_tool_id][
|
|
||||||
"arguments"
|
|
||||||
] = parsed_args
|
|
||||||
except json.JSONDecodeError:
|
|
||||||
pass
|
|
||||||
|
|
||||||
# Find the end of the current tool call and remove only that part from buffer
|
# Section finalized -- advance buffer and prepare next call.
|
||||||
tool_call_end_pattern = (
|
self._buffer = buffer[end_idx + len(self.tool_call_end_token) :]
|
||||||
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 = ""
|
|
||||||
|
|
||||||
result = StreamingParseResult(normal_text="", calls=calls)
|
|
||||||
self.current_tool_id += 1
|
self.current_tool_id += 1
|
||||||
self._last_arguments = ""
|
self._reset_inflight_call_state()
|
||||||
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:
|
except Exception as e:
|
||||||
logger.error("Error in parse_streaming_increment: %s", e, exc_info=True)
|
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:
|
def structure_info(self) -> _GetInfoFunc:
|
||||||
"""Return function that creates StructureInfo for guided generation."""
|
"""Return function that creates StructureInfo for guided generation."""
|
||||||
|
|||||||
@@ -111,6 +111,32 @@ class TestKimiK2DetectorBasic(unittest.TestCase):
|
|||||||
self.assertEqual(result.calls[1].name, "get_weather")
|
self.assertEqual(result.calls[1].name, "get_weather")
|
||||||
self.assertEqual(result.calls[1].parameters, '{"city": "Tokyo"}')
|
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):
|
def test_normal_text_before_tool_call(self):
|
||||||
"""Normal text before tool call markers is preserved."""
|
"""Normal text before tool call markers is preserved."""
|
||||||
text = (
|
text = (
|
||||||
@@ -183,6 +209,15 @@ class TestKimiK2DetectorHyphenatedNames(unittest.TestCase):
|
|||||||
class TestKimiK2DetectorStreaming(unittest.TestCase):
|
class TestKimiK2DetectorStreaming(unittest.TestCase):
|
||||||
"""Streaming incremental parsing tests for KimiK2Detector."""
|
"""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):
|
def setUp(self):
|
||||||
self.tools = [
|
self.tools = [
|
||||||
_make_tool("ReadFile"),
|
_make_tool("ReadFile"),
|
||||||
@@ -627,6 +662,409 @@ class TestKimiK2EndToEnd(unittest.TestCase):
|
|||||||
self.assertEqual(len(name_calls), 1)
|
self.assertEqual(len(name_calls), 1)
|
||||||
self.assertEqual(name_calls[0].name, "get_weather")
|
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):
|
def test_e2e_multiple_tool_calls_without_think_close(self):
|
||||||
"""Multiple tool calls inside <think> without </think>."""
|
"""Multiple tool calls inside <think> without </think>."""
|
||||||
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
|
reasoning_det = KimiK2ReasoningDetector(stream_reasoning=True)
|
||||||
@@ -664,6 +1102,155 @@ class TestKimiK2EndToEnd(unittest.TestCase):
|
|||||||
self.assertEqual(name_calls[1].name, "ReadFile")
|
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
|
# Part 3: Bare-counter tool call ID parsing
|
||||||
# ============================================================
|
# ============================================================
|
||||||
|
|||||||
Reference in New Issue
Block a user