From 3257c3df172977c97fdfe3740e3a5edeb812e0b5 Mon Sep 17 00:00:00 2001
From: Nic Davidson <59121543+nicdavidson@users.noreply.github.com>
Date: Sat, 4 Apr 2026 18:21:45 -0600
Subject: [PATCH] Add Gemma 4 tool call parser (#1105)

---
 mlx_lm/tokenizer_utils.py     |  2 ++
 mlx_lm/tool_parsers/gemma4.py | 45 +++++++++++++++++++++++++++++++++++
 tests/test_tool_parsing.py    | 31 ++++++++++++++++++++++++
 3 files changed, 78 insertions(+)
 create mode 100644 mlx_lm/tool_parsers/gemma4.py

diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py
index 4e11e81..b340d3c 100644
--- a/mlx_lm/tokenizer_utils.py
+++ b/mlx_lm/tokenizer_utils.py
@@ -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:
diff --git a/mlx_lm/tool_parsers/gemma4.py b/mlx_lm/tool_parsers/gemma4.py
new file mode 100644
index 0000000..90257a5
--- /dev/null
+++ b/mlx_lm/tool_parsers/gemma4.py
@@ -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|>"
diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py
index 08d452d..e8c4169 100644
--- a/tests/test_tool_parsing.py
+++ b/tests/test_tool_parsing.py
@@ -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 = (