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.
pull/615/head
Sambhav Dixit 2 months ago committed by GitHub
parent 8900d645c8
commit d72d2e74ff
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -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)

Loading…
Cancel
Save