Update Gradio_UI.py

Gradio_UI.py

@@ -140,12 +140,14 @@ def stream_to_gradio(
     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"):
+        # Patch : ne compte les tokens que si la valeur n'est pas None
+        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
-            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
+        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(
@@ -175,6 +177,7 @@ def stream_to_gradio(
         yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
 
 
+
 class GradioUI:
     """A one-line interface to launch your agent in Gradio"""