Skip to content

Conversation

kalanyuz
Copy link
Contributor

With response API you get additional usage token information in the form of litellm's InputTokenDetails and OutputTokenDetails.

Because these are objects the current _merge_usage_entries cannot handle them and thus cause the entire request to fail.
This PR fixes that.

This pull request improves the aggregation logic for usage tracking and adds comprehensive tests to ensure correct merging of token usage details, including handling of custom objects and None values.

Improvements to usage aggregation logic

  • Enhanced the _merge_usage_entries method in dspy/utils/usage_tracker.py to correctly merge usage entries when values are custom objects (with a __dict__), ensuring that details from types like InputTokensDetails and OutputTokensDetails are properly aggregated.
  • Updated the method signature for _merge_usage_entries for better readability and consistency.

Expanded test coverage

  • Added new tests in tests/utils/test_usage_tracker.py to verify that input and output token details are accurately aggregated across multiple usage entries, including cases with missing or None values.
  • Included imports for InputTokensDetails and OutputTokensDetails from litellm.types.llms.openai to support new test scenarios.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant