From ff58f8dc9cac38cfe8e31f776ac4f607de297095 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 7 Aug 2026 17:52:55 -0700 Subject: [PATCH] chore: conform project budget code to staging lint gates Staging's Final-annotation pass landed after this branch forked, so the merged project-budget code breached three ratcheted ceilings: - LIT010: annotate the new single-assignment locals Final, and rewrite _team_member_budget_check_for_key so the membership lookup and the cross-pod spend read are conditional expressions rather than rebinds - ANN202: annotate _update_project_db's return type - reportGeneralTypeIssues: staging declares user_api_key_cache Final in update_spend, so the project cache invalidation binds its own name --- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 40 +++++++++++-------- litellm/proxy/db/db_spend_update_writer.py | 10 ++--- litellm/proxy/db/spend_counter_reseed.py | 3 +- .../spend_tracking/budget_reservation.py | 10 ++--- 5 files changed, 36 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3bf3193c783..acebcf40ecc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4192,7 +4192,7 @@ async def _project_max_budget_check( from litellm.proxy.proxy_server import get_current_spend - project_spend = await get_current_spend( + project_spend: Final = await get_current_spend( counter_key=f"spend:project:{project_object.project_id}", fallback_spend=project_object.spend or 0.0, max_budget=max_budget, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ca1d1abce31..938e557ae22 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1054,44 +1054,50 @@ async def _team_member_budget_check_for_key( if prisma_client is None: return - _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" + _cache_key: Final = f"{valid_token.team_id}_{valid_token.user_id}" - team_member_info = await user_api_key_cache.async_get_cache( + cached_member: Final = await user_api_key_cache.async_get_cache( key=_cache_key, model_type=LiteLLM_TeamMembership, ) - if team_member_info is None and valid_token.user_id is not None and valid_token.team_id is not None: - _db_member = await TeamMembershipRepository(prisma_client).table.find_first( + _db_member: Final = ( + await TeamMembershipRepository(prisma_client).table.find_first( where={ # mutable-ok: prisma query argument shape "user_id": valid_token.user_id, "team_id": valid_token.team_id, }, include={"litellm_budget_table": True}, # mutable-ok: prisma query argument shape ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) + if cached_member is None and valid_token.user_id is not None and valid_token.team_id is not None + else None + ) + fresh_member: Final = LiteLLM_TeamMembership(**_db_member.dict()) if _db_member is not None else None + if fresh_member is not None: + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=fresh_member, + model_type=LiteLLM_TeamMembership, + ttl=5, + ) + team_member_info: Final = cached_member if cached_member is not None else fresh_member if team_member_info is None or team_member_info.litellm_budget_table is None: return - team_member_budget = team_member_info.litellm_budget_table.max_budget + team_member_budget: Final = team_member_info.litellm_budget_table.max_budget if team_member_budget is None or team_member_budget <= 0: return from litellm.proxy.proxy_server import get_current_spend - team_member_spend = valid_token.team_member_spend - if valid_token.user_id is not None and valid_token.team_id is not None: - team_member_spend = await get_current_spend( + team_member_spend: Final = ( + await get_current_spend( counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", - fallback_spend=team_member_spend, + fallback_spend=valid_token.team_member_spend, max_budget=team_member_budget, ) + if valid_token.user_id is not None and valid_token.team_id is not None + else valid_token.team_member_spend + ) if team_member_spend > team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d6349b89725..ca502f2e24b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -720,7 +720,7 @@ class DBSpendUpdateWriter: response_cost: float | None, project_id: str | None, prisma_client: PrismaClient | None, - ): + ) -> None: try: if project_id is None or prisma_client is None: return @@ -1390,7 +1390,7 @@ class DBSpendUpdateWriter: ) ### UPDATE PROJECT TABLE ### - project_list_transactions = db_spend_update_transactions.get("project_list_transactions") + project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Project", transactions=project_list_transactions, @@ -1401,10 +1401,10 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) if project_list_transactions and proxy_logging_obj is not None: - user_api_key_cache = proxy_logging_obj.call_details.get("user_api_key_cache") - if user_api_key_cache is not None: + project_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache") + if project_api_key_cache is not None: for project_id in project_list_transactions: - await user_api_key_cache.async_delete_cache(key=f"project_id:{project_id}") + await project_api_key_cache.async_delete_cache(key=f"project_id:{project_id}") ### UPDATE TAG TABLE ### tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"] diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 577933499db..783dea9dd17 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -114,7 +114,8 @@ class SpendCounterReseed: org_id: Final = counter_key[len("spend:org:") :] row = await OrganizationRepository(prisma_client).table.find_unique(where={"organization_id": org_id}) elif counter_key.startswith("spend:project:"): - project_id = counter_key[len("spend:project:") :] + project_id: Final = counter_key[len("spend:project:") :] + # rebind-ok: one repository lookup per mutually exclusive counter-key branch row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id}) else: return None diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 27ab59b6d30..39e3c9273af 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -437,7 +437,7 @@ async def _get_budget_counters( if org_counter is not None: counters.append(org_counter) - project_counter = await _get_project_budget_counter( + project_counter: Final = await _get_project_budget_counter( valid_token=valid_token, user_api_key_cache=user_api_key_cache, ) @@ -583,13 +583,13 @@ async def _get_project_budget_counter( if valid_token.project_id is None: return None - source_cache_key = f"project_id:{valid_token.project_id}" - project_object = await user_api_key_cache.async_get_cache(key=source_cache_key) + source_cache_key: Final = f"project_id:{valid_token.project_id}" + project_object: Final = await user_api_key_cache.async_get_cache(key=source_cache_key) if project_object is None: return None - budget_table = _get_value(project_object, "litellm_budget_table") - max_budget = _to_float(_get_value(budget_table, "max_budget")) + budget_table: Final = _get_value(project_object, "litellm_budget_table") + max_budget: Final = _to_float(_get_value(budget_table, "max_budget")) if max_budget is None or max_budget <= 0: return None