Update Gradio_UI.py

Gradio_UI.py

@@ -136,18 +136,24 @@ def stream_to_gradio(
         )
     import gradio as gr
 
+    # Initialize token counters
     total_input_tokens = 0
     total_output_tokens = 0
-
+    
+    # Iterate over the agent run logs
     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"):
+        if hasattr(agent.model, "last_input_token_count") and agent.model.last_input_token_count is not None:
             total_input_tokens += 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.input_token_count = agent.model.last_input_token_count
-                step_log.output_token_count = agent.model.last_output_token_count
+    
+        # Store token counts in step logs if it's an ActionStep
+        if isinstance(step_log, ActionStep):
+            step_log.input_token_count = agent.model.last_input_token_count if hasattr(agent.model, "last_input_token_count") else 0
+            step_log.output_token_count = agent.model.last_output_token_count if hasattr(agent.model, "last_output_token_count") else 0
 
+        
         for message in pull_messages_from_step(
             step_log,
         ):