Add Gemma 4 tool call parser (#1105)
This commit is contained in:
@@ -473,6 +473,8 @@ def _infer_tool_parser(chat_template):
|
||||
return None
|
||||
elif "<minimax:tool_call>" in chat_template:
|
||||
return "minimax_m2"
|
||||
elif "<|tool_call>" in chat_template and "<tool_call|>" in chat_template:
|
||||
return "gemma4"
|
||||
elif "<start_function_call>" in chat_template:
|
||||
return "function_gemma"
|
||||
elif "<longcat_tool_call>" in chat_template:
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
# Copyright © 2025 Apple Inc.
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Optional
|
||||
|
||||
_tool_call_regex = re.compile(r"call:(\w+)(\{.*\})", re.DOTALL)
|
||||
|
||||
|
||||
def _gemma4_args_to_json(text: str) -> str:
|
||||
"""Convert Gemma 4 tool call args to valid JSON.
|
||||
|
||||
Gemma 4 uses unquoted keys and <|"|> as string delimiters
|
||||
instead of standard double quotes.
|
||||
"""
|
||||
strings = []
|
||||
|
||||
def _capture(m):
|
||||
strings.append(m.group(1))
|
||||
return f"\x00{len(strings) - 1}\x00"
|
||||
|
||||
# Extract <|"|>-delimited strings and replace with placeholders
|
||||
text = re.sub(r'<\|"\|>(.*?)<\|"\|>', _capture, text, flags=re.DOTALL)
|
||||
# Quote bare keys
|
||||
text = re.sub(r"(?<=[{,])(\w+):", r'"\1":', text)
|
||||
# Restore captured strings as properly escaped JSON strings
|
||||
for i, s in enumerate(strings):
|
||||
text = text.replace(f"\x00{i}\x00", json.dumps(s))
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def parse_tool_call(text: str, _: Optional[Any] = None):
|
||||
match = _tool_call_regex.search(text)
|
||||
if not match:
|
||||
raise ValueError("No function provided.")
|
||||
func_name = match.group(1)
|
||||
args_str = match.group(2)
|
||||
json_str = _gemma4_args_to_json(args_str)
|
||||
arguments = json.loads(json_str)
|
||||
return dict(name=func_name, arguments=arguments)
|
||||
|
||||
|
||||
tool_call_start = "<|tool_call>"
|
||||
tool_call_end = "<tool_call|>"
|
||||
@@ -3,6 +3,7 @@ from pathlib import Path
|
||||
|
||||
from mlx_lm.tool_parsers import (
|
||||
function_gemma,
|
||||
gemma4,
|
||||
glm47,
|
||||
json_tools,
|
||||
kimi_k2,
|
||||
@@ -18,6 +19,7 @@ class TestToolParsing(unittest.TestCase):
|
||||
def test_parsers(self):
|
||||
test_cases = [
|
||||
("call:multiply{a:12234585,b:48838483920}", function_gemma),
|
||||
("call:multiply{a:12234585,b:48838483920}", gemma4),
|
||||
(
|
||||
'{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}',
|
||||
glm47,
|
||||
@@ -89,6 +91,10 @@ class TestToolParsing(unittest.TestCase):
|
||||
"call:get_current_temperature{location:<escape>London<escape>}",
|
||||
function_gemma,
|
||||
),
|
||||
(
|
||||
'call:get_current_temperature{location:<|"|>London<|"|>}',
|
||||
gemma4,
|
||||
),
|
||||
(
|
||||
'get_current_temperature<arg_key>location</arg_key><arg_value>"London"</arg_value>',
|
||||
glm47,
|
||||
@@ -191,6 +197,31 @@ class TestToolParsing(unittest.TestCase):
|
||||
self.assertEqual(tool_call["arguments"]["filters"], {"category": "books"})
|
||||
self.assertEqual(tool_call["arguments"]["tags"], ["fiction", "new"])
|
||||
|
||||
def test_gemma4(self):
|
||||
# Nested object
|
||||
test_case = 'call:configure{settings:{enabled:true,name:<|"|>test<|"|>}}'
|
||||
tool_call = gemma4.parse_tool_call(test_case, None)
|
||||
self.assertEqual(tool_call["name"], "configure")
|
||||
self.assertEqual(
|
||||
tool_call["arguments"],
|
||||
{"settings": {"enabled": True, "name": "test"}},
|
||||
)
|
||||
|
||||
# Array of strings
|
||||
test_case = 'call:tag{items:[<|"|>foo<|"|>,<|"|>bar<|"|>]}'
|
||||
tool_call = gemma4.parse_tool_call(test_case, None)
|
||||
self.assertEqual(tool_call["name"], "tag")
|
||||
self.assertEqual(tool_call["arguments"], {"items": ["foo", "bar"]})
|
||||
|
||||
# Mixed types
|
||||
test_case = 'call:search{query:<|"|>hello world<|"|>,limit:10,verbose:false}'
|
||||
tool_call = gemma4.parse_tool_call(test_case, None)
|
||||
self.assertEqual(tool_call["name"], "search")
|
||||
self.assertEqual(
|
||||
tool_call["arguments"],
|
||||
{"query": "hello world", "limit": 10, "verbose": False},
|
||||
)
|
||||
|
||||
def test_kimi_k2(self):
|
||||
# Single tool call
|
||||
test_case = (
|
||||
|
||||
Reference in New Issue
Block a user