Update Gradio_UI.py

Gradio_UI.py

@@ -130,24 +130,26 @@ def stream_to_gradio(
     additional_args: Optional[dict] = None,
 ):
     """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
-    if not _is_package_available("gradio"):
+    if not is_package_available("gradio"):
         raise ModuleNotFoundError(
             "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
         )
     import gradio as gr
-
     total_input_tokens = 0
     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"):
-            total_input_tokens += agent.model.last_input_token_count
-            total_output_tokens += agent.model.last_output_token_count
+            # Check if token counts are None before adding
+            input_tokens = agent.model.last_input_token_count or 0
+            output_tokens = agent.model.last_output_token_count or 0
+            
+            total_input_tokens += input_tokens
+            total_output_tokens += output_tokens
+            
             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
-
+                step_log.input_token_count = input_tokens
+                step_log.output_token_count = output_tokens
         for message in pull_messages_from_step(
             step_log,
         ):