Spaces:
Sleeping
Sleeping
| 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 |