Updated Gradio

Gradio_UI.py

@@ -136,8 +136,6 @@ def stream_to_gradio(
         )
     import gradio as gr
 
-    total_input_tokens = 0
-    total_output_tokens = 0
     final_answer_step = None
 
     for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
@@ -145,10 +143,8 @@ def stream_to_gradio(
             final_answer_step = step_log
             continue # Don't display the final answer step itself, process it at the end
 
-        # Track tokens if model provides them
-        if hasattr(agent.model, "last_input_token_count"):
-            total_input_tokens += agent.model.last_input_token_count
-            total_output_tokens += agent.model.last_output_token_count
+        # FIX: Check if token counts exist before trying to use them
+        if hasattr(agent.model, "last_input_token_count") and agent.model.last_input_token_count is not None:
             if isinstance(step_log, ActionStep):
                 step_log.input_token_count = agent.model.last_input_token_count
                 step_log.output_token_count = agent.model.last_output_token_count
@@ -158,41 +154,20 @@ def stream_to_gradio(
         ):
             yield message
     
-    # FIX: Process the final answer correctly after the loop
+    # Process the final answer correctly after the loop
     if final_answer_step:
         # Extract the actual value from the FinalAnswerStep object
         final_answer_value = getattr(final_answer_step, 'final_answer', final_answer_step)
         
-        # Check if the final answer is a string that looks like an image file path
-        if isinstance(final_answer_value, str) and final_answer_value.endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp')):
-            # It's an image file path, create an AgentImage object
-            from smolagents.agent_types import AgentImage
-            processed_answer = AgentImage(value=final_answer_value)
-        else:
-            # Try the default conversion
-            processed_answer = handle_agent_output_types(final_answer_value)
-
-        if isinstance(processed_answer, AgentText):
-            yield gr.ChatMessage(
-                role="assistant",
-                content=f"**Final answer:**\n{processed_answer.to_string()}\n",
-            )
-        elif isinstance(processed_answer, AgentImage):
-            # For images, pass the file path directly - Gradio will handle it
-            image_path = processed_answer.to_string()
-            yield gr.ChatMessage(
-                role="assistant",
-                content=f"**Final answer:**\n{image_path}",  # Just pass the path as a string
-            )
-        elif isinstance(processed_answer, AgentAudio):
-            # Same for audio files
+        # Directly check if the final answer is a path to an image file
+        if isinstance(final_answer_value, str) and final_answer_value.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
             yield gr.ChatMessage(
                 role="assistant",
-                content=f"**Final answer:**\n{processed_answer.to_string()}",
+                content=final_answer_value, # Pass the file path directly to Gradio
             )
         else:
-            # Fallback for any other type
-            yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(processed_answer)}")
+            # Fallback for text or any other type of answer
+            yield gr.ChatMessage(role="assistant", content=f"**Final answer:**\n{str(final_answer_value)}")
 
 
 class GradioUI:
@@ -286,7 +261,7 @@ class GradioUI:
             file_uploads_log = gr.State([])
             chatbot = gr.Chatbot(
                 label="Agent",
-                type="messages", # FIX: Re-added this crucial parameter
+                type="messages",
                 render_markdown=True,
                 avatar_images=(
                     None,