Fix last step missing duration

huggingface · Nov 11, 2024 · c73eeb7 · c73eeb7
1 parent a131fd0
commit c73eeb7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/transformers/agents/agents.py b/src/transformers/agents/agents.py
@@ -834,6 +834,7 @@ def stream_run(self, task: str):
             self.logger.error(error_message, exc_info=1)
             final_answer = self.provide_final_answer(task)
             final_step_log["final_answer"] = final_answer
+            final_step_log["step_duration"] = 0
             for callback in self.step_callbacks:
                 callback(final_step_log)
             yield final_step_log
@@ -874,6 +875,7 @@ def direct_run(self, task: str):
             self.logger.error(error_message, exc_info=1)
             final_answer = self.provide_final_answer(task)
             final_step_log["final_answer"] = final_answer
+            final_step_log["step_duration"] = 0
             for callback in self.step_callbacks:
                 callback(final_step_log)
 

diff --git a/src/transformers/agents/monitoring.py b/src/transformers/agents/monitoring.py
@@ -87,7 +87,7 @@ def __init__(self, tracked_llm_engine):
             self.total_output_token_count = 0
 
     def update_metrics(self, step_log):
-        step_duration = step_log.get("step_duration", None)
+        step_duration = step_log["step_duration"]
         self.step_durations.append(step_duration)
         logger.info(f"Step {len(self.step_durations)}:")
         logger.info(f"- Time taken: {step_duration:.2f} seconds (valid only if step succeeded)")