From 4b8368a27d99f4732905bdbe56a784f027a61259 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 19 Aug 2026 15:52:24 -0700 Subject: [PATCH] refactor(key management): tighten the key budgets response types entity_type was declared str, so the generated client saw a bare string even though the resolver already carried a Litellm_EntityType and only called .value at the boundary. Declaring the enum keeps the wire format identical and lets a caller match a BudgetExceededError's entity against a row without comparing loose strings. key is populated on every 200, since a request that resolves no key 404s before the response is built, so declaring it optional understated the contract. --- .../key_budget_resolver.py | 2 +- .../key_management_endpoints.py | 29 +++++++++++-------- .../key_management_endpoints.py | 6 ++-- 3 files changed, 22 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_budget_resolver.py b/litellm/proxy/management_endpoints/key_budget_resolver.py index bbe1494ec16..d5e6ba58c0e 100644 --- a/litellm/proxy/management_endpoints/key_budget_resolver.py +++ b/litellm/proxy/management_endpoints/key_budget_resolver.py @@ -390,7 +390,7 @@ def _to_entry(plan: _PlannedBudget, spend: float | None) -> KeyBudgetEntry: ) return KeyBudgetEntry( scope=plan.scope, - entity_type=_ENTITY_TYPE_BY_SCOPE[plan.scope].value, + entity_type=_ENTITY_TYPE_BY_SCOPE[plan.scope], entity_id=plan.entity_id, entity_label=plan.entity_label, enforcement=plan.enforcement, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3913777895b..229a8140389 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3753,6 +3753,15 @@ async def info_key_fn( raise handle_exception_on_proxy(e) +def _key_not_found_error() -> ProxyException: + return ProxyException( + message="Key not found in database", + type=ProxyErrorTypes.not_found_error, + param="key", + code=status.HTTP_404_NOT_FOUND, + ) + + @router.get( "/key/{key_id}/budgets", tags=("key management",), @@ -3798,8 +3807,8 @@ async def key_budgets_fn( - budgets: list - One entry per applicable budget - scope: str - `proxy`, `key`, `key_window`, `key_model`, `team`, `team_window`, `team_member`, `user`, `organization`, `project`, `tag`, `end_user` or `end_user_model` - - entity_type: str - The `Litellm_EntityType` a `BudgetExceededError` from this scope - carries, so a denial message maps back to a row here + - entity_type: Litellm_EntityType - The entity a `BudgetExceededError` from this scope + names, so a denial message maps back to a row here - entity_id / entity_label: str | None - Which entity is limited, and its human-facing alias - enforcement: str - `hard` blocks the request, `soft` only raises an alert - max_budget: float | None - The limit in effect. `null` means this scope applies to the key @@ -3839,17 +3848,13 @@ async def key_budgets_fn( ) key: Final = key_id or user_api_key_dict.api_key - hashed_key: Final = _hash_token_if_needed(token=key) if key is not None else None - key_info: Final = ( - await VerificationTokenRepository(prisma_client).find_by_id(hashed_key) if hashed_key is not None else None - ) + if key is None: + raise _key_not_found_error() + + hashed_key: Final = _hash_token_if_needed(token=key) + key_info: Final = await VerificationTokenRepository(prisma_client).find_by_id(hashed_key) if key_info is None: - raise ProxyException( - message="Key not found in database", - type=ProxyErrorTypes.not_found_error, - param="key", - code=status.HTTP_404_NOT_FOUND, - ) + raise _key_not_found_error() if ( await _can_user_query_key_info( diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 69c38873c1e..23fcee1f26d 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -3,6 +3,8 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, model_validator +from litellm.proxy._types import Litellm_EntityType + class BulkUpdateKeyRequestItem(BaseModel): """Individual key update request item""" @@ -135,7 +137,7 @@ class KeyBudgetEntry(BaseModel): """One budget that can gate requests made with a key, with its live spend.""" scope: BudgetScope - entity_type: str + entity_type: Litellm_EntityType entity_id: str | None = None entity_label: str | None = None enforcement: BudgetEnforcement @@ -154,5 +156,5 @@ class KeyBudgetEntry(BaseModel): class KeyBudgetsResponse(BaseModel): """Every budget that applies to one key, including the ones left unconfigured.""" - key: str | None = None + key: str budgets: tuple[KeyBudgetEntry, ...]