Update Gradio_UI.py

Gradio_UI.py

@@ -166,29 +166,26 @@ def stream_to_gradio(
         print(f"DEBUG: Inner final_answer: {final_answer.final_answer}")
         print(f"DEBUG: Inner final_answer type: {type(final_answer.final_answer)}")
 
-    if isinstance(final_answer, AgentText):
+    # Extract the actual final answer from FinalAnswerStep if needed
+    actual_final_answer = final_answer.final_answer if hasattr(final_answer, 'final_answer') else final_answer
+    
+    if isinstance(actual_final_answer, AgentText):
         yield gr.ChatMessage(
             role="assistant",
-            content=f"**Final answer:**\n{final_answer.to_string()}\n",
+            content=f"**Final answer:**\n{actual_final_answer.to_string()}\n",
         )
-    elif isinstance(final_answer, AgentImage):
-        print(f"DEBUG: Got AgentImage object: {final_answer}")
-        print(f"DEBUG: AgentImage type: {type(final_answer)}")
-        print(f"DEBUG: AgentImage attributes: {dir(final_answer)}")
-        image_path = final_answer.to_string()
-        print(f"DEBUG: Image path from to_string(): {image_path}")
-        print(f"DEBUG: Does path exist? {os.path.exists(image_path) if image_path else 'No path'}")
+    elif isinstance(actual_final_answer, AgentImage):
         yield gr.ChatMessage(
             role="assistant",
-            content={"path": image_path, "mime_type": "image/png"},
+            content={"path": actual_final_answer.to_string(), "mime_type": "image/png"},
         )
-    elif isinstance(final_answer, AgentAudio):
+    elif isinstance(actual_final_answer, AgentAudio):
         yield gr.ChatMessage(
             role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
+            content={"path": actual_final_answer.to_string(), "mime_type": "audio/wav"},
         )
     else:
-        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
+        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(actual_final_answer)}")
 
 
 class GradioUI: