Add image handling in FinalAnswerTool to support saving and validating AgentImage instances

tools/final_answer.py

@@ -1,5 +1,7 @@
-from typing import Any, Optional
+from typing import Any
 from smolagents.tools import Tool
+from smolagents.agent_types import AgentImage
+import tempfile
 
 class FinalAnswerTool(Tool):
     name = "final_answer"
@@ -8,7 +10,26 @@ class FinalAnswerTool(Tool):
     output_type = "any"
 
     def forward(self, answer: Any) -> Any:
+        # Handle PIL Images by converting them to AgentImage with saved file
+        if hasattr(answer, 'save') and hasattr(answer, 'mode'):  # PIL Image
+            # Create a temporary file to save the image
+            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
+                answer.save(tmp_file.name, format='PNG')
+                # Create AgentImage from the saved file
+                return AgentImage(tmp_file.name)
+        
+        # If it's already an AgentImage, ensure it has valid data
+        if isinstance(answer, AgentImage):
+            # If the AgentImage doesn't have a valid path, try to save it
+            if not hasattr(answer, '_path') or not answer._path:
+                if hasattr(answer, '_value') and answer._value is not None:
+                    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
+                        answer._value.save(tmp_file.name, format='PNG')
+                        answer._path = tmp_file.name
+            return answer
+            
         return answer
 
     def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         self.is_initialized = False