Update Gradio_UI.py

Gradio_UI.py

@@ -153,16 +153,47 @@ def stream_to_gradio(
         ):
             yield message
 
-    final_step = step_log  # last step returned by agent.run()
-    final_output = None
-    
-    if hasattr(final_step, "tool_calls"):
+    final_step = step_log  # Final ActionStep
+
+    # Debugging tool calls to inspect final_answer structure
+    print("DEBUG TOOL CALLS:", final_step.tool_calls)
+
+    # Attempt to get the direct tool_output
+    final_output = getattr(final_step, "tool_output", None)
+
+    # Fallback if tool_output is missing
+    if not final_output and getattr(final_step, "tool_calls", None):
         for call in final_step.tool_calls:
             if call.name == "final_answer":
-                # Prefer the actual output if available, fallback to arguments
-                final_output = call.output if call.output is not None else call.arguments.get("answer")
+                print("DEBUG FINAL ANSWER TOOL CALL:", call)
+                # Try the `output` field if present (some agents set this explicitly)
+                final_output = getattr(call, "output", None)
+                if final_output is None:
+                    # Fallback to arguments dict
+                    final_output = call.arguments.get("answer")
                 break
 
+    final_output = handle_agent_output_types(final_output)
+
+    if isinstance(final_output, AgentText):
+        yield gr.ChatMessage(
+            role="assistant",
+            content=final_output.to_string().strip(),
+        )
+    elif isinstance(final_output, AgentImage):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_output.to_string(), "mime_type": "image/png"},
+        )
+    elif isinstance(final_output, AgentAudio):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_output.to_string(), "mime_type": "audio/wav"},
+        )
+    else:
+        yield gr.ChatMessage(role="assistant", content=str(final_output).strip())
+
+
 
     
     final_output = handle_agent_output_types(final_output)