Update Gradio_UI.py

Gradio_UI.py

@@ -150,7 +150,12 @@ def stream_to_gradio(
                 # Handle the case where it's None, for example, by assuming a default value of 0
                 total_input_tokens += 0
 
-            total_output_tokens += agent.model.last_output_token_count
+            # total_output_tokens += agent.model.last_output_token_count
+            if agent.model.last_output_token_count is not None:
+                total_output_tokens += agent.model.last_output_token_count
+            else:
+                # Handle the case where it's None, for example, by assuming a default value of 0
+                total_output_tokens += 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