ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix _reset_budget_for_key/team/user

Browse Source
This commit is contained in:
Ishaan Jaff 2025-03-17 18:56:24 -07:00
parent cf86ce3908
commit 91f4d4d865
1 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
litellm/proxy/common_utils/reset_budget_job.py
View File

@ -319,7 +319,7 @@ class ResetBudgetJob:
item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken],
current_time: datetime,
item_type: str,
) -> Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken]:
):
"""
Common logic for resetting budget for a team, user, or key
"""
@ -339,19 +339,19 @@ class ResetBudgetJob:
async def _reset_budget_for_team(
team: LiteLLM_TeamTable, current_time: datetime
) -> Optional[LiteLLM_TeamTable]:
result = await ResetBudgetJob._reset_budget_common(team, current_time, "team")
return result if isinstance(result, LiteLLM_TeamTable) else None
await ResetBudgetJob._reset_budget_common(team, current_time, "team")
return team
@staticmethod
async def _reset_budget_for_user(
user: LiteLLM_UserTable, current_time: datetime
) -> Optional[LiteLLM_UserTable]:
result = await ResetBudgetJob._reset_budget_common(user, current_time, "user")
return result if isinstance(result, LiteLLM_UserTable) else None
await ResetBudgetJob._reset_budget_common(user, current_time, "user")
return user
@staticmethod
async def _reset_budget_for_key(
key: LiteLLM_VerificationToken, current_time: datetime
) -> Optional[LiteLLM_VerificationToken]:
result = await ResetBudgetJob._reset_budget_common(key, current_time, "key")
return result if isinstance(result, LiteLLM_VerificationToken) else None
await ResetBudgetJob._reset_budget_common(key, current_time, "key")
return key