Skip to content

Add type hints for UsageTracker #8636

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Aug 14, 2025

Conversation

TomeHirata
Copy link
Collaborator

Add type hints to UsageTracker for internal consistency

Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR adds type hints to the UsageTracker class to improve internal consistency and type safety. The changes focus on adding parameter and return type annotations to methods and the context manager function.

  • Added type hints to method parameters and return types in the UsageTracker class
  • Updated the track_usage() context manager function with proper Generator type annotation
  • Refactored conditional statements to use walrus operator for improved readability

Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.

@@ -20,16 +20,16 @@ def __init__(self):
# }
self.usage_data = defaultdict(list)

def _flatten_usage_entry(self, usage_entry) -> dict[str, dict[str, Any]]:
def _flatten_usage_entry(self, usage_entry: dict[str, Any]) -> dict[str, Any]:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

usage_entry is not necessarily a dict as i remember, it could be the litellm usage object so I skipped type hints for this method (and the others)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is a dict when add_usage is called. We convert the litellm usage object to a dict before calling it. That's probably why the type signature of add_usage has been dict

settings.usage_tracker.add_usage(self.model, dict(results.usage))

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for causing confusion, here is the more detailed context: internally we have control over what we pass to UsageTracker, but this UsageTracker can also be used in custom usage tracking, and the input can be any dict-like types, e.g., litellm usage or openai usage instance.

Given a second thought, it's not bad to force users to call dict() over their dict-like type though, and keeping consistency of using /not using type hints inside the same API is important. I am going ahead approving this PR.

@okhat
Copy link
Collaborator

okhat commented Aug 13, 2025

Thanks so much @TomeHirata ! Feel free to merge whenever concerns are discussed/resolved!

@TomeHirata TomeHirata merged commit 801680d into stanfordnlp:main Aug 14, 2025
10 checks passed
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.

3 participants