From a2e92c74a9eb88b42087fc8e9554c0a50d0db75e Mon Sep 17 00:00:00 2001
From: Robert Hodgson <84035529+Strings-RH@users.noreply.github.com>
Date: Fri, 14 Feb 2025 11:21:41 +0000
Subject: [PATCH] Allow for last_input_token_count=None (#644)

---
 src/smolagents/gradio_ui.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/smolagents/gradio_ui.py b/src/smolagents/gradio_ui.py
index 4425f1e..a0f04b9 100644
--- a/src/smolagents/gradio_ui.py
+++ b/src/smolagents/gradio_ui.py
@@ -141,7 +141,7 @@ 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 getattr(agent.model, "last_input_token_count", None):
+        if getattr(agent.model, "last_input_token_count", None) 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):