Skip to content

Commit 3257c3d

Browse files
authored
Add Gemma 4 tool call parser (ml-explore#1105)
1 parent d4eb136 commit 3257c3d

3 files changed

Lines changed: 78 additions & 0 deletions

File tree

mlx_lm/tokenizer_utils.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -473,6 +473,8 @@ def _infer_tool_parser(chat_template):
473473
return None
474474
elif "<minimax:tool_call>" in chat_template:
475475
return "minimax_m2"
476+
elif "<|tool_call>" in chat_template and "<tool_call|>" in chat_template:
477+
return "gemma4"
476478
elif "<start_function_call>" in chat_template:
477479
return "function_gemma"
478480
elif "<longcat_tool_call>" in chat_template:

mlx_lm/tool_parsers/gemma4.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# Copyright © 2025 Apple Inc.
2+
3+
import json
4+
import re
5+
from typing import Any, Optional
6+
7+
_tool_call_regex = re.compile(r"call:(\w+)(\{.*\})", re.DOTALL)
8+
9+
10+
def _gemma4_args_to_json(text: str) -> str:
11+
"""Convert Gemma 4 tool call args to valid JSON.
12+
13+
Gemma 4 uses unquoted keys and <|"|> as string delimiters
14+
instead of standard double quotes.
15+
"""
16+
strings = []
17+
18+
def _capture(m):
19+
strings.append(m.group(1))
20+
return f"\x00{len(strings) - 1}\x00"
21+
22+
# Extract <|"|>-delimited strings and replace with placeholders
23+
text = re.sub(r'<\|"\|>(.*?)<\|"\|>', _capture, text, flags=re.DOTALL)
24+
# Quote bare keys
25+
text = re.sub(r"(?<=[{,])(\w+):", r'"\1":', text)
26+
# Restore captured strings as properly escaped JSON strings
27+
for i, s in enumerate(strings):
28+
text = text.replace(f"\x00{i}\x00", json.dumps(s))
29+
30+
return text
31+
32+
33+
def parse_tool_call(text: str, _: Optional[Any] = None):
34+
match = _tool_call_regex.search(text)
35+
if not match:
36+
raise ValueError("No function provided.")
37+
func_name = match.group(1)
38+
args_str = match.group(2)
39+
json_str = _gemma4_args_to_json(args_str)
40+
arguments = json.loads(json_str)
41+
return dict(name=func_name, arguments=arguments)
42+
43+
44+
tool_call_start = "<|tool_call>"
45+
tool_call_end = "<tool_call|>"

tests/test_tool_parsing.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
from mlx_lm.tool_parsers import (
55
function_gemma,
6+
gemma4,
67
glm47,
78
json_tools,
89
kimi_k2,
@@ -18,6 +19,7 @@ class TestToolParsing(unittest.TestCase):
1819
def test_parsers(self):
1920
test_cases = [
2021
("call:multiply{a:12234585,b:48838483920}", function_gemma),
22+
("call:multiply{a:12234585,b:48838483920}", gemma4),
2123
(
2224
'{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}',
2325
glm47,
@@ -89,6 +91,10 @@ def test_parsers(self):
8991
"call:get_current_temperature{location:<escape>London<escape>}",
9092
function_gemma,
9193
),
94+
(
95+
'call:get_current_temperature{location:<|"|>London<|"|>}',
96+
gemma4,
97+
),
9298
(
9399
'get_current_temperature<arg_key>location</arg_key><arg_value>"London"</arg_value>',
94100
glm47,
@@ -191,6 +197,31 @@ def test_qwen3_coder_single_quoted_params(self):
191197
self.assertEqual(tool_call["arguments"]["filters"], {"category": "books"})
192198
self.assertEqual(tool_call["arguments"]["tags"], ["fiction", "new"])
193199

200+
def test_gemma4(self):
201+
# Nested object
202+
test_case = 'call:configure{settings:{enabled:true,name:<|"|>test<|"|>}}'
203+
tool_call = gemma4.parse_tool_call(test_case, None)
204+
self.assertEqual(tool_call["name"], "configure")
205+
self.assertEqual(
206+
tool_call["arguments"],
207+
{"settings": {"enabled": True, "name": "test"}},
208+
)
209+
210+
# Array of strings
211+
test_case = 'call:tag{items:[<|"|>foo<|"|>,<|"|>bar<|"|>]}'
212+
tool_call = gemma4.parse_tool_call(test_case, None)
213+
self.assertEqual(tool_call["name"], "tag")
214+
self.assertEqual(tool_call["arguments"], {"items": ["foo", "bar"]})
215+
216+
# Mixed types
217+
test_case = 'call:search{query:<|"|>hello world<|"|>,limit:10,verbose:false}'
218+
tool_call = gemma4.parse_tool_call(test_case, None)
219+
self.assertEqual(tool_call["name"], "search")
220+
self.assertEqual(
221+
tool_call["arguments"],
222+
{"query": "hello world", "limit": 10, "verbose": False},
223+
)
224+
194225
def test_kimi_k2(self):
195226
# Single tool call
196227
test_case = (

0 commit comments

Comments
 (0)