Fix Gemma4 tool parser: support hyphenated names and braces in strings (#1150)

Author: AkashKhamkar

mlx_lm/server.py

@@ -78,7 +78,14 @@ def __call__(self, tool_calls):
 
         result = []
         for tool_text in tool_calls:
-            parsed = self._tool_parser(tool_text, self._tools)
+            try:
+                parsed = self._tool_parser(tool_text, self._tools)
+            except (ValueError, json.JSONDecodeError) as e:
+                logging.warning(
+                    f"Failed to parse tool call ({type(e).__name__}: {e}) — "
+                    f"tool text was likely truncated mid-generation."
+                )
+                continue
             if not isinstance(parsed, list):
                 parsed = [parsed]
             result.extend(self._format(tc) for tc in parsed)

mlx_lm/tool_parsers/gemma4.py

@@ -5,10 +5,19 @@
 
 import regex as re
 
+# Matches <|"|>...<|"|> string literals (Gemma 4's string delimiter).
+_GEMMA4_STR = r'<\|"\|>(?:(?!<\|"\|>)[\s\S])*?<\|"\|>'
+
 # Matches call:name{...} with balanced braces via the regex module's
-# recursive (?R)-style support.  (\{(?:[^{}]|(?2))*\}) recurses on the
-# second capture group so nested objects like {a:{b:1}} are captured whole.
-_tool_call_regex = re.compile(r"call:(\w+)(\{(?:[^{}]|(?2))*\})", re.DOTALL)
+# recursive (?R)-style support.  The inner alternatives handle:
+#   [^{}<]          – any char that is not a brace or start of <|"|>
+#   <(?!\|"\|>)     – a lone '<' that is NOT the start of <|"|>
+#   <|"|>...<|"|>   – a complete string literal (braces inside are ignored)
+#   (?2)            – recursively balanced nested brace group
+_tool_call_regex = re.compile(
+    r"call:([\w-]+)(\{(?:[^{}<]|<(?!\|\"\|>)|" + _GEMMA4_STR + r"|(?2))*\})",
+    re.DOTALL,
+)
 
 
 def _gemma4_args_to_json(text: str) -> str:

tests/test_tool_parsing.py

@@ -254,6 +254,27 @@ def test_gemma4(self):
             {"settings": {"enabled": True, "name": "test"}},
         )
 
+        # Hyphenated function name (e.g. manim-video)
+        test_case = (
+            'call:manim-video{mode:<|"|>plan<|"|>,prompt:<|"|>explain KV caching<|"|>}'
+        )
+        tool_call = gemma4.parse_tool_call(test_case, None)
+        self.assertEqual(tool_call["name"], "manim-video")
+        self.assertEqual(
+            tool_call["arguments"],
+            {"mode": "plan", "prompt": "explain KV caching"},
+        )
+
+        # Braces inside a string argument (e.g. code snippets or markdown in content)
+        test_case = (
+            'call:skill_manage{action:<|"|>create<|"|>,'
+            'content:<|"|>use a dict like {key: value} in your code<|"|>}'
+        )
+        tool_call = gemma4.parse_tool_call(test_case, None)
+        self.assertEqual(tool_call["name"], "skill_manage")
+        self.assertEqual(tool_call["arguments"]["action"], "create")
+        self.assertIn("{", tool_call["arguments"]["content"])
+
     def test_kimi_k2(self):
         # Single tool call
         test_case = (