Update Gradio_UI.py

Gradio_UI.py

@@ -140,18 +140,27 @@ 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 (some backends return None)
-        in_tok = getattr(agent.model, "last_input_token_count", None)
-        out_tok = getattr(agent.model, "last_output_token_count", None)
+        # # Track tokens if model provides them (some backends return None)
+        # in_tok = getattr(agent.model, "last_input_token_count", None)
+        # out_tok = getattr(agent.model, "last_output_token_count", None)
         
-        if isinstance(in_tok, int):
-            total_input_tokens += in_tok
-        if isinstance(out_tok, int):
-            total_output_tokens += out_tok
+        # if isinstance(in_tok, int):
+        #     total_input_tokens += in_tok
+        # if isinstance(out_tok, int):
+        #     total_output_tokens += out_tok
         
-        if isinstance(step_log, ActionStep):
-            step_log.input_token_count = in_tok if isinstance(in_tok, int) else 0
-            step_log.output_token_count = out_tok if isinstance(out_tok, int) else 0
+        # if isinstance(step_log, ActionStep):
+        #     step_log.input_token_count = in_tok if isinstance(in_tok, int) else 0
+        #     step_log.output_token_count = out_tok if isinstance(out_tok, int) else 0
+
+        # 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
+            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
+
 
         for message in pull_messages_from_step(
             step_log,