Update tools/final_answer.py

tools/final_answer.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Optional
 from smolagents.tools import Tool
 from PIL import Image
 import os
@@ -6,53 +6,51 @@ import datetime
 
 class FinalAnswerTool(Tool):
     name = "final_answer"
-    description = "Provides a final answer to the given problem. Can include a textual response and/or a path to a generated image."
+    description = "Provides a final answer. Can include text and/or a generated image."
     
-    # CORRECTED: Add "nullable": True for both 'text' and 'image_path'
     inputs = {
-        'text': {'type': 'string', 'description': 'The final textual answer to the problem', 'nullable': True},
-        'image_path': {'type': 'string', 'description': 'Optional path to a generated image file (e.g., .webp, .png, .jpg) to be displayed.', 'nullable': True}
+        '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" # Output will be a Markdown string
+    output_type = "string"
 
     def forward(self, text: Optional[str] = None, image_path: Optional[str] = None) -> str:
-        """
-        Processes the final answer, potentially embedding an image from a given path.
-        """
-        output_content = []
-
+        # If an image exists, prefer displaying it
         if image_path and os.path.exists(image_path):
             image_display_path = image_path
             if image_path.lower().endswith(".webp"):
                 try:
                     img = Image.open(image_path)
-                    temp_dir = os.path.dirname(image_path) if os.path.dirname(image_path) else "/tmp"
+                    temp_dir = os.path.dirname(image_path) or "/tmp"
                     os.makedirs(temp_dir, exist_ok=True)
-                    
+
                     base_name = os.path.basename(image_path)
                     name_without_ext = os.path.splitext(base_name)[0]
                     temp_png_path = os.path.join(
                         temp_dir,
                         f"{name_without_ext}_{datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')}.png"
                     )
-                    
+
                     img.save(temp_png_path, "PNG")
                     image_display_path = temp_png_path
-                    print(f"[DEBUG] Converted {image_path} to {temp_png_path} for display.")
                 except Exception as e:
-                    print(f"[ERROR] Failed to convert image {image_path} to PNG: {e}. Attempting to use original path.")
-        
-            output_content.append(f"![Generated Image]({image_display_path})")
+                    print(f"[ERROR] Failed to convert image {image_path} to PNG: {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:
-            output_content.append(text)
+            return text
         
-        if not output_content:
-            # If both are None, provide a default message
-            return "No final answer provided."
+        return "No final answer provided."
 
-        return "\n\n".join(output_content)
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)