From af2a2ac618394a3c2d36c194f6b4ff346758ea7c Mon Sep 17 00:00:00 2001 From: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Date: Thu, 7 May 2026 22:20:02 +0100 Subject: [PATCH] fix(function_call): handle Kimi-K2.5 bare numeric tool call IDs (#23950) --- .../srt/function_call/kimik2_detector.py | 123 ++++++++++++--- .../function_call/test_kimik2_detector.py | 143 ++++++++++++++++++ 2 files changed, 244 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py index 113031c8e..da2c76fd0 100644 --- a/python/sglang/srt/function_call/kimik2_detector.py +++ b/python/sglang/srt/function_call/kimik2_detector.py @@ -35,13 +35,18 @@ class KimiK2Detector(BaseFormatDetector): """ Detector for Kimi K2 / K2.5 model function call format. - Format Structure: + Format Structure (standard): ``` <|tool_calls_section_begin|> <|tool_call_begin|>functions.{func_name}:{index}<|tool_call_argument_begin|>{json_args}<|tool_call_end|> <|tool_calls_section_end|> ``` + Format Structure (bare counter — model omits function name): + ``` + <|tool_call_begin|>{counter}<|tool_call_argument_begin|>{json_args}<|tool_call_end|> + ``` + Reference: https://huggingface.co/moonshotai/Kimi-K2-Instruct/blob/main/docs/tool_call_guidance.md """ @@ -55,23 +60,96 @@ class KimiK2Detector(BaseFormatDetector): self.tool_call_end_token: str = "<|tool_call_end|>" self.tool_call_argument_begin_token: str = "<|tool_call_argument_begin|>" - # Support hyphenated function names (common in MCP tools, e.g. mcp__portal__search-documents) + # Capture tool_call_id broadly: the model may emit standard IDs + # like "functions.ReadFile:0" or bare call counters like "3". self.tool_call_regex = re.compile( - r"<\|tool_call_begin\|>\s*(?P[\w.\-]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*?\})\s*<\|tool_call_end\|>", + r"<\|tool_call_begin\|>\s*(?P[^\s<|]+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*?\})\s*<\|tool_call_end\|>", re.DOTALL, ) self.stream_tool_call_portion_regex = re.compile( - r"<\|tool_call_begin\|>\s*(?P[\w.\-]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*)", + r"<\|tool_call_begin\|>\s*(?P[^\s<|]+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*)", re.DOTALL, ) self._last_arguments = "" + self._current_stream_function_name: str | None = None - # Robust parser for ids like "functions.search:0", "functions.mcp__search-docs:0", or fallback "search:0" + # Standard ID: "functions.search:0", "search:0" self.tool_call_id_regex = re.compile( r"^(?:functions\.)?(?P[\w.\-]+):(?P\d+)$" ) + # Bare call counter: "0", "3" (model uses auto-incrementing counter) + self.tool_call_id_counter_regex = re.compile(r"^\d+$") + + def _parse_tool_call_id( + self, function_id: str, tools: List[Tool], function_args: str = None + ): + """Parse a tool call ID into (function_name, call_index). + + Standard format: "functions.ReadFile:0" → ("ReadFile", 0) + Bare counter: "3" → call_index=3, infer name from arguments. + + The bare counter is a conversation-level auto-increment, NOT an index + into the tools list. The function name is inferred by matching argument + keys against tool parameter schemas. + """ + m = self.tool_call_id_regex.match(function_id) + if m: + return m.group("name"), int(m.group("index")) + + if self.tool_call_id_counter_regex.match(function_id): + call_index = int(function_id) + name = self._infer_tool_name(tools, function_args) + if name: + return name, call_index + return None, call_index + + logger.warning("Unexpected tool_call_id format: %s", function_id) + return None, 0 + + def _infer_tool_name(self, tools: List[Tool], function_args: str = None): + """Infer function name when the model omits it (bare counter ID). + + Matches argument keys against tool parameter schemas, preferring the + tool whose declared properties best match the actual arguments. + """ + if not tools: + return None + if len(tools) == 1: + return tools[0].function.name + + if not function_args: + logger.debug( + "No function_args for tool name inference with %d tools", len(tools) + ) + return None + + try: + arg_keys = set(json.loads(function_args).keys()) + except (json.JSONDecodeError, TypeError): + logger.debug( + "Could not parse function_args for tool name inference " + "(may be partial JSON in streaming)" + ) + return None + + # Pick the tool whose properties best match the argument keys. + best_name = None + best_score = -1 + for tool in tools: + params = tool.function.parameters or {} + props = set(params.get("properties", {}).keys()) + if not props: + continue + overlap = len(arg_keys & props) + extra = len(arg_keys - props) + score = overlap - extra + if score > best_score: + best_score = score + best_name = tool.function.name + + return best_name def has_tool_call(self, text: str) -> bool: """Check if the text contains a KimiK2 format tool call.""" @@ -83,15 +161,11 @@ class KimiK2Detector(BaseFormatDetector): :param text: The complete text to parse. :param tools: List of available tools. - :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. + :return: StreamingParseResult with normal_text (content before tool calls) and calls (parsed items). """ if self.bot_token not in text: return StreamingParseResult(normal_text=text, calls=[]) try: - # there are two possible captures - between tags, or between a - # tag and end-of-string so the result of - # findall is an array of tuples where one is a function call and - # the other is None function_call_tuples = self.tool_call_regex.findall(text) logger.debug("function_call_tuples: %s", function_call_tuples) @@ -99,12 +173,11 @@ class KimiK2Detector(BaseFormatDetector): tool_calls = [] for match in function_call_tuples: function_id, function_args = match - m = self.tool_call_id_regex.match(function_id) - if not m: - logger.warning("Unexpected tool_call_id format: %s", function_id) + function_name, function_idx = self._parse_tool_call_id( + function_id, tools, function_args + ) + if function_name is None: continue - function_name = m.group("name") - function_idx = int(m.group("index")) logger.debug(f"function_name {function_name}") @@ -120,8 +193,7 @@ class KimiK2Detector(BaseFormatDetector): return StreamingParseResult(normal_text=content, calls=tool_calls) except Exception as e: - logger.error(f"Error in detect_and_parse: {e}") - # return the normal text if parsing fails + logger.error("Error in detect_and_parse: %s", e, exc_info=True) return StreamingParseResult(normal_text=text) def parse_streaming_increment( @@ -153,11 +225,16 @@ class KimiK2Detector(BaseFormatDetector): function_id = match.group("tool_call_id") function_args = match.group("function_arguments") - m = self.tool_call_id_regex.match(function_id) - if not m: - logger.warning("Unexpected tool_call_id format: %s", function_id) + # 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 + ) + if function_name is None: return StreamingParseResult(normal_text="", calls=calls) - function_name = m.group("name") # Initialize state if this is the first tool call if self.current_tool_id == -1: @@ -180,6 +257,7 @@ class KimiK2Detector(BaseFormatDetector): ) ) 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": {}, @@ -234,12 +312,13 @@ class KimiK2Detector(BaseFormatDetector): 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) except Exception as e: - logger.error(f"Error in parse_streaming_increment: {e}") + logger.error("Error in parse_streaming_increment: %s", e, exc_info=True) return StreamingParseResult(normal_text=_strip_special_tokens(current_text)) def structure_info(self) -> _GetInfoFunc: diff --git a/test/registered/function_call/test_kimik2_detector.py b/test/registered/function_call/test_kimik2_detector.py index 68b7c98a9..be0cd83c1 100644 --- a/test/registered/function_call/test_kimik2_detector.py +++ b/test/registered/function_call/test_kimik2_detector.py @@ -663,5 +663,148 @@ class TestKimiK2EndToEnd(unittest.TestCase): self.assertEqual(name_calls[1].name, "ReadFile") +# ============================================================ +# Part 3: Bare-counter tool call ID parsing +# ============================================================ + + +class TestKimiK2BareCounterParsing(unittest.TestCase): + """Tests for bare numeric tool_call_id format (e.g., '3' instead of 'functions.ReadFile:0').""" + + def setUp(self): + self.detector = KimiK2FuncDetector() + self.tools = [ + _make_tool("ReadFile"), + _make_tool( + "get_weather", + { + "type": "object", + "properties": { + "city": {"type": "string"}, + "unit": {"type": "string"}, + }, + "required": ["city"], + }, + ), + ] + + # --- _parse_tool_call_id --- + + def test_standard_format_with_functions_prefix(self): + name, idx = self.detector._parse_tool_call_id( + "functions.ReadFile:0", self.tools + ) + self.assertEqual(name, "ReadFile") + self.assertEqual(idx, 0) + + def test_standard_format_without_functions_prefix(self): + name, idx = self.detector._parse_tool_call_id("ReadFile:1", self.tools) + self.assertEqual(name, "ReadFile") + self.assertEqual(idx, 1) + + def test_bare_counter_single_tool(self): + single_tool = [_make_tool("search")] + name, idx = self.detector._parse_tool_call_id( + "3", single_tool, '{"query": "test"}' + ) + self.assertEqual(name, "search") + self.assertEqual(idx, 3) + + def test_bare_counter_infers_by_args(self): + name, idx = self.detector._parse_tool_call_id( + "0", self.tools, '{"city": "Tokyo"}' + ) + self.assertEqual(name, "get_weather") + self.assertEqual(idx, 0) + + def test_bare_counter_no_tools_returns_none(self): + name, idx = self.detector._parse_tool_call_id("5", [], '{"x": 1}') + self.assertIsNone(name) + self.assertEqual(idx, 5) + + def test_bare_counter_no_args_multiple_tools_returns_none(self): + name, idx = self.detector._parse_tool_call_id("2", self.tools, None) + self.assertIsNone(name) + self.assertEqual(idx, 2) + + def test_unexpected_format_returns_none(self): + name, idx = self.detector._parse_tool_call_id("some_garbage", self.tools) + self.assertIsNone(name) + self.assertEqual(idx, 0) + + # --- _infer_tool_name --- + + def test_infer_no_tools(self): + self.assertIsNone(self.detector._infer_tool_name([], '{"x": 1}')) + + def test_infer_single_tool(self): + result = self.detector._infer_tool_name([_make_tool("only_one")], '{"x": 1}') + self.assertEqual(result, "only_one") + + def test_infer_by_argument_overlap(self): + result = self.detector._infer_tool_name( + self.tools, '{"city": "Paris", "unit": "celsius"}' + ) + self.assertEqual(result, "get_weather") + + def test_infer_malformed_json_returns_none(self): + result = self.detector._infer_tool_name(self.tools, '{"city": "Par') + self.assertIsNone(result) + + def test_infer_empty_args_returns_none(self): + self.assertIsNone(self.detector._infer_tool_name(self.tools, None)) + self.assertIsNone(self.detector._infer_tool_name(self.tools, "")) + + def test_infer_no_matching_props_returns_none(self): + tools_no_props = [ + _make_tool("a", {"type": "object"}), + _make_tool("b", {"type": "object"}), + ] + result = self.detector._infer_tool_name(tools_no_props, '{"x": 1}') + self.assertIsNone(result) + + # --- detect_and_parse with bare counter (end-to-end) --- + + def test_detect_and_parse_bare_counter(self): + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>0" + '<|tool_call_argument_begin|>{"city": "Tokyo"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + self.assertEqual(result.calls[0].parameters, '{"city": "Tokyo"}') + + def test_detect_and_parse_bare_counter_skips_unknown(self): + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>0" + '<|tool_call_argument_begin|>{"unknown_key": "value"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = self.detector.detect_and_parse(text, self.tools) + # No tool props match, _infer_tool_name returns None, call is skipped + self.assertEqual(len(result.calls), 0) + + def test_streaming_bare_counter_single_tool(self): + detector = KimiK2FuncDetector() + single_tool = [_make_tool("search")] + chunks = [ + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>0" + '<|tool_call_argument_begin|>{"path', + '": "/test"}', + "<|tool_call_end|>", + "<|tool_calls_section_end|>", + ] + tool_calls, _ = _collect_streaming_tool_calls(detector, chunks, single_tool) + self.assertEqual(len(tool_calls), 1) + self.assertEqual(tool_calls[0]["name"], "search") + + if __name__ == "__main__": unittest.main()