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.

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,

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(

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, ...]