diff --git a/mlx_lm/server.py b/mlx_lm/server.py index c74942c..685ae9a 100644 --- a/mlx_lm/server.py +++ b/mlx_lm/server.py @@ -78,7 +78,14 @@ class ToolCallFormatter: 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) diff --git a/mlx_lm/tool_parsers/gemma4.py b/mlx_lm/tool_parsers/gemma4.py index 1df4d29..c491487 100644 --- a/mlx_lm/tool_parsers/gemma4.py +++ b/mlx_lm/tool_parsers/gemma4.py @@ -5,10 +5,19 @@ from typing import Any, Optional 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: diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index b19e1d5..52892b7 100644 --- a/tests/test_tool_parsing.py +++ b/tests/test_tool_parsing.py @@ -254,6 +254,27 @@ class TestToolParsing(unittest.TestCase): {"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 = (