make final answer work with gardio ui

tools/final_answer.py

@@ -15,9 +15,7 @@
 
 from typing import Any, Optional
 from smolagents.tools import Tool
-import base64
-from pathlib import Path
-import mimetypes
+from smolagents.agent_types import AgentImage  # Import this class
 
 class FinalAnswerTool(Tool):
     name = "final_answer"
@@ -26,29 +24,10 @@ class FinalAnswerTool(Tool):
     output_type = "any"
 
     def forward(self, answer: Any) -> Any:
-        # Check if the answer is an image path (specific to gradio temp directories)
-        if isinstance(answer, str) and '/tmp/gradio/' in answer:
-            try:
-                # Read the image file as binary
-                img_path = Path(answer)
-                if img_path.exists():
-                    with open(img_path, 'rb') as img_file:
-                        img_data = img_file.read()
-                    
-                    # Determine the MIME type
-                    mime_type = mimetypes.guess_type(answer)[0] or 'image/webp'
-                    
-                    # Return both the original path and the image data
-                    return {
-                        "type": "image",
-                        "path": answer,
-                        "mime_type": mime_type,
-                        "data": img_data
-                    }
-                else:
-                    return f"Image was generated but could not be accessed at: {answer}"
-            except Exception as e:
-                return f"Error processing image: {str(e)}"
+        # Check if the answer is an image path
+        if isinstance(answer, str) and ('/tmp/gradio/' in answer or answer.endswith(('.png', '.jpg', '.jpeg', '.webp'))):
+            # Return as AgentImage for proper handling by stream_to_gradio
+            return AgentImage(answer)
         
         # Return the original answer for non-image responses
         return answer