From 06d3b15e0108e301f1df6bebfa7ffd6522173fc6 Mon Sep 17 00:00:00 2001
From: diego-coder <39010417+diego-coder@users.noreply.github.com>
Date: Fri, 12 Sep 2025 16:29:01 -0700
Subject: [PATCH] feat: improve JSON arg parsing in ChatMLX

---
 .../langchain_community/chat_models/mlx.py    | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/libs/community/langchain_community/chat_models/mlx.py b/libs/community/langchain_community/chat_models/mlx.py
index 44e76087..4d5abfe3 100644
--- a/libs/community/langchain_community/chat_models/mlx.py
+++ b/libs/community/langchain_community/chat_models/mlx.py
@@ -1,5 +1,7 @@
 """MLX Chat Wrapper."""
 
+import json
+import logging
 from typing import (
     Any,
     Callable,
@@ -38,6 +40,8 @@
 
 from langchain_community.llms.mlx_pipeline import MLXPipeline
 
+logger = logging.getLogger(__name__)
+
 DEFAULT_SYSTEM_PROMPT = """You are a helpful, respectful, and honest assistant."""
 
 
@@ -69,6 +73,25 @@ def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
         self.tokenizer = self.llm.tokenizer
 
+    def _parse_tool_args(self, arg_text: str) -> Dict[str, Any]:
+        """Parse the arguments for a tool call.
+
+        Args:
+            arg_text: JSON string representation of the tool arguments.
+
+        Returns:
+            Parsed arguments dictionary. If parsing fails, returns a dict with
+            the original text under the ``input`` key.
+        """
+        try:
+            args = json.loads(arg_text)
+        except json.JSONDecodeError:
+            args = {"input": arg_text}
+        except Exception as e:  # pragma: no cover - defensive
+            logger.warning("Unexpected error during tool argument parsing: %s", e)
+            args = {"input": arg_text}
+        return args
+
     def _generate(
         self,
         messages: List[BaseMessage],