Update Gradio_UI.py

fix the streaming err

Gradio_UI.py

@@ -122,7 +122,6 @@ def pull_messages_from_step(
         yield gr.ChatMessage(role="assistant", content=f"{step_footnote}")
         yield gr.ChatMessage(role="assistant", content="-----")
 
-
 def stream_to_gradio(
     agent,
     task: str,
@@ -139,6 +138,62 @@ def stream_to_gradio(
     total_input_tokens = 0
     total_output_tokens = 0
 
+    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
+        # Track tokens if model provides them
+        if hasattr(agent.model, "last_input_token_count"):
+            # Check if token counts are not None before adding
+            if agent.model.last_input_token_count is not None:
+                total_input_tokens += agent.model.last_input_token_count
+                if isinstance(step_log, ActionStep):
+                    step_log.input_token_count = agent.model.last_input_token_count
+            
+            if hasattr(agent.model, "last_output_token_count") and agent.model.last_output_token_count is not None:
+                total_output_tokens += agent.model.last_output_token_count
+                if isinstance(step_log, ActionStep):
+                    step_log.output_token_count = agent.model.last_output_token_count
+
+        for message in pull_messages_from_step(
+            step_log,
+        ):
+            yield message
+
+    final_answer = step_log  # Last log is the run's final_answer
+    final_answer = handle_agent_output_types(final_answer)
+
+    if isinstance(final_answer, AgentText):
+        yield gr.ChatMessage(
+            role="assistant",
+            content=f"**Final answer:**\n{final_answer.to_string()}\n",
+        )
+    elif isinstance(final_answer, AgentImage):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_answer.to_string(), "mime_type": "image/png"},
+        )
+    elif isinstance(final_answer, AgentAudio):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
+        )
+    else:
+        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
+        
+def stream_to_gradio_bk(
+    agent,
+    task: str,
+    reset_agent_memory: bool = False,
+    additional_args: Optional[dict] = None,
+):
+    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
+    if not _is_package_available("gradio"):
+        raise ModuleNotFoundError(
+            "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
+        )
+    import gradio as gr
+
+    total_input_tokens = 0
+    total_output_tokens = 0
+
     for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
         # Track tokens if model provides them
         if hasattr(agent.model, "last_input_token_count"):