Better tool UI

app.py

@@ -142,13 +142,11 @@ async def start():
 async def execute_tool_call(tool_call: Any):
     """
     Helper function to find the correct MCP session, execute a tool call,
-    and robustly serialize the result to a JSON string.
+    and robustly serialize the result for both the UI and the LLM.
     """
     tool_name = tool_call.function.name
     tool_args_str = tool_call.function.arguments
 
-    log.debug(f"Executing tool call: {tool_name} with args: {tool_args_str}")
-
     history = cl.user_session.get("message_history")
 
     mcp_connection_name = None
@@ -173,7 +171,7 @@ async def execute_tool_call(tool_call: Any):
         return
 
     async with cl.Step(type="tool", name=tool_name) as step:
-        step.input = tool_args_str
+        step.input = json.loads(tool_args_str)
 
         try:
             mcp_session, _ = cl.context.session.mcp_sessions.get(mcp_connection_name)
@@ -181,56 +179,47 @@ async def execute_tool_call(tool_call: Any):
             tool_result = await mcp_session.call_tool(tool_name, tool_args)
 
             content = tool_result.content
-            final_content_str = ""
+            output_for_step = content  # Default to showing raw content
+            output_for_llm = str(content)  # Default to string for LLM
 
-            if isinstance(content, str):
-                final_content_str = content
-            elif (
+            if (
                 isinstance(content, list)
-                and len(content) == 1
+                and len(content) > 0
                 and hasattr(content[0], "text")
-                and not hasattr(content[0], "data")
             ):
-                final_content_str = content[0].text
-            else:
-
-                def to_serializable(obj: Any) -> Any:
-                    if hasattr(obj, "model_dump"):
-                        return obj.model_dump(exclude_none=True)
-                    if hasattr(obj, "dict"):
-                        return obj.dict(exclude_none=True)
-                    if hasattr(obj, "to_dict"):
-                        return obj.to_dict()
-                    if isinstance(obj, list):
-                        return [to_serializable(item) for item in obj]
-                    return obj
-
-                serializable_content = to_serializable(content)
-                final_content_str = json.dumps(serializable_content, indent=2)
-
-            step.output = final_content_str
-            log.debug(f"Tool {tool_name} executed successfully")
+                text_content = content[0].text
+                output_for_llm = text_content  # LLM always gets the original string
+
+                # ✅ **THIS IS THE FIX**
+                # Try to parse the text as JSON for a nice UI display.
+                try:
+                    output_for_step = json.loads(text_content)
+                except json.JSONDecodeError:
+                    # If it's not valid JSON, it's just plain text.
+                    output_for_step = text_content
 
+            # Set the step output for the UI
+            step.output = output_for_step
+
+            # Append the string version to history for the LLM
             history.append(
                 {
                     "tool_call_id": tool_call.id,
                     "role": "tool",
                     "name": tool_name,
-                    "content": final_content_str,
+                    "content": output_for_llm,
                 }
             )
 
         except Exception as e:
-            error_msg = f"Error executing tool {tool_name}: {str(e)}"
-            log.error(error_msg)
-            log.error(f"Full traceback: {traceback.format_exc()}")
+            error_msg = f"Error executing tool `{tool_name}`: {traceback.format_exc()}"
             step.error = str(e)
             history.append(
                 {
                     "tool_call_id": tool_call.id,
                     "role": "tool",
                     "name": tool_name,
-                    "content": f"Error executing tool {tool_name}: {traceback.format_exc()}",
+                    "content": error_msg,
                 }
             )