Update tools/final_answer.py

tools/final_answer.py

@@ -1,14 +1,29 @@
-from typing import Any, Optional
-from smolagents.tools import Tool
+from typing import Any
+from smolagents.tools import Tool, AgentImage
 
 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"
+    description = "Provides the final formatted response to the user."
+    inputs = {
+        'answer': {'type': 'any', 'description': 'The final answer to the problem'},
+        'image': {'type': 'any', 'description': 'Optional image result to be displayed'}
+    }
+    output_type = "string"
 
-    def forward(self, answer: Any) -> Any:
-        return answer
+    def forward(self, answer: Any = "", image: Any = None) -> str:
+        # Handle AgentImage rendering
+        if isinstance(image, AgentImage):
+            try:
+                image_html = image.display()  # or image.to_html()
+            except:
+                image_html = "[Image could not be displayed]"
+            # Combine text + image if both are provided
+            if answer:
+                return f"{answer}\n\n{image_html}"
+            else:
+                return image_html
+        else:
+            return str(answer)
 
     def __init__(self, *args, **kwargs):
-        self.is_initialized = False
+        super().__init__(*args, **kwargs)