From fe5796abeb88ce35f5228841dde771608896405b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Hawrylak?= <phawrylak@outlook.com>
Date: Sat, 20 Sep 2025 15:42:00 +0200
Subject: [PATCH] Fix tool call argument serialization issue (#39)

* Tool calling arguments JSON fix

* Extracted duplicated code to helper function

* Update utils.py

---------

Co-authored-by: Game_Time <108236317+RayBytes@users.noreply.github.com>
---
 chatmock/utils.py | 38 ++++++++++++++++++++++++++------------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/chatmock/utils.py b/chatmock/utils.py
index 6c344be..a4ada3f 100644
--- a/chatmock/utils.py
+++ b/chatmock/utils.py
@@ -254,6 +254,30 @@ def sse_translate_chat(
     ws_index: dict[str, int] = {}
     ws_next_index: int = 0
     
+    def _serialize_tool_args(eff_args: Any) -> str:
+        """
+        Serialize tool call arguments with proper JSON handling.
+        
+        Args:
+            eff_args: Arguments to serialize (dict, list, str, or other)
+            
+        Returns:
+            JSON string representation of the arguments
+        """
+        if isinstance(eff_args, (dict, list)):
+            return json.dumps(eff_args)
+        elif isinstance(eff_args, str):
+            try:
+                parsed = json.loads(eff_args)
+                if isinstance(parsed, (dict, list)):
+                    return json.dumps(parsed) 
+                else:
+                    return json.dumps({"query": eff_args})  
+            except (json.JSONDecodeError, ValueError):
+                return json.dumps({"query": eff_args})
+        else:
+            return "{}"
+    
     def _extract_usage(evt: Dict[str, Any]) -> Dict[str, int] | None:
         try:
             usage = (evt.get("response") or {}).get("usage")
@@ -320,12 +344,7 @@ def sse_translate_chat(
                         except Exception:
                             pass
                     eff_params = ws_state.get(call_id, params if isinstance(params, (dict, list, str)) else {})
-                    if isinstance(eff_params, (dict, list)):
-                        args_str = json.dumps(eff_params)
-                    elif isinstance(eff_params, str):
-                        args_str = json.dumps({"query": eff_params})
-                    else:
-                        args_str = "{}"
+                    args_str = _serialize_tool_args(eff_params)
                     if call_id not in ws_index:
                         ws_index[call_id] = ws_next_index
                         ws_next_index += 1
@@ -402,12 +421,7 @@ def sse_translate_chat(
                             pass
                     eff_args = ws_state.get(call_id, raw_args if isinstance(raw_args, (dict, list, str)) else {})
                     try:
-                        if isinstance(eff_args, (dict, list)):
-                            args = json.dumps(eff_args)
-                        elif isinstance(eff_args, str):
-                            args = json.dumps({"query": eff_args})
-                        else:
-                            args = "{}"
+                        args = _serialize_tool_args(eff_args)
                     except Exception:
                         args = "{}"
                     if item.get("type") == "web_search_call" and verbose and vlog: