Update tools/final_answer.py

tools/final_answer.py

@@ -1,5 +1,8 @@
-from typing import Any, Optional
+from typing import Any
 from smolagents.tools import Tool
+from smolagents.agent_types import AgentImage
+import os
+from uuid import uuid4
 
 class FinalAnswerTool(Tool):
     name = "final_answer"
@@ -8,7 +11,15 @@ class FinalAnswerTool(Tool):
     output_type = "any"
 
     def forward(self, answer: Any) -> Any:
+        # If it's an AgentImage, save it to file and return an updated object
+        if isinstance(answer, AgentImage):
+            image = answer
+            # Save to temp folder inside the Space
+            output_dir = "/tmp/final_outputs"
+            os.makedirs(output_dir, exist_ok=True)
+            path = os.path.join(output_dir, f"{uuid4().hex}.png")
+            image.save(path)
+            # Patch the object so `to_string()` will work
+            image.path = path
+            return image
         return answer
-
-    def __init__(self, *args, **kwargs):
-        self.is_initialized = False