Update Gradio_UI.py

Gradio_UI.py

@@ -142,8 +142,9 @@ def stream_to_gradio(
     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"):
-            total_input_tokens += agent.model.last_input_token_count
-            total_output_tokens += agent.model.last_output_token_count
+            # Ensure that None is handled by defaulting to 0
+            total_input_tokens += agent.model.last_input_token_count or 0
+            total_output_tokens += agent.model.last_output_token_count or 0
             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
@@ -293,4 +294,4 @@ class GradioUI:
         demo.launch(debug=True, share=True, **kwargs)
 
 
-__all__ = ["stream_to_gradio", "GradioUI"]
+__all__ = ["stream_to_gradio", "GradioUI"]