File size: 3,231 Bytes
fd18d93 8fe992b d24c217 8fe992b d24c217 8fe992b d24c217 8fe992b d24c217 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | from typing import Any, Optional
from smolagents.tools import Tool
from PIL import Image # Add this import
import os # Add this import
import datetime # Add this import
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."
# IMPORTANT: Update the inputs to explicitly allow 'text' and 'image_path'
inputs = {
'text': {'type': 'str', 'description': 'The final textual answer to the problem'},
'image_path': {'type': 'str', 'description': 'Optional path to a generated image file (e.g., .webp, .png, .jpg) to be displayed.'}
}
output_type = "str" # The output will be a Markdown 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 image_path and os.path.exists(image_path):
# Gradio should typically handle .webp directly, but converting to PNG
# is a robust fallback for broader compatibility or specific display quirks.
image_display_path = image_path
if image_path.lower().endswith(".webp"):
try:
img = Image.open(image_path)
# Create a temporary filename for the PNG
# Use a more robust temporary path strategy
temp_dir = os.path.dirname(image_path) if os.path.dirname(image_path) else "/tmp"
os.makedirs(temp_dir, exist_ok=True) # Ensure temp_dir exists
base_name = os.path.basename(image_path)
name_without_ext = os.path.splitext(base_name)[0]
# Add microsecond precision to avoid collisions if multiple images are generated rapidly
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.")
# Fallback to original path if conversion fails
# Embed the image using Markdown syntax, which Gradio will render
output_content.append(f"")
if text:
output_content.append(text)
if not output_content:
return "No final answer provided."
return "\n\n".join(output_content)
def __init__(self, *args, **kwargs):
# Important: Call the superclass __init__ to properly initialize the Tool
super().__init__(*args, **kwargs)
self.is_initialized = False # You can keep this if it's used elsewhere |