07b615e96e
- Added new environments for reinforcement learning, including `HermesSweEnv` for software engineering tasks and `TerminalTestEnv` for inline testing. - Introduced `ToolContext` for unrestricted access to tools during reward computation. - Updated `.gitignore` to exclude `wandb/` directory. - Enhanced `README.md` with detailed architecture and usage instructions for Atropos environments. - Added configuration files for SWE and terminal test environments to streamline setup. - Removed unnecessary compiled Python files from `__pycache__`.
94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
"""
|
|
Kimi K2 tool call parser.
|
|
|
|
Format:
|
|
<|tool_calls_section_begin|>
|
|
<|tool_call_begin|>function_id:0<|tool_call_argument_begin|>{"arg": "val"}<|tool_call_end|>
|
|
<|tool_calls_section_end|>
|
|
|
|
The function_id format is typically "functions.func_name:index" or "func_name:index".
|
|
|
|
Based on VLLM's KimiK2ToolParser.extract_tool_calls()
|
|
"""
|
|
|
|
import re
|
|
import uuid
|
|
from typing import List, Optional
|
|
|
|
from openai.types.chat.chat_completion_message_tool_call import (
|
|
ChatCompletionMessageToolCall,
|
|
Function,
|
|
)
|
|
|
|
from environments.tool_call_parsers import ParseResult, ToolCallParser, register_parser
|
|
|
|
|
|
@register_parser("kimi_k2")
|
|
class KimiK2ToolCallParser(ToolCallParser):
|
|
"""
|
|
Parser for Kimi K2 tool calls.
|
|
|
|
Uses section begin/end tokens wrapping individual tool call begin/end tokens.
|
|
The tool_call_id contains the function name (after last dot, before colon).
|
|
"""
|
|
|
|
# Support both singular and plural variants
|
|
START_TOKENS = [
|
|
"<|tool_calls_section_begin|>",
|
|
"<|tool_call_section_begin|>",
|
|
]
|
|
|
|
# Regex captures: tool_call_id (e.g., "functions.get_weather:0"), function_arguments
|
|
PATTERN = re.compile(
|
|
r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[^<]+:\d+)\s*"
|
|
r"<\|tool_call_argument_begin\|>\s*"
|
|
r"(?P<function_arguments>(?:(?!<\|tool_call_begin\|>).)*?)\s*"
|
|
r"<\|tool_call_end\|>",
|
|
re.DOTALL,
|
|
)
|
|
|
|
def parse(self, text: str) -> ParseResult:
|
|
# Check for any variant of the start token
|
|
has_start = any(token in text for token in self.START_TOKENS)
|
|
if not has_start:
|
|
return text, None
|
|
|
|
try:
|
|
matches = self.PATTERN.findall(text)
|
|
if not matches:
|
|
return text, None
|
|
|
|
tool_calls: List[ChatCompletionMessageToolCall] = []
|
|
for match in matches:
|
|
function_id, function_args = match
|
|
|
|
# Extract function name from ID format: "functions.get_weather:0" -> "get_weather"
|
|
function_name = function_id.split(":")[0].split(".")[-1]
|
|
|
|
tool_calls.append(
|
|
ChatCompletionMessageToolCall(
|
|
id=function_id, # Preserve the original ID format
|
|
type="function",
|
|
function=Function(
|
|
name=function_name,
|
|
arguments=function_args.strip(),
|
|
),
|
|
)
|
|
)
|
|
|
|
if not tool_calls:
|
|
return text, None
|
|
|
|
# Content is everything before the tool calls section
|
|
earliest_start = len(text)
|
|
for token in self.START_TOKENS:
|
|
idx = text.find(token)
|
|
if idx >= 0 and idx < earliest_start:
|
|
earliest_start = idx
|
|
|
|
content = text[:earliest_start].strip()
|
|
return content if content else None, tool_calls
|
|
|
|
except Exception:
|
|
return text, None
|