diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py index 611f760..492fbb8 100644 --- a/mlx_lm/tokenizer_utils.py +++ b/mlx_lm/tokenizer_utils.py @@ -291,7 +291,10 @@ class TokenizerWrapper: self._tool_call_end = tool_call_end vocab = tokenizer.get_vocab() - THINK_TOKENS = [("", "")] + THINK_TOKENS = [ + ("", ""), + ("", ""), + ] for think_start, think_end in THINK_TOKENS: if think_start in vocab and think_end in vocab: self._think_start = think_start @@ -472,6 +475,8 @@ def _infer_tool_parser(chat_template): return "minimax_m2" elif "" in chat_template: return "function_gemma" + elif "" in chat_template: + return "longcat" elif "" in chat_template: return "glm47" elif "\n", re.DOTALL) +_func_arg_regex = re.compile( + r"(.*?)(?:\\n|\s)*(.*?)", + re.DOTALL, +) + +tool_call_start = "" +tool_call_end = "" + + +def _is_string_type( + tool_name: str, + arg_name: str, + tools: list[Any] | None, +) -> bool: + if tools is None: + return False + for tool in tools: + func = tool["function"] + if func["name"] == tool_name: + params = func["parameters"] + if params is None: + return False + arg_type = params.get("properties", {}).get(arg_name, {}).get("type", None) + return arg_type == "string" + return False + + +def _deserialize(value: str) -> Any: + try: + return json.loads(value) + except Exception: + pass + + try: + return ast.literal_eval(value) + except Exception: + pass + return value + + +def parse_tool_call(text: str, tools: list[Any] | None = None): + text = text.strip() + + if text.startswith("{"): + try: + return json.loads(text) + except json.JSONDecodeError: + pass + + func_name = _func_name_regex.search(text).group(1).strip() + pairs = _func_arg_regex.findall(text) + arg_dct = {} + for key, value in pairs: + arg_key = key.strip() + arg_val = value.strip() + if not _is_string_type(func_name, arg_key, tools): + arg_val = _deserialize(arg_val) + arg_dct[arg_key] = arg_val + return dict(name=func_name, arguments=arg_dct) diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index d72c8cc..bcef5c7 100644 --- a/tests/test_tool_parsing.py +++ b/tests/test_tool_parsing.py @@ -6,6 +6,7 @@ from mlx_lm.tool_parsers import ( glm47, json_tools, kimi_k2, + longcat, minimax_m2, qwen3_coder, ) @@ -14,12 +15,23 @@ from mlx_lm.tool_parsers import ( class TestToolParsing(unittest.TestCase): def test_parsers(self): - parsers = [function_gemma, glm47, json_tools, kimi_k2, minimax_m2, qwen3_coder] + parsers = [ + function_gemma, + glm47, + json_tools, + longcat, + longcat, + kimi_k2, + minimax_m2, + qwen3_coder, + ] test_cases = [ "call:multiply{a:12234585,b:48838483920}", "multiplya12234585b48838483920", '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}', + "multiplya\n12234585\nb\n48838483920", + '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}', '<|tool_call_begin|>functions.multiply:0<|tool_call_argument_begin|>{"a": 12234585, "b": 48838483920}<|tool_call_end|>', '\n12234585\n48838483920\n', "\n\n12234585\n\n\n48838483920\n\n", @@ -55,6 +67,8 @@ class TestToolParsing(unittest.TestCase): "call:get_current_temperature{location:London}", 'get_current_temperaturelocation"London"', '{"name": "get_current_temperature", "arguments": {"location": "London"}}', + "get_current_temperaturelocation\nLondon", + '{"name": "get_current_temperature", "arguments": {"location": "London"}}', '<|tool_call_begin|>functions.get_current_temperature:0<|tool_call_argument_begin|>{"location": "London"}<|tool_call_end|>', '\nLondon\n', "\n\nLondon\n\n",