From d72d2e74ff047daa6bf359b952f3414655d3affe Mon Sep 17 00:00:00 2001 From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com> Date: Sun, 27 Oct 2024 19:14:23 +0530 Subject: [PATCH] Add test for agent output updating in log_step_metadata - Implemented test_agent_output_updating to verify that logging step metadata correctly updates the total token count and ensures that the agent's output steps are properly tracked confirming only one step is recorded. --- tests/agents/test_agent_logging.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py index 98b0f8a5..2fe13ea9 100644 --- a/tests/agents/test_agent_logging.py +++ b/tests/agents/test_agent_logging.py @@ -47,3 +47,14 @@ class TestAgentLogging(unittest.TestCase): log_result = self.agent.log_step_metadata(1, "prompt", "response") self.assertEqual(log_result['tokens']['total'], 400) + + def test_agent_output_updating(self): + initial_total_tokens = sum(step['tokens']['total'] for step in self.agent.agent_output.steps) + self.agent.log_step_metadata(1, "prompt", "response") + + final_total_tokens = sum(step['tokens']['total'] for step in self.agent.agent_output.steps) + self.assertEqual( + final_total_tokens - initial_total_tokens, + 200 + ) + self.assertEqual(len(self.agent.agent_output.steps), 1)