Add Gemma 4 tool call parser (#1105)

This commit is contained in:
Nic Davidson
2026-04-04 18:21:45 -06:00
committed by GitHub
parent d4eb136d44
commit 3257c3df17
3 changed files with 78 additions and 0 deletions
+2
View File
@@ -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:
+45
View File
@@ -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|>"
+31
View File
@@ -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 = (