Update tools/final_answer.py

tools/final_answer.py

@@ -1,27 +1,31 @@
-from typing import Optional
-from smolagents.tools import Tool
-from PIL import Image
 import os
 import datetime
+from typing import Optional, Union, Tuple
+from PIL import Image
+from smolagents import Tool
 
 class FinalAnswerTool(Tool):
     name = "final_answer"
     description = "Provides a final answer. Can include text and/or a generated image."
-    
+
     inputs = {
         'text': {'type': 'string', 'description': 'Final textual answer', 'nullable': True},
         'image_path': {'type': 'string', 'description': 'Path to a generated image file', 'nullable': True}
     }
-    
-    output_type = "string"
 
-    def forward(self, text: Optional[str] = None, image_path: Optional[str] = None) -> str:
-        # If an image exists, prefer displaying it
+    output_type = "any"  # <-- can now return text, image, or both
+
+    def forward(
+        self, 
+        text: Optional[str] = None, 
+        image_path: Optional[str] = None
+    ) -> Union[str, Image.Image, Tuple[Image.Image, str]]:
+        
         if image_path and os.path.exists(image_path):
-            image_display_path = image_path
+            # Convert .webp to .png (since some frontends don't support webp well)
             if image_path.lower().endswith(".webp"):
                 try:
-                    img = Image.open(image_path)
+                    img = Image.open(image_path).convert("RGB")
                     temp_dir = os.path.dirname(image_path) or "/tmp"
                     os.makedirs(temp_dir, exist_ok=True)
 
@@ -33,25 +37,18 @@ class FinalAnswerTool(Tool):
                     )
 
                     img.save(temp_png_path, "PNG")
-                    image_display_path = temp_png_path
+                    image_path = temp_png_path
                 except Exception as e:
-                    print(f"[ERROR] Failed to convert image {image_path} to PNG: {e}")
+                    print(f"[ERROR] Failed to convert image {image_path}: {e}")
 
-            # Embed image in Markdown
-            output = f"![Generated Image]({image_display_path})"
-            if text:
-                output += f"\n\n{text}"
-            return output
-
-        # Fallback to text only
-        if text:
-            return text
-        
-        return "No final answer provided."
+            img = Image.open(image_path)
 
-    def __call__(self, *args, **kwargs):
-        return self.forward(*args, **kwargs)
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.is_initialized = False
+            if text:
+                # Return tuple if both text + image
+                return (img, text)
+            else:
+                # Return only image
+                return img
+
+        # Only text case
+        return text