-
Notifications
You must be signed in to change notification settings - Fork 2.1k
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
Conversation
There was a problem hiding this 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 properGenerator
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.
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
@@ -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]: |
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
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
Line 173 in 34090a2
settings.usage_tracker.add_usage(self.model, dict(results.usage)) |
There was a problem hiding this comment.
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.
Thanks so much @TomeHirata ! Feel free to merge whenever concerns are discussed/resolved! |
Add type hints to UsageTracker for internal consistency