Fixed final_answer function

Gradio_UI.py

@@ -160,38 +160,39 @@ def stream_to_gradio(
     
     # FIX: Process the final answer correctly after the loop
     if final_answer_step:
-        # Extract the actual value from the FinalAnswerStep object
-        final_answer_value = getattr(final_answer_step, 'final_answer', final_answer_step)
-        
-        # Check if the final answer is a string that looks like an image file path
-        if isinstance(final_answer_value, str) and final_answer_value.endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp')):
-            # It's an image file path, create an AgentImage object
-            from smolagents.agent_types import AgentImage
-            processed_answer = AgentImage(value=final_answer_value)
-        else:
-            # Try the default conversion
-            processed_answer = handle_agent_output_types(final_answer_value)
-
-        if isinstance(processed_answer, AgentText):
-            yield gr.ChatMessage(
-                role="assistant",
-                content=f"**Final answer:**\n{processed_answer.to_string()}\n",
-            )
-        elif isinstance(processed_answer, AgentImage):
-            # For images, Gradio needs the file path directly
-            image_path = processed_answer.to_string()
-            yield gr.ChatMessage(
-                role="assistant",
-                content={"file": image_path, "alt_text": "Generated image"},
-            )
-        elif isinstance(processed_answer, AgentAudio):
-            yield gr.ChatMessage(
-                role="assistant",
-                content={"file": processed_answer.to_string(), "alt_text": "Generated audio"},
-            )
-        else:
-            # Fallback for any other type
-            yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(processed_answer)}")
+    # Extract the actual value from the FinalAnswerStep object
+    final_answer_value = getattr(final_answer_step, 'final_answer', final_answer_step)
+    
+    # Check if the final answer is a string that looks like an image file path
+    if isinstance(final_answer_value, str) and final_answer_value.endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp')):
+        # It's an image file path, create an AgentImage object
+        from smolagents.agent_types import AgentImage
+        processed_answer = AgentImage(value=final_answer_value)
+    else:
+        # Try the default conversion
+        processed_answer = handle_agent_output_types(final_answer_value)
+
+    if isinstance(processed_answer, AgentText):
+        yield gr.ChatMessage(
+            role="assistant",
+            content=f"**Final answer:**\n{processed_answer.to_string()}\n",
+        )
+    elif isinstance(processed_answer, AgentImage):
+        # For images, pass the file path directly - Gradio will handle it
+        image_path = processed_answer.to_string()
+        yield gr.ChatMessage(
+            role="assistant",
+            content=f"**Final answer:**\n{image_path}",  # Just pass the path as a string
+        )
+    elif isinstance(processed_answer, AgentAudio):
+        # Same for audio files
+        yield gr.ChatMessage(
+            role="assistant",
+            content=f"**Final answer:**\n{processed_answer.to_string()}",
+        )
+    else:
+        # Fallback for any other type
+        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(processed_answer)}")
 
 
 class GradioUI: