slackroo commited on
Commit
7aa811c
·
verified ·
1 Parent(s): d62215b

Update Gradio_UI.py

Browse files
Files changed (1) hide show
  1. Gradio_UI.py +5 -6
Gradio_UI.py CHANGED
@@ -142,12 +142,11 @@ 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
- if agent.model.last_input_token_count:
146
- total_input_tokens += agent.model.last_input_token_count
147
- total_output_tokens += agent.model.last_output_token_count
148
- if isinstance(step_log, ActionStep):
149
- step_log.input_token_count = agent.model.last_input_token_count
150
- step_log.output_token_count = agent.model.last_output_token_count
151
 
152
  for message in pull_messages_from_step(
153
  step_log,
 
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
149
+ step_log.output_token_count = agent.model.last_output_token_count
 
150
 
151
  for message in pull_messages_from_step(
152
  step_log,