diff --git a/mlx_lm/tool_parsers/json_tools.py b/mlx_lm/tool_parsers/json_tools.py index 27a9caa44..5099f221a 100644 --- a/mlx_lm/tool_parsers/json_tools.py +++ b/mlx_lm/tool_parsers/json_tools.py @@ -2,10 +2,46 @@ import json -tool_call_start = "" +# The start marker intentionally omits the closing ">". Tool-call markers are matched as +# exact token-id sequences, and "" encodes to a run ending in a standalone ">" +# token. Many tokenizers merge that ">" with the following byte (e.g. "\n" -> a +# single ">\n" token), so the full-marker token run never appears in the generated stream +# and the call is never captured. Matching the stable "" before the JSON. +tool_call_start = ""/newline left by the prefix start marker and to a trailing + "" if the end marker likewise merges. Output is identical to + json.loads(text.strip()) for a clean JSON-only segment. + """ + start = text.find("{") + if start == -1: + raise ValueError("no JSON object in tool call segment") + depth = 0 + in_str = False + esc = False + for i in range(start, len(text)): + c = text[i] + if in_str: + if esc: + esc = False + elif c == "\\": + esc = True + elif c == '"': + in_str = False + continue + if c == '"': + in_str = True + elif c == "{": + depth += 1 + elif c == "}": + depth -= 1 + if depth == 0: + return json.loads(text[start : i + 1]) + raise ValueError("unbalanced JSON in tool call segment") diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index 52892b7ff..686646687 100644 --- a/tests/test_tool_parsing.py +++ b/tests/test_tool_parsing.py @@ -207,6 +207,28 @@ def test_gemma4(self): {"settings": {"enabled": True, "name": "test"}}, ) + def test_json_tools_marker_merge(self): + # The start marker is "", which many tokenizers + # merge with the next byte). The server therefore captures the tool segment starting + # at the leftover ">", and the segment may also include a trailing "" if + # that marker's tokens merge too. parse_tool_call must extract the JSON regardless. + expected = {"name": "get_weather", "arguments": {"city": "Paris"}} + variants = [ + '{"name": "get_weather", "arguments": {"city": "Paris"}}', # clean (backward-compat) + '>\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n', # leftover ">" + '>\n{"name": "get_weather", "arguments": {"city": "Paris"}}\n', # + end marker + ] + for text in variants: + with self.subTest(text=text): + self.assertEqual(json_tools.parse_tool_call(text, None), expected) + + # nested braces and a "}" inside a string value must not truncate extraction + text = '{"name": "f", "arguments": {"expr": "a{b}c", "n": {"k": 2}}}' + self.assertEqual( + json_tools.parse_tool_call(text, None), + {"name": "f", "arguments": {"expr": "a{b}c", "n": {"k": 2}}}, + ) + # Array of strings test_case = 'call:tag{items:[<|"|>foo<|"|>,<|"|>bar<|"|>]}' tool_call = gemma4.parse_tool_call(test_case, None)