eli-the-pooh commited on
Commit
20933fc
·
verified ·
1 Parent(s): 1818bab

Update Gradio_UI.py

Browse files
Files changed (1) hide show
  1. Gradio_UI.py +6 -1
Gradio_UI.py CHANGED
@@ -142,7 +142,12 @@ def stream_to_gradio(
142
  for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
143
  # Track tokens if model provides them
144
  if hasattr(agent.model, "last_input_token_count"):
145
- total_input_tokens += agent.model.last_input_token_count
 
 
 
 
 
146
  total_output_tokens += agent.model.last_output_token_count
147
  if isinstance(step_log, ActionStep):
148
  step_log.input_token_count = agent.model.last_input_token_count
 
142
  for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
143
  # Track tokens if model provides them
144
  if hasattr(agent.model, "last_input_token_count"):
145
+ # Correction suggested by ChtGPT. Removed:
146
+ # total_input_tokens += agent.model.last_input_token_count
147
+ # if last_input_token_count is None, treat it as 0
148
+ input_tokens = agent.model.last_input_token_count or 0
149
+ total_input_tokens += input_tokens
150
+
151
  total_output_tokens += agent.model.last_output_token_count
152
  if isinstance(step_log, ActionStep):
153
  step_log.input_token_count = agent.model.last_input_token_count