chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[feat] Enhance Kimi-K2/K2.5 function call and reasoning detection (#19552)

Browse Source 
Co-authored-by: alfredyyang <alfredyyang@tencent.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
AlfredYong
2026-03-20 03:57:57 +08:00
committed by GitHub
parent 2ee7f41e25
commit c562e0d13b
2 changed files with 700 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
python/sglang/srt/function_call/kimik2_detector.py
View File

@@ -15,10 +15,25 @@ from sglang.srt.function_call.utils import _is_complete_json
logger = logging.getLogger(__name__)
_KIMI_K2_SPECIAL_TOKENS = [
"<|tool_calls_section_begin|>",
"<|tool_calls_section_end|>",
"<|tool_call_begin|>",
"<|tool_call_end|>",
"<|tool_call_argument_begin|>",
]
def _strip_special_tokens(text: str) -> str:
"""Remove all Kimi-K2 tool-call special tokens from text."""
for token in _KIMI_K2_SPECIAL_TOKENS:
text = text.replace(token, "")
return text
class KimiK2Detector(BaseFormatDetector):
"""
Detector for Kimi K2 model function call format.
Detector for Kimi K2 / K2.5 model function call format.
Format Structure:
```
@@ -38,22 +53,24 @@ class KimiK2Detector(BaseFormatDetector):
self.tool_call_start_token: str = "<|tool_call_begin|>"
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)
self.tool_call_regex = re.compile(
r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*?\})\s*<\|tool_call_end\|>",
r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w.\-]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*?\})\s*<\|tool_call_end\|>",
re.DOTALL,
)
self.stream_tool_call_portion_regex = re.compile(
r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*)",
r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[\w.\-]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>\{.*)",
re.DOTALL,
)
self._last_arguments = ""
# Robust parser for ids like "functions.search:0" or fallback "search:0"
# Robust parser for ids like "functions.search:0", "functions.mcp__search-docs:0", or fallback "search:0"
self.tool_call_id_regex = re.compile(
r"^(?:functions\.)?(?P<name>[\w\.]+):(?P<index>\d+)$"
r"^(?:functions\.)?(?P<name>[\w.\-]+):(?P<index>\d+)$"
)
def has_tool_call(self, text: str) -> bool:
@@ -123,10 +140,8 @@ class KimiK2Detector(BaseFormatDetector):
if not has_tool_call:
self._buffer = ""
for e_token in [self.eot_token, self.tool_call_end_token]:
if e_token in new_text:
new_text = new_text.replace(e_token, "")
return StreamingParseResult(normal_text=new_text)
normal_text = _strip_special_tokens(new_text)
return StreamingParseResult(normal_text=normal_text)
if not hasattr(self, "_tool_indices"):
self._tool_indices = self._get_tool_indices(tools)
@@ -165,7 +180,6 @@ class KimiK2Detector(BaseFormatDetector):
)
)
self.current_tool_name_sent = True
# Store the tool call info for serving layer completions endpoint
self.prev_tool_call_arr[self.current_tool_id] = {
"name": function_name,
"arguments": {},
@@ -177,10 +191,11 @@ class KimiK2Detector(BaseFormatDetector):
else function_args
)
parsed_args_diff = argument_diff.split("<|tool_call_end|>", 1)[0]
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,
@@ -188,12 +203,12 @@ class KimiK2Detector(BaseFormatDetector):
parameters=parsed_args_diff,
)
)
self._last_arguments += argument_diff
self._last_arguments += parsed_args_diff
self.streamed_args_for_tool[
self.current_tool_id
] += parsed_args_diff
parsed_args = function_args.split("<|tool_call_end|>", 1)[0]
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)
@@ -207,12 +222,11 @@ class KimiK2Detector(BaseFormatDetector):
tool_call_end_pattern = (
r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>"
)
match = re.search(
end_match = re.search(
tool_call_end_pattern, current_text, re.DOTALL
)
if match:
# Remove the completed tool call from buffer, keep any remaining content
self._buffer = current_text[match.end() :]
if end_match:
self._buffer = current_text[end_match.end() :]
else:
self._buffer = ""
@@ -226,7 +240,7 @@ class KimiK2Detector(BaseFormatDetector):
except Exception as e:
logger.error(f"Error in parse_streaming_increment: {e}")
return StreamingParseResult(normal_text=current_text)
return StreamingParseResult(normal_text=_strip_special_tokens(current_text))
def structure_info(self) -> _GetInfoFunc:
"""Return function that creates StructureInfo for guided generation."""