mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(proxy): add key-level budget_fallbacks to reroute requests when a per-model budget is exceeded (#31783)
This commit is contained in:
parent
70785b9b6a
commit
28ddad271e
17 changed files with 705 additions and 9 deletions
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "budget_fallbacks" JSONB NOT NULL DEFAULT '{}';
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "budget_fallbacks" JSONB NOT NULL DEFAULT '{}';
|
||||
|
|
@ -419,6 +419,7 @@ model LiteLLM_VerificationToken {
|
|||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
budget_id String?
|
||||
organization_id String?
|
||||
object_permission_id String?
|
||||
|
|
@ -512,6 +513,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
router_settings Json? @default("{}")
|
||||
budget_id String?
|
||||
organization_id String?
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
|
|||
permissions: Dict = {}
|
||||
model_spend: Dict = {}
|
||||
model_max_budget: Dict = {}
|
||||
budget_fallbacks: dict[str, list[str]] = {}
|
||||
soft_budget_cooldown: bool = False
|
||||
blocked: Optional[bool] = None
|
||||
litellm_budget_table: Optional[dict] = None
|
||||
|
|
|
|||
|
|
@ -1039,6 +1039,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
|
|||
config: Optional[dict] = {}
|
||||
permissions: Optional[dict] = {}
|
||||
model_max_budget: Optional[dict] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
|
||||
budget_fallbacks: Optional[dict[str, list[str]]] = None
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
model_rpm_limit: Optional[dict] = None
|
||||
|
|
@ -1136,6 +1137,7 @@ class GenerateKeyResponse(KeyRequestBase):
|
|||
"config",
|
||||
"permissions",
|
||||
"model_max_budget",
|
||||
"budget_fallbacks",
|
||||
"router_settings",
|
||||
"budget_limits",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -11,8 +11,10 @@ import asyncio
|
|||
import fnmatch
|
||||
import re
|
||||
import secrets
|
||||
|
||||
import orjson
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, Iterator, NamedTuple, List, Optional, Tuple, Union, cast
|
||||
from typing import Any, Dict, Iterator, NamedTuple, List, Optional, Protocol, Tuple, Union, cast
|
||||
|
||||
import fastapi
|
||||
from fastapi import HTTPException, Request, WebSocket, status
|
||||
|
|
@ -30,6 +32,7 @@ from litellm.proxy._types import *
|
|||
from litellm.proxy.auth.auth_checks import (
|
||||
ExperimentalUIJWTToken,
|
||||
_cache_key_object,
|
||||
_can_object_call_model,
|
||||
_check_end_user_budget,
|
||||
_delete_cache_key_object,
|
||||
_get_user_role,
|
||||
|
|
@ -73,6 +76,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
|
|||
_read_request_body,
|
||||
_safe_get_request_headers,
|
||||
_safe_get_request_query_params,
|
||||
_safe_set_request_parsed_body,
|
||||
populate_request_with_path_params,
|
||||
)
|
||||
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
|
||||
|
|
@ -171,6 +175,87 @@ def _get_model_names_for_budget_checks(
|
|||
return model
|
||||
|
||||
|
||||
class _KeyModelBudgetLimiter(Protocol):
|
||||
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
|
||||
|
||||
async def get_fallback_model_within_budget(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, model: str
|
||||
) -> Optional[str]: ...
|
||||
|
||||
|
||||
async def _check_key_model_budget_with_fallback(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
model_max_budget_limiter: _KeyModelBudgetLimiter,
|
||||
model_name: str,
|
||||
request_data: dict,
|
||||
request: Request,
|
||||
llm_model_list: Optional[list] = None,
|
||||
llm_router: Optional[litellm.Router] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Enforce the key's per-model budget for `model_name`. If exceeded and the
|
||||
key has a `budget_fallbacks` chain configured for `model_name`, reroute
|
||||
the request to the first fallback model still within its own budget
|
||||
instead of rejecting the request.
|
||||
|
||||
The selected fallback is validated against the key's model-access
|
||||
allowlist and the team's model restrictions so that budget_fallbacks
|
||||
cannot bypass model authorization. The rewrite is persisted to the
|
||||
parsed-body cache, Starlette's JSON cache (``request._json``), and
|
||||
path parameters so that downstream handlers see the final model
|
||||
regardless of whether they consume ``_read_request_body()``,
|
||||
``request.json()``, or the path ``model`` parameter.
|
||||
|
||||
Fallback is only attempted when ``model_name`` matches the top-level
|
||||
``request_data["model"]``; models extracted from nested fields
|
||||
(``session.model``, ``completion.model``, etc.) are not rewritable
|
||||
and raise immediately.
|
||||
|
||||
Raises:
|
||||
BudgetExceededError: if `model_name` is over budget and no configured
|
||||
fallback is within budget either (or the fallback is not authorized).
|
||||
"""
|
||||
try:
|
||||
await model_max_budget_limiter.is_key_within_model_budget(
|
||||
user_api_key_dict=valid_token,
|
||||
model=model_name,
|
||||
)
|
||||
except litellm.BudgetExceededError as e:
|
||||
if request_data.get("model") != model_name:
|
||||
raise e
|
||||
fallback_model = await model_max_budget_limiter.get_fallback_model_within_budget(
|
||||
user_api_key_dict=valid_token,
|
||||
model=model_name,
|
||||
)
|
||||
if fallback_model is None:
|
||||
raise e
|
||||
try:
|
||||
await can_key_call_model(
|
||||
model=fallback_model,
|
||||
llm_model_list=llm_model_list,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
if valid_token.team_models:
|
||||
_can_object_call_model(
|
||||
model=fallback_model,
|
||||
llm_router=llm_router,
|
||||
models=valid_token.team_models,
|
||||
team_model_aliases=valid_token.team_model_aliases,
|
||||
team_id=valid_token.team_id,
|
||||
object_type="team",
|
||||
)
|
||||
except ProxyException:
|
||||
raise e
|
||||
request_data["model"] = fallback_model
|
||||
_safe_set_request_parsed_body(request=request, parsed_body=request_data)
|
||||
request._json = request_data # type: ignore[attr-defined]
|
||||
request._body = orjson.dumps(request_data) # type: ignore[attr-defined]
|
||||
path_params = request.scope.get("path_params")
|
||||
if isinstance(path_params, dict) and "model" in path_params:
|
||||
path_params["model"] = fallback_model
|
||||
|
||||
|
||||
def _get_bearer_token_or_received_api_key(api_key: str) -> str:
|
||||
if api_key.startswith("Bearer "): # ensure Bearer token passed in
|
||||
api_key = api_key.replace("Bearer ", "") # extract the token
|
||||
|
|
@ -1779,11 +1864,26 @@ async def _user_api_key_auth_builder(
|
|||
):
|
||||
## GET THE SPEND FOR THIS MODEL
|
||||
for model_name in current_models:
|
||||
await model_max_budget_limiter.is_key_within_model_budget(
|
||||
user_api_key_dict=valid_token,
|
||||
model=model_name,
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
model_name=model_name,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
# Recompute after a potential budget-fallback rewrite so
|
||||
# the end-user check below validates the final model
|
||||
current_model = _get_model_from_request_context(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
current_models = _get_model_names_for_budget_checks(model=current_model)
|
||||
|
||||
# Check 5b. End-user model max budget
|
||||
end_user_mmb = valid_token.end_user_model_max_budget
|
||||
if (
|
||||
|
|
@ -2838,11 +2938,26 @@ async def _run_post_custom_auth_checks(
|
|||
and valid_token.token is not None
|
||||
):
|
||||
for model_name in current_models:
|
||||
await model_max_budget_limiter.is_key_within_model_budget(
|
||||
user_api_key_dict=valid_token,
|
||||
model=model_name,
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
model_name=model_name,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
# Recompute after a potential budget-fallback rewrite so
|
||||
# the end-user check below validates the final model
|
||||
current_model = _get_model_from_request_context(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
current_models = _get_model_names_for_budget_checks(model=current_model)
|
||||
|
||||
# 4. Check end-user model_max_budget
|
||||
end_user_mmb = valid_token.end_user_model_max_budget
|
||||
if (
|
||||
|
|
|
|||
|
|
@ -80,6 +80,20 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
|
||||
return True
|
||||
|
||||
async def get_fallback_model_within_budget(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
model: str,
|
||||
) -> Optional[str]:
|
||||
budget_fallbacks: dict[str, list[str]] = user_api_key_dict.budget_fallbacks or {}
|
||||
for fallback_model in budget_fallbacks.get(model, []):
|
||||
try:
|
||||
await self.is_key_within_model_budget(user_api_key_dict=user_api_key_dict, model=fallback_model)
|
||||
return fallback_model
|
||||
except litellm.BudgetExceededError:
|
||||
continue
|
||||
return None
|
||||
|
||||
async def is_end_user_within_model_budget(
|
||||
self,
|
||||
end_user_id: str,
|
||||
|
|
|
|||
|
|
@ -374,6 +374,7 @@ async def new_user(
|
|||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
- model_tpm_limit: Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
|
|
@ -1375,6 +1376,7 @@ async def user_update(
|
|||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
|
||||
- model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
- mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
- model_tpm_limit: Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
|
|
|
|||
|
|
@ -1244,7 +1244,7 @@ async def _check_team_key_limits(
|
|||
)
|
||||
# Exclude the key being updated to avoid double-counting its limits.
|
||||
# data.key may be a raw key (sk-...) or a pre-hashed token_id.
|
||||
if isinstance(data, UpdateKeyRequest):
|
||||
if isinstance(data, UpdateKeyRequest) and data.key is not None:
|
||||
hashed_key = _hash_token_if_needed(data.key)
|
||||
keys = [key for key in keys if key.token != hashed_key]
|
||||
check_team_key_model_specific_limits(
|
||||
|
|
@ -1426,7 +1426,7 @@ async def _check_org_key_limits(
|
|||
)
|
||||
# Exclude the key being updated to avoid double-counting its limits.
|
||||
# data.key may be a raw key (sk-...) or a pre-hashed token_id.
|
||||
if isinstance(data, UpdateKeyRequest):
|
||||
if isinstance(data, UpdateKeyRequest) and data.key is not None:
|
||||
hashed_key = _hash_token_if_needed(data.key)
|
||||
keys = [key for key in keys if key.token != hashed_key]
|
||||
check_org_key_model_specific_limits(
|
||||
|
|
@ -1485,6 +1485,7 @@ async def generate_key_fn(
|
|||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
- mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit.
|
||||
|
|
@ -1691,6 +1692,7 @@ async def generate_service_account_key_fn(
|
|||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
- mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit.
|
||||
|
|
@ -2486,6 +2488,7 @@ async def update_key_fn(
|
|||
- spend: Optional[float] - Amount spent by key
|
||||
- max_budget: Optional[float] - Max budget for key
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
- soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
- max_parallel_requests: Optional[int] - Rate limit for parallel requests
|
||||
|
|
@ -3526,6 +3529,7 @@ async def generate_key_helper_fn(
|
|||
allowed_cache_controls: Optional[list] = [],
|
||||
permissions: Optional[dict] = {},
|
||||
model_max_budget: Optional[dict] = {},
|
||||
budget_fallbacks: Optional[dict] = None,
|
||||
model_rpm_limit: Optional[dict] = None,
|
||||
model_tpm_limit: Optional[dict] = None,
|
||||
mcp_rpm_limit: Optional[dict] = None,
|
||||
|
|
@ -3616,6 +3620,7 @@ async def generate_key_helper_fn(
|
|||
metadata_json = json.dumps(metadata)
|
||||
validate_model_max_budget(model_max_budget)
|
||||
model_max_budget_json = json.dumps(model_max_budget)
|
||||
budget_fallbacks_json = json.dumps(budget_fallbacks or {})
|
||||
user_role = user_role
|
||||
tpm_limit = tpm_limit
|
||||
rpm_limit = rpm_limit
|
||||
|
|
@ -3668,6 +3673,7 @@ async def generate_key_helper_fn(
|
|||
"allowed_cache_controls": allowed_cache_controls,
|
||||
"permissions": permissions_json,
|
||||
"model_max_budget": model_max_budget_json,
|
||||
"budget_fallbacks": budget_fallbacks_json,
|
||||
"organization_id": organization_id,
|
||||
"budget_id": budget_id,
|
||||
"blocked": blocked,
|
||||
|
|
@ -4018,6 +4024,7 @@ def _transform_verification_tokens_to_deleted_records(
|
|||
"metadata",
|
||||
"model_spend",
|
||||
"model_max_budget",
|
||||
"budget_fallbacks",
|
||||
"router_settings",
|
||||
]:
|
||||
if json_field in record and record[json_field] is not None:
|
||||
|
|
@ -4519,6 +4526,7 @@ async def regenerate_key_fn(
|
|||
- spend: Optional[float] - Amount spent by key
|
||||
- max_budget: Optional[float] - Max budget for key
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
- max_parallel_requests: Optional[int] - Rate limit for parallel requests
|
||||
|
|
|
|||
|
|
@ -419,6 +419,7 @@ model LiteLLM_VerificationToken {
|
|||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
budget_id String?
|
||||
organization_id String?
|
||||
object_permission_id String?
|
||||
|
|
@ -512,6 +513,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
router_settings Json? @default("{}")
|
||||
budget_id String?
|
||||
organization_id String?
|
||||
|
|
|
|||
|
|
@ -419,6 +419,7 @@ model LiteLLM_VerificationToken {
|
|||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
budget_id String?
|
||||
organization_id String?
|
||||
object_permission_id String?
|
||||
|
|
@ -512,6 +513,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
access_group_ids String[] @default([])
|
||||
model_spend Json @default("{}")
|
||||
model_max_budget Json @default("{}")
|
||||
budget_fallbacks Json @default("{}")
|
||||
router_settings Json? @default("{}")
|
||||
budget_id String?
|
||||
organization_id String?
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from pydantic import AliasPath, BaseModel, Field, RootModel
|
|||
from e2e_gateway import Gateway, build_gateway
|
||||
from e2e_http import NoBody, StreamingResponse, Success, unwrap
|
||||
from models import (
|
||||
AnthropicMessagesBody,
|
||||
BudgetWindow,
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
|
|
@ -171,6 +172,7 @@ class BudgetClient:
|
|||
user_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
model_max_budget: dict[str, ModelBudgetEntry] | None = None,
|
||||
budget_fallbacks: dict[str, list[str]] | None = None,
|
||||
budget_limits: list[BudgetWindow] | None = None,
|
||||
) -> str:
|
||||
return self.gateway.generate_key(
|
||||
|
|
@ -183,6 +185,7 @@ class BudgetClient:
|
|||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
model_max_budget=model_max_budget,
|
||||
budget_fallbacks=budget_fallbacks,
|
||||
budget_limits=budget_limits,
|
||||
)
|
||||
)
|
||||
|
|
@ -217,6 +220,24 @@ class BudgetClient:
|
|||
),
|
||||
)
|
||||
|
||||
def messages(
|
||||
self,
|
||||
key: str,
|
||||
model: str,
|
||||
content: str,
|
||||
*,
|
||||
max_tokens: int = 16,
|
||||
) -> StreamingResponse:
|
||||
return self.gateway.transport.send(
|
||||
"/v1/messages",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
max_tokens=max_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
# ---- internal user --------------------------------------------------
|
||||
|
||||
def create_user(self, *, max_budget: float) -> str:
|
||||
|
|
|
|||
60
tests/e2e/budgets/test_budget_fallback_e2e.py
Normal file
60
tests/e2e/budgets/test_budget_fallback_e2e.py
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
"""Live e2e: a virtual key's per-model `budget_fallbacks` reroutes `/v1/messages`
|
||||
transparently from an exhausted Anthropic model to an OpenAI model, instead of
|
||||
blocking the caller with a `budget_exceeded` error. Coverage for the
|
||||
budget_fallbacks feature in litellm/proxy/hooks/model_max_budget_limiter.py.
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from budget_client import BudgetClient, model_budget
|
||||
from e2e_config import unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
PRIMARY_MODEL = "claude-haiku-4-5"
|
||||
FALLBACK_MODEL = "gpt-5.5"
|
||||
|
||||
|
||||
def test_budget_fallback_reroutes_anthropic_messages_to_openai(
|
||||
client: BudgetClient, resources: ResourceManager
|
||||
) -> None:
|
||||
key = client.generate_key(
|
||||
model_max_budget=model_budget(PRIMARY_MODEL, 1e-6),
|
||||
budget_fallbacks={PRIMARY_MODEL: [FALLBACK_MODEL]},
|
||||
)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
# Exhaust the primary model's near-zero budget. Once exceeded, every
|
||||
# subsequent /v1/messages call for this key must reroute to the fallback
|
||||
# instead of surfacing a budget_exceeded block.
|
||||
served_by = None
|
||||
deadline = time.monotonic() + 60
|
||||
while time.monotonic() < deadline:
|
||||
result = client.messages(
|
||||
key, PRIMARY_MODEL, f"hi {unique_marker()}", max_tokens=16
|
||||
)
|
||||
if not result.ok:
|
||||
pytest.fail(
|
||||
"budget_fallbacks must reroute transparently, never surface a "
|
||||
f"block; status={result.status_code} body={result.body[:300]}"
|
||||
)
|
||||
served_by = json.loads(result.body)["model"]
|
||||
if FALLBACK_MODEL in served_by:
|
||||
break
|
||||
time.sleep(1)
|
||||
assert served_by is not None and FALLBACK_MODEL in served_by, (
|
||||
f"{PRIMARY_MODEL}'s budget_fallbacks never rerouted to {FALLBACK_MODEL}"
|
||||
)
|
||||
|
||||
# The rerouted call must be recorded under the fallback model, not the
|
||||
# exhausted primary - proving spend tracking followed the reroute.
|
||||
rows = client.gateway.poll_logs_for_key(
|
||||
key, predicate=lambda rows: any(FALLBACK_MODEL in (r.model or "") for r in rows)
|
||||
)
|
||||
assert any(FALLBACK_MODEL in (r.model or "") for r in rows), (
|
||||
f"no spend log recorded against {FALLBACK_MODEL} after the reroute"
|
||||
)
|
||||
|
|
@ -31,6 +31,7 @@ class KeyGenerateBody(BaseModel):
|
|||
team_id: str | None = None
|
||||
budget_id: str | None = None
|
||||
model_max_budget: dict[str, ModelBudgetEntry] | None = None
|
||||
budget_fallbacks: dict[str, list[str]] | None = None
|
||||
budget_limits: list[BudgetWindow] | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
|
|
@ -95,6 +96,12 @@ class ChatBody(BaseModel):
|
|||
metadata: ChatMetadata | None = None
|
||||
|
||||
|
||||
class AnthropicMessagesBody(BaseModel):
|
||||
model: str
|
||||
messages: list[ChatMessage]
|
||||
max_tokens: int
|
||||
|
||||
|
||||
class OutMessage(BaseModel):
|
||||
content: str | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -452,6 +452,83 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config
|
|||
mock_push.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fallback_model_within_budget_returns_none_without_fallbacks(
|
||||
budget_limiter,
|
||||
):
|
||||
user_api_key = UserAPIKeyAuth(token="test-key", budget_fallbacks={})
|
||||
assert (
|
||||
await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4")
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fallback_model_within_budget_returns_first_within_budget(
|
||||
budget_limiter,
|
||||
):
|
||||
user_api_key = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
model_max_budget={"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}},
|
||||
budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]},
|
||||
)
|
||||
with patch.object(
|
||||
budget_limiter, "_get_virtual_key_spend_for_model", return_value=1.0
|
||||
):
|
||||
result = await budget_limiter.get_fallback_model_within_budget(
|
||||
user_api_key, "gpt-4"
|
||||
)
|
||||
assert result == "gpt-4o-mini"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fallback_model_within_budget_skips_exhausted_fallback(
|
||||
budget_limiter,
|
||||
):
|
||||
user_api_key = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
model_max_budget={
|
||||
"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"},
|
||||
"claude-haiku": {"budget_limit": 100.0, "time_period": "1d"},
|
||||
},
|
||||
budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]},
|
||||
)
|
||||
|
||||
async def _spend_for_model(user_api_key_hash, model, key_budget_config):
|
||||
return 150.0 if model == "gpt-4o-mini" else 1.0
|
||||
|
||||
with patch.object(
|
||||
budget_limiter,
|
||||
"_get_virtual_key_spend_for_model",
|
||||
side_effect=_spend_for_model,
|
||||
):
|
||||
result = await budget_limiter.get_fallback_model_within_budget(
|
||||
user_api_key, "gpt-4"
|
||||
)
|
||||
assert result == "claude-haiku"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_fallback_model_within_budget_returns_none_when_chain_exhausted(
|
||||
budget_limiter,
|
||||
):
|
||||
user_api_key = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
model_max_budget={
|
||||
"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"},
|
||||
"claude-haiku": {"budget_limit": 100.0, "time_period": "1d"},
|
||||
},
|
||||
budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]},
|
||||
)
|
||||
with patch.object(
|
||||
budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0
|
||||
):
|
||||
result = await budget_limiter.get_fallback_model_within_budget(
|
||||
user_api_key, "gpt-4"
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_log_success_event_skips_redis_push_without_redis(budget_limiter):
|
||||
"""When dual_cache has no Redis backend, do not await _push_in_memory_increments_to_redis."""
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.proxy.auth.handle_jwt import JWTHandler
|
|||
from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_check_key_model_budget_with_fallback,
|
||||
_PendingAutoRegister,
|
||||
_matches_routing_override,
|
||||
_reserve_budget_after_common_checks,
|
||||
|
|
@ -4084,3 +4085,265 @@ async def test_auth_path_caches_team_object_under_canonical_team_id_key():
|
|||
assert served is not None and served.team_id == team_id
|
||||
assert cache.get_cache(key=team_id) is None
|
||||
assert cache.get_cache(key=None) is None
|
||||
|
||||
|
||||
class TestCheckKeyModelBudgetWithFallback:
|
||||
"""`_check_key_model_budget_with_fallback` must reroute a request to the
|
||||
first configured `budget_fallbacks` entry still within its own budget,
|
||||
and only raise `BudgetExceededError` when no fallback is available."""
|
||||
|
||||
def _make_request(self):
|
||||
request = MagicMock()
|
||||
request.scope = {}
|
||||
return request
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_within_budget_does_not_reroute(self):
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
limiter.is_key_within_model_budget.return_value = True
|
||||
request_data = {"model": "gpt-4o"}
|
||||
request = self._make_request()
|
||||
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert request_data["model"] == "gpt-4o"
|
||||
limiter.get_fallback_model_within_budget.assert_not_awaited()
|
||||
assert "parsed_body" not in request.scope
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exceeded_budget_reroutes_to_fallback(self):
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
budget_fallbacks={"gpt-4o": ["gpt-4o-mini", "claude-haiku"]},
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(
|
||||
current_cost=10, max_budget=5
|
||||
)
|
||||
limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini"
|
||||
request_data = {"model": "gpt-4o"}
|
||||
request = self._make_request()
|
||||
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert request_data["model"] == "gpt-4o-mini"
|
||||
limiter.get_fallback_model_within_budget.assert_awaited_once_with(
|
||||
user_api_key_dict=valid_token, model="gpt-4o"
|
||||
)
|
||||
# the rerouted model must be visible to a later, separate
|
||||
# `_read_request_body` call on the same `request` (route handlers
|
||||
# re-parse the body from this cache instead of reusing the dict).
|
||||
cached_keys, cached_body = request.scope["parsed_body"]
|
||||
assert cached_body["model"] == "gpt-4o-mini"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_every_fallback_also_exceeded(self):
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5)
|
||||
limiter.is_key_within_model_budget.side_effect = original_error
|
||||
limiter.get_fallback_model_within_budget.return_value = None
|
||||
request_data = {"model": "gpt-4o"}
|
||||
request = self._make_request()
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert exc_info.value is original_error
|
||||
assert request_data["model"] == "gpt-4o"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_fallback_not_authorized(self):
|
||||
"""If the fallback model is within budget but the key is not allowed
|
||||
to call it, the original BudgetExceededError must be raised instead
|
||||
of rerouting to an unauthorized model."""
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
models=["gpt-4o"],
|
||||
budget_fallbacks={"gpt-4o": ["restricted-model"]},
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5)
|
||||
limiter.is_key_within_model_budget.side_effect = original_error
|
||||
limiter.get_fallback_model_within_budget.return_value = "restricted-model"
|
||||
request_data = {"model": "gpt-4o"}
|
||||
request = self._make_request()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.can_key_call_model",
|
||||
side_effect=ProxyException(
|
||||
message="model not allowed",
|
||||
type=ProxyErrorTypes.budget_exceeded,
|
||||
param="model",
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
),
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
assert exc_info.value is original_error
|
||||
assert request_data["model"] == "gpt-4o"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reroute_succeeds_when_fallback_is_authorized(self):
|
||||
"""If the fallback model is within budget AND authorized, the request
|
||||
must be rerouted to it."""
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
models=["gpt-4o", "gpt-4o-mini"],
|
||||
budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]},
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(
|
||||
current_cost=10, max_budget=5
|
||||
)
|
||||
limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini"
|
||||
request_data = {"model": "gpt-4o"}
|
||||
request = self._make_request()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.can_key_call_model",
|
||||
return_value=True,
|
||||
):
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
assert request_data["model"] == "gpt-4o-mini"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_fallback_blocked_by_team_models(self):
|
||||
"""If the key allows the fallback but the team does not, the original
|
||||
BudgetExceededError must be raised."""
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
models=["gpt-4o", "restricted-model"],
|
||||
team_id="team-1",
|
||||
team_models=["gpt-4o"],
|
||||
budget_fallbacks={"gpt-4o": ["restricted-model"]},
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5)
|
||||
limiter.is_key_within_model_budget.side_effect = original_error
|
||||
limiter.get_fallback_model_within_budget.return_value = "restricted-model"
|
||||
request_data = {"model": "gpt-4o"}
|
||||
request = self._make_request()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.can_key_call_model",
|
||||
return_value=True,
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
assert exc_info.value is original_error
|
||||
assert request_data["model"] == "gpt-4o"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reroute_updates_path_params_model(self):
|
||||
"""On path-model routes (/openai/deployments/{model}/...) the fallback
|
||||
must also update path_params so downstream logic does not revert to the
|
||||
original path model."""
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
models=["gpt-4o", "gpt-4o-mini"],
|
||||
budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]},
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(
|
||||
current_cost=10, max_budget=5
|
||||
)
|
||||
limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini"
|
||||
request_data = {"model": "gpt-4o"}
|
||||
request = self._make_request()
|
||||
request.scope["path_params"] = {"model": "gpt-4o"}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.can_key_call_model",
|
||||
return_value=True,
|
||||
):
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
assert request_data["model"] == "gpt-4o-mini"
|
||||
assert request.scope["path_params"]["model"] == "gpt-4o-mini"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_model_from_nested_field(self):
|
||||
"""Budget fallback must not attempt rewrite when the checked model
|
||||
came from a nested field (e.g. session.model) rather than the
|
||||
top-level request_data['model']."""
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-key",
|
||||
models=["gpt-4o", "gpt-4o-mini"],
|
||||
budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]},
|
||||
)
|
||||
limiter = AsyncMock()
|
||||
original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5)
|
||||
limiter.is_key_within_model_budget.side_effect = original_error
|
||||
request_data = {"session": {"model": "gpt-4o"}}
|
||||
request = self._make_request()
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_key_model_budget_with_fallback(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=limiter,
|
||||
model_name="gpt-4o",
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert exc_info.value is original_error
|
||||
assert "model" not in request_data
|
||||
|
|
|
|||
|
|
@ -762,6 +762,58 @@ async def test_generate_key_helper_fn_with_access_group_ids(monkeypatch):
|
|||
assert captured_key_data.get("access_group_ids") == ["ag-1", "ag-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_key_helper_fn_with_budget_fallbacks(monkeypatch):
|
||||
"""Regression: /key/generate must accept `budget_fallbacks` end-to-end.
|
||||
|
||||
generate_key_helper_fn previously had no `budget_fallbacks` parameter, so
|
||||
passing it via /key/generate (which unpacks the full request body as
|
||||
kwargs) raised "unexpected keyword argument" before ever reaching the DB.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.jsonify_object = lambda data: data # type: ignore
|
||||
mock_prisma_client.db = MagicMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable = MagicMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock(
|
||||
return_value=MagicMock(object_permission_id=None)
|
||||
)
|
||||
|
||||
captured_key_data = {}
|
||||
|
||||
async def _insert_data_side_effect(*args, **kwargs):
|
||||
table_name = kwargs.get("table_name")
|
||||
if table_name == "user":
|
||||
return MagicMock(models=[], spend=0)
|
||||
elif table_name == "key":
|
||||
captured_key_data.update(kwargs.get("data", {}))
|
||||
return MagicMock(
|
||||
token="hashed_token_budget_fallbacks",
|
||||
litellm_budget_table=None,
|
||||
object_permission=None,
|
||||
created_at=None,
|
||||
updated_at=None,
|
||||
)
|
||||
return MagicMock()
|
||||
|
||||
mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
generate_key_helper_fn,
|
||||
)
|
||||
|
||||
await generate_key_helper_fn(
|
||||
request_type="key",
|
||||
table_name="key",
|
||||
user_id="test-user",
|
||||
budget_fallbacks={"anthropic-haiku-4-5": ["gpt-5.5"]},
|
||||
)
|
||||
|
||||
assert json.loads(captured_key_data["budget_fallbacks"]) == {
|
||||
"anthropic-haiku-4-5": ["gpt-5.5"]
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_generation_with_mcp_tool_permissions(monkeypatch):
|
||||
"""
|
||||
|
|
@ -4222,6 +4274,7 @@ def test_transform_verification_tokens_to_deleted_records():
|
|||
permissions={"permission": True},
|
||||
metadata={},
|
||||
model_max_budget={"gpt-4": {"budget_limit": 100.0}},
|
||||
budget_fallbacks={"gpt-4": ["gpt-4o-mini"]},
|
||||
model_spend={},
|
||||
soft_budget_cooldown=False,
|
||||
allowed_routes=[],
|
||||
|
|
@ -4260,6 +4313,8 @@ def test_transform_verification_tokens_to_deleted_records():
|
|||
record2 = records[1]
|
||||
assert record2["token"] == "hashed-token-2"
|
||||
assert isinstance(record2["model_max_budget"], str)
|
||||
assert isinstance(record2["budget_fallbacks"], str)
|
||||
assert json.loads(record2["budget_fallbacks"]) == {"gpt-4": ["gpt-4o-mini"]}
|
||||
|
||||
|
||||
def test_transform_verification_tokens_to_deleted_records_empty_list():
|
||||
|
|
|
|||
60
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
60
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -6504,6 +6504,7 @@ export interface paths {
|
|||
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||||
* - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
* - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
* - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit.
|
||||
|
|
@ -6710,6 +6711,7 @@ export interface paths {
|
|||
* - spend: Optional[float] - Amount spent by key
|
||||
* - max_budget: Optional[float] - Max budget for key
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
* - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
* - max_parallel_requests: Optional[int] - Rate limit for parallel requests
|
||||
|
|
@ -6785,6 +6787,7 @@ export interface paths {
|
|||
* - guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
* - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||||
* - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||||
* - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit.
|
||||
|
|
@ -6883,6 +6886,7 @@ export interface paths {
|
|||
* - spend: Optional[float] - Amount spent by key
|
||||
* - max_budget: Optional[float] - Max budget for key
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
* - soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
* - max_parallel_requests: Optional[int] - Rate limit for parallel requests
|
||||
|
|
@ -6964,6 +6968,7 @@ export interface paths {
|
|||
* - spend: Optional[float] - Amount spent by key
|
||||
* - max_budget: Optional[float] - Max budget for key
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||||
* - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||||
* - max_parallel_requests: Optional[int] - Rate limit for parallel requests
|
||||
|
|
@ -14613,6 +14618,7 @@ export interface paths {
|
|||
* - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
* - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
|
||||
* - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
* - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
* - model_tpm_limit: Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
|
|
@ -14693,6 +14699,7 @@ export interface paths {
|
|||
* - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
* - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
|
||||
* - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
* - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
* - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user.
|
||||
* - model_tpm_limit: Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
|
||||
|
|
@ -23579,6 +23586,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
@ -23719,6 +23730,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
@ -24702,6 +24717,13 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/**
|
||||
* Budget Fallbacks
|
||||
* @default {}
|
||||
*/
|
||||
budget_fallbacks: {
|
||||
[key: string]: string[];
|
||||
};
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
@ -26087,6 +26109,13 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/**
|
||||
* Budget Fallbacks
|
||||
* @default {}
|
||||
*/
|
||||
budget_fallbacks: {
|
||||
[key: string]: string[];
|
||||
};
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
@ -27895,6 +27924,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Limits */
|
||||
budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null;
|
||||
/**
|
||||
|
|
@ -28029,6 +28062,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
@ -29647,6 +29684,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
@ -31606,6 +31647,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
@ -32069,6 +32114,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Limits */
|
||||
budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null;
|
||||
/**
|
||||
|
|
@ -32167,6 +32216,10 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Fallbacks */
|
||||
budget_fallbacks?: {
|
||||
[key: string]: string[];
|
||||
} | null;
|
||||
/** Budget Limits */
|
||||
budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null;
|
||||
/**
|
||||
|
|
@ -32394,6 +32447,13 @@ export interface components {
|
|||
blocked?: boolean | null;
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/**
|
||||
* Budget Fallbacks
|
||||
* @default {}
|
||||
*/
|
||||
budget_fallbacks: {
|
||||
[key: string]: string[];
|
||||
};
|
||||
/** Budget Id */
|
||||
budget_id?: string | null;
|
||||
/** Budget Limits */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue