Update Gradio UI streaming to use agent.run() with improved error handling

- Refactored stream_to_gradio function to use agent.run() with streaming
- Added Gradio-specific error message handling
- Improved token tracking and error logging
- Imported gradio explicitly for error message generation

Gradio_UI.py

@@ -29,6 +29,7 @@ from smolagents.agents import ActionStep, MultiStepAgent
 from smolagents.memory import MemoryStep
 from smolagents.utils import _is_package_available
 import logging
+
 logger = logging.getLogger(__name__)
 
 
@@ -160,26 +161,31 @@ def pull_messages_from_step(
 
 def stream_to_gradio(agent, task: str, reset_agent_memory: bool = True):
     """Stream agent responses to Gradio interface with better error handling"""
+    import gradio as gr
+
     total_input_tokens = 0
     total_output_tokens = 0
 
     try:
-        for msg in agent.chat(task, reset_memory=reset_agent_memory):
-            # Safely handle token counting
-            if hasattr(agent.model, "last_input_token_count"):
-                input_tokens = agent.model.last_input_token_count or 0
-                total_input_tokens += input_tokens
+        # Use agent.run() instead of agent.chat()
+        for step_log in agent.run(task, stream=True, reset=reset_agent_memory):
+            # Extract messages from the step
+            for msg in pull_messages_from_step(step_log):
+                # Safely handle token counting
+                if hasattr(agent.model, "last_input_token_count"):
+                    input_tokens = agent.model.last_input_token_count or 0
+                    total_input_tokens += input_tokens
 
-            if hasattr(agent.model, "last_output_token_count"):
-                output_tokens = agent.model.last_output_token_count or 0
-                total_output_tokens += output_tokens
+                if hasattr(agent.model, "last_output_token_count"):
+                    output_tokens = agent.model.last_output_token_count or 0
+                    total_output_tokens += output_tokens
 
-            yield msg
+                yield msg
 
     except Exception as e:
         error_msg = f"Error during chat: {str(e)}"
         logger.error(error_msg)
-        yield f"⚠️ {error_msg}"
+        yield gr.ChatMessage(role="assistant", content=f"⚠️ {error_msg}")
 
 
 class GradioUI: