final answer to display image

tools/final_answer.py

@@ -1,5 +1,23 @@
+# from typing import Any, Optional
+# from smolagents.tools import Tool
+
+# class FinalAnswerTool(Tool):
+#     name = "final_answer"
+#     description = "Provides a final answer to the given problem."
+#     inputs = {'answer': {'type': 'any', 'description': 'The final answer to the problem'}}
+#     output_type = "any"
+
+#     def forward(self, answer: Any) -> Any:
+#         return answer
+
+#     def __init__(self, *args, **kwargs):
+#         self.is_initialized = False
+
 from typing import Any, Optional
 from smolagents.tools import Tool
+import base64
+from pathlib import Path
+import mimetypes
 
 class FinalAnswerTool(Tool):
     name = "final_answer"
@@ -8,7 +26,32 @@ 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)}"
+        
+        # Return the original answer for non-image responses
         return answer
 
     def __init__(self, *args, **kwargs):
-        self.is_initialized = False
+        self.is_initialized = False