mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
feat(proxy): per-model team member budgets with router fallback (#43887)
Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
7383f59ad2
commit
5325c4a236
41 changed files with 5039 additions and 245 deletions
|
|
@ -84,6 +84,7 @@ _PROPAGATED_METADATA_KEYS: Final = (
|
|||
"user_api_key_model_max_budget",
|
||||
"user_api_key_team_model_max_budget",
|
||||
"user_api_key_user_model_max_budget",
|
||||
"user_api_key_team_member_model_max_budget",
|
||||
"user_api_key_end_user_model_max_budget",
|
||||
"litellm_call_id",
|
||||
"litellm_parent_otel_span",
|
||||
|
|
@ -125,6 +126,23 @@ class _SummaryAcompletion(Protocol):
|
|||
) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ...
|
||||
|
||||
|
||||
class _SummaryModelBudgetLimiter(Protocol):
|
||||
async def is_team_member_within_model_budget(
|
||||
self,
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
team_member_model_max_budget: Mapping[str, object],
|
||||
model: str,
|
||||
) -> bool: ...
|
||||
|
||||
async def is_end_user_within_model_budget(
|
||||
self,
|
||||
end_user_id: str,
|
||||
end_user_model_max_budget: Mapping[str, object],
|
||||
model: str,
|
||||
) -> bool: ...
|
||||
|
||||
|
||||
class _CreateRateLimitDescriptors(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
|
|
@ -392,6 +410,42 @@ async def _check_summary_model_access(
|
|||
return True
|
||||
|
||||
|
||||
async def _check_summary_team_member_model_budget(
|
||||
user_api_key_auth: "UserAPIKeyAuth",
|
||||
model_max_budget_limiter: _SummaryModelBudgetLimiter,
|
||||
summary_model: str,
|
||||
) -> bool:
|
||||
team_member_model_max_budget: Final = user_api_key_auth.team_member_model_max_budget
|
||||
user_id: Final = user_api_key_auth.user_id
|
||||
team_id: Final = user_api_key_auth.team_id
|
||||
if (
|
||||
not isinstance(team_member_model_max_budget, Mapping)
|
||||
or not team_member_model_max_budget
|
||||
or not isinstance(user_id, str)
|
||||
or not isinstance(team_id, str)
|
||||
):
|
||||
return True
|
||||
|
||||
try:
|
||||
await model_max_budget_limiter.is_team_member_within_model_budget(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
team_member_model_max_budget=team_member_model_max_budget,
|
||||
model=summary_model,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return False
|
||||
except Exception as e: # noqa: BLE001 # unexpected budget-check errors must deny compaction
|
||||
verbose_logger.warning(
|
||||
"compact_20260112: unexpected error during team member model-budget "
|
||||
"check for summary_model=%s; denying: %s",
|
||||
summary_model,
|
||||
e,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _check_summary_model_budget(
|
||||
user_api_key_auth: Optional["UserAPIKeyAuth"],
|
||||
summary_model: str,
|
||||
|
|
@ -407,9 +461,9 @@ async def _check_summary_model_budget(
|
|||
per-model budget is configured.
|
||||
|
||||
Every scope is checked because the summary's spend is charged to every
|
||||
scope: this file propagates the key, team, user and end-user budgets into the
|
||||
subrequest's metadata, so skipping one of them would let compaction
|
||||
increment a counter it can never be refused by.
|
||||
scope: this file propagates the key, team, user, team-member and end-user
|
||||
budgets into the subrequest's metadata, so skipping one of them would let
|
||||
compaction increment a counter it can never be refused by.
|
||||
"""
|
||||
if user_api_key_auth is None:
|
||||
return True
|
||||
|
|
@ -497,7 +551,11 @@ async def _check_summary_model_budget(
|
|||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
return await _check_summary_team_member_model_budget(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
summary_model=summary_model,
|
||||
)
|
||||
|
||||
|
||||
def _without_parallel_request_gauges(
|
||||
|
|
|
|||
|
|
@ -2106,6 +2106,7 @@ class NewTeamRequest(TeamBase):
|
|||
team_member_tpm_limit: int | None = None # allow user to set TPM limit for all team members
|
||||
team_member_key_duration: str | None = None # e.g. "1d", "1w", "1m"
|
||||
team_member_budget_duration: str | None = None # e.g. "30d", "1mo"
|
||||
team_member_model_max_budget: GenericBudgetConfigType | None = None
|
||||
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
|
||||
enforced_batch_output_expires_after: dict | None = None
|
||||
enforced_file_expires_after: dict | None = None
|
||||
|
|
@ -2165,6 +2166,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
|
|||
require_trace_id: bool | None = None
|
||||
team_member_budget: float | None = None
|
||||
team_member_budget_duration: str | None = None
|
||||
team_member_model_max_budget: GenericBudgetConfigType | None = None
|
||||
team_member_rpm_limit: int | None = None
|
||||
team_member_tpm_limit: int | None = None
|
||||
team_member_key_duration: str | None = None
|
||||
|
|
@ -3363,6 +3365,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
# and validating it here would make one malformed row fail auth outright.
|
||||
# resolve_model_budget validates the single entry a request actually needs.
|
||||
user_model_max_budget: Mapping[str, object] | None = None
|
||||
team_member_model_max_budget: Mapping[str, object] | None = Field(default=None, exclude=True)
|
||||
request_route: str | None = None
|
||||
is_session_token: bool = False
|
||||
# Server-only marker set exclusively by the MCP gateway admission path
|
||||
|
|
@ -4651,6 +4654,13 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
|
|||
default=None,
|
||||
description="UTC expiry for temp_budget_increase",
|
||||
)
|
||||
model_max_budget: GenericBudgetConfigType | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Per-model spend caps for this team member, each with its own budget_duration. "
|
||||
"Overrides the team's default per-model member budget. Pass an empty dict to fall back to the team default."
|
||||
),
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_temp_budget(self) -> "TeamMemberUpdateRequest":
|
||||
|
|
@ -4669,6 +4679,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse):
|
|||
allowed_models: list[str] | None = None
|
||||
temp_budget_increase: float | None = None
|
||||
temp_budget_expiry: datetime | None = None
|
||||
model_max_budget: GenericBudgetConfigType | None = None
|
||||
|
||||
|
||||
class TeamModelAddRequest(LiteLLMBaseModel):
|
||||
|
|
|
|||
|
|
@ -1739,6 +1739,7 @@ async def get_team_member_default_budget(
|
|||
budget_id: str,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
raise_on_lookup_error: bool = False,
|
||||
) -> LiteLLM_BudgetTable | None:
|
||||
"""
|
||||
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
|
||||
|
|
@ -1773,6 +1774,8 @@ async def get_team_member_default_budget(
|
|||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception("Error fetching team-default member budget %s", budget_id)
|
||||
if raise_on_lookup_error:
|
||||
raise
|
||||
return None
|
||||
|
||||
if budget_record is None:
|
||||
|
|
|
|||
|
|
@ -11,13 +11,13 @@ attempted, mirroring `fallback_model_access.py`. It deliberately leaves the prim
|
|||
a zero-cost model is never blocked by budget, and only the paid fallback is refused. On by default;
|
||||
set `general_settings.enforce_fallback_budget: false` to restore the unguarded behaviour.
|
||||
|
||||
Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation
|
||||
path before it can be: team, team-member, end-user, org, global and per-model budgets, whose
|
||||
auth-path functions enforce rather than report (they raise), so reusing them would fire threshold
|
||||
alerts and take spend reservations for a target that is then skipped; and the key's rolling
|
||||
`budget_limits` windows, whose accumulated spend lives only in per-window counters
|
||||
(`spend:key:{token}:window:{budget_duration}`), so enforcing them means more counter reads on the
|
||||
fallback path rather than reusing state auth already loaded.
|
||||
Scope: the key's and the user's `max_budget`, plus the team member's per-model caps. Not covered yet,
|
||||
and each needs a read-only evaluation path before it can be: team and team-member overall, end-user,
|
||||
org, global and other per-model budgets, whose auth-path functions enforce rather than report (they
|
||||
raise), so reusing them would fire threshold alerts and take spend reservations for a target that is
|
||||
then skipped; and the key's rolling `budget_limits` windows, whose accumulated spend lives only in
|
||||
per-window counters (`spend:key:{token}:window:{budget_duration}`), so enforcing them means more
|
||||
counter reads on the fallback path rather than reusing state auth already loaded.
|
||||
|
||||
Two known limitations of that narrow scope, both shared with `fallback_model_access.py`:
|
||||
|
||||
|
|
@ -33,10 +33,11 @@ Two known limitations of that narrow scope, both shared with `fallback_model_acc
|
|||
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module exports
|
||||
|
|
@ -55,6 +56,16 @@ class _FallbackBudgetSettings(LiteLLMBaseModel):
|
|||
enforce_fallback_budget: bool = True
|
||||
|
||||
|
||||
class _TeamMemberModelBudgetLimiter(Protocol):
|
||||
async def is_team_member_within_model_budget(
|
||||
self,
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
team_member_model_max_budget: Mapping[str, object],
|
||||
model: str,
|
||||
) -> bool: ...
|
||||
|
||||
|
||||
def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None:
|
||||
try:
|
||||
return _RequestMetadata.model_validate(metadata).user_api_key_auth
|
||||
|
|
@ -103,7 +114,13 @@ async def _counter_spend(counter_key: str, fallback_spend: float, max_budget: fl
|
|||
)
|
||||
|
||||
|
||||
async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool:
|
||||
async def is_token_within_budget_for_model(
|
||||
*,
|
||||
model: str,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: Router,
|
||||
team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
True when the key and the user behind it can still pay for `model`.
|
||||
|
||||
|
|
@ -139,6 +156,23 @@ async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKe
|
|||
if user_spend >= user_budget:
|
||||
return False
|
||||
|
||||
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
|
||||
if (
|
||||
team_member_model_budget_limiter is not None
|
||||
and team_member_model_max_budget
|
||||
and isinstance(valid_token.user_id, str)
|
||||
and isinstance(valid_token.team_id, str)
|
||||
):
|
||||
try:
|
||||
await team_member_model_budget_limiter.is_team_member_within_model_budget(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=valid_token.team_id,
|
||||
team_member_model_max_budget=team_member_model_max_budget,
|
||||
model=model,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
|
|
@ -151,6 +185,7 @@ class RouterFallbackBudgetCheck:
|
|||
"""
|
||||
|
||||
is_enforced: Callable[[], bool]
|
||||
team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None
|
||||
|
||||
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool:
|
||||
if not self.is_enforced():
|
||||
|
|
@ -158,11 +193,20 @@ class RouterFallbackBudgetCheck:
|
|||
valid_token: Final = _user_api_key_auth_from_request(request_kwargs)
|
||||
if valid_token is None:
|
||||
return True
|
||||
try:
|
||||
return await is_token_within_budget_for_model(model=model, valid_token=valid_token, llm_router=llm_router)
|
||||
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller
|
||||
verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
|
||||
return False
|
||||
budget_models: Final = tuple(dict.fromkeys((model, llm_router.routable_model_group(model))))
|
||||
for budget_model in budget_models:
|
||||
try:
|
||||
if not await is_token_within_budget_for_model(
|
||||
model=budget_model,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
team_member_model_budget_limiter=self.team_member_model_budget_limiter,
|
||||
):
|
||||
return False
|
||||
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller
|
||||
verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
router_fallback_budget_check: Final = RouterFallbackBudgetCheck(is_enforced=_enforced_by_general_settings)
|
||||
|
|
|
|||
|
|
@ -11,12 +11,11 @@ import asyncio
|
|||
import fnmatch
|
||||
import re
|
||||
import secrets
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final, NamedTuple, Protocol, Union, cast
|
||||
|
||||
import fastapi
|
||||
import orjson
|
||||
from fastapi import HTTPException, Request, WebSocket, status
|
||||
from fastapi.security.api_key import APIKeyHeader
|
||||
from starlette.exceptions import WebSocketException
|
||||
|
|
@ -51,13 +50,18 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
|
|||
_get_user_role, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_is_user_proxy_admin, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # reuse common-check group budget enforcement
|
||||
_team_member_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # reuse common-check pricing enforcement
|
||||
_virtual_key_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_virtual_key_max_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_virtual_key_soft_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
cache_key_object,
|
||||
can_customer_access_model,
|
||||
can_key_call_model,
|
||||
can_object_call_model,
|
||||
can_project_access_model,
|
||||
can_team_access_model,
|
||||
check_end_user_budget,
|
||||
common_checks,
|
||||
delete_cache_key_object,
|
||||
|
|
@ -67,6 +71,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
|
|||
get_object_permission,
|
||||
get_org_object_for_request,
|
||||
get_project_object,
|
||||
get_team_member_default_budget,
|
||||
get_team_membership,
|
||||
get_team_object,
|
||||
get_user_object,
|
||||
|
|
@ -76,9 +81,12 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
|
|||
is_user_proxy_admin,
|
||||
is_valid_fallback_model,
|
||||
jwt_key_mapping_cache_key,
|
||||
key_access_group_grants_model,
|
||||
key_model_aliases_for_auth_check,
|
||||
resolve_and_validate_end_user_id,
|
||||
resolve_default_end_user_budget,
|
||||
route_skips_budget_checks,
|
||||
stamp_matched_model_access_groups,
|
||||
team_member_max_budget_alert_check,
|
||||
virtual_key_max_budget_alert_check,
|
||||
virtual_key_max_budget_check,
|
||||
|
|
@ -169,7 +177,13 @@ from litellm.proxy.utils import (
|
|||
)
|
||||
from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
PROVIDER_SCOPED_RESOURCE_KEYS,
|
||||
fallback_lookup_groups,
|
||||
fallbacks_disabled_for_request,
|
||||
get_fallback_model_group_for_lookup_groups,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
|
|
@ -332,6 +346,18 @@ def _get_model_names_for_budget_checks(
|
|||
return model
|
||||
|
||||
|
||||
def _budget_fallback_target_model(target: object) -> str | None:
|
||||
if isinstance(target, str):
|
||||
return target
|
||||
if not isinstance(target, Mapping):
|
||||
return None
|
||||
target_mapping: Final = cast(Mapping[str, object], target) # cast-ok: isinstance cannot type the container's items
|
||||
if set(target_mapping) != {"model"}:
|
||||
return None
|
||||
model: Final = target_mapping.get("model")
|
||||
return model if isinstance(model, str) else None
|
||||
|
||||
|
||||
class _KeyModelBudgetLimiter(Protocol):
|
||||
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
|
||||
|
||||
|
|
@ -354,6 +380,34 @@ class _TeamModelBudgetLimiter(Protocol):
|
|||
) -> bool: ...
|
||||
|
||||
|
||||
class _EveryModelBudgetLimiter(Protocol):
|
||||
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
|
||||
|
||||
async def is_user_within_model_budget(
|
||||
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
|
||||
) -> bool: ...
|
||||
|
||||
async def is_end_user_within_model_budget(
|
||||
self, end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str
|
||||
) -> bool: ...
|
||||
|
||||
async def is_team_within_model_budget(
|
||||
self,
|
||||
team_id: str,
|
||||
team_model_max_budget: Mapping[str, object],
|
||||
key_model_max_budget: Mapping[str, object] | None,
|
||||
model: str,
|
||||
) -> bool: ...
|
||||
|
||||
async def is_team_member_within_model_budget(
|
||||
self,
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
team_member_model_max_budget: Mapping[str, object],
|
||||
model: str,
|
||||
) -> bool: ...
|
||||
|
||||
|
||||
class _TokenTeamModels(Protocol):
|
||||
@property
|
||||
def team_models(self) -> list[str]: ...
|
||||
|
|
@ -374,6 +428,30 @@ def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
|
|||
return valid_token.team_models
|
||||
|
||||
|
||||
class _TokenModelBudgetFields(Protocol):
|
||||
@property
|
||||
def model_max_budget(self) -> Mapping[str, object] | None: ...
|
||||
|
||||
@property
|
||||
def end_user_model_max_budget(self) -> Mapping[str, object] | None: ...
|
||||
|
||||
@property
|
||||
def team_model_aliases(self) -> dict[str, str] | None: ... # mutable-ok: mirrors _can_object_call_model's contract
|
||||
|
||||
|
||||
def _token_model_budget_fields(valid_token: _TokenModelBudgetFields) -> _TokenModelBudgetFields:
|
||||
return valid_token
|
||||
|
||||
|
||||
class _RouterFallbacks(Protocol):
|
||||
@property
|
||||
def fallbacks(self) -> Sequence[object] | None: ...
|
||||
|
||||
|
||||
def _router_fallbacks(router: _RouterFallbacks) -> Sequence[object] | None:
|
||||
return router.fallbacks
|
||||
|
||||
|
||||
async def _read_user_model_max_budget(
|
||||
user_id: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
|
|
@ -443,13 +521,455 @@ async def _check_team_model_budget(
|
|||
)
|
||||
|
||||
|
||||
async def _is_within_every_model_budget(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
model_max_budget_limiter: _EveryModelBudgetLimiter,
|
||||
model: str,
|
||||
) -> bool:
|
||||
budget_fields: Final = _token_model_budget_fields(valid_token)
|
||||
key_model_max_budget: Final = budget_fields.model_max_budget
|
||||
user_model_max_budget: Final = valid_token.user_model_max_budget
|
||||
end_user_model_max_budget: Final = budget_fields.end_user_model_max_budget
|
||||
team_model_max_budget: Final = valid_token.team_model_max_budget
|
||||
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
|
||||
try:
|
||||
if valid_token.token is not None and isinstance(key_model_max_budget, Mapping) and key_model_max_budget:
|
||||
await model_max_budget_limiter.is_key_within_model_budget(
|
||||
user_api_key_dict=valid_token,
|
||||
model=model,
|
||||
)
|
||||
if valid_token.user_id is not None and isinstance(user_model_max_budget, Mapping) and user_model_max_budget:
|
||||
await model_max_budget_limiter.is_user_within_model_budget(
|
||||
user_id=valid_token.user_id,
|
||||
user_model_max_budget=user_model_max_budget,
|
||||
model=model,
|
||||
)
|
||||
if (
|
||||
valid_token.end_user_id is not None
|
||||
and isinstance(end_user_model_max_budget, Mapping)
|
||||
and end_user_model_max_budget
|
||||
):
|
||||
await model_max_budget_limiter.is_end_user_within_model_budget(
|
||||
end_user_id=valid_token.end_user_id,
|
||||
end_user_model_max_budget=end_user_model_max_budget,
|
||||
model=model,
|
||||
)
|
||||
if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget:
|
||||
await model_max_budget_limiter.is_team_within_model_budget(
|
||||
team_id=valid_token.team_id,
|
||||
team_model_max_budget=team_model_max_budget,
|
||||
key_model_max_budget=key_model_max_budget,
|
||||
model=model,
|
||||
)
|
||||
if (
|
||||
valid_token.user_id is not None
|
||||
and valid_token.team_id is not None
|
||||
and isinstance(team_member_model_max_budget, Mapping)
|
||||
and team_member_model_max_budget
|
||||
):
|
||||
await model_max_budget_limiter.is_team_member_within_model_budget(
|
||||
user_id=valid_token.user_id,
|
||||
team_id=valid_token.team_id,
|
||||
team_member_model_max_budget=team_member_model_max_budget,
|
||||
model=model,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _is_fallback_model_authorized(
|
||||
model: str,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
llm_router: litellm.Router | None,
|
||||
prisma_client: PrismaClient | None = None,
|
||||
team_membership: LiteLLM_TeamMembership | None = None,
|
||||
project_object: LiteLLM_ProjectTableCachedObj | None = None,
|
||||
refuse_agent_keys: bool = False,
|
||||
team_object: LiteLLM_TeamTable | None = None,
|
||||
) -> bool:
|
||||
if refuse_agent_keys and valid_token.agent_id is not None:
|
||||
return False
|
||||
try:
|
||||
await can_key_call_model(
|
||||
model=model,
|
||||
llm_model_list=llm_model_list,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if team_object is not None:
|
||||
try:
|
||||
await can_team_access_model(
|
||||
model=model,
|
||||
team_object=team_object,
|
||||
llm_router=llm_router,
|
||||
team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases,
|
||||
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
except ProxyException as team_denial:
|
||||
if team_denial.type != ProxyErrorTypes.team_model_access_denied:
|
||||
raise
|
||||
if not await key_access_group_grants_model(
|
||||
model=model,
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
llm_router=llm_router,
|
||||
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
|
||||
):
|
||||
raise
|
||||
else:
|
||||
team_models: Final = _token_team_models(valid_token)
|
||||
if team_models:
|
||||
can_object_call_model(
|
||||
model=model,
|
||||
llm_router=llm_router,
|
||||
models=team_models,
|
||||
team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases,
|
||||
team_id=valid_token.team_id,
|
||||
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
|
||||
object_type="team",
|
||||
)
|
||||
member_allowed_models: Final = (
|
||||
team_membership.litellm_budget_table.allowed_models
|
||||
if team_membership is not None and team_membership.litellm_budget_table is not None
|
||||
else None
|
||||
)
|
||||
if member_allowed_models:
|
||||
can_object_call_model(
|
||||
model=model,
|
||||
llm_router=llm_router,
|
||||
models=member_allowed_models,
|
||||
team_id=valid_token.team_id,
|
||||
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
|
||||
object_type="team",
|
||||
)
|
||||
if project_object is not None and len(project_object.models) > 0:
|
||||
can_project_access_model(
|
||||
model=model,
|
||||
project_object=project_object,
|
||||
llm_router=llm_router,
|
||||
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
|
||||
)
|
||||
except ProxyException:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _is_within_all_fallback_model_budgets(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
model_max_budget_limiter: _EveryModelBudgetLimiter,
|
||||
models: Sequence[str],
|
||||
) -> bool:
|
||||
for model in models:
|
||||
if not await _is_within_every_model_budget(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
model=model,
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _check_team_member_model_budget_with_fallback(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
model_max_budget_limiter: _EveryModelBudgetLimiter,
|
||||
models: Sequence[str],
|
||||
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
|
||||
request: Request,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
llm_router: litellm.Router | None,
|
||||
project_object: LiteLLM_ProjectTableCachedObj | None = None,
|
||||
team_object: LiteLLM_TeamTable | None = None,
|
||||
end_user_object: LiteLLM_EndUserTable | None = None,
|
||||
route: str = "",
|
||||
prisma_client: PrismaClient | None = None,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> None:
|
||||
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
|
||||
if (
|
||||
valid_token.user_id is None
|
||||
or valid_token.team_id is None
|
||||
or not isinstance(team_member_model_max_budget, Mapping)
|
||||
or not team_member_model_max_budget
|
||||
):
|
||||
return
|
||||
|
||||
user_id: Final = valid_token.user_id
|
||||
team_id: Final = valid_token.team_id
|
||||
for model_name in models:
|
||||
await _check_team_member_model_budget_for_model(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
team_member_model_max_budget=team_member_model_max_budget,
|
||||
model_name=model_name,
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_router=llm_router,
|
||||
project_object=project_object,
|
||||
team_object=team_object,
|
||||
end_user_object=end_user_object,
|
||||
route=route,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
|
||||
async def _check_team_member_model_budget_for_model(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
team_member_model_max_budget: Mapping[str, object],
|
||||
model_name: str,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
model_max_budget_limiter: _EveryModelBudgetLimiter,
|
||||
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
|
||||
request: Request,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
llm_router: litellm.Router | None,
|
||||
project_object: LiteLLM_ProjectTableCachedObj | None = None,
|
||||
team_object: LiteLLM_TeamTable | None = None,
|
||||
end_user_object: LiteLLM_EndUserTable | None = None,
|
||||
route: str = "",
|
||||
prisma_client: PrismaClient | None = None,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
await model_max_budget_limiter.is_team_member_within_model_budget(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
team_member_model_max_budget=team_member_model_max_budget,
|
||||
model=model_name,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
if (
|
||||
request_data.get("model") != model_name
|
||||
or llm_router is None
|
||||
or fallbacks_disabled_for_request(kwargs=request_data)
|
||||
):
|
||||
raise
|
||||
fallback_succeeded: Final = await _try_team_member_model_budget_fallback(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
model_name=model_name,
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_router=llm_router,
|
||||
project_object=project_object,
|
||||
team_object=team_object,
|
||||
end_user_object=end_user_object,
|
||||
route=route,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not fallback_succeeded:
|
||||
raise
|
||||
|
||||
|
||||
async def _try_team_member_model_budget_fallback(
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
model_name: str,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
model_max_budget_limiter: _EveryModelBudgetLimiter,
|
||||
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
|
||||
request: Request,
|
||||
llm_model_list: Sequence[object] | None,
|
||||
llm_router: litellm.Router,
|
||||
project_object: LiteLLM_ProjectTableCachedObj | None = None,
|
||||
team_object: LiteLLM_TeamTable | None = None,
|
||||
end_user_object: LiteLLM_EndUserTable | None = None,
|
||||
route: str = "",
|
||||
prisma_client: PrismaClient | None = None,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
) -> bool:
|
||||
if request.headers.get("content-type", "").lower().startswith("multipart/form-data") or any(
|
||||
request_data.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS
|
||||
):
|
||||
return False
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
resolved_prisma_client: Final = prisma_client if prisma_client is not None else proxy_server.prisma_client
|
||||
resolved_user_api_key_cache: Final = (
|
||||
user_api_key_cache if user_api_key_cache is not None else proxy_server.user_api_key_cache
|
||||
)
|
||||
resolved_proxy_logging_obj: Final = (
|
||||
proxy_logging_obj if proxy_logging_obj is not None else proxy_server.proxy_logging_obj
|
||||
)
|
||||
|
||||
try:
|
||||
router_settings: Final[dict[str, object] | None] = (
|
||||
cast( # cast-ok: proxy settings return an unparameterized dict
|
||||
dict[str, object] | None,
|
||||
await proxy_config.get_hierarchical_router_settings(
|
||||
user_api_key_dict=valid_token,
|
||||
prisma_client=resolved_prisma_client,
|
||||
proxy_logging_obj=resolved_proxy_logging_obj,
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback.
|
||||
return False
|
||||
|
||||
configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None
|
||||
router_fallbacks: Final[Sequence[object] | None] = (
|
||||
cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below
|
||||
if isinstance(configured_fallbacks, list)
|
||||
else _router_fallbacks(llm_router)
|
||||
)
|
||||
if router_fallbacks is None:
|
||||
return False
|
||||
|
||||
fallback_models, _ = get_fallback_model_group_for_lookup_groups(
|
||||
fallbacks=list(router_fallbacks),
|
||||
lookup_groups=fallback_lookup_groups(kwargs=request_data, model_group=model_name),
|
||||
)
|
||||
fallback_targets: Final[Sequence[object]] = fallback_models or ()
|
||||
fallback_candidates: Final[tuple[str, ...]] = tuple(
|
||||
model for target in fallback_targets if (model := _budget_fallback_target_model(target)) is not None
|
||||
)
|
||||
if not fallback_candidates:
|
||||
return False
|
||||
|
||||
model_group_aliases: Final[object] = (
|
||||
router_settings.get("model_group_alias") if router_settings is not None else None
|
||||
)
|
||||
key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime
|
||||
object, valid_token.aliases
|
||||
)
|
||||
resolved_fallback_candidates: Final[tuple[tuple[str, str], ...]] = tuple(
|
||||
(
|
||||
candidate,
|
||||
target if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None else candidate,
|
||||
)
|
||||
for candidate in fallback_candidates
|
||||
)
|
||||
fallback_candidates_with_served_models: Final = (
|
||||
(
|
||||
fallback_candidate,
|
||||
fallback_model,
|
||||
resolve_served_model(
|
||||
fallback_model,
|
||||
litellm.model_alias_map,
|
||||
key_aliases,
|
||||
model_group_aliases,
|
||||
),
|
||||
)
|
||||
for fallback_candidate, fallback_model in resolved_fallback_candidates
|
||||
)
|
||||
|
||||
try:
|
||||
team_membership: Final = await get_team_membership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
prisma_client=resolved_prisma_client,
|
||||
user_api_key_cache=resolved_user_api_key_cache,
|
||||
proxy_logging_obj=resolved_proxy_logging_obj,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # membership lookup failures must not authorize a fallback
|
||||
return False
|
||||
|
||||
for fallback_candidate, fallback_model, served_model in fallback_candidates_with_served_models:
|
||||
if fallback_model == model_name:
|
||||
continue
|
||||
try:
|
||||
if (
|
||||
litellm.block_requests_for_models_without_pricing
|
||||
and RouteChecks.is_llm_api_route(route=route)
|
||||
and _unpriced_models_in_request(model=served_model, llm_router=llm_router)
|
||||
):
|
||||
continue
|
||||
if end_user_object is not None and end_user_object.models:
|
||||
can_customer_access_model(
|
||||
model=fallback_model,
|
||||
end_user_object=end_user_object,
|
||||
llm_router=llm_router,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
if not is_model_cost_zero(
|
||||
model=served_model, llm_router=llm_router
|
||||
) and not await _is_within_all_fallback_model_budgets(
|
||||
valid_token=valid_token,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
models=tuple(
|
||||
dict.fromkeys(
|
||||
(
|
||||
fallback_candidate,
|
||||
fallback_model,
|
||||
served_model,
|
||||
llm_router.routable_model_group(served_model),
|
||||
)
|
||||
)
|
||||
),
|
||||
):
|
||||
continue
|
||||
if not await _is_fallback_model_authorized(
|
||||
model=fallback_model,
|
||||
valid_token=valid_token,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_router=llm_router,
|
||||
prisma_client=prisma_client,
|
||||
team_object=team_object,
|
||||
team_membership=team_membership,
|
||||
project_object=project_object,
|
||||
refuse_agent_keys=True,
|
||||
):
|
||||
continue
|
||||
valid_token.matched_model_access_groups = [] # rebind-ok: clear prior groups before evaluating this candidate
|
||||
matched_groups = await stamp_matched_model_access_groups(
|
||||
model=fallback_model,
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
project_object=project_object,
|
||||
llm_router=llm_router,
|
||||
prisma_client=resolved_prisma_client,
|
||||
user_api_key_cache=resolved_user_api_key_cache,
|
||||
proxy_logging_obj=resolved_proxy_logging_obj,
|
||||
team_membership=team_membership,
|
||||
team_membership_loaded=True,
|
||||
)
|
||||
valid_token.matched_model_access_groups = list(matched_groups) # rebind-ok: spend uses the candidate groups
|
||||
if matched_groups:
|
||||
await _model_access_group_max_budget_check(
|
||||
matched_model_access_groups=list(matched_groups),
|
||||
prisma_client=resolved_prisma_client,
|
||||
user_api_key_cache=resolved_user_api_key_cache,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # A candidate check failure must not escape the fail-closed walk.
|
||||
verbose_proxy_logger.warning(
|
||||
"Skipping fallback candidate %s after a budget or authorization check error",
|
||||
fallback_model.replace("\r", "").replace("\n", ""),
|
||||
exc_info=True,
|
||||
)
|
||||
continue
|
||||
rewrite_request_model(request_data=request_data, request=request, model=fallback_model)
|
||||
if "path_params" in request.scope:
|
||||
request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def _check_key_model_budget_with_fallback(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
model_max_budget_limiter: _KeyModelBudgetLimiter,
|
||||
model_name: str,
|
||||
request_data: dict,
|
||||
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
|
||||
request: Request,
|
||||
llm_model_list: list | None = None,
|
||||
llm_model_list: Sequence[object] | None = None,
|
||||
llm_router: litellm.Router | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -489,32 +1009,16 @@ async def _check_key_model_budget_with_fallback(
|
|||
)
|
||||
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,
|
||||
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
|
||||
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 # pyright: ignore[reportPrivateUsage] # Starlette JSON cache
|
||||
request._body = orjson.dumps(request_data) # pyright: ignore[reportPrivateUsage] # Starlette body cache
|
||||
path_params: Final = request.scope.get("path_params")
|
||||
if isinstance(path_params, dict) and "model" in path_params:
|
||||
path_params["model"] = fallback_model
|
||||
if not await _is_fallback_model_authorized(
|
||||
model=fallback_model,
|
||||
valid_token=valid_token,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_router=llm_router,
|
||||
):
|
||||
raise
|
||||
rewrite_request_model(request_data=request_data, request=request, model=fallback_model)
|
||||
if "path_params" in request.scope:
|
||||
request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model
|
||||
|
||||
|
||||
def _get_bearer_token_or_received_api_key(api_key: str) -> str:
|
||||
|
|
@ -2543,7 +3047,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
|
|||
model_name=model_name,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
|
||||
Sequence[object] | None, llm_model_list
|
||||
),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
|
|
@ -2836,6 +3342,79 @@ def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, o
|
|||
)
|
||||
|
||||
|
||||
async def _load_team_member_default_model_budget(
|
||||
user_api_key_auth_obj: UserAPIKeyAuth,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
skip_budget_checks: bool,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
) -> None:
|
||||
team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None
|
||||
team_metadata: Final[Mapping[str, object] | None] = (
|
||||
cast( # cast-ok: team metadata is a string-keyed JSON object
|
||||
Mapping[str, object], team_metadata_value
|
||||
)
|
||||
if isinstance(team_metadata_value, Mapping)
|
||||
else None
|
||||
)
|
||||
team_member_budget_id: Final[object | None] = (
|
||||
team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None
|
||||
)
|
||||
if (
|
||||
not isinstance(user_api_key_auth_obj.user_id, str)
|
||||
or not isinstance(user_api_key_auth_obj.team_id, str)
|
||||
or prisma_client is None
|
||||
):
|
||||
return
|
||||
|
||||
try:
|
||||
membership: Final = await get_team_membership(
|
||||
user_id=user_api_key_auth_obj.user_id,
|
||||
team_id=user_api_key_auth_obj.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
member_budget_row: Final = membership.litellm_budget_table if membership is not None else None
|
||||
member_budget_value: Final[object] = (
|
||||
getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None
|
||||
)
|
||||
if (
|
||||
membership is not None
|
||||
and membership.budget_id is not None
|
||||
and membership.budget_id != team_member_budget_id
|
||||
and isinstance(member_budget_value, Mapping)
|
||||
and member_budget_value
|
||||
):
|
||||
user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream
|
||||
member_budget_value
|
||||
)
|
||||
return
|
||||
if not isinstance(team_member_budget_id, str):
|
||||
return
|
||||
team_member_default_budget: Final = await get_team_member_default_budget(
|
||||
budget_id=team_member_budget_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
raise_on_lookup_error=True,
|
||||
)
|
||||
except Exception:
|
||||
if not skip_budget_checks:
|
||||
raise
|
||||
verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True)
|
||||
return
|
||||
|
||||
default_budget_value: Final[object] = (
|
||||
getattr(team_member_default_budget, "model_max_budget", None)
|
||||
if team_member_default_budget is not None
|
||||
else None
|
||||
)
|
||||
user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream
|
||||
default_budget_value if isinstance(default_budget_value, Mapping) else None
|
||||
)
|
||||
|
||||
|
||||
@tracer.wrap()
|
||||
async def run_centralized_common_checks(
|
||||
user_api_key_auth_obj: UserAPIKeyAuth,
|
||||
|
|
@ -2864,6 +3443,7 @@ async def run_centralized_common_checks(
|
|||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
litellm_proxy_admin_name,
|
||||
llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
|
||||
llm_router,
|
||||
master_key,
|
||||
model_max_budget_limiter,
|
||||
|
|
@ -3124,6 +3704,17 @@ async def run_centralized_common_checks(
|
|||
llm_router=llm_router,
|
||||
team_id=user_api_key_auth_obj.team_id,
|
||||
)
|
||||
typed_team_object: Final[LiteLLM_TeamTable | None] = (
|
||||
team_object if isinstance(team_object, LiteLLM_TeamTable) else None
|
||||
)
|
||||
await _load_team_member_default_model_budget(
|
||||
user_api_key_auth_obj=user_api_key_auth_obj,
|
||||
team_object=typed_team_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
skip_budget_checks=skip_budget_checks or route_skips_budget_checks(route=route),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Pin the metadata variable name (litellm_metadata vs metadata) before
|
||||
# any tag merge runs. Without this, header tags from
|
||||
|
|
@ -3164,19 +3755,42 @@ async def run_centralized_common_checks(
|
|||
project_object=project_object,
|
||||
)
|
||||
if not skip_budget_checks:
|
||||
models: Final = _get_model_names_for_budget_checks(
|
||||
model=_get_model_from_request_context(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=user_api_key_auth_obj.team_id,
|
||||
)
|
||||
)
|
||||
await _check_team_model_budget(
|
||||
valid_token=user_api_key_auth_obj,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
models=_get_model_names_for_budget_checks(
|
||||
model=_get_model_from_request_context(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=user_api_key_auth_obj.team_id,
|
||||
)
|
||||
),
|
||||
models=models,
|
||||
)
|
||||
typed_project_object: Final[LiteLLM_ProjectTableCachedObj | None] = (
|
||||
project_object if isinstance(project_object, LiteLLM_ProjectTableCachedObj) else None
|
||||
)
|
||||
if not route_skips_budget_checks(route=route):
|
||||
await _check_team_member_model_budget_with_fallback(
|
||||
valid_token=user_api_key_auth_obj,
|
||||
model_max_budget_limiter=model_max_budget_limiter,
|
||||
models=models,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
|
||||
Sequence[object] | None, llm_model_list
|
||||
),
|
||||
llm_router=llm_router,
|
||||
project_object=typed_project_object,
|
||||
team_object=typed_team_object,
|
||||
end_user_object=end_user_object if isinstance(end_user_object, LiteLLM_EndUserTable) else None,
|
||||
route=route,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
await _reserve_budget_after_common_checks(
|
||||
user_api_key_auth_obj=user_api_key_auth_obj,
|
||||
|
|
@ -3184,7 +3798,10 @@ async def run_centralized_common_checks(
|
|||
request_data=request_data,
|
||||
route=route,
|
||||
llm_router=llm_router,
|
||||
team_object=team_object,
|
||||
team_object=cast( # cast-ok: preserve plain team-table objects through reservation
|
||||
LiteLLM_TeamTableCachedObj | None,
|
||||
team_object,
|
||||
),
|
||||
user_object=user_object,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
|
|
@ -3855,13 +4472,13 @@ async def enforce_key_and_fallback_model_access(
|
|||
)
|
||||
|
||||
|
||||
_enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access
|
||||
_enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access # pyright: ignore[reportUnknownVariableType] # existing tests use this alias
|
||||
|
||||
|
||||
async def _run_post_custom_auth_checks(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
request: Request,
|
||||
request_data: dict,
|
||||
request_data: dict[str, object],
|
||||
route: str,
|
||||
parent_otel_span: Span | None,
|
||||
) -> UserAPIKeyAuth:
|
||||
|
|
@ -3956,7 +4573,9 @@ async def _run_post_custom_auth_checks(
|
|||
model_name=model_name,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
llm_model_list=llm_model_list,
|
||||
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
|
||||
Sequence[object] | None, llm_model_list
|
||||
),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -83,6 +83,7 @@ from litellm.litellm_core_utils.streaming_handler import (
|
|||
from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
can_key_call_resolved_model,
|
||||
is_model_cost_zero,
|
||||
request_skips_budget_checks,
|
||||
tag_max_budget_check_for_tags,
|
||||
)
|
||||
|
|
@ -127,7 +128,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
|
|||
)
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
from litellm.types.router_weights import validate_router_weights
|
||||
|
|
@ -250,6 +251,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
add_litellm_data_to_request,
|
||||
refresh_proxy_server_request_body_snapshot,
|
||||
reject_url_valued_destination,
|
||||
sanitize_for_log,
|
||||
)
|
||||
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -737,6 +739,7 @@ async def _resolve_per_request_model_group_alias(
|
|||
|
||||
|
||||
_REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
|
||||
_ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
|
||||
|
||||
def _request_model(data: Mapping[str, object]) -> str | list[str] | None:
|
||||
|
|
@ -1733,6 +1736,69 @@ def _timing_values(
|
|||
return getattr(logging_obj, "response_timing_metrics", None) or {}
|
||||
|
||||
|
||||
async def _team_member_within_fallback_model_budget(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
model: str,
|
||||
llm_router: Router | None,
|
||||
proxy_config: ProxyConfig | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> bool:
|
||||
member_caps: Final = user_api_key_dict.team_member_model_max_budget
|
||||
user_id: Final = user_api_key_dict.user_id
|
||||
team_id: Final = user_api_key_dict.team_id
|
||||
if not isinstance(member_caps, Mapping) or not member_caps or user_id is None or team_id is None:
|
||||
return True
|
||||
from litellm.proxy.proxy_server import model_max_budget_limiter
|
||||
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
router_settings: Final = (
|
||||
_ROUTER_SETTINGS_ADAPTER.validate_python(
|
||||
await proxy_config.get_hierarchical_router_settings( # pyright: ignore[reportUnknownMemberType] # legacy bare-dict return
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
)
|
||||
if proxy_config is not None and llm_router is not None
|
||||
else None
|
||||
)
|
||||
raw_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime
|
||||
object, user_api_key_dict.aliases
|
||||
)
|
||||
served_model: Final = resolve_served_model(
|
||||
model,
|
||||
litellm.model_alias_map,
|
||||
raw_aliases,
|
||||
router_settings.get("model_group_alias") if router_settings is not None else None,
|
||||
)
|
||||
if is_model_cost_zero(model=served_model, llm_router=llm_router):
|
||||
return True
|
||||
budget_models: Final = (
|
||||
tuple(dict.fromkeys((model, served_model, llm_router.routable_model_group(served_model))))
|
||||
if llm_router is not None
|
||||
else tuple(dict.fromkeys((model, served_model)))
|
||||
)
|
||||
for budget_model in budget_models:
|
||||
await model_max_budget_limiter.is_team_member_within_model_budget(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
team_member_model_max_budget=member_caps,
|
||||
model=budget_model,
|
||||
)
|
||||
except litellm.BudgetExceededError:
|
||||
return False
|
||||
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the member
|
||||
verbose_proxy_logger.warning(
|
||||
"Skipping rate limit fallback to model=%s: member budget lookup failed: %s",
|
||||
sanitize_for_log(model),
|
||||
sanitize_for_log(e),
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class ProxyBaseLLMRequestProcessing:
|
||||
def __init__(self, data: dict):
|
||||
self.data = data
|
||||
|
|
@ -2371,6 +2437,14 @@ class ProxyBaseLLMRequestProcessing:
|
|||
for fallback_model in fallback_models:
|
||||
if fallback_model == original_model:
|
||||
continue
|
||||
if not await _team_member_within_fallback_model_budget(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
model=fallback_model,
|
||||
llm_router=llm_router,
|
||||
proxy_config=proxy_config,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
):
|
||||
continue
|
||||
self.data = independent_snapshot(pristine)
|
||||
self.data["model"] = fallback_model
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import time
|
|||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import Final, cast # noqa: TID251 # untyped callback payload requires explicit narrowing
|
||||
|
||||
from openai.types import Batch
|
||||
|
||||
|
|
@ -25,6 +25,7 @@ VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend"
|
|||
END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend"
|
||||
USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend"
|
||||
TEAM_SPEND_CACHE_KEY_PREFIX: Final = "team_model_spend"
|
||||
TEAM_MEMBER_SPEND_CACHE_KEY_PREFIX: Final = "team_member_model_spend"
|
||||
|
||||
_SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
|
||||
{
|
||||
|
|
@ -32,6 +33,7 @@ _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
|
|||
Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX,
|
||||
Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX,
|
||||
Litellm_EntityType.TEAM: TEAM_SPEND_CACHE_KEY_PREFIX,
|
||||
Litellm_EntityType.TEAM_MEMBER: TEAM_MEMBER_SPEND_CACHE_KEY_PREFIX,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -45,6 +47,7 @@ _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType(
|
|||
Litellm_EntityType.USER: "user_model_budget_start_time",
|
||||
Litellm_EntityType.END_USER: "end_user_budget_start_time",
|
||||
Litellm_EntityType.TEAM: "team_model_budget_start_time",
|
||||
Litellm_EntityType.TEAM_MEMBER: "team_member_model_budget_start_time",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -63,6 +66,10 @@ class ResolvedModelBudget:
|
|||
budget_config: BudgetConfig
|
||||
|
||||
|
||||
def team_member_budget_entity_id(user_id: str, team_id: str) -> str:
|
||||
return f"{len(user_id)}:{user_id}:{team_id}"
|
||||
|
||||
|
||||
def model_budget_spend_cache_key(
|
||||
entity_type: Litellm_EntityType,
|
||||
entity_id: str | None,
|
||||
|
|
@ -290,6 +297,18 @@ def _resolve_entity_model_budgets(
|
|||
)
|
||||
|
||||
|
||||
def _metadata_value(metadata: object, key: str) -> object | None:
|
||||
metadata_mapping: Final[Mapping[str, object]] = cast( # cast-ok: callback metadata has string keys
|
||||
Mapping[str, object], metadata
|
||||
)
|
||||
return metadata_mapping.get(key)
|
||||
|
||||
|
||||
def _metadata_string_value(metadata: object, key: str) -> str | None:
|
||||
value: Final = _metadata_value(metadata, key)
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
||||
"""
|
||||
Handles budgets for model + virtual key
|
||||
|
|
@ -410,6 +429,24 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}",
|
||||
)
|
||||
|
||||
@with_service_target("model_budgets")
|
||||
async def is_team_member_within_model_budget(
|
||||
self,
|
||||
user_id: str,
|
||||
team_id: str,
|
||||
team_member_model_max_budget: Mapping[str, object],
|
||||
model: str,
|
||||
) -> bool:
|
||||
return await self._is_entity_within_model_budget(
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER,
|
||||
entity_id=team_member_budget_entity_id(user_id=user_id, team_id=team_id),
|
||||
model_max_budget=team_member_model_max_budget,
|
||||
model=model,
|
||||
exceeded_message=(
|
||||
f"LiteLLM Team Member: user={user_id}, team={team_id}, exceeded budget for model={model}"
|
||||
),
|
||||
)
|
||||
|
||||
async def _is_entity_within_model_budget(
|
||||
self,
|
||||
entity_type: Litellm_EntityType,
|
||||
|
|
@ -523,6 +560,19 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
|
||||
response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
|
||||
key_model_max_budget: Final = _metadata.get("user_api_key_model_max_budget")
|
||||
team_member_user_id: Final = _metadata_string_value(
|
||||
payload_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
|
||||
"user_api_key_user_id",
|
||||
)
|
||||
team_member_team_id: Final = _metadata_string_value(
|
||||
payload_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
|
||||
"user_api_key_team_id",
|
||||
)
|
||||
team_member_entity_id: Final = (
|
||||
team_member_budget_entity_id(user_id=team_member_user_id, team_id=team_member_team_id)
|
||||
if team_member_user_id is not None and team_member_team_id is not None
|
||||
else None
|
||||
)
|
||||
entity_budgets: Final = (
|
||||
(
|
||||
Litellm_EntityType.KEY,
|
||||
|
|
@ -553,6 +603,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
standard_logging_payload.get("end_user") or payload_metadata.get("user_api_key_end_user_id"),
|
||||
_metadata.get("user_api_key_end_user_model_max_budget"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TEAM_MEMBER,
|
||||
team_member_entity_id,
|
||||
_metadata_value(
|
||||
_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
|
||||
"user_api_key_team_member_model_max_budget",
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
resolved_budgets: Final = _resolve_entity_model_budgets(model=model, entity_budgets=entity_budgets)
|
||||
|
|
|
|||
|
|
@ -2508,6 +2508,9 @@ async def add_litellm_data_to_request(
|
|||
data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget
|
||||
user_model_budget: Final = user_api_key_dict.user_model_max_budget
|
||||
data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param
|
||||
data[_metadata_variable_name][ # rebind-ok: adds member budget for spend tracking
|
||||
"user_api_key_team_member_model_max_budget"
|
||||
] = user_api_key_dict.team_member_model_max_budget
|
||||
data[_metadata_variable_name].update(carried_budget_metadata(user_api_key_dict))
|
||||
|
||||
data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@ from typing import Final
|
|||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from litellm._internal_context import service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
|
@ -141,7 +143,7 @@ async def update_budget(
|
|||
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
|
||||
- budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
|
||||
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -199,6 +201,11 @@ async def update_budget(
|
|||
where={"budget_id": budget_obj.budget_id},
|
||||
data=budget_obj_jsonified,
|
||||
)
|
||||
with service_target("team_member_budgets"):
|
||||
try:
|
||||
await user_api_key_cache.async_delete_cache(key=f"team_member_default_budget:{budget_obj.budget_id}")
|
||||
except Exception: # noqa: BLE001 # Cache failure must not fail a committed budget update.
|
||||
verbose_proxy_logger.warning("Failed to evict the team member default budget cache", exc_info=True)
|
||||
|
||||
return response
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import json
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
|
|
@ -580,21 +581,31 @@ MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType(
|
|||
"allowed_models": "allowed_models",
|
||||
"temp_budget_increase": "temp_budget_increase",
|
||||
"temp_budget_expiry": "temp_budget_expiry",
|
||||
"model_max_budget": "model_max_budget",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _prisma_value(value: object) -> object:
|
||||
_JSON_BUDGET_COLUMNS: Final = frozenset({"model_max_budget"})
|
||||
|
||||
|
||||
def _prisma_value(column: str, value: object) -> object:
|
||||
if column in _JSON_BUDGET_COLUMNS:
|
||||
return json.dumps(value if value is not None else {})
|
||||
return list(value) if isinstance(value, tuple) else value
|
||||
|
||||
|
||||
def _prisma_payload(data: Mapping[str, object]) -> dict[str, object]:
|
||||
return {column: _prisma_value(column, value) for column, value in data.items()}
|
||||
|
||||
|
||||
def member_budget_patch(source: BaseModel) -> Mapping[str, object]:
|
||||
"""Map the per-member limit fields a request actually set to their budget-table
|
||||
columns (merge-patch: a sent value updates, an explicit null clears, an absent
|
||||
field is left untouched)."""
|
||||
provided: Final = source.model_dump(exclude_unset=True)
|
||||
return {
|
||||
column: _prisma_value(provided[request_field])
|
||||
column: provided[request_field]
|
||||
for request_field, column in MEMBER_BUDGET_PATCH_FIELDS.items()
|
||||
if request_field in provided
|
||||
}
|
||||
|
|
@ -603,7 +614,7 @@ def member_budget_patch(source: BaseModel) -> Mapping[str, object]:
|
|||
def _is_set_budget_value(value: object) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
if isinstance(value, list) and len(value) == 0:
|
||||
if isinstance(value, (list, dict)) and len(value) == 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
|
@ -645,10 +656,26 @@ async def upsert_budget_and_membership(
|
|||
A patch that only touches the temporary budget pair never copies permanent
|
||||
limits into a new row, so the member keeps inheriting the live team default.
|
||||
"""
|
||||
if not budget_patch:
|
||||
ignore_default_model_budget_clear: Final = (
|
||||
existing_budget_id is not None
|
||||
and existing_budget_id == team_default_budget_id
|
||||
and "model_max_budget" in budget_patch
|
||||
and (
|
||||
budget_patch["model_max_budget"] is None
|
||||
or (isinstance(budget_patch["model_max_budget"], Mapping) and not budget_patch["model_max_budget"])
|
||||
)
|
||||
)
|
||||
effective_patch: Final = MappingProxyType(
|
||||
{
|
||||
field: value
|
||||
for field, value in budget_patch.items()
|
||||
if not (ignore_default_model_budget_clear and field == "model_max_budget")
|
||||
}
|
||||
)
|
||||
if not effective_patch:
|
||||
return
|
||||
|
||||
write_data: Final = dict(budget_patch)
|
||||
write_data: Final = dict(effective_patch)
|
||||
if "budget_duration" in write_data:
|
||||
duration: Final = write_data["budget_duration"]
|
||||
write_data["budget_reset_at"] = (
|
||||
|
|
@ -675,7 +702,7 @@ async def upsert_budget_and_membership(
|
|||
return
|
||||
await tx.litellm_budgettable.update(
|
||||
where={"budget_id": existing_budget_id},
|
||||
data={"updated_by": user_api_key_dict.user_id or "", **write_data},
|
||||
data={"updated_by": user_api_key_dict.user_id or "", **_prisma_payload(write_data)},
|
||||
)
|
||||
return
|
||||
|
||||
|
|
@ -689,14 +716,24 @@ async def upsert_budget_and_membership(
|
|||
create_data: Final[dict[str, object]] = { # mutable-ok: Prisma create payloads are dict-shaped
|
||||
"created_by": user_api_key_dict.user_id or "",
|
||||
"updated_by": user_api_key_dict.user_id or "",
|
||||
"model_max_budget": {},
|
||||
**MappingProxyType(
|
||||
{f: source[f] for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS if _is_set_budget_value(source.get(f))}
|
||||
{
|
||||
f: source[f]
|
||||
for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS
|
||||
if (
|
||||
f != "model_max_budget"
|
||||
or team_default_budget_id is None
|
||||
or source.get("budget_id") != team_default_budget_id
|
||||
)
|
||||
and _is_set_budget_value(source.get(f))
|
||||
}
|
||||
),
|
||||
**write_data,
|
||||
}
|
||||
|
||||
# Restarting an inherited window on an unrelated edit hands the member a free period.
|
||||
carried: Final = source.get("budget_reset_at") if "budget_duration" not in budget_patch else None
|
||||
carried: Final = source.get("budget_reset_at") if "budget_duration" not in effective_patch else None
|
||||
if carried is not None:
|
||||
create_data["budget_reset_at"] = carried
|
||||
if create_data.get("budget_reset_at") is None:
|
||||
|
|
@ -708,7 +745,7 @@ async def upsert_budget_and_membership(
|
|||
return
|
||||
|
||||
new_budget: Final = await tx.litellm_budgettable.create(
|
||||
data=create_data,
|
||||
data=_prisma_payload(create_data),
|
||||
include={"team_membership": True},
|
||||
)
|
||||
await tx.litellm_teammembership.upsert(
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from pydantic import BaseModel, Field, JsonValue, TypeAdapter, ValidationError
|
|||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
|
@ -228,7 +229,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
|
|||
TeamUserSpendRow,
|
||||
UpdateTeamMemberPermissionsRequest,
|
||||
)
|
||||
from litellm.types.utils import BudgetConfig
|
||||
from litellm.types.utils import BudgetConfig, GenericBudgetConfigType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
|
|
@ -547,6 +548,7 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit: int | None = None,
|
||||
team_member_tpm_limit: int | None = None,
|
||||
team_member_budget_duration: str | None = None,
|
||||
team_member_model_max_budget: GenericBudgetConfigType | None = None,
|
||||
) -> bool:
|
||||
"""Check if any team member limits are provided"""
|
||||
return any(
|
||||
|
|
@ -555,9 +557,22 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit is not None,
|
||||
team_member_tpm_limit is not None,
|
||||
team_member_budget_duration is not None,
|
||||
team_member_model_max_budget is not None,
|
||||
]
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@with_service_target("team_member_budgets")
|
||||
async def _evict_default_budget_cache(
|
||||
budget_id: str,
|
||||
user_api_key_cache: UserApiKeyCache | None,
|
||||
) -> None:
|
||||
if user_api_key_cache is not None:
|
||||
try:
|
||||
await user_api_key_cache.async_delete_cache(key=f"team_member_default_budget:{budget_id}")
|
||||
except Exception: # noqa: BLE001 # Cache failure must not fail a committed budget update.
|
||||
verbose_proxy_logger.warning("Failed to evict the team member default budget cache", exc_info=True)
|
||||
|
||||
@staticmethod
|
||||
async def create_team_member_budget_table(
|
||||
data: NewTeamRequest | _TeamAliasBudgetRow,
|
||||
|
|
@ -567,6 +582,8 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit: int | None = None,
|
||||
team_member_tpm_limit: int | None = None,
|
||||
team_member_budget_duration: str | None = None,
|
||||
team_member_model_max_budget: GenericBudgetConfigType | None = None,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
explicitly_set_fields: AbstractSet[str] = frozenset(),
|
||||
) -> dict:
|
||||
"""Create team member budget table with provided limits.
|
||||
|
|
@ -602,11 +619,17 @@ class TeamMemberBudgetHandler:
|
|||
budget_request.tpm_limit = team_member_tpm_limit
|
||||
if team_member_budget_duration is not None:
|
||||
budget_request.budget_duration = team_member_budget_duration
|
||||
if team_member_model_max_budget is not None:
|
||||
budget_request.model_max_budget = team_member_model_max_budget
|
||||
|
||||
team_member_budget_table: Final = await _as_budget_write(new_budget)(
|
||||
budget_obj=budget_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
await TeamMemberBudgetHandler._evict_default_budget_cache(
|
||||
budget_id=team_member_budget_table.budget_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
# Add team_member_budget_id as metadata field to team table
|
||||
if new_team_data_json.get("metadata") is None:
|
||||
|
|
@ -627,6 +650,8 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit: int | None = None,
|
||||
team_member_tpm_limit: int | None = None,
|
||||
team_member_budget_duration: str | None = None,
|
||||
team_member_model_max_budget: GenericBudgetConfigType | None = None,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
explicitly_set_fields: AbstractSet[str] = frozenset(),
|
||||
) -> dict:
|
||||
"""Upsert team member budget table with provided limits.
|
||||
|
|
@ -657,11 +682,19 @@ class TeamMemberBudgetHandler:
|
|||
budget_request.budget_duration = team_member_budget_duration
|
||||
if team_member_budget_duration is None:
|
||||
budget_request.budget_reset_at = None
|
||||
if team_member_model_max_budget is not None or "team_member_model_max_budget" in explicitly_set_fields:
|
||||
budget_request.model_max_budget = (
|
||||
{} if team_member_model_max_budget is None else team_member_model_max_budget
|
||||
)
|
||||
|
||||
budget_row: Final = await _as_budget_write(update_budget)(
|
||||
budget_obj=budget_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
await TeamMemberBudgetHandler._evict_default_budget_cache(
|
||||
budget_id=team_member_budget_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"Updated team member budget table: %s, with team_member_budget=%s, team_member_rpm_limit=%s, team_member_tpm_limit=%s",
|
||||
budget_row.budget_id,
|
||||
|
|
@ -682,6 +715,8 @@ class TeamMemberBudgetHandler:
|
|||
team_member_rpm_limit=team_member_rpm_limit,
|
||||
team_member_tpm_limit=team_member_tpm_limit,
|
||||
team_member_budget_duration=team_member_budget_duration,
|
||||
team_member_model_max_budget=team_member_model_max_budget,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
explicitly_set_fields=explicitly_set_fields,
|
||||
)
|
||||
|
||||
|
|
@ -696,6 +731,9 @@ class TeamMemberBudgetHandler:
|
|||
data_dict.pop("team_member_budget_duration", None)
|
||||
data_dict.pop("team_member_rpm_limit", None)
|
||||
data_dict.pop("team_member_tpm_limit", None)
|
||||
data_dict.pop( # pyright: ignore[reportUnknownMemberType] # legacy raw dict lacks key typing
|
||||
"team_member_model_max_budget", None
|
||||
)
|
||||
|
||||
_clean_team_member_fields = clean_team_member_fields
|
||||
|
||||
|
|
@ -705,6 +743,7 @@ class TeamMemberBudgetHandler:
|
|||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
updated_kv: dict,
|
||||
explicitly_set_fields: set,
|
||||
user_api_key_cache: UserApiKeyCache | None = None,
|
||||
) -> dict:
|
||||
"""Clear explicitly-nulled fields on the team member budget row."""
|
||||
from litellm.proxy._types import BudgetNewRequest
|
||||
|
|
@ -727,10 +766,16 @@ class TeamMemberBudgetHandler:
|
|||
budget_request.rpm_limit = None
|
||||
if "team_member_tpm_limit" in explicitly_set_fields:
|
||||
budget_request.tpm_limit = None
|
||||
if "team_member_model_max_budget" in explicitly_set_fields:
|
||||
budget_request.model_max_budget = {}
|
||||
await update_budget(
|
||||
budget_obj=budget_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
await TeamMemberBudgetHandler._evict_default_budget_cache(
|
||||
budget_id=team_member_budget_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
TeamMemberBudgetHandler.clean_team_member_fields(updated_kv)
|
||||
return updated_kv
|
||||
|
|
@ -1420,6 +1465,7 @@ async def new_team(
|
|||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
|
||||
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
|
||||
- team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
|
||||
- team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
|
||||
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
|
||||
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
|
||||
|
|
@ -1508,6 +1554,10 @@ async def new_team(
|
|||
validate_budget_duration(data.budget_duration)
|
||||
validate_budget_duration(data.team_member_budget_duration)
|
||||
validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user)
|
||||
validate_team_model_max_budget(
|
||||
model_max_budget=data.team_member_model_max_budget,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
if data.soft_budget is not None:
|
||||
if data.max_budget is not None:
|
||||
|
|
@ -1702,6 +1752,7 @@ async def new_team(
|
|||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
team_member_model_max_budget=data.team_member_model_max_budget,
|
||||
):
|
||||
data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
|
||||
data=data,
|
||||
|
|
@ -1711,6 +1762,8 @@ async def new_team(
|
|||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
team_member_model_max_budget=data.team_member_model_max_budget,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
explicitly_set_fields=data.model_fields_set,
|
||||
)
|
||||
|
||||
|
|
@ -2190,6 +2243,7 @@ async def update_team(
|
|||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
|
||||
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
|
||||
- team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
|
||||
- team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
|
||||
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
|
||||
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
|
||||
|
|
@ -2276,6 +2330,10 @@ async def update_team(
|
|||
validate_budget_duration(data.budget_duration)
|
||||
validate_budget_duration(data.team_member_budget_duration)
|
||||
validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user)
|
||||
validate_team_model_max_budget(
|
||||
model_max_budget=data.team_member_model_max_budget,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique(
|
||||
where={"team_id": data.team_id}
|
||||
|
|
@ -2495,6 +2553,7 @@ async def update_team(
|
|||
"team_member_rpm_limit",
|
||||
"team_member_tpm_limit",
|
||||
"team_member_budget_duration",
|
||||
"team_member_model_max_budget",
|
||||
]
|
||||
if field in updated_kv
|
||||
}
|
||||
|
|
@ -2524,6 +2583,7 @@ async def update_team(
|
|||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
team_member_model_max_budget=data.team_member_model_max_budget,
|
||||
):
|
||||
updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table(
|
||||
team_table=existing_team_row,
|
||||
|
|
@ -2533,6 +2593,8 @@ async def update_team(
|
|||
team_member_rpm_limit=data.team_member_rpm_limit,
|
||||
team_member_tpm_limit=data.team_member_tpm_limit,
|
||||
team_member_budget_duration=data.team_member_budget_duration,
|
||||
team_member_model_max_budget=data.team_member_model_max_budget,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
explicitly_set_fields=_team_member_fields_in_request,
|
||||
)
|
||||
# Backfill team_memberships for members who joined before the
|
||||
|
|
@ -2556,6 +2618,7 @@ async def update_team(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
updated_kv=updated_kv,
|
||||
explicitly_set_fields=_team_member_fields_in_request,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
else:
|
||||
TeamMemberBudgetHandler.clean_team_member_fields(updated_kv)
|
||||
|
|
@ -3866,6 +3929,7 @@ async def team_member_update(
|
|||
)
|
||||
|
||||
validate_budget_duration(data.budget_duration)
|
||||
validate_team_model_max_budget(data.model_max_budget, premium_user=premium_user)
|
||||
|
||||
_existing_team_row: Final = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
|
||||
|
||||
|
|
@ -3982,6 +4046,7 @@ async def team_member_update(
|
|||
allowed_models=data.allowed_models,
|
||||
temp_budget_increase=data.temp_budget_increase,
|
||||
temp_budget_expiry=data.temp_budget_expiry,
|
||||
model_max_budget=data.model_max_budget,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -677,6 +677,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
_metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget
|
||||
_metadata["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget
|
||||
_metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget
|
||||
_metadata["user_api_key_team_member_model_max_budget"] = user_api_key_dict.team_member_model_max_budget
|
||||
_metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget
|
||||
else:
|
||||
for field in (
|
||||
|
|
|
|||
|
|
@ -359,7 +359,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
log_once_if_budget_reservation_disabled,
|
||||
warn_once_if_custom_auth_skips_common_checks,
|
||||
)
|
||||
from litellm.proxy.auth.fallback_budget import router_fallback_budget_check
|
||||
from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck, router_fallback_budget_check
|
||||
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
|
||||
|
|
@ -7196,7 +7196,10 @@ class ProxyConfig:
|
|||
),
|
||||
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
fallback_budget_check=router_fallback_budget_check,
|
||||
fallback_budget_check=RouterFallbackBudgetCheck(
|
||||
is_enforced=router_fallback_budget_check.is_enforced,
|
||||
team_member_model_budget_limiter=model_max_budget_limiter,
|
||||
),
|
||||
auto_router_capability_limit=_license_check.auto_router_capability_limit,
|
||||
)
|
||||
|
||||
|
|
@ -7691,7 +7694,10 @@ class ProxyConfig:
|
|||
search_tools=search_tools,
|
||||
ignore_invalid_deployments=True,
|
||||
fallback_access_check=router_fallback_access_check,
|
||||
fallback_budget_check=router_fallback_budget_check,
|
||||
fallback_budget_check=RouterFallbackBudgetCheck(
|
||||
is_enforced=router_fallback_budget_check.is_enforced,
|
||||
team_member_model_budget_limiter=model_max_budget_limiter,
|
||||
),
|
||||
auto_router_capability_limit=_license_check.auto_router_capability_limit,
|
||||
)
|
||||
verbose_proxy_logger.debug("updated llm_router: %s", llm_router)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ from collections.abc import Mapping, Sequence
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeVar
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.llms.openai import OpenAIFileObject
|
||||
|
||||
|
|
@ -19,6 +21,7 @@ from litellm.types.utils import LlmProviders
|
|||
_V = TypeVar("_V")
|
||||
|
||||
ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors"})
|
||||
_OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object])
|
||||
|
||||
|
||||
def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]:
|
||||
|
|
@ -72,6 +75,30 @@ def resolve_model_group_alias(model_group_alias: object, model: str) -> str | No
|
|||
return target
|
||||
|
||||
|
||||
def resolve_served_model(
|
||||
model: str,
|
||||
model_alias_map: object,
|
||||
key_aliases: object,
|
||||
model_group_alias: object,
|
||||
) -> str:
|
||||
"""Model the request preparation serves for `model`.
|
||||
|
||||
Apply the global alias map, then key aliases, then `router_settings.model_group_alias`.
|
||||
"""
|
||||
global_aliases: Final = (
|
||||
_OBJECT_MAPPING_ADAPTER.validate_python(model_alias_map) if isinstance(model_alias_map, Mapping) else None
|
||||
)
|
||||
global_target: Final = global_aliases.get(model) if global_aliases is not None else None
|
||||
after_global: Final = global_target if isinstance(global_target, str) else model
|
||||
key_alias_mapping: Final = (
|
||||
_OBJECT_MAPPING_ADAPTER.validate_python(key_aliases) if isinstance(key_aliases, Mapping) else None
|
||||
)
|
||||
key_target: Final = key_alias_mapping.get(after_global) if key_alias_mapping is not None else None
|
||||
after_key: Final = key_target if isinstance(key_target, str) else after_global
|
||||
router_target: Final = resolve_model_group_alias(model_group_alias, after_key)
|
||||
return router_target if router_target is not None else after_key
|
||||
|
||||
|
||||
def truncate_fallback_error_detail(detail: str) -> str:
|
||||
"""
|
||||
Bound a fallback failure detail before it is logged or appended to an exception message.
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ Coverage:
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List
|
||||
from typing import Any, Dict, Final, List
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -1198,6 +1198,7 @@ def _fake_user_api_key_auth(
|
|||
end_user_model_max_budget=None,
|
||||
end_user_id=None,
|
||||
user_model_max_budget=None,
|
||||
team_member_model_max_budget=None,
|
||||
user_id=None,
|
||||
token=None,
|
||||
):
|
||||
|
|
@ -1219,6 +1220,7 @@ def _fake_user_api_key_auth(
|
|||
auth.end_user_model_max_budget = end_user_model_max_budget
|
||||
auth.end_user_id = end_user_id
|
||||
auth.user_model_max_budget = user_model_max_budget
|
||||
auth.team_member_model_max_budget = team_member_model_max_budget
|
||||
auth.user_id = user_id
|
||||
auth.token = token
|
||||
return auth
|
||||
|
|
@ -1659,6 +1661,120 @@ async def test_summary_model_denied_when_user_over_model_budget():
|
|||
assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts"
|
||||
|
||||
|
||||
async def test_summary_model_denied_when_team_member_over_model_budget(monkeypatch: pytest.MonkeyPatch):
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.llms.anthropic.pass_through.context_management.constants import COMPACT_SUMMARY_MODEL_SETTING_KEY
|
||||
from litellm.models.team_membership import LiteLLM_TeamMembership
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import get_team_member_default_budget
|
||||
from litellm.proxy.auth.user_api_key_auth import _load_team_member_default_model_budget
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
budget_id: Final = "summary-budget"
|
||||
member_budget: Final = {"claude-haiku-4-5": {"max_budget": 0.0, "budget_duration": "1d"}}
|
||||
user_id: Final = "member-over-budget"
|
||||
team_id: Final = "team-1"
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
prisma_client: Final = MagicMock()
|
||||
team_object: Final = LiteLLM_TeamTableCachedObj(
|
||||
team_id=team_id,
|
||||
models=["all-proxy-models"],
|
||||
metadata={"team_member_budget_id": budget_id},
|
||||
)
|
||||
membership: Final = LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
budget_id=budget_id,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(budget_id=budget_id, model_max_budget={}),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
|
||||
value=membership,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=f"team_member_default_budget:{budget_id}",
|
||||
value=LiteLLM_BudgetTable(budget_id=budget_id, model_max_budget=member_budget),
|
||||
model_type=LiteLLM_BudgetTable,
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=f"team_id:{team_id}",
|
||||
value=team_object,
|
||||
model_type=LiteLLM_TeamTableCachedObj,
|
||||
)
|
||||
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
token="hashed-token",
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
models=["all-proxy-models"],
|
||||
)
|
||||
await _load_team_member_default_model_budget(
|
||||
user_api_key_auth_obj=auth,
|
||||
team_object=team_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
skip_budget_checks=False,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
assert auth.team_member_model_max_budget == member_budget
|
||||
assert (
|
||||
await get_team_member_default_budget(
|
||||
budget_id=budget_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
raise_on_lookup_error=True,
|
||||
)
|
||||
).model_max_budget == member_budget
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{"model_name": MODEL, "litellm_params": {"model": MODEL}},
|
||||
{
|
||||
"model_name": "claude-haiku-4-5",
|
||||
"litellm_params": {
|
||||
"model": "claude-haiku-4-5",
|
||||
"mock_response": "<summary>x</summary>",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "model_max_budget_limiter", limiter)
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setitem(general_settings, COMPACT_SUMMARY_MODEL_SETTING_KEY, "claude-haiku-4-5")
|
||||
messages: Final = [{"role": "user", "content": "hello " * 60_000}]
|
||||
|
||||
result: Final = await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec={"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}},
|
||||
user_api_key_auth=auth,
|
||||
llm_router=router,
|
||||
)
|
||||
|
||||
assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded"
|
||||
assert result.compaction_block is None
|
||||
|
||||
|
||||
async def test_summary_model_denied_when_end_user_over_model_budget():
|
||||
"""End-user per-model budget is enforced for the summary subrequest too."""
|
||||
import litellm
|
||||
|
|
@ -2022,9 +2138,7 @@ async def test_summary_model_denied_when_team_over_model_budget():
|
|||
limiter = MagicMock()
|
||||
limiter.is_key_within_model_budget = AsyncMock(return_value=True)
|
||||
limiter.is_team_within_model_budget = AsyncMock(
|
||||
side_effect=litellm.BudgetExceededError(
|
||||
message="over budget", current_cost=10, max_budget=5
|
||||
)
|
||||
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
|
||||
)
|
||||
|
||||
with (
|
||||
|
|
@ -2064,9 +2178,7 @@ async def test_summary_model_denied_when_team_over_model_budget():
|
|||
PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
)
|
||||
|
||||
real_params = inspect.signature(
|
||||
PROXY_VirtualKeyModelMaxBudgetLimiter.is_team_within_model_budget
|
||||
).parameters
|
||||
real_params = inspect.signature(PROXY_VirtualKeyModelMaxBudgetLimiter.is_team_within_model_budget).parameters
|
||||
for kwarg in ("team_id", "team_model_max_budget", "key_model_max_budget", "model"):
|
||||
assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter does not accept"
|
||||
|
||||
|
|
@ -2585,6 +2697,9 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata()
|
|||
"user_api_key": "sk-parent",
|
||||
"user_api_key_team_id": "team-abc",
|
||||
"user_api_key_user_id": "user-xyz",
|
||||
"user_api_key_team_member_model_max_budget": {
|
||||
"claude-haiku-4-5": {"budget_limit": 4.0, "time_period": "1d"}
|
||||
},
|
||||
"litellm_call_id": "call-1",
|
||||
},
|
||||
additional_drop_params=None,
|
||||
|
|
@ -2595,6 +2710,9 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata()
|
|||
assert captured_summary_metadata.get("user_api_key") == "sk-parent"
|
||||
assert captured_summary_metadata.get("user_api_key_team_id") == "team-abc"
|
||||
assert captured_summary_metadata.get("user_api_key_user_id") == "user-xyz"
|
||||
assert captured_summary_metadata.get("user_api_key_team_member_model_max_budget") == {
|
||||
"claude-haiku-4-5": {"budget_limit": 4.0, "time_period": "1d"}
|
||||
}
|
||||
assert captured_summary_metadata.get("litellm_call_id") == "call-1"
|
||||
# Anthropic-shape ``metadata.user_id`` must not leak in as a propagated field.
|
||||
assert "user_id" not in captured_summary_metadata
|
||||
|
|
|
|||
|
|
@ -7703,6 +7703,47 @@ async def test_get_team_member_default_budget_caches_json_safe_payload():
|
|||
mock_prisma_client.db.litellm_budgettable.find_unique.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_team_member_default_budget_reraises_lookup_error_when_requested() -> None:
|
||||
from litellm.proxy.auth.auth_checks import get_team_member_default_budget
|
||||
|
||||
lookup_error: Final = RuntimeError("budget database unavailable")
|
||||
cache: Final = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(side_effect=lookup_error)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
await get_team_member_default_budget(
|
||||
budget_id="default-budget",
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
raise_on_lookup_error=True,
|
||||
)
|
||||
|
||||
assert exc_info.value is lookup_error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_team_member_default_budget_returns_none_on_lookup_error_by_default() -> None:
|
||||
from litellm.proxy.auth.auth_checks import get_team_member_default_budget
|
||||
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
|
||||
side_effect=RuntimeError("budget database unavailable")
|
||||
)
|
||||
cache: Final = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
|
||||
budget: Final = await get_team_member_default_budget(
|
||||
budget_id="default-budget",
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
|
||||
assert budget is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_end_user_object_db_fetch_returns_validated_end_user():
|
||||
from litellm.proxy.auth.auth_checks import get_end_user_object
|
||||
|
|
|
|||
|
|
@ -1,12 +1,23 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.fallback_budget import (
|
||||
RouterFallbackBudgetCheck,
|
||||
is_token_within_budget_for_model,
|
||||
router_fallback_budget_check,
|
||||
)
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
model_budget_spend_cache_key,
|
||||
team_member_budget_entity_id,
|
||||
)
|
||||
|
||||
FREE_MODEL = {
|
||||
"model_name": "free-model",
|
||||
|
|
@ -48,6 +59,41 @@ def _token(**overrides) -> UserAPIKeyAuth:
|
|||
return UserAPIKeyAuth(**fields)
|
||||
|
||||
|
||||
async def _team_member_limiter(
|
||||
spend_by_model: Mapping[str, float] | None = None,
|
||||
) -> tuple[PROXY_VirtualKeyModelMaxBudgetLimiter, MagicMock]:
|
||||
entity_id: Final = team_member_budget_entity_id(user_id="u1", team_id="t1")
|
||||
spend_by_key: Final = {
|
||||
model_budget_spend_cache_key(
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER,
|
||||
entity_id=entity_id,
|
||||
budget_model=model,
|
||||
budget_duration="1d",
|
||||
): spend
|
||||
for model, spend in (spend_by_model or {}).items()
|
||||
}
|
||||
redis_cache: Final = MagicMock()
|
||||
|
||||
async def get_spend(*, key: str) -> object | None:
|
||||
return spend_by_key.get(key)
|
||||
|
||||
redis_cache.async_get_cache = AsyncMock(side_effect=get_spend)
|
||||
return PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=redis_cache)), redis_cache
|
||||
|
||||
|
||||
def _member_budget_spend_keys(models: Sequence[str]) -> tuple[str, ...]:
|
||||
entity_id: Final = team_member_budget_entity_id(user_id="u1", team_id="t1")
|
||||
return tuple(
|
||||
model_budget_spend_cache_key(
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER,
|
||||
entity_id=entity_id,
|
||||
budget_model=model,
|
||||
budget_duration="1d",
|
||||
)
|
||||
for model in models
|
||||
)
|
||||
|
||||
|
||||
ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: True)
|
||||
NOT_ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: False)
|
||||
|
||||
|
|
@ -70,6 +116,201 @@ async def test_paid_target_refused_when_over_user_budget():
|
|||
assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_refuses_an_exhausted_paid_target() -> None:
|
||||
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
|
||||
token: Final = _token(
|
||||
user_id="u1",
|
||||
team_id="t1",
|
||||
team_member_model_max_budget=member_caps,
|
||||
)
|
||||
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 10.0})
|
||||
|
||||
within_budget: Final = await is_token_within_budget_for_model(
|
||||
model="paid-model",
|
||||
valid_token=token,
|
||||
llm_router=_router(),
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
|
||||
assert within_budget is False
|
||||
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
|
||||
_member_budget_spend_keys(("paid-model",))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_fallback_budget_check_resolves_model_group_alias_before_member_cap_check() -> None:
|
||||
expensive_model: Final = {
|
||||
**PAID_MODEL,
|
||||
"model_name": "expensive",
|
||||
"model_info": {"id": "expensive-model-id"},
|
||||
}
|
||||
router: Final = Router(
|
||||
model_list=[FREE_MODEL, expensive_model],
|
||||
model_group_alias={"backup": "expensive"},
|
||||
)
|
||||
member_caps: Final = {
|
||||
"backup": {"max_budget": 100.0, "budget_duration": "1d"},
|
||||
"expensive": {"max_budget": 5.0, "budget_duration": "1d"},
|
||||
}
|
||||
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
|
||||
limiter, redis_cache = await _team_member_limiter(spend_by_model={"backup": 1.0, "expensive": 10.0})
|
||||
budget_check: Final = RouterFallbackBudgetCheck(
|
||||
is_enforced=lambda: True,
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
|
||||
within_budget: Final = await budget_check(
|
||||
model="backup",
|
||||
request_kwargs={"metadata": {"user_api_key_auth": token}},
|
||||
llm_router=router,
|
||||
)
|
||||
|
||||
assert within_budget is False
|
||||
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
|
||||
_member_budget_spend_keys(("backup", "expensive"))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_fallback_budget_check_refuses_an_exhausted_alias_cap() -> None:
|
||||
expensive_model: Final = {
|
||||
**PAID_MODEL,
|
||||
"model_name": "expensive",
|
||||
"model_info": {"id": "expensive-model-id"},
|
||||
}
|
||||
router: Final = Router(
|
||||
model_list=[FREE_MODEL, expensive_model],
|
||||
model_group_alias={"backup": "expensive"},
|
||||
)
|
||||
member_caps: Final = {"backup": {"max_budget": 5.0, "budget_duration": "1d"}}
|
||||
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
|
||||
limiter, redis_cache = await _team_member_limiter(spend_by_model={"backup": 10.0})
|
||||
budget_check: Final = RouterFallbackBudgetCheck(
|
||||
is_enforced=lambda: True,
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
|
||||
within_budget: Final = await budget_check(
|
||||
model="backup",
|
||||
request_kwargs={"metadata": {"user_api_key_auth": token}},
|
||||
llm_router=router,
|
||||
)
|
||||
|
||||
assert within_budget is False
|
||||
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
|
||||
_member_budget_spend_keys(("backup",))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_fallback_budget_check_allows_alias_when_both_names_are_within_cap() -> None:
|
||||
expensive_model: Final = {
|
||||
**PAID_MODEL,
|
||||
"model_name": "expensive",
|
||||
"model_info": {"id": "expensive-model-id"},
|
||||
}
|
||||
router: Final = Router(
|
||||
model_list=[FREE_MODEL, expensive_model],
|
||||
model_group_alias={"backup": "expensive"},
|
||||
)
|
||||
member_caps: Final = {
|
||||
"backup": {"max_budget": 5.0, "budget_duration": "1d"},
|
||||
"expensive": {"max_budget": 5.0, "budget_duration": "1d"},
|
||||
}
|
||||
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
|
||||
limiter, redis_cache = await _team_member_limiter(spend_by_model={"backup": 1.0, "expensive": 1.0})
|
||||
budget_check: Final = RouterFallbackBudgetCheck(
|
||||
is_enforced=lambda: True,
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
|
||||
within_budget: Final = await budget_check(
|
||||
model="backup",
|
||||
request_kwargs={"metadata": {"user_api_key_auth": token}},
|
||||
llm_router=router,
|
||||
)
|
||||
|
||||
assert within_budget is True
|
||||
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
|
||||
_member_budget_spend_keys(("backup", "expensive"))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_fallback_budget_check_deduplicates_non_alias_target() -> None:
|
||||
router: Final = _router()
|
||||
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
|
||||
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
|
||||
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 1.0})
|
||||
budget_check: Final = RouterFallbackBudgetCheck(
|
||||
is_enforced=lambda: True,
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
|
||||
within_budget: Final = await budget_check(
|
||||
model="paid-model",
|
||||
request_kwargs={"metadata": {"user_api_key_auth": token}},
|
||||
llm_router=router,
|
||||
)
|
||||
|
||||
assert within_budget is True
|
||||
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
|
||||
_member_budget_spend_keys(("paid-model",))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_allows_a_paid_target_when_under_cap() -> None:
|
||||
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
|
||||
token: Final = _token(
|
||||
user_id="u1",
|
||||
team_id="t1",
|
||||
team_member_model_max_budget=member_caps,
|
||||
)
|
||||
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 1.0})
|
||||
|
||||
within_budget: Final = await is_token_within_budget_for_model(
|
||||
model="paid-model",
|
||||
valid_token=token,
|
||||
llm_router=_router(),
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
|
||||
assert within_budget is True
|
||||
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
|
||||
_member_budget_spend_keys(("paid-model",))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_caps_are_not_evaluated_without_a_limiter() -> None:
|
||||
token: Final = _token(
|
||||
user_id="u1",
|
||||
team_id="t1",
|
||||
team_member_model_max_budget={"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}},
|
||||
)
|
||||
|
||||
assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_limiter_is_not_called_without_member_caps() -> None:
|
||||
token: Final = _token(user_id="u1", team_id="t1")
|
||||
limiter, redis_cache = await _team_member_limiter()
|
||||
|
||||
within_budget: Final = await is_token_within_budget_for_model(
|
||||
model="paid-model",
|
||||
valid_token=token,
|
||||
llm_router=_router(),
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
|
||||
assert within_budget is True
|
||||
assert redis_cache.async_get_cache.call_args_list == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_cost_target_allowed_even_when_over_budget():
|
||||
"""Refusing a free target would deny a request on spend some other model accrued."""
|
||||
|
|
@ -123,6 +364,33 @@ async def test_enforced_check_reads_the_key_from_request_metadata(metadata_field
|
|||
assert await ENFORCED(model="paid-model", request_kwargs=under, llm_router=_router()) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_check_passes_member_limiter_from_request_metadata() -> None:
|
||||
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
|
||||
token: Final = _token(
|
||||
user_id="u1",
|
||||
team_id="t1",
|
||||
team_member_model_max_budget=member_caps,
|
||||
)
|
||||
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 10.0})
|
||||
budget_check: Final = RouterFallbackBudgetCheck(
|
||||
is_enforced=lambda: True,
|
||||
team_member_model_budget_limiter=limiter,
|
||||
)
|
||||
request_kwargs: Final = {"metadata": {"user_api_key_auth": token}}
|
||||
|
||||
within_budget: Final = await budget_check(
|
||||
model="paid-model",
|
||||
request_kwargs=request_kwargs,
|
||||
llm_router=_router(),
|
||||
)
|
||||
|
||||
assert within_budget is False
|
||||
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
|
||||
_member_budget_spend_keys(("paid-model",))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_stale_low_counter_still_refuses_a_paid_target(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1,6 +1,8 @@
|
|||
# tests/litellm/proxy/common_utils/test_upsert_budget_membership.py
|
||||
import json
|
||||
import types
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -230,7 +232,11 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user):
|
|||
)
|
||||
|
||||
mock_tx.litellm_budgettable.create.assert_awaited_once()
|
||||
assert stored_budget_row(mock_tx) == {"temp_budget_increase": 5.0, "temp_budget_expiry": expiry}
|
||||
assert stored_budget_row(mock_tx) == {
|
||||
"model_max_budget": "{}",
|
||||
"temp_budget_increase": 5.0,
|
||||
"temp_budget_expiry": expiry,
|
||||
}
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
mock_tx.litellm_teammembership.update.assert_not_called()
|
||||
|
||||
|
|
@ -252,7 +258,11 @@ async def test_create_from_temp_pair_never_snapshots_team_default(mock_tx, fake_
|
|||
)
|
||||
|
||||
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
|
||||
assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}
|
||||
assert stored_budget_row(mock_tx) == {
|
||||
"model_max_budget": "{}",
|
||||
"temp_budget_increase": 1.0,
|
||||
"temp_budget_expiry": expiry,
|
||||
}
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
|
|
@ -274,7 +284,11 @@ async def test_temp_pair_on_shared_default_member_creates_bare_row(mock_tx, fake
|
|||
|
||||
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
|
||||
mock_tx.litellm_budgettable.update.assert_not_called()
|
||||
assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}
|
||||
assert stored_budget_row(mock_tx) == {
|
||||
"model_max_budget": "{}",
|
||||
"temp_budget_increase": 1.0,
|
||||
"temp_budget_expiry": expiry,
|
||||
}
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
|
|
@ -296,6 +310,58 @@ async def test_clearing_temp_pair_on_shared_default_member_is_noop(mock_tx, fake
|
|||
mock_tx.litellm_teammembership.upsert.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clearing_model_budget_to_none_on_shared_default_is_noop(mock_tx, fake_user):
|
||||
await upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-default",
|
||||
user_id="user-on-default",
|
||||
existing_budget_id="team-default-budget-1",
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"model_max_budget": None},
|
||||
team_default_budget_id="team-default-budget-1",
|
||||
)
|
||||
|
||||
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
|
||||
mock_tx.litellm_budgettable.create.assert_not_called()
|
||||
mock_tx.litellm_budgettable.update.assert_not_called()
|
||||
mock_tx.litellm_teammembership.update.assert_not_called()
|
||||
mock_tx.litellm_teammembership.upsert.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clearing_model_budget_to_none_with_tpm_clones_shared_default(mock_tx, fake_user):
|
||||
shared_default_id: Final = "team-default-budget-1"
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(
|
||||
budget_id=shared_default_id,
|
||||
max_budget=100.0,
|
||||
model_max_budget={"claude-sonnet-4-6": {"max_budget": 5.0}},
|
||||
)
|
||||
)
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-default",
|
||||
user_id="user-on-default",
|
||||
existing_budget_id=shared_default_id,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"model_max_budget": None, "tpm_limit": 5},
|
||||
team_default_budget_id=shared_default_id,
|
||||
)
|
||||
|
||||
create_data: Final = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
|
||||
assert create_data == {
|
||||
"created_by": fake_user.user_id,
|
||||
"updated_by": fake_user.user_id,
|
||||
"model_max_budget": "{}",
|
||||
"max_budget": 100.0,
|
||||
"tpm_limit": 5,
|
||||
}
|
||||
mock_tx.litellm_budgettable.update.assert_not_called()
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temp_pair_with_permanent_field_still_clones_shared_default(mock_tx, fake_user):
|
||||
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
|
||||
|
|
@ -337,7 +403,7 @@ async def test_create_from_plain_patch_does_not_snapshot_team_default(mock_tx, f
|
|||
)
|
||||
|
||||
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
|
||||
assert stored_budget_row(mock_tx) == {"tpm_limit": 500}
|
||||
assert stored_budget_row(mock_tx) == {"model_max_budget": "{}", "tpm_limit": 500}
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
|
|
@ -380,6 +446,7 @@ async def test_clone_on_write_from_shared_default(mock_tx, fake_user):
|
|||
assert create_data == {
|
||||
"created_by": fake_user.user_id,
|
||||
"updated_by": fake_user.user_id,
|
||||
"model_max_budget": "{}",
|
||||
"max_budget": 50.0, # caller wins
|
||||
"tpm_limit": 500, # cloned from default
|
||||
"budget_duration": "1d", # cloned from default
|
||||
|
|
@ -401,6 +468,69 @@ async def test_clone_on_write_from_shared_default(mock_tx, fake_user):
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clone_on_write_preserves_model_caps_from_shared_nondefault_row(mock_tx, fake_user):
|
||||
shared_budget_id: Final = "shared-member-budget"
|
||||
model_caps: Final = {"gpt-4o": {"max_budget": 4.0, "budget_duration": "1d"}}
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(
|
||||
budget_id=shared_budget_id,
|
||||
max_budget=100.0,
|
||||
model_max_budget=model_caps,
|
||||
)
|
||||
)
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-shared",
|
||||
user_id="user-shared",
|
||||
existing_budget_id=shared_budget_id,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"tpm_limit": 1000},
|
||||
team_default_budget_id="team-default-budget",
|
||||
shared_budget_ids=frozenset({shared_budget_id}),
|
||||
)
|
||||
|
||||
create_data: Final = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
|
||||
persisted_model_caps: Final = create_data["model_max_budget"]
|
||||
assert (
|
||||
json.loads(persisted_model_caps) if isinstance(persisted_model_caps, str) else persisted_model_caps
|
||||
) == model_caps
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clearing_model_budget_to_none_clones_shared_nondefault_row(mock_tx, fake_user):
|
||||
shared_budget_id: Final = "shared-member-budget"
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(
|
||||
budget_id=shared_budget_id,
|
||||
max_budget=100.0,
|
||||
model_max_budget={"gpt-4o": {"max_budget": 4.0}},
|
||||
)
|
||||
)
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-shared",
|
||||
user_id="user-shared",
|
||||
existing_budget_id=shared_budget_id,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"model_max_budget": None},
|
||||
team_default_budget_id="team-default-budget",
|
||||
shared_budget_ids=frozenset({shared_budget_id}),
|
||||
)
|
||||
|
||||
create_data: Final = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
|
||||
assert create_data == {
|
||||
"created_by": fake_user.user_id,
|
||||
"updated_by": fake_user.user_id,
|
||||
"model_max_budget": "{}",
|
||||
"max_budget": 100.0,
|
||||
}
|
||||
mock_tx.litellm_budgettable.update.assert_not_called()
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
# TEST: forking the shared default while clearing its duration must drop the
|
||||
# duration (and not carry a reset time) on the new private budget.
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -431,6 +561,7 @@ async def test_clone_on_write_clears_duration(mock_tx, fake_user):
|
|||
assert create_data == {
|
||||
"created_by": fake_user.user_id,
|
||||
"updated_by": fake_user.user_id,
|
||||
"model_max_budget": "{}",
|
||||
"max_budget": 200.0,
|
||||
"tpm_limit": 500,
|
||||
"budget_duration": None,
|
||||
|
|
@ -460,3 +591,31 @@ async def test_private_budget_updates_in_place(mock_tx, fake_user):
|
|||
)
|
||||
mock_tx.litellm_budgettable.create.assert_not_called()
|
||||
mock_tx.litellm_teammembership.upsert.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clearing_model_budget_to_none_updates_private_row(mock_tx, fake_user):
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(
|
||||
max_budget=10.0,
|
||||
model_max_budget={"gpt-4o": {"max_budget": 4.0}},
|
||||
)
|
||||
)
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-mixed",
|
||||
user_id="user-private",
|
||||
existing_budget_id="private-budget-xyz",
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"model_max_budget": None},
|
||||
team_default_budget_id="team-default-budget-1",
|
||||
)
|
||||
|
||||
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
|
||||
where={"budget_id": "private-budget-xyz"},
|
||||
data={"updated_by": fake_user.user_id, "model_max_budget": "{}"},
|
||||
)
|
||||
mock_tx.litellm_budgettable.create.assert_not_called()
|
||||
mock_tx.litellm_teammembership.update.assert_not_called()
|
||||
mock_tx.litellm_teammembership.upsert.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -6,15 +6,20 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
model_budget_spend_cache_key,
|
||||
team_member_budget_entity_id,
|
||||
)
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.utils import LiteLLMBatch, Usage
|
||||
|
||||
KEY_HASH: Final = "key-hash-batch"
|
||||
USER_ID: Final = "user-batch"
|
||||
TEAM_ID: Final = "team-batch"
|
||||
MODEL_GROUP: Final = "batch-qa-primary"
|
||||
BATCH_COST: Final = 2.925e-05
|
||||
CHAT_COST: Final = 0.001
|
||||
|
|
@ -35,7 +40,13 @@ def _batch(batch_id: str, status: str) -> LiteLLMBatch:
|
|||
)
|
||||
|
||||
|
||||
def _event(call_type: str, response_cost: float) -> dict[str, object]:
|
||||
def _event(
|
||||
call_type: str,
|
||||
response_cost: float,
|
||||
team_member_model_max_budget: Mapping[str, object] | None = None,
|
||||
user_id: str = USER_ID,
|
||||
team_id: str = TEAM_ID,
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"call_type": call_type,
|
||||
"standard_logging_object": {
|
||||
|
|
@ -43,12 +54,21 @@ def _event(call_type: str, response_cost: float) -> dict[str, object]:
|
|||
"response_cost": response_cost,
|
||||
"model": "openai/gpt-5.4-mini",
|
||||
"model_group": MODEL_GROUP,
|
||||
"metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_user_id": USER_ID},
|
||||
"metadata": {
|
||||
"user_api_key_hash": KEY_HASH,
|
||||
"user_api_key_user_id": user_id,
|
||||
"user_api_key_team_id": team_id,
|
||||
},
|
||||
},
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}},
|
||||
"user_api_key_user_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}},
|
||||
**(
|
||||
{"user_api_key_team_member_model_max_budget": team_member_model_max_budget}
|
||||
if team_member_model_max_budget is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
},
|
||||
}
|
||||
|
|
@ -60,9 +80,23 @@ async def _poll(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMBa
|
|||
)
|
||||
|
||||
|
||||
async def _chat(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter) -> None:
|
||||
async def _chat(
|
||||
limiter: PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
team_member_model_max_budget: Mapping[str, object] | None = None,
|
||||
user_id: str = USER_ID,
|
||||
team_id: str = TEAM_ID,
|
||||
) -> None:
|
||||
await limiter.async_log_success_event(
|
||||
_event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None
|
||||
_event(
|
||||
"acompletion",
|
||||
CHAT_COST,
|
||||
team_member_model_max_budget=team_member_model_max_budget,
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
),
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -74,6 +108,87 @@ def _local_spend(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str)
|
|||
return limiter.dual_cache.in_memory_cache.get_cache(key=spend_key) or 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_is_charged_and_enforced_after_success():
|
||||
budget: Final = {MODEL_GROUP: {"max_budget": 0.0005, "budget_duration": "1d"}}
|
||||
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
|
||||
await _chat(limiter, team_member_model_max_budget=budget)
|
||||
|
||||
entity_id: Final = team_member_budget_entity_id(user_id=USER_ID, team_id=TEAM_ID)
|
||||
spend_key: Final = model_budget_spend_cache_key(
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER,
|
||||
entity_id=entity_id,
|
||||
budget_model=MODEL_GROUP,
|
||||
budget_duration="1d",
|
||||
)
|
||||
assert await _spend(limiter, spend_key) == CHAT_COST
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError, match="LiteLLM Team Member"):
|
||||
await limiter.is_team_member_within_model_budget(
|
||||
user_id=USER_ID,
|
||||
team_id=TEAM_ID,
|
||||
team_member_model_max_budget=budget,
|
||||
model=MODEL_GROUP,
|
||||
)
|
||||
|
||||
|
||||
def test_team_member_budget_entity_id_is_unambiguous():
|
||||
first_entity_id: Final = team_member_budget_entity_id(user_id="a:b", team_id="c")
|
||||
second_entity_id: Final = team_member_budget_entity_id(user_id="a", team_id="b:c")
|
||||
|
||||
assert first_entity_id == "3:a:b:c"
|
||||
assert second_entity_id == "1:a:b:c"
|
||||
|
||||
first_spend_key: Final = model_budget_spend_cache_key(
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER,
|
||||
entity_id=first_entity_id,
|
||||
budget_model=MODEL_GROUP,
|
||||
budget_duration="1d",
|
||||
)
|
||||
second_spend_key: Final = model_budget_spend_cache_key(
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER,
|
||||
entity_id=second_entity_id,
|
||||
budget_model=MODEL_GROUP,
|
||||
budget_duration="1d",
|
||||
)
|
||||
assert first_spend_key != second_spend_key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_spend_isolated_by_user_and_team():
|
||||
budget: Final = {MODEL_GROUP: {"max_budget": CHAT_COST * 2, "budget_duration": "1d"}}
|
||||
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
|
||||
await _chat(limiter, team_member_model_max_budget=budget)
|
||||
|
||||
assert await limiter.is_team_member_within_model_budget(
|
||||
user_id="different-user",
|
||||
team_id=TEAM_ID,
|
||||
team_member_model_max_budget=budget,
|
||||
model=MODEL_GROUP,
|
||||
)
|
||||
assert await limiter.is_team_member_within_model_budget(
|
||||
user_id=USER_ID,
|
||||
team_id="different-team",
|
||||
team_member_model_max_budget=budget,
|
||||
model=MODEL_GROUP,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_team_member_model_budget_rejects_first_request():
|
||||
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError, match="LiteLLM Team Member"):
|
||||
await limiter.is_team_member_within_model_budget(
|
||||
user_id=USER_ID,
|
||||
team_id=TEAM_ID,
|
||||
team_member_model_max_budget={MODEL_GROUP: {"max_budget": 0, "budget_duration": "1d"}},
|
||||
model=MODEL_GROUP,
|
||||
)
|
||||
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.seconds = 0.0
|
||||
|
|
|
|||
|
|
@ -85,6 +85,10 @@ def _matches(row: Mapping[str, object], where: Mapping[str, object]) -> bool:
|
|||
return all((wanted := _wanted(where, field)) is not None and row.get(field) in wanted for field in where)
|
||||
|
||||
|
||||
def _as_stored(data: Mapping[str, object]) -> dict[str, object]:
|
||||
return {k: json.loads(v) if k == "model_max_budget" and isinstance(v, str) else v for k, v in data.items()}
|
||||
|
||||
|
||||
class _BudgetTable:
|
||||
def __init__(self, budgets: Sequence[_BudgetRow]) -> None:
|
||||
self.rows: dict[str, _BudgetRow] = {b.budget_id: b for b in budgets}
|
||||
|
|
@ -94,13 +98,13 @@ class _BudgetTable:
|
|||
|
||||
async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> _BudgetRow:
|
||||
row: Final = self.rows[where["budget_id"]]
|
||||
updated: Final = row.model_copy(update=dict(data))
|
||||
updated: Final = row.model_copy(update=_as_stored(data))
|
||||
self.rows[row.budget_id] = updated
|
||||
return updated
|
||||
|
||||
async def create(self, data: Mapping[str, object], include: Mapping[str, bool] | None = None) -> _BudgetRow:
|
||||
budget_id: Final = f"new-budget-{len(self.rows) + 1}"
|
||||
row: Final = _BudgetRow.model_validate({**data, "budget_id": budget_id})
|
||||
row: Final = _BudgetRow.model_validate({**_as_stored(data), "budget_id": budget_id})
|
||||
self.rows[budget_id] = row
|
||||
return row
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ from fastapi.testclient import TestClient
|
|||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, CommonProxyErrors
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
|
||||
|
||||
|
|
@ -102,6 +104,25 @@ async def test_update_budget_rejects_a_duration_that_never_advances(client_and_m
|
|||
mock_table.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_budget_evicts_local_cache_when_redis_delete_fails(
|
||||
client_and_mocks, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
client, _, _ = client_and_mocks
|
||||
cache_key = "team_member_default_budget:budget_456"
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache())
|
||||
cache.set_cache(cache_key, {"budget_id": "budget_456"}, local_only=True)
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis unavailable"))
|
||||
cache.redis_cache = redis_cache
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", cache)
|
||||
|
||||
response = client.post("/budget/update", json={"budget_id": "budget_456", "max_budget": 10.0})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert cache.get_cache(cache_key, local_only=True) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_budget_db_not_connected(client_and_mocks, monkeypatch):
|
||||
client, mock_prisma, mock_table = client_and_mocks
|
||||
|
|
@ -136,6 +157,26 @@ async def test_update_budget_success(client_and_mocks, monkeypatch):
|
|||
assert body["updated_by"] == "test_user"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_budget_evicts_team_member_default_budget_cache(client_and_mocks, monkeypatch):
|
||||
client, _, _ = client_and_mocks
|
||||
cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache())
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", cache)
|
||||
cache_key: Final = "team_member_default_budget:budget_456"
|
||||
await cache.async_set_cache(key=cache_key, value={"budget_id": "budget_456"})
|
||||
assert await cache.async_get_cache(key=cache_key) == {"budget_id": "budget_456"}
|
||||
|
||||
try:
|
||||
response: Final = client.post(
|
||||
"/budget/update",
|
||||
json={"budget_id": "budget_456", "max_budget": 10.0},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert await cache.async_get_cache(key=cache_key) is None
|
||||
finally:
|
||||
await cache.async_delete_cache(key=cache_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_and_update_budget_persist_tpd_limit(client_and_mocks):
|
||||
client, _, mock_table = client_and_mocks
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from pydantic import BaseModel
|
|||
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_OrganizationMembershipTable,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
|
|
@ -35,25 +36,132 @@ from litellm.proxy.management_endpoints.common_utils import (
|
|||
update_metadata_fields,
|
||||
user_api_key_has_admin_view,
|
||||
user_has_admin_privileges,
|
||||
upsert_budget_and_membership,
|
||||
)
|
||||
from litellm.types.utils import BudgetConfig
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"budget_patch",
|
||||
[
|
||||
{"tpm_limit": 200},
|
||||
{"model_max_budget": {}, "tpm_limit": 5},
|
||||
],
|
||||
)
|
||||
async def test_shared_default_clone_keeps_inheriting_updated_team_model_caps(
|
||||
budget_patch: dict[str, object],
|
||||
) -> None:
|
||||
from litellm.proxy.auth.user_api_key_auth import _load_team_member_default_model_budget
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, team_membership_reservation_cache_key
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership
|
||||
|
||||
source_caps = {"claude-sonnet-4-6": {"max_budget": 3.0, "budget_duration": "1d"}}
|
||||
default_caps = {"claude-sonnet-4-6": {"max_budget": 1.0, "budget_duration": "1d"}}
|
||||
budget_table = SimpleNamespace(
|
||||
find_unique=AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
model_dump=MagicMock(
|
||||
return_value={
|
||||
"budget_id": "default-budget",
|
||||
"model_max_budget": source_caps,
|
||||
"tpm_limit": 100,
|
||||
}
|
||||
)
|
||||
)
|
||||
),
|
||||
create=AsyncMock(return_value=SimpleNamespace(budget_id="private-budget")),
|
||||
)
|
||||
membership_table = SimpleNamespace(upsert=AsyncMock())
|
||||
tx = SimpleNamespace(litellm_budgettable=budget_table, litellm_teammembership=membership_table)
|
||||
token = UserAPIKeyAuth(user_id="member-1", team_id="team-1")
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
tx,
|
||||
team_id="team-1",
|
||||
user_id="member-1",
|
||||
existing_budget_id="default-budget",
|
||||
user_api_key_dict=token,
|
||||
budget_patch=budget_patch,
|
||||
team_default_budget_id="default-budget",
|
||||
)
|
||||
|
||||
budget_table.create.assert_awaited_once()
|
||||
membership_table.upsert.assert_awaited_once()
|
||||
assert budget_table.create.await_args.kwargs["data"]["tpm_limit"] == budget_patch["tpm_limit"]
|
||||
model_max_budget: Final = budget_table.create.await_args.kwargs["data"]["model_max_budget"]
|
||||
assert isinstance(model_max_budget, str)
|
||||
assert model_max_budget == "{}"
|
||||
membership = LiteLLM_TeamMembership(
|
||||
user_id="member-1",
|
||||
team_id="team-1",
|
||||
budget_id="private-budget",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(budget_id="private-budget", model_max_budget={}),
|
||||
)
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=team_membership_reservation_cache_key(user_id="member-1", team_id="team-1"),
|
||||
value=membership,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="team_member_default_budget:default-budget",
|
||||
value=LiteLLM_BudgetTable(budget_id="default-budget", model_max_budget=default_caps),
|
||||
model_type=LiteLLM_BudgetTable,
|
||||
)
|
||||
team_object = SimpleNamespace(
|
||||
team_id="team-1",
|
||||
metadata={"team_member_budget_id": "default-budget"},
|
||||
)
|
||||
|
||||
await _load_team_member_default_model_budget(
|
||||
user_api_key_auth_obj=token,
|
||||
team_object=team_object,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
skip_budget_checks=False,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
assert token.team_member_model_max_budget == default_caps
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_model_budget_patch_on_team_default_keeps_membership_attached() -> None:
|
||||
tx: Final = MagicMock()
|
||||
tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
model_dump=MagicMock(
|
||||
return_value={
|
||||
"budget_id": "default-budget",
|
||||
"max_budget": 10.0,
|
||||
"model_max_budget": {"claude-sonnet-4-6": {"max_budget": 1.0}},
|
||||
}
|
||||
)
|
||||
)
|
||||
)
|
||||
tx.litellm_budgettable.create = AsyncMock()
|
||||
tx.litellm_budgettable.update = AsyncMock()
|
||||
tx.litellm_teammembership.upsert = AsyncMock()
|
||||
tx.litellm_teammembership.update = AsyncMock()
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
tx,
|
||||
team_id="team-1",
|
||||
user_id="member-1",
|
||||
existing_budget_id="default-budget",
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
|
||||
budget_patch={"model_max_budget": {}},
|
||||
team_default_budget_id="default-budget",
|
||||
)
|
||||
|
||||
tx.litellm_budgettable.create.assert_not_awaited()
|
||||
tx.litellm_budgettable.update.assert_not_awaited()
|
||||
tx.litellm_teammembership.upsert.assert_not_awaited()
|
||||
tx.litellm_teammembership.update.assert_not_awaited()
|
||||
|
||||
|
||||
class TestUpdateMetadataFieldsEmptyCollections:
|
||||
"""
|
||||
Regression tests for issue #20304.
|
||||
|
||||
The UI sends empty arrays (`[]`) for enterprise-only fields like
|
||||
guardrails, policies, and logging even when the user hasn't configured
|
||||
these features. The backend must not treat empty collections as an
|
||||
intent to use the feature, and therefore must not trigger the premium
|
||||
license check.
|
||||
|
||||
However, empty collections must still be written into metadata so that
|
||||
users can intentionally clear a previously-set field (e.g. removing all
|
||||
guardrails by sending `guardrails: []`).
|
||||
"""
|
||||
|
||||
@patch("litellm.proxy.management_endpoints.common_utils.premium_user_check")
|
||||
def test_empty_list_does_not_trigger_premium_check(self, mock_premium_check):
|
||||
"""Empty lists for premium fields must not trigger the premium check."""
|
||||
|
|
|
|||
|
|
@ -2591,6 +2591,8 @@ async def test_update_team_team_member_budget_not_passed_to_db(
|
|||
team_member_rpm_limit=None,
|
||||
team_member_tpm_limit=None,
|
||||
team_member_budget_duration=None,
|
||||
team_member_model_max_budget=None,
|
||||
user_api_key_cache=None,
|
||||
explicitly_set_fields=frozenset(),
|
||||
):
|
||||
# Remove team_member_budget from updated_kv as the real function does
|
||||
|
|
@ -3159,6 +3161,8 @@ async def test_update_team_with_team_member_budget_duration(
|
|||
team_member_rpm_limit=None,
|
||||
team_member_tpm_limit=None,
|
||||
team_member_budget_duration=None,
|
||||
team_member_model_max_budget=None,
|
||||
user_api_key_cache=None,
|
||||
explicitly_set_fields=frozenset(),
|
||||
):
|
||||
result_kv = updated_kv.copy()
|
||||
|
|
@ -10891,6 +10895,7 @@ async def test_create_team_member_budget_table_with_duration():
|
|||
team_alias="test-team",
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
team_member_model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}},
|
||||
)
|
||||
|
||||
with patch(
|
||||
|
|
@ -10904,15 +10909,181 @@ async def test_create_team_member_budget_table_with_duration():
|
|||
user_api_key_dict=mock_admin,
|
||||
team_member_budget=20.0,
|
||||
team_member_budget_duration="30d",
|
||||
team_member_model_max_budget=data.team_member_model_max_budget,
|
||||
)
|
||||
|
||||
mock_new_budget.assert_awaited_once()
|
||||
budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
|
||||
assert budget_request.budget_duration == "30d"
|
||||
assert budget_request.max_budget == 20.0
|
||||
assert budget_request.model_max_budget == data.team_member_model_max_budget
|
||||
assert result["metadata"]["team_member_budget_id"] == "budget-abc"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_team_member_model_max_budget_updates_default_budget_and_evicts_cache():
|
||||
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
|
||||
|
||||
model_budget: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}
|
||||
team_table: Final = LiteLLM_TeamTable(
|
||||
team_id="test-team",
|
||||
metadata={"team_member_budget_id": "member-budget-1"},
|
||||
members_with_roles=[],
|
||||
)
|
||||
cache: Final = MagicMock()
|
||||
cache.async_delete_cache = AsyncMock()
|
||||
budget_cache: Final = MagicMock()
|
||||
budget_cache.async_delete_cache = AsyncMock()
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_budgettable.update = AsyncMock(return_value=MagicMock(budget_id="member-budget-1"))
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", budget_cache),
|
||||
patch("litellm.proxy.proxy_server.premium_user", True),
|
||||
):
|
||||
result: Final = await TeamMemberBudgetHandler.upsert_team_member_budget_table(
|
||||
team_table=team_table,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
|
||||
updated_kv={"team_id": "test-team", "team_member_model_max_budget": model_budget},
|
||||
team_member_model_max_budget=model_budget,
|
||||
user_api_key_cache=cache,
|
||||
explicitly_set_fields={"team_member_model_max_budget"},
|
||||
)
|
||||
|
||||
budget_update: Final = prisma.db.litellm_budgettable.update.call_args.kwargs
|
||||
assert budget_update["where"]["budget_id"] == "member-budget-1"
|
||||
assert budget_update["data"]["budget_id"] == "member-budget-1"
|
||||
assert json.loads(budget_update["data"]["model_max_budget"]) == model_budget
|
||||
assert result["metadata"]["team_member_budget_id"] == "member-budget-1"
|
||||
assert "team_member_model_max_budget" not in result
|
||||
cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:member-budget-1")
|
||||
budget_cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:member-budget-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_clears_explicitly_null_member_model_budget_and_evicts_cache(
|
||||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import UpdateTeamRequest
|
||||
from litellm.proxy.management_endpoints.team_endpoints import update_team
|
||||
|
||||
existing: Final = _existing_team_with_model_caps({})
|
||||
existing.metadata = {"team_member_budget_id": "member-budget-1"}
|
||||
existing.members_with_roles = []
|
||||
existing.model_dump.return_value["metadata"] = existing.metadata
|
||||
existing.model_dump.return_value["members_with_roles"] = []
|
||||
cache: Final = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
cache.async_set_cache = AsyncMock()
|
||||
cache.async_delete_cache = AsyncMock()
|
||||
updated: Final = _existing_team_with_model_caps({})
|
||||
updated.metadata = {"team_member_budget_id": "member-budget-1"}
|
||||
updated.model_dump.return_value.update({"metadata": updated.metadata, "members_with_roles": []})
|
||||
updated.litellm_model_table = None
|
||||
update_budget_response: Final = MagicMock(budget_id="member-budget-1")
|
||||
prisma: Final = MagicMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
patch("litellm.proxy.proxy_server.premium_user", True),
|
||||
):
|
||||
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
|
||||
prisma.jsonify_team_object = lambda db_data: db_data
|
||||
prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated)
|
||||
prisma.db.litellm_budgettable.update = AsyncMock(return_value=update_budget_response)
|
||||
|
||||
request_data: Final = UpdateTeamRequest(team_id="standalone-team-123", team_member_model_max_budget=None)
|
||||
await update_team(
|
||||
data=request_data,
|
||||
http_request=MagicMock(spec=Request),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
|
||||
)
|
||||
|
||||
budget_update: Final = prisma.db.litellm_budgettable.update.call_args.kwargs
|
||||
assert budget_update["where"]["budget_id"] == "member-budget-1"
|
||||
assert budget_update["data"]["budget_id"] == "member-budget-1"
|
||||
assert budget_update["data"]["model_max_budget"] == "{}"
|
||||
assert "team_member_model_max_budget" in request_data.model_fields_set
|
||||
assert "model_max_budget" in budget_update["data"]
|
||||
assert cache.async_delete_cache.await_args_list == [
|
||||
call(key="team_member_default_budget:member-budget-1"),
|
||||
call(key="team_member_default_budget:member-budget-1"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_persists_member_model_budget_to_default_budget_row(
|
||||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import UpdateTeamRequest
|
||||
from litellm.proxy.management_endpoints.team_endpoints import update_team
|
||||
|
||||
member_caps: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}
|
||||
existing: Final = _existing_team_with_model_caps({})
|
||||
existing.metadata = {"team_member_budget_id": "member-budget-1"}
|
||||
existing.members_with_roles = []
|
||||
existing.model_dump.return_value["metadata"] = existing.metadata
|
||||
existing.model_dump.return_value["members_with_roles"] = []
|
||||
cache: Final = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
cache.async_set_cache = AsyncMock()
|
||||
cache.async_delete_cache = AsyncMock()
|
||||
updated: Final = _existing_team_with_model_caps({})
|
||||
updated.metadata = {"team_member_budget_id": "member-budget-1"}
|
||||
updated.model_dump.return_value.update({"metadata": updated.metadata, "members_with_roles": []})
|
||||
updated.litellm_model_table = None
|
||||
update_budget_response: Final = MagicMock(budget_id="member-budget-1")
|
||||
prisma: Final = MagicMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
patch("litellm.proxy.proxy_server.premium_user", True),
|
||||
):
|
||||
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
|
||||
prisma.jsonify_team_object = lambda db_data: db_data
|
||||
prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated)
|
||||
prisma.db.litellm_budgettable.update = AsyncMock(return_value=update_budget_response)
|
||||
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="standalone-team-123", team_member_model_max_budget=member_caps),
|
||||
http_request=MagicMock(spec=Request),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
|
||||
)
|
||||
|
||||
budget_update: Final = prisma.db.litellm_budgettable.update.call_args.kwargs
|
||||
assert budget_update["where"]["budget_id"] == "member-budget-1"
|
||||
assert budget_update["data"]["budget_id"] == "member-budget-1"
|
||||
stored_member_model_caps: Final = json.loads(budget_update["data"]["model_max_budget"])
|
||||
assert {
|
||||
model: {
|
||||
**budget,
|
||||
"tpm_limit": budget.get("tpm_limit"),
|
||||
"rpm_limit": budget.get("rpm_limit"),
|
||||
}
|
||||
for model, budget in stored_member_model_caps.items()
|
||||
} == {
|
||||
"gpt-4o": {
|
||||
"max_budget": 5.0,
|
||||
"budget_duration": "1d",
|
||||
"tpm_limit": None,
|
||||
"rpm_limit": None,
|
||||
}
|
||||
}
|
||||
assert cache.async_delete_cache.await_args_list == [
|
||||
call(key="team_member_default_budget:member-budget-1"),
|
||||
call(key="team_member_default_budget:member-budget-1"),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for _batch_resolve_access_group_resources
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -12519,11 +12690,13 @@ async def _drive_team_write(
|
|||
**(existing_kwargs or {}),
|
||||
)
|
||||
auth = user or UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u")
|
||||
user_api_key_cache: Final = MagicMock()
|
||||
user_api_key_cache.async_delete_cache = AsyncMock()
|
||||
|
||||
with (
|
||||
_patch("litellm.proxy.proxy_server.prisma_client") as pc,
|
||||
_patch("litellm.proxy.proxy_server.llm_router", None),
|
||||
_patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||
_patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),
|
||||
_patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
|
||||
_patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
_patch(
|
||||
|
|
@ -16245,13 +16418,24 @@ def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_a
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth):
|
||||
@pytest.mark.parametrize(
|
||||
("member_caps", "member_budget_duration"),
|
||||
[
|
||||
({"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, "30d"),
|
||||
({}, None),
|
||||
],
|
||||
)
|
||||
async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth, member_caps, member_budget_duration):
|
||||
mock_db_client.jsonify_team_object = lambda db_data: db_data
|
||||
mock_db_client.get_data = AsyncMock(return_value=None)
|
||||
mock_db_client.update_data = AsyncMock(return_value=MagicMock())
|
||||
mock_db_client.db = MagicMock()
|
||||
mock_db_client.db.litellm_modeltable = MagicMock()
|
||||
mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123"))
|
||||
mock_db_client.db.litellm_budgettable = MagicMock()
|
||||
mock_db_client.db.litellm_budgettable.create = AsyncMock(
|
||||
side_effect=lambda *, data: MagicMock(budget_id=data["budget_id"])
|
||||
)
|
||||
|
||||
team_create_result = MagicMock(team_id="team-model-caps")
|
||||
team_create_result.model_dump.return_value = {"team_id": "team-model-caps"}
|
||||
|
|
@ -16269,11 +16453,18 @@ async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_aut
|
|||
from litellm.proxy._types import NewTeamRequest
|
||||
from litellm.proxy.management_endpoints.team_endpoints import new_team
|
||||
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
member_budget_cache: Final = MagicMock()
|
||||
member_budget_cache.async_delete_cache = AsyncMock()
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.premium_user", True),
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", member_budget_cache),
|
||||
):
|
||||
await new_team(
|
||||
data=NewTeamRequest(
|
||||
team_alias="model-caps",
|
||||
model_max_budget={"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}},
|
||||
team_member_budget_duration=member_budget_duration,
|
||||
team_member_model_max_budget=member_caps,
|
||||
),
|
||||
http_request=MagicMock(spec=Request),
|
||||
user_api_key_dict=mock_admin_auth,
|
||||
|
|
@ -16283,6 +16474,75 @@ async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_aut
|
|||
assert team_data["model_max_budget"] == {
|
||||
"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d", "tpm_limit": None, "rpm_limit": None}
|
||||
}
|
||||
member_budget_row: Final = mock_db_client.db.litellm_budgettable.create.call_args.kwargs["data"]
|
||||
member_budget_id: Final = member_budget_row["budget_id"]
|
||||
assert team_data["metadata"]["team_member_budget_id"] == member_budget_id
|
||||
assert member_budget_row.get("budget_duration") == member_budget_duration
|
||||
stored_member_budget_model_caps: Final = json.loads(member_budget_row["model_max_budget"])
|
||||
expected_member_model_caps: Final = (
|
||||
{
|
||||
"gpt-4o": {
|
||||
"max_budget": 5.0,
|
||||
"budget_duration": "1d",
|
||||
"tpm_limit": None,
|
||||
"rpm_limit": None,
|
||||
}
|
||||
}
|
||||
if member_caps
|
||||
else {}
|
||||
)
|
||||
assert {
|
||||
model: {
|
||||
**budget,
|
||||
"tpm_limit": budget.get("tpm_limit"),
|
||||
"rpm_limit": budget.get("rpm_limit"),
|
||||
}
|
||||
for model, budget in stored_member_budget_model_caps.items()
|
||||
} == expected_member_model_caps
|
||||
member_budget_cache.async_delete_cache.assert_awaited_once_with(
|
||||
key=f"team_member_default_budget:{member_budget_id}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("endpoint_name", ["new", "update"])
|
||||
@pytest.mark.parametrize(
|
||||
("premium_user", "model_budget", "expected_status"),
|
||||
[
|
||||
(False, {"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}}, 403),
|
||||
(True, {"gpt-4o": {"max_budget": 10.0}}, 400),
|
||||
],
|
||||
)
|
||||
async def test_team_member_model_budget_validation(endpoint_name, premium_user, model_budget, expected_status):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import NewTeamRequest, ProxyException, UpdateTeamRequest
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
new_team,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
update_team as update_team_endpoint,
|
||||
)
|
||||
|
||||
endpoint: Final = new_team if endpoint_name == "new" else update_team_endpoint
|
||||
data: Final = (
|
||||
NewTeamRequest(team_alias="member-model-validation", team_member_model_max_budget=model_budget)
|
||||
if endpoint_name == "new"
|
||||
else UpdateTeamRequest(team_id="member-model-validation", team_member_model_max_budget=model_budget)
|
||||
)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
|
||||
patch("litellm.proxy.proxy_server.premium_user", premium_user),
|
||||
pytest.raises(ProxyException) as exc_info,
|
||||
):
|
||||
await endpoint(
|
||||
data=data,
|
||||
http_request=MagicMock(spec=Request),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
|
||||
)
|
||||
|
||||
assert exc_info.value.code == str(expected_status)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -17098,6 +17358,87 @@ def test_member_budget_patch_maps_temp_budget_fields() -> None:
|
|||
}
|
||||
|
||||
|
||||
def test_member_budget_patch_stores_model_max_budget_as_json_and_clears_on_empty() -> None:
|
||||
from litellm.proxy.management_endpoints.common_utils import member_budget_patch
|
||||
|
||||
request: Final = TeamMemberUpdateRequest(
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
model_max_budget={"gpt-4o": {"budget_limit": 5, "time_period": "30d"}},
|
||||
)
|
||||
assert member_budget_patch(request) == {"model_max_budget": {"gpt-4o": {"max_budget": 5.0, "budget_duration": "30d"}}}
|
||||
|
||||
cleared: Final = TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", model_max_budget={})
|
||||
assert member_budget_patch(cleared) == {"model_max_budget": {}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_budget_and_membership_clears_own_model_max_budget_to_empty_json() -> None:
|
||||
from litellm.proxy.management_endpoints.common_utils import upsert_budget_and_membership
|
||||
|
||||
own_row: Final = SimpleNamespace(
|
||||
model_dump=lambda: {
|
||||
"budget_id": "member-b",
|
||||
"max_budget": 50.0,
|
||||
"model_max_budget": {"gpt-4o": {"max_budget": 2.5, "budget_duration": "7d"}},
|
||||
}
|
||||
)
|
||||
tx: Final = MagicMock()
|
||||
tx.litellm_budgettable.find_unique = AsyncMock(return_value=own_row)
|
||||
tx.litellm_budgettable.update = AsyncMock()
|
||||
tx.litellm_teammembership.update = AsyncMock()
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
tx,
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
existing_budget_id="member-b",
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
|
||||
budget_patch={"model_max_budget": None},
|
||||
team_default_budget_id="default-b",
|
||||
)
|
||||
|
||||
assert tx.litellm_budgettable.update.await_args.kwargs["data"]["model_max_budget"] == "{}"
|
||||
tx.litellm_teammembership.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_budget_and_membership_clone_keeps_shared_default_model_max_budget_as_json() -> None:
|
||||
from litellm.proxy.management_endpoints.common_utils import upsert_budget_and_membership
|
||||
|
||||
default_model_budget: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "30d"}}
|
||||
default_row: Final = SimpleNamespace(
|
||||
model_dump=lambda: {
|
||||
"budget_id": "default-b",
|
||||
"max_budget": 10.0,
|
||||
"model_max_budget": default_model_budget,
|
||||
"budget_reset_at": None,
|
||||
}
|
||||
)
|
||||
tx: Final = MagicMock()
|
||||
tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row)
|
||||
tx.litellm_budgettable.create = AsyncMock(return_value=SimpleNamespace(budget_id="member-b"))
|
||||
tx.litellm_teammembership.upsert = AsyncMock()
|
||||
|
||||
await upsert_budget_and_membership(
|
||||
tx,
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
existing_budget_id="default-b",
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
|
||||
budget_patch={"tpm_limit": 7},
|
||||
team_default_budget_id="default-b",
|
||||
)
|
||||
|
||||
create_data: Final = tx.litellm_budgettable.create.await_args.kwargs["data"]
|
||||
assert create_data["tpm_limit"] == 7
|
||||
assert create_data["max_budget"] == 10.0
|
||||
assert json.loads(create_data["model_max_budget"]) == {}
|
||||
assert tx.litellm_teammembership.upsert.await_args.kwargs["data"]["update"] == {
|
||||
"litellm_budget_table": {"connect": {"budget_id": "member-b"}}
|
||||
}
|
||||
|
||||
|
||||
def test_team_member_update_request_temp_budget_fields_must_be_set_together() -> None:
|
||||
with pytest.raises(ValidationError, match="temp_budget_increase and temp_budget_expiry must be set together"):
|
||||
TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", temp_budget_increase=50.0)
|
||||
|
|
@ -17211,3 +17552,18 @@ def test_aggregated_date_range_error_accepts_canonical_dates_and_keeps_range_che
|
|||
assert aggregated_date_range_error("2026-09-24", "2026-09-26") is None
|
||||
assert aggregated_date_range_error("2026-09-26", "2026-09-24") == "end_date must be on or after start_date"
|
||||
assert aggregated_date_range_error("2020-01-01", "2026-12-31") == "Date range must be at most 400 days"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_budget_cache_eviction_does_not_fail_after_redis_error() -> None:
|
||||
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis unavailable"))
|
||||
|
||||
await TeamMemberBudgetHandler._evict_default_budget_cache(
|
||||
budget_id="budget_456",
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
user_api_key_cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:budget_456")
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
|
|
@ -817,6 +818,10 @@ async def test_add_new_member_with_user_email_links_default_budget():
|
|||
mock_prisma_client.db.litellm_budgettable.create.assert_not_called()
|
||||
|
||||
|
||||
def _as_stored(data: Mapping[str, object]) -> dict[str, object]:
|
||||
return {k: json.loads(v) if k == "model_max_budget" and isinstance(v, str) else v for k, v in data.items()}
|
||||
|
||||
|
||||
class _FakeBudgetTable:
|
||||
def __init__(self) -> None:
|
||||
self.rows: dict[str, dict[str, object]] = {}
|
||||
|
|
@ -829,14 +834,14 @@ class _FakeBudgetTable:
|
|||
self, *, data: Mapping[str, object], include: Mapping[str, bool] | None = None
|
||||
) -> LiteLLM_BudgetTable:
|
||||
budget_id: Final = str(data.get("budget_id") or uuid.uuid4())
|
||||
self.rows[budget_id] = {**data, "budget_id": budget_id}
|
||||
self.rows[budget_id] = {**_as_stored(data), "budget_id": budget_id}
|
||||
return self._record(budget_id)
|
||||
|
||||
async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_BudgetTable | None:
|
||||
return self._record(where["budget_id"]) if where["budget_id"] in self.rows else None
|
||||
|
||||
async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_BudgetTable:
|
||||
self.rows[where["budget_id"]] = {**self.rows[where["budget_id"]], **data}
|
||||
self.rows[where["budget_id"]] = {**self.rows[where["budget_id"]], **_as_stored(data)}
|
||||
return self._record(where["budget_id"])
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -7918,14 +7918,17 @@ def test_passthrough_carries_the_per_model_budgets():
|
|||
key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}
|
||||
user_budget = {"claude-opus-4-8": {"budget_limit": 2.0, "time_period": "1mo"}}
|
||||
end_user_budget = {"claude-opus-4-8": {"budget_limit": 3.0, "time_period": "1d"}}
|
||||
team_member_budget = {"claude-opus-4-8": {"budget_limit": 4.0, "time_period": "1d"}}
|
||||
|
||||
kwargs = _passthrough_kwargs_for_reservation(
|
||||
UserAPIKeyAuth(
|
||||
token="hash",
|
||||
user_id="u-1",
|
||||
team_id="team-1",
|
||||
model_max_budget=key_budget,
|
||||
user_model_max_budget=user_budget,
|
||||
end_user_model_max_budget=end_user_budget,
|
||||
team_member_model_max_budget=team_member_budget,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -7933,6 +7936,7 @@ def test_passthrough_carries_the_per_model_budgets():
|
|||
assert metadata["user_api_key_model_max_budget"] == key_budget
|
||||
assert metadata["user_api_key_user_model_max_budget"] == user_budget
|
||||
assert metadata["user_api_key_end_user_model_max_budget"] == end_user_budget
|
||||
assert metadata["user_api_key_team_member_model_max_budget"] == team_member_budget
|
||||
|
||||
|
||||
def test_passthrough_budget_metadata_cannot_be_forged_by_the_request_body():
|
||||
|
|
@ -7981,7 +7985,10 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata
|
|||
budget: Final = {"managed-model": {"budget_limit": 0.1, "time_period": "1d"}}
|
||||
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(DualCache())
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
api_key="custom-key", token="custom-key", team_id="shared-team", team_model_max_budget=budget,
|
||||
api_key="custom-key",
|
||||
token="custom-key",
|
||||
team_id="shared-team",
|
||||
team_model_max_budget=budget,
|
||||
)
|
||||
endpoint: Final = create_pass_through_route(
|
||||
endpoint="/custom-budget-test",
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
|
|||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.bug_report import (
|
||||
DISABLE_ENV_VAR,
|
||||
ISSUE_URL_BASE,
|
||||
|
|
@ -59,11 +60,18 @@ from litellm.proxy.common_request_processing import (
|
|||
sse_error_payload,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK
|
||||
from litellm.proxy.dd_span_tagger import DDSpanTagger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException
|
||||
from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth
|
||||
from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
model_budget_spend_cache_key,
|
||||
team_member_budget_entity_id,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment
|
||||
|
|
@ -7097,6 +7105,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
fallbacks: list[dict[str, list[str]]],
|
||||
model_guardrails: dict[str, list[str]] | None = None,
|
||||
untagged_rate_limited_models: frozenset[str] = frozenset(),
|
||||
fallback_budget_check: RouterFallbackBudgetCheck | None = None,
|
||||
free_models: frozenset[str] = frozenset(),
|
||||
extra_models: frozenset[str] = frozenset(),
|
||||
) -> tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]]:
|
||||
"""Real v3 limiter (the default ``parallel_request_limiter``) wired in through the
|
||||
``proxy_logging_obj`` seam, so ``common_processing_pre_call_logic`` runs for real:
|
||||
|
|
@ -7145,11 +7156,17 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
"api_key": "fake",
|
||||
**({"guardrails": guardrails_by_group[group]} if group in guardrails_by_group else {}),
|
||||
},
|
||||
**(
|
||||
{"model_info": {"input_cost_per_token": 0, "output_cost_per_token": 0}}
|
||||
if group in free_models
|
||||
else {}
|
||||
),
|
||||
}
|
||||
for chain in fallbacks
|
||||
for group in (*chain.keys(), *(m for models in chain.values() for m in models))
|
||||
for group in (*chain.keys(), *(m for models in chain.values() for m in models), *sorted(extra_models))
|
||||
],
|
||||
fallbacks=fallbacks,
|
||||
fallback_budget_check=fallback_budget_check,
|
||||
)
|
||||
return proxy_logging_obj, router, proxy_server.ProxyConfig(), limiter_models
|
||||
|
||||
|
|
@ -7636,6 +7653,138 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
assert data["metadata"]["guardrails"] == [structured_guardrail]
|
||||
assert rig[3] == [primary_model, primary_model, fallback_model]
|
||||
|
||||
@staticmethod
|
||||
async def _install_member_spend(monkeypatch: pytest.MonkeyPatch, model: str, spend: float) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
member_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
await member_limiter.dual_cache.async_set_cache(
|
||||
key=model_budget_spend_cache_key(
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER,
|
||||
entity_id=team_member_budget_entity_id(user_id="user-1", team_id="team-1"),
|
||||
budget_model=model,
|
||||
budget_duration="1d",
|
||||
),
|
||||
value=spend,
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "model_max_budget_limiter", member_limiter)
|
||||
|
||||
def _member_key(self, primary_model: str, fallback_model: str) -> ProxyUserAPIKeyAuth:
|
||||
return self._otel_key(model_rpm_limit={primary_model: 1}).model_copy(
|
||||
update={
|
||||
"user_id": "user-1",
|
||||
"team_id": "team-1",
|
||||
"team_member_model_max_budget": {fallback_model: {"max_budget": 5.0, "budget_duration": "1d"}},
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _client_request(model: str) -> dict[str, object]:
|
||||
return {"model": model, "messages": [{"role": "user", "content": "hi"}]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_fallback_skips_backup_over_member_model_budget(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
primary_model: Final = "gpt-4.1"
|
||||
fallback_model: Final = "gpt-4.1-mini"
|
||||
await self._install_member_spend(monkeypatch, fallback_model, spend=5.0)
|
||||
key: Final = self._member_key(primary_model, fallback_model)
|
||||
rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
|
||||
|
||||
_, (first_data, _) = await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=self._client_request(primary_model))
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await self._run(processor, key, rig)
|
||||
|
||||
assert first_data["model"] == primary_model
|
||||
assert rig[3] == [primary_model, primary_model]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_fallback_skips_backup_aliased_to_model_over_member_budget(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
primary_model: Final = "gpt-4.1"
|
||||
fallback_model: Final = "gpt-4.1-mini"
|
||||
resolved_model: Final = "expensive"
|
||||
await self._install_member_spend(monkeypatch, resolved_model, spend=5.0)
|
||||
key: Final = self._member_key(primary_model, fallback_model).model_copy(
|
||||
update={
|
||||
"router_settings": {"model_group_alias": {fallback_model: resolved_model}},
|
||||
"team_member_model_max_budget": {resolved_model: {"max_budget": 5.0, "budget_duration": "1d"}},
|
||||
}
|
||||
)
|
||||
rig: Final = self._v3_limiter_rig(
|
||||
monkeypatch,
|
||||
key,
|
||||
[{primary_model: [fallback_model]}],
|
||||
extra_models=frozenset({resolved_model}),
|
||||
)
|
||||
|
||||
await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=self._client_request(primary_model))
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await self._run(processor, key, rig)
|
||||
|
||||
assert processor.data["model"] == primary_model
|
||||
assert rig[3] == [primary_model, primary_model]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_fallback_uses_backup_within_member_model_budget(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
primary_model: Final = "gpt-4.1"
|
||||
fallback_model: Final = "gpt-4.1-mini"
|
||||
await self._install_member_spend(monkeypatch, fallback_model, spend=4.0)
|
||||
key: Final = self._member_key(primary_model, fallback_model)
|
||||
rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
|
||||
|
||||
await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
_, (data, logging_obj) = await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
|
||||
assert data["model"] == fallback_model
|
||||
assert logging_obj.model == fallback_model
|
||||
assert rig[3] == [primary_model, primary_model, fallback_model]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_fallback_uses_free_backup_with_spent_member_model_budget(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
primary_model: Final = "gpt-4.1"
|
||||
fallback_model: Final = "gpt-4.1-mini"
|
||||
await self._install_member_spend(monkeypatch, fallback_model, spend=5.0)
|
||||
key: Final = self._member_key(primary_model, fallback_model)
|
||||
rig: Final = self._v3_limiter_rig(
|
||||
monkeypatch, key, [{primary_model: [fallback_model]}], free_models=frozenset({fallback_model})
|
||||
)
|
||||
|
||||
await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
_, (data, _) = await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
|
||||
assert data["model"] == fallback_model
|
||||
assert rig[3] == [primary_model, primary_model, fallback_model]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_fallback_ignores_key_budget_already_admitted_at_auth(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
primary_model: Final = "gpt-4.1"
|
||||
fallback_model: Final = "gpt-4.1-mini"
|
||||
key: Final = self._otel_key(model_rpm_limit={primary_model: 1}).model_copy(
|
||||
update={"token": "hashed-key", "max_budget": 1.0, "spend": 2.0}
|
||||
)
|
||||
rig: Final = self._v3_limiter_rig(
|
||||
monkeypatch,
|
||||
key,
|
||||
[{primary_model: [fallback_model]}],
|
||||
fallback_budget_check=RouterFallbackBudgetCheck(is_enforced=lambda: True),
|
||||
)
|
||||
|
||||
await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
_, (data, _) = await self._pre_call(self._client_request(primary_model), key, rig)
|
||||
|
||||
assert data["model"] == fallback_model
|
||||
assert rig[3] == [primary_model, primary_model, fallback_model]
|
||||
|
||||
class _RecordingSuccessLogger(CustomLogger):
|
||||
def __init__(self):
|
||||
|
|
|
|||
|
|
@ -2016,6 +2016,8 @@ async def test_add_litellm_data_to_request_user_spend_and_budget():
|
|||
api_key="hashed-key",
|
||||
metadata={},
|
||||
team_metadata={},
|
||||
team_id="team-1",
|
||||
team_member_model_max_budget={"gpt-3.5-turbo": {"max_budget": 15.0, "budget_duration": "1d"}},
|
||||
user_spend=150.0,
|
||||
user_max_budget=500.0,
|
||||
)
|
||||
|
|
@ -2032,6 +2034,9 @@ async def test_add_litellm_data_to_request_user_spend_and_budget():
|
|||
metadata = updated_data.get("metadata", {})
|
||||
assert metadata["user_api_key_user_spend"] == 150.0
|
||||
assert metadata["user_api_key_user_max_budget"] == 500.0
|
||||
assert metadata["user_api_key_team_member_model_max_budget"] == {
|
||||
"gpt-3.5-turbo": {"max_budget": 15.0, "budget_duration": "1d"}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm.router_utils.common_utils import (
|
|||
filter_web_search_deployments,
|
||||
provider_for_generic_call,
|
||||
resolve_model_group_alias,
|
||||
resolve_served_model,
|
||||
truncate_fallback_error_detail,
|
||||
PROVIDER_SCOPED_CREDENTIAL_PARAMS,
|
||||
warn_on_provider_credential_mismatch,
|
||||
|
|
@ -566,6 +567,25 @@ class TestResolveModelGroupAlias:
|
|||
assert router.get_model_from_alias("group-b") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, model_alias_map, key_aliases, model_group_alias, expected",
|
||||
[
|
||||
pytest.param("a", {"a": "b"}, {"b": "c"}, {"c": "d"}, "d", id="global-key-router-order"),
|
||||
pytest.param("a", {}, ["a", "b"], {}, "a", id="non-mapping-key-aliases-ignored"),
|
||||
pytest.param("a", {}, {}, {"a": {"model": "x"}}, "x", id="router-item-entry"),
|
||||
pytest.param("a", {}, {}, {}, "a", id="no-alias"),
|
||||
],
|
||||
)
|
||||
def test_resolves_served_model_aliases_in_order(
|
||||
model: str,
|
||||
model_alias_map: object,
|
||||
key_aliases: object,
|
||||
model_group_alias: object,
|
||||
expected: str,
|
||||
) -> None:
|
||||
assert resolve_served_model(model, model_alias_map, key_aliases, model_group_alias) == expected
|
||||
|
||||
|
||||
class TestTruncateFallbackErrorDetail:
|
||||
def test_short_detail_is_returned_unchanged(self):
|
||||
assert truncate_fallback_error_detail("boom") == "boom"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,83 @@
|
|||
import { fireEvent, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders, testQueryClient } from "../../tests/test-utils";
|
||||
import Teams from "./Teams";
|
||||
|
||||
const sessionCookie = () => {
|
||||
const encode = (value: object) =>
|
||||
btoa(JSON.stringify(value)).replaceAll("=", "").replaceAll("+", "-").replaceAll("/", "_");
|
||||
const claims = {
|
||||
key: "sk-session-test",
|
||||
user_id: "test-admin",
|
||||
user_role: "proxy_admin",
|
||||
premium_user: true,
|
||||
exp: Date.now() / 1000 + 3600,
|
||||
};
|
||||
document.cookie = `token=${encode({ alg: "none" })}.${encode(claims)}.test; Path=/`;
|
||||
};
|
||||
|
||||
const responseFor = (url: string): object | object[] => {
|
||||
if (url.includes("/v2/team/list")) return { teams: [], total: 0, page: 1, page_size: 50, total_pages: 1 };
|
||||
if (url.includes("/models")) return { data: [{ id: "gpt-4" }] };
|
||||
if (url.includes("/v1/access_group")) return [];
|
||||
if (url.includes("/v1/mcp/access_groups")) return { access_groups: [] };
|
||||
if (url.includes("/guardrails/list")) return { guardrails: [] };
|
||||
if (url.includes("/policies/list")) return { policies: [] };
|
||||
if (url.includes("/organization/list")) return { organizations: [] };
|
||||
if (url.includes("/team/new")) return { team_id: "new-team-1" };
|
||||
if (url.includes("/default")) return { values: {} };
|
||||
if (url.includes("/litellm-ui-config")) return { server_root_path: "", proxy_base_url: "" };
|
||||
return {};
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
testQueryClient.clear();
|
||||
sessionCookie();
|
||||
vi.stubGlobal(
|
||||
"fetch",
|
||||
vi.fn(async (input: RequestInfo | URL) => Response.json(responseFor(String(input)))),
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
testQueryClient.clear();
|
||||
vi.unstubAllGlobals();
|
||||
document.cookie = "token=; Max-Age=0; Path=/";
|
||||
});
|
||||
|
||||
describe("Teams creation integration", () => {
|
||||
it("sends a configured default member per-model budget", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
renderWithProviders(<Teams accessToken="sk-session-test" userID="test-admin" userRole="Admin" premiumUser />);
|
||||
|
||||
await user.click(screen.getAllByRole("button", { name: /create team/i })[0]);
|
||||
await screen.findByLabelText(/team name/i);
|
||||
await user.click(screen.getByText("Additional Settings"));
|
||||
await screen.findByText(/Team Member Key Duration/);
|
||||
|
||||
const addBudgetButtons = screen.getAllByRole("button", { name: /Add Model Budget/i });
|
||||
await user.click(addBudgetButtons[1]);
|
||||
const modelSelectors = screen.getAllByPlaceholderText("Select model");
|
||||
await user.click(modelSelectors[0]);
|
||||
await user.click(await screen.findByRole("option", { name: "gpt-4" }));
|
||||
fireEvent.change(screen.getAllByPlaceholderText("Max spend ($)")[0], { target: { value: "4" } });
|
||||
await user.click(screen.getByText("Additional Settings"));
|
||||
|
||||
fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Budget Team" } });
|
||||
await user.click(screen.getAllByRole("button", { name: /create team/i }).at(-1)!);
|
||||
await waitFor(() => {
|
||||
expect(vi.mocked(fetch)).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
const createRequest = vi.mocked(fetch).mock.calls.find(([input]) => String(input).includes("/team/new"));
|
||||
expect(createRequest).toBeDefined();
|
||||
const payload = JSON.parse(String(createRequest?.[1]?.body)) as Record<string, unknown>;
|
||||
expect(payload.team_member_model_max_budget).toStrictEqual({
|
||||
"gpt-4": { budget_limit: 4, time_period: "30d" },
|
||||
});
|
||||
expect(JSON.parse(JSON.stringify(payload)).team_member_model_max_budget).toStrictEqual({
|
||||
"gpt-4": { budget_limit: 4, time_period: "30d" },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -51,6 +51,7 @@ import BudgetDurationDropdown, {
|
|||
import { Organization, getDefaultTeamSettings, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking";
|
||||
import NumericalInput from "./shared/numerical_input";
|
||||
import { ModelMaxBudget, ModelMaxBudgetField } from "./key_team_helpers/ModelMaxBudgetEditor";
|
||||
import { modelMaxBudgetUpdate } from "./key_team_helpers/modelMaxBudgetPayload";
|
||||
import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
|
||||
import SearchToolSelector from "./search_tools/SearchToolSelector";
|
||||
import SkillSelector from "./skills/SkillSelector";
|
||||
|
|
@ -280,6 +281,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
|
|||
const [loggingSettings, setLoggingSettings] = useState<any[]>([]);
|
||||
const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({});
|
||||
const [modelMaxBudget, setModelMaxBudget] = useState<ModelMaxBudget>({});
|
||||
const [teamMemberModelMaxBudget, setTeamMemberModelMaxBudget] = useState<ModelMaxBudget>({});
|
||||
const [routerSettings, setRouterSettings] = useState<RouterSettingsAccordionValue | null>(null);
|
||||
const [routerSettingsKey, setRouterSettingsKey] = useState<number>(0);
|
||||
|
||||
|
|
@ -358,6 +360,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
|
|||
setLoggingSettings([]);
|
||||
setModelAliases({});
|
||||
setModelMaxBudget({});
|
||||
setTeamMemberModelMaxBudget({});
|
||||
setRouterSettings(null);
|
||||
setRouterSettingsKey((prev) => prev + 1);
|
||||
};
|
||||
|
|
@ -551,9 +554,11 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
|
|||
}
|
||||
|
||||
const { require_trace_id: requireTraceId, ...teamCreateValues } = formValues;
|
||||
const teamMemberModelBudgets = modelMaxBudgetUpdate(teamMemberModelMaxBudget, {});
|
||||
await teamCreateCall(accessToken, {
|
||||
...teamCreateValues,
|
||||
...(requireTraceId === undefined ? {} : { require_trace_id: requireTraceId }),
|
||||
...(teamMemberModelBudgets !== undefined ? { team_member_model_max_budget: teamMemberModelBudgets } : {}),
|
||||
models: normalizeTeamModelSelection(formValues.models),
|
||||
});
|
||||
toast.success("Team created");
|
||||
|
|
@ -912,6 +917,15 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
|
|||
/>
|
||||
)}
|
||||
</FormField>
|
||||
<ModelMaxBudgetField
|
||||
key={`team-member-model-max-budget-${routerSettingsKey}`}
|
||||
label="Default Per-Model Budget"
|
||||
premiumUser={premiumUser}
|
||||
value={teamMemberModelMaxBudget}
|
||||
onChange={setTeamMemberModelMaxBudget}
|
||||
availableModels={userModels}
|
||||
hint="Set a default per-model spend cap for each team member, with an independent reset window."
|
||||
/>
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="team_member_key_duration"
|
||||
|
|
|
|||
|
|
@ -92,6 +92,13 @@ export function ModelMaxBudgetEditor({
|
|||
}: ModelMaxBudgetEditorProps) {
|
||||
const [entries, setEntries] = useState<ModelBudgetEntry[]>(() => modelMaxBudgetToEntries(value));
|
||||
|
||||
const incoming = JSON.stringify(value ?? {});
|
||||
const [syncedValue, setSyncedValue] = useState(incoming);
|
||||
if (incoming !== syncedValue) {
|
||||
setSyncedValue(incoming);
|
||||
if (incoming !== JSON.stringify(entriesToModelMaxBudget(entries))) setEntries(modelMaxBudgetToEntries(value));
|
||||
}
|
||||
|
||||
const emitChange = (updated: ModelBudgetEntry[]) => {
|
||||
setEntries(updated);
|
||||
onChange(entriesToModelMaxBudget(updated));
|
||||
|
|
@ -219,14 +226,15 @@ export function ModelMaxBudgetEditor({
|
|||
|
||||
interface ModelMaxBudgetFieldProps extends ModelMaxBudgetEditorProps {
|
||||
hint: string;
|
||||
label?: string;
|
||||
}
|
||||
|
||||
/** The editor with its label, so every form that offers it presents it the same way. */
|
||||
export function ModelMaxBudgetField({ hint, ...editorProps }: ModelMaxBudgetFieldProps) {
|
||||
export function ModelMaxBudgetField({ hint, label = "Per-Model Budgets", ...editorProps }: ModelMaxBudgetFieldProps) {
|
||||
return (
|
||||
<Field>
|
||||
<FieldLabel>
|
||||
<span title={hint}>Per-Model Budgets</span>
|
||||
<span title={hint}>{label}</span>
|
||||
</FieldLabel>
|
||||
<ModelMaxBudgetEditor {...editorProps} />
|
||||
</Field>
|
||||
|
|
|
|||
|
|
@ -94,6 +94,7 @@ import { Team } from "./key_team_helpers/key_list";
|
|||
import { EmailEventSettingsResponse, EmailEventSettingsUpdateRequest } from "./email_events/types";
|
||||
import type { ListPluginsResponse, SkillRegisterRequest } from "./claude_code_plugins/types";
|
||||
import type { ModelBudgetUsage, ModelMaxBudget } from "./key_team_helpers/ModelMaxBudgetEditor";
|
||||
import type { StoredModelMaxBudget } from "./key_team_helpers/modelMaxBudgetPayload";
|
||||
import type { ObjectPermission } from "./object_permission_types";
|
||||
import type { components } from "@/lib/http/schema";
|
||||
import { fetchClient } from "@/lib/http/api";
|
||||
|
|
@ -2831,6 +2832,7 @@ export interface Member {
|
|||
allowed_models?: string[] | null;
|
||||
temp_budget_increase?: number | null;
|
||||
temp_budget_expiry?: string | null;
|
||||
model_max_budget?: StoredModelMaxBudget | null;
|
||||
}
|
||||
|
||||
export const teamMemberAddCall = async (accessToken: string, teamId: string, formValues: Member) => {
|
||||
|
|
@ -2971,6 +2973,9 @@ export const teamMemberUpdateCall = async (
|
|||
if ("temp_budget_expiry" in formValues) {
|
||||
requestBody.temp_budget_expiry = orNull(formValues.temp_budget_expiry);
|
||||
}
|
||||
if (formValues.model_max_budget !== undefined) {
|
||||
requestBody.model_max_budget = formValues.model_max_budget;
|
||||
}
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import { useState } from "react";
|
||||
import { fireEvent, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders } from "../../../tests/test-utils";
|
||||
import EditMembership from "./EditMembership";
|
||||
|
|
@ -104,6 +106,79 @@ describe("EditMembership submit payload", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("seeds the per-model budget editor from the member's stored budget and submits the edited cap", async () => {
|
||||
const modelBudgetConfig = {
|
||||
...teamMemberConfig,
|
||||
additionalFields: [
|
||||
...additionalFields,
|
||||
{
|
||||
name: "model_max_budget",
|
||||
label: "Per-Model Budgets",
|
||||
type: "model-max-budget" as const,
|
||||
availableModels: ["gpt-4o"],
|
||||
premiumUser: true,
|
||||
},
|
||||
],
|
||||
};
|
||||
const member = {
|
||||
user_id: "u1",
|
||||
user_email: "a@b.com",
|
||||
role: "user",
|
||||
model_max_budget: { "gpt-4o": { max_budget: 5, budget_duration: "30d" } },
|
||||
};
|
||||
renderEdit(modelBudgetConfig, member);
|
||||
|
||||
const cap = await screen.findByPlaceholderText("Max spend ($)");
|
||||
expect(cap).toHaveValue(5);
|
||||
fireEvent.change(cap, { target: { value: "9" } });
|
||||
save();
|
||||
|
||||
await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce());
|
||||
expect(submitted().model_max_budget).toStrictEqual({ "gpt-4o": { budget_limit: 9, time_period: "30d" } });
|
||||
});
|
||||
|
||||
it("saves a role-only edit for a member with stored model caps", async () => {
|
||||
const modelBudgetConfig = {
|
||||
...teamMemberConfig,
|
||||
additionalFields: [
|
||||
...additionalFields,
|
||||
{
|
||||
name: "model_max_budget",
|
||||
label: "Per-Model Budgets",
|
||||
type: "model-max-budget" as const,
|
||||
availableModels: ["gpt-4o"],
|
||||
premiumUser: true,
|
||||
},
|
||||
],
|
||||
};
|
||||
const member = {
|
||||
user_id: "u1",
|
||||
user_email: "a@b.com",
|
||||
role: "user",
|
||||
model_max_budget: { "gpt-4o": { max_budget: 5, budget_duration: "30d" } },
|
||||
};
|
||||
renderEdit(modelBudgetConfig, member);
|
||||
|
||||
const user = userEvent.setup();
|
||||
await user.click(screen.getByRole("combobox", { name: "Role(Current: User)" }));
|
||||
await user.click(await screen.findByRole("option", { name: "Admin" }));
|
||||
save();
|
||||
|
||||
await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce());
|
||||
const expectedPayload = {
|
||||
user_email: "a@b.com",
|
||||
user_id: "u1",
|
||||
role: "admin",
|
||||
max_budget_in_team: null,
|
||||
budget_duration: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
allowed_models: [],
|
||||
model_max_budget: { "gpt-4o": { budget_limit: 5, time_period: "30d" } },
|
||||
};
|
||||
expect(submitted()).toStrictEqual(expectedPayload);
|
||||
});
|
||||
|
||||
it("omits a field the config hides even when the member record carries it", async () => {
|
||||
renderEdit({ ...orgMemberConfig, showUserId: false }, { user_id: "u1", user_email: "a@b.com", role: "admin" });
|
||||
|
||||
|
|
@ -342,3 +417,57 @@ describe("EditMembership submit payload", () => {
|
|||
expect(screen.getByLabelText("Email")).toHaveValue("");
|
||||
});
|
||||
});
|
||||
|
||||
it("discards an unfinished per-model budget row when an edit session is reopened", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const modelBudgetConfig = {
|
||||
...teamMemberConfig,
|
||||
additionalFields: [
|
||||
...additionalFields,
|
||||
{
|
||||
name: "model_max_budget",
|
||||
label: "Per-Model Budgets",
|
||||
type: "model-max-budget" as const,
|
||||
availableModels: ["gpt-4o"],
|
||||
premiumUser: true,
|
||||
},
|
||||
],
|
||||
};
|
||||
const member = { user_id: "u1", user_email: "a@b.com", role: "user", model_max_budget: {} };
|
||||
const otherMember = { ...member, user_id: "u2", user_email: "c@d.com" };
|
||||
const Harness = ({ initialData }: { initialData: typeof member }) => {
|
||||
const [visible, setVisible] = useState(false);
|
||||
return (
|
||||
<>
|
||||
<button type="button" onClick={() => setVisible(true)}>
|
||||
Open Edit Member
|
||||
</button>
|
||||
<EditMembership
|
||||
visible={visible}
|
||||
onCancel={() => setVisible(false)}
|
||||
onSubmit={vi.fn()}
|
||||
mode="edit"
|
||||
config={modelBudgetConfig as never}
|
||||
initialData={initialData as never}
|
||||
/>
|
||||
</>
|
||||
);
|
||||
};
|
||||
|
||||
const { rerender } = renderWithProviders(<Harness initialData={member} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Open Edit Member" }));
|
||||
await user.click(screen.getByRole("button", { name: /Add Model Budget/i }));
|
||||
expect(screen.getByPlaceholderText("Select model")).toBeInTheDocument();
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Cancel" }));
|
||||
await user.click(screen.getByRole("button", { name: "Open Edit Member" }));
|
||||
|
||||
expect(screen.queryByPlaceholderText("Select model")).not.toBeInTheDocument();
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /Add Model Budget/i }));
|
||||
expect(screen.getByPlaceholderText("Select model")).toBeInTheDocument();
|
||||
|
||||
rerender(<Harness initialData={otherMember} />);
|
||||
expect(screen.queryByPlaceholderText("Select model")).not.toBeInTheDocument();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -12,11 +12,13 @@ import { Input } from "@/components/ui/input";
|
|||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
|
||||
import { UtcDateTimeInput } from "@/components/shared/form/UtcDateTimeInput";
|
||||
import { ModelMaxBudgetEditor } from "../key_team_helpers/ModelMaxBudgetEditor";
|
||||
import { useZodForm } from "@/lib/forms/useZodForm";
|
||||
import {
|
||||
buildMemberFormData,
|
||||
buildMemberFormValues,
|
||||
emptyMemberFormValues,
|
||||
isModelMaxBudget,
|
||||
TEMP_BUDGET_PAIR_MESSAGE,
|
||||
tempBudgetPairError,
|
||||
type MemberAdditionalField,
|
||||
|
|
@ -51,7 +53,11 @@ const ROLE_REQUIRED_MESSAGE = "Please select a role!";
|
|||
const isEmailish = (value: string): boolean =>
|
||||
value === "" || z.email({ pattern: z.regexes.html5Email }).safeParse(value).success;
|
||||
|
||||
const memberFieldSchema = z.union([z.string(), z.number(), z.null(), z.array(z.string())]).optional();
|
||||
const modelMaxBudgetSchema = z.record(z.string(), z.looseObject({ budget_limit: z.number(), time_period: z.string() }));
|
||||
|
||||
const memberFieldSchema = z
|
||||
.union([z.string(), z.number(), z.null(), z.array(z.string()), modelMaxBudgetSchema])
|
||||
.optional();
|
||||
|
||||
const buildMemberSchema = (config: ModalConfig): z.ZodType<MemberFormValues, MemberFormValues> => {
|
||||
const shape = {
|
||||
|
|
@ -171,6 +177,16 @@ const MemberModal = <T extends BaseMember>({
|
|||
onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)}
|
||||
/>
|
||||
);
|
||||
case "model-max-budget":
|
||||
return (
|
||||
<ModelMaxBudgetEditor
|
||||
key={`${visible}-${initialData?.user_id ?? initialData?.user_email ?? ""}`}
|
||||
value={isModelMaxBudget(value) ? value : {}}
|
||||
onChange={onChange}
|
||||
availableModels={field.availableModels ?? []}
|
||||
premiumUser={field.premiumUser ?? false}
|
||||
/>
|
||||
);
|
||||
case "utc-datetime":
|
||||
return (
|
||||
<UtcDateTimeInput
|
||||
|
|
|
|||
|
|
@ -721,4 +721,98 @@ describe("TeamInfoView - member budget reset prompt", () => {
|
|||
expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument();
|
||||
expect(bulkUpdatePOST).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("omits the default member model budget when it has not changed", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const storedBudget = { "gpt-4": { budget_limit: 3, time_period: "30d" } };
|
||||
const data = createMockTeamData({
|
||||
models: ["gpt-4"],
|
||||
team_member_budget_table: {
|
||||
max_budget: 10,
|
||||
budget_duration: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
model_max_budget: storedBudget,
|
||||
},
|
||||
});
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(data);
|
||||
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" });
|
||||
|
||||
renderWithProviders(<TeamInfoView {...props} premiumUser />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await user.click(screen.getByText("Team Member Settings"));
|
||||
await screen.findByPlaceholderText("Max spend ($)");
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).not.toHaveProperty("team_member_model_max_budget");
|
||||
});
|
||||
|
||||
it("sends the default member model budget after an edit", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
const storedBudget = { "gpt-4": { budget_limit: 3, time_period: "30d" } };
|
||||
const data = createMockTeamData({
|
||||
models: ["gpt-4"],
|
||||
team_member_budget_table: {
|
||||
max_budget: 10,
|
||||
budget_duration: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
model_max_budget: storedBudget,
|
||||
},
|
||||
});
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(data);
|
||||
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" });
|
||||
|
||||
renderWithProviders(<TeamInfoView {...props} premiumUser />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await user.click(screen.getByText("Team Member Settings"));
|
||||
fireEvent.change(await screen.findByPlaceholderText("Max spend ($)"), { target: { value: "7" } });
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_model_max_budget).toStrictEqual({
|
||||
"gpt-4": { budget_limit: 7, time_period: "30d" },
|
||||
});
|
||||
});
|
||||
|
||||
it("offers access-group models in the default member budget picker", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(
|
||||
createMockTeamData({
|
||||
models: ["gpt-4o"],
|
||||
access_group_models: ["claude-haiku-4-5"],
|
||||
}) as TeamData,
|
||||
);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...props} userModels={["gpt-4o", "claude-haiku-4-5"]} premiumUser />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await user.click(screen.getByText("Team Member Settings"));
|
||||
await user.click(screen.getAllByRole("button", { name: /add model budget/i })[0]);
|
||||
await user.click(screen.getByPlaceholderText("Select model"));
|
||||
|
||||
expect(await screen.findByRole("option", { name: "claude-haiku-4-5" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("offers user models in the default member budget picker when the team model list is empty", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: [] }) as TeamData);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...props} userModels={["gpt-4o", "claude-haiku-4-5"]} premiumUser />);
|
||||
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
|
||||
await user.click(screen.getByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
await user.click(screen.getByText("Team Member Settings"));
|
||||
await user.click(screen.getAllByRole("button", { name: /add model budget/i })[0]);
|
||||
await user.click(screen.getByPlaceholderText("Select model"));
|
||||
|
||||
expect(await screen.findByRole("option", { name: "gpt-4o" })).toBeInTheDocument();
|
||||
expect(screen.getByRole("option", { name: "claude-haiku-4-5" })).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -286,7 +286,7 @@ export interface TeamMembership {
|
|||
max_parallel_requests: number | null;
|
||||
tpm_limit: number | null;
|
||||
rpm_limit: number | null;
|
||||
model_max_budget: Record<string, number> | null;
|
||||
model_max_budget: StoredModelMaxBudget | null;
|
||||
budget_duration: string | null;
|
||||
budget_reset_at: string | null;
|
||||
allowed_models?: string[] | null;
|
||||
|
|
@ -339,6 +339,7 @@ export interface TeamData {
|
|||
budget_duration: string | null;
|
||||
tpm_limit: number | null;
|
||||
rpm_limit: number | null;
|
||||
model_max_budget?: StoredModelMaxBudget | null;
|
||||
} | null;
|
||||
};
|
||||
keys: any[];
|
||||
|
|
@ -631,6 +632,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
const [isTeamSaving, setIsTeamSaving] = useState(false);
|
||||
const [teamModelAliases, setTeamModelAliases] = useState<Record<string, string>>({});
|
||||
const [teamModelMaxBudget, setTeamModelMaxBudget] = useState<ModelMaxBudget>({});
|
||||
const [teamMemberModelMaxBudget, setTeamMemberModelMaxBudget] = useState<ModelMaxBudget>({});
|
||||
const routerSettingsRef = React.useRef<RouterSettingsAccordionRef>(null);
|
||||
const [organization, setOrganization] = useState<Organization | null>(null);
|
||||
const { userRole, isViewOnly } = useAuthorized();
|
||||
|
|
@ -661,12 +663,20 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
] as const
|
||||
).find(([failed]) => failed)?.[1] ?? null;
|
||||
const availableRateLimitModels = useMemo(() => {
|
||||
const selected = watchedModels ?? teamData?.team_info?.models ?? [];
|
||||
const selected = (isEditing ? watchedModels : undefined) ?? teamData?.team_info?.models ?? [];
|
||||
if (selected.includes("all-proxy-models") || selected.includes("all-team-models")) {
|
||||
return userModels;
|
||||
}
|
||||
return unfurlWildcardModelsInList(selected, userModels);
|
||||
}, [watchedModels, teamData, userModels]);
|
||||
}, [isEditing, watchedModels, teamData, userModels]);
|
||||
const availableModelBudgetModels = useMemo(() => {
|
||||
const selected = (isEditing ? watchedModels : undefined) ?? teamData?.team_info?.models ?? [];
|
||||
if (selected.length === 0 || selected.includes("all-proxy-models") || selected.includes("all-team-models")) {
|
||||
return userModels;
|
||||
}
|
||||
const granted = [...selected, ...(teamData?.team_info?.access_group_models ?? [])];
|
||||
return unfurlWildcardModelsInList(Array.from(new Set(granted)), userModels);
|
||||
}, [isEditing, watchedModels, teamData, userModels]);
|
||||
|
||||
const teamEditAccess = useMemo(() => parseTeamEditAccess(teamData?.team_info?.caller_edit_access), [teamData]);
|
||||
const canEditTeam = is_team_admin || is_proxy_admin || teamEditAccess.kind !== "none";
|
||||
|
|
@ -684,6 +694,9 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
const startEditing = () => {
|
||||
form.reset(teamFormValues());
|
||||
setTeamModelMaxBudget((teamData?.team_info?.model_max_budget ?? {}) as ModelMaxBudget);
|
||||
setTeamMemberModelMaxBudget(
|
||||
(teamData?.team_info?.team_member_budget_table?.model_max_budget ?? {}) as ModelMaxBudget,
|
||||
);
|
||||
setTeamMemberSettingsOpen(false);
|
||||
setSearchToolSettingsOpen(false);
|
||||
setIsEditing(true);
|
||||
|
|
@ -866,6 +879,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
allowed_models: values.allowed_models,
|
||||
temp_budget_increase: values.temp_budget_increase,
|
||||
temp_budget_expiry: values.temp_budget_expiry,
|
||||
model_max_budget: modelMaxBudgetUpdate(values.model_max_budget ?? {}, selectedEditMember?.model_max_budget),
|
||||
};
|
||||
toast.dismiss(); // Remove all existing toasts
|
||||
|
||||
|
|
@ -1223,6 +1237,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
updateData.model_max_budget = modelBudgets;
|
||||
}
|
||||
|
||||
const teamMemberModelBudgets = modelMaxBudgetUpdate(
|
||||
teamMemberModelMaxBudget,
|
||||
info.team_member_budget_table?.model_max_budget,
|
||||
);
|
||||
if (teamMemberModelBudgets !== undefined) {
|
||||
updateData.team_member_model_max_budget = teamMemberModelBudgets;
|
||||
}
|
||||
|
||||
// Handle router_settings - read fresh values from DOM at save time.
|
||||
const currentRouterSettings = routerSettingsRef.current?.getValue();
|
||||
if (currentRouterSettings?.router_settings) {
|
||||
|
|
@ -1658,6 +1680,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
/>
|
||||
)}
|
||||
</FormField>
|
||||
<ModelMaxBudgetField
|
||||
label="Default Per-Model Budget"
|
||||
premiumUser={premiumUser}
|
||||
value={teamMemberModelMaxBudget}
|
||||
onChange={setTeamMemberModelMaxBudget}
|
||||
availableModels={availableModelBudgetModels}
|
||||
hint="Set a default per-model spend cap for each team member, with an independent reset window."
|
||||
/>
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="team_member_key_duration"
|
||||
|
|
@ -2347,6 +2377,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
</p>
|
||||
<div>Max Budget: {info.team_member_budget_table?.max_budget ?? "No Limit"}</div>
|
||||
<div>Budget Duration: {info.team_member_budget_table?.budget_duration || "No Limit"}</div>
|
||||
{modelMaxBudgetToEntries(
|
||||
info.team_member_budget_table?.model_max_budget as ModelMaxBudget | null | undefined,
|
||||
).map(({ model, budgetLimit, timePeriod }) => (
|
||||
<div key={model}>
|
||||
Default Per-Model Budget ({model}):{" "}
|
||||
{budgetLimit !== null ? `$${formatNumberWithCommas(budgetLimit, 4)}` : "No Limit"} per {timePeriod}
|
||||
</div>
|
||||
))}
|
||||
<div>Key Duration: {info.metadata?.team_member_key_duration || "No Limit"}</div>
|
||||
<div>TPM Limit: {info.team_member_budget_table?.tpm_limit ?? "No Limit"}</div>
|
||||
<div>RPM Limit: {info.team_member_budget_table?.rpm_limit ?? "No Limit"}</div>
|
||||
|
|
@ -2554,6 +2592,20 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
),
|
||||
type: "utc-datetime" as const,
|
||||
},
|
||||
{
|
||||
name: "model_max_budget",
|
||||
label: (
|
||||
<span>
|
||||
Per-Model Budgets{" "}
|
||||
<SimpleTooltip content="Spend cap per model for this member, each with its own reset period. Starts from the team's default per-model budget; saving a change stores a per-model budget for this member only.">
|
||||
<Info className="ml-1 inline size-3.5 align-text-bottom" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
),
|
||||
type: "model-max-budget" as const,
|
||||
availableModels: availableModelBudgetModels,
|
||||
premiumUser,
|
||||
},
|
||||
{
|
||||
name: "tpm_limit",
|
||||
label: (
|
||||
|
|
|
|||
|
|
@ -432,10 +432,31 @@ describe("TeamMembersComponent", () => {
|
|||
allowed_models: [],
|
||||
temp_budget_increase: 0,
|
||||
temp_budget_expiry: "2030-01-02T03:04:00Z",
|
||||
model_max_budget: null,
|
||||
};
|
||||
expect(seedMemberBudgetFields({ user_id: "user1@test.com", role: "member" }, budget)).toStrictEqual(seeded);
|
||||
});
|
||||
|
||||
it("seeds the team's default per-model budget when the member has no per-model override", () => {
|
||||
const teamDefault = { "gpt-4o": { max_budget: 10, budget_duration: "30d" } };
|
||||
const ownRowWithoutModelCaps = {
|
||||
...createMockTeamData().team_memberships[0].litellm_budget_table,
|
||||
model_max_budget: {},
|
||||
};
|
||||
const memberOverride = { "gpt-4o": { max_budget: 2.5, budget_duration: "7d" } };
|
||||
const ownRowWithModelCaps = { ...ownRowWithoutModelCaps, model_max_budget: memberOverride };
|
||||
|
||||
expect(
|
||||
seedMemberBudgetFields({ user_id: "u", role: "user" }, ownRowWithoutModelCaps, teamDefault).model_max_budget,
|
||||
).toBe(teamDefault);
|
||||
expect(seedMemberBudgetFields({ user_id: "u", role: "user" }, undefined, teamDefault).model_max_budget).toBe(
|
||||
teamDefault,
|
||||
);
|
||||
expect(
|
||||
seedMemberBudgetFields({ user_id: "u", role: "user" }, ownRowWithModelCaps, teamDefault).model_max_budget,
|
||||
).toBe(memberOverride);
|
||||
});
|
||||
|
||||
it("seeds null temporary budget fields for a member without a budget row", () => {
|
||||
expect(seedMemberBudgetFields({ user_id: "user2@test.com", role: "admin" }, undefined)).toMatchObject({
|
||||
max_budget_in_team: null,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import { Button } from "@/components/ui/button";
|
|||
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import MemberTable from "@/components/common_components/MemberTable";
|
||||
import type { StoredModelMaxBudget } from "@/components/key_team_helpers/modelMaxBudgetPayload";
|
||||
import { Member } from "@/components/networking";
|
||||
import { parseErrorMessage } from "@/components/shared/errorUtils";
|
||||
import { DateCell, MoneyCell } from "@/components/shared/table_cells";
|
||||
|
|
@ -25,9 +26,13 @@ const BUDGET_SOURCE_LABELS: Record<Exclude<TeamMemberBudgetSource, "none">, stri
|
|||
const formatBudget = (value: number | null): string =>
|
||||
value === null ? "Unlimited" : `$${formatNumberWithCommas(value, 2)}`;
|
||||
|
||||
const hasModelBudgets = (budget: StoredModelMaxBudget | null | undefined): budget is StoredModelMaxBudget =>
|
||||
budget != null && Object.keys(budget).length > 0;
|
||||
|
||||
export const seedMemberBudgetFields = (
|
||||
record: Member,
|
||||
budget: TeamMembership["litellm_budget_table"] | undefined,
|
||||
teamDefaultModelMaxBudget?: StoredModelMaxBudget | null,
|
||||
): Member => ({
|
||||
...record,
|
||||
max_budget_in_team: budget?.max_budget ?? null,
|
||||
|
|
@ -37,6 +42,9 @@ export const seedMemberBudgetFields = (
|
|||
allowed_models: budget?.allowed_models || [],
|
||||
temp_budget_increase: budget?.temp_budget_increase ?? null,
|
||||
temp_budget_expiry: budget?.temp_budget_expiry ?? null,
|
||||
model_max_budget: hasModelBudgets(budget?.model_max_budget)
|
||||
? budget.model_max_budget
|
||||
: teamDefaultModelMaxBudget ?? null,
|
||||
});
|
||||
|
||||
interface TeamMemberTabProps {
|
||||
|
|
@ -305,7 +313,13 @@ export default function TeamMemberTab({
|
|||
canEdit={canEditTeam}
|
||||
onEdit={(record) => {
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === record.user_id);
|
||||
setSelectedEditMember(seedMemberBudgetFields(record, membership?.litellm_budget_table));
|
||||
setSelectedEditMember(
|
||||
seedMemberBudgetFields(
|
||||
record,
|
||||
membership?.litellm_budget_table,
|
||||
teamData.team_info.team_member_budget_table?.model_max_budget,
|
||||
),
|
||||
);
|
||||
setIsEditMemberModalVisible(true);
|
||||
}}
|
||||
onDelete={handleMemberDelete}
|
||||
|
|
|
|||
|
|
@ -1,8 +1,21 @@
|
|||
export type MemberFieldValue = string | number | null | undefined | string[];
|
||||
import {
|
||||
entriesToModelMaxBudget,
|
||||
modelMaxBudgetToEntries,
|
||||
type ModelMaxBudget,
|
||||
} from "../key_team_helpers/ModelMaxBudgetEditor";
|
||||
|
||||
export type MemberFieldValue = string | number | null | undefined | string[] | ModelMaxBudget;
|
||||
|
||||
export type MemberFormValues = Record<string, MemberFieldValue>;
|
||||
|
||||
export type MemberFieldType = "input" | "select" | "numerical" | "multi-select" | "budget-duration" | "utc-datetime";
|
||||
export type MemberFieldType =
|
||||
| "input"
|
||||
| "select"
|
||||
| "numerical"
|
||||
| "multi-select"
|
||||
| "budget-duration"
|
||||
| "utc-datetime"
|
||||
| "model-max-budget";
|
||||
|
||||
export interface MemberAdditionalField {
|
||||
name: string;
|
||||
|
|
@ -12,6 +25,8 @@ export interface MemberAdditionalField {
|
|||
step?: number;
|
||||
min?: number;
|
||||
placeholder?: string;
|
||||
availableModels?: string[];
|
||||
premiumUser?: boolean;
|
||||
}
|
||||
|
||||
export interface MemberFieldsConfig {
|
||||
|
|
@ -33,6 +48,9 @@ export const TEMP_BUDGET_PAIR_MESSAGE = "Set both a temporary budget increase an
|
|||
|
||||
const isUnset = (value: MemberFieldValue): boolean => value === null || value === undefined || value === "";
|
||||
|
||||
export const isModelMaxBudget = (value: MemberFieldValue): value is ModelMaxBudget =>
|
||||
typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
|
||||
export const tempBudgetPairError = (values: MemberFormValues): "temp_budget_increase" | "temp_budget_expiry" | null => {
|
||||
const increaseUnset = isUnset(values.temp_budget_increase);
|
||||
const expiryUnset = isUnset(values.temp_budget_expiry);
|
||||
|
|
@ -66,6 +84,9 @@ export const buildMemberFormValues = (
|
|||
allowed_models: initialData.allowed_models || [],
|
||||
temp_budget_increase: initialData.temp_budget_increase ?? null,
|
||||
temp_budget_expiry: initialData.temp_budget_expiry || null,
|
||||
model_max_budget: isModelMaxBudget(initialData.model_max_budget)
|
||||
? entriesToModelMaxBudget(modelMaxBudgetToEntries(initialData.model_max_budget))
|
||||
: {},
|
||||
};
|
||||
|
||||
return pickFieldNames(config, seeded);
|
||||
|
|
@ -78,6 +99,8 @@ const emptyValueForType = (type: MemberFieldType | undefined): MemberFieldValue
|
|||
switch (type) {
|
||||
case "multi-select":
|
||||
return [];
|
||||
case "model-max-budget":
|
||||
return {};
|
||||
case "numerical":
|
||||
case "budget-duration":
|
||||
case "utc-datetime":
|
||||
|
|
|
|||
25
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
25
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -16882,6 +16882,7 @@ export interface paths {
|
|||
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
* - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
|
||||
* - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
|
||||
* - team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
|
||||
* - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
|
||||
* - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
|
||||
* - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
|
||||
|
|
@ -17111,6 +17112,7 @@ export interface paths {
|
|||
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||||
* - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
|
||||
* - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
|
||||
* - team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
|
||||
* - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
|
||||
* - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
|
||||
* - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
|
||||
|
|
@ -38499,6 +38501,10 @@ export interface components {
|
|||
team_member_budget_duration?: string | null;
|
||||
/** Team Member Key Duration */
|
||||
team_member_key_duration?: string | null;
|
||||
/** Team Member Model Max Budget */
|
||||
team_member_model_max_budget?: {
|
||||
[key: string]: components["schemas"]["BudgetConfig"];
|
||||
} | null;
|
||||
/** Team Member Permissions */
|
||||
team_member_permissions?: string[] | null;
|
||||
/** Team Member Rpm Limit */
|
||||
|
|
@ -39584,6 +39590,10 @@ export interface components {
|
|||
team_member_budget_duration?: string | null;
|
||||
/** Team Member Key Duration */
|
||||
team_member_key_duration?: string | null;
|
||||
/** Team Member Model Max Budget */
|
||||
team_member_model_max_budget?: {
|
||||
[key: string]: components["schemas"]["BudgetConfig"];
|
||||
} | null;
|
||||
/** Team Member Rpm Limit */
|
||||
team_member_rpm_limit?: number | null;
|
||||
/** Team Member Tpm Limit */
|
||||
|
|
@ -45580,6 +45590,13 @@ export interface components {
|
|||
budget_duration?: string | null;
|
||||
/** Max Budget In Team */
|
||||
max_budget_in_team?: number | null;
|
||||
/**
|
||||
* Model Max Budget
|
||||
* @description Per-model spend caps for this team member, each with its own budget_duration. Overrides the team's default per-model member budget. Pass an empty dict to fall back to the team default.
|
||||
*/
|
||||
model_max_budget?: {
|
||||
[key: string]: components["schemas"]["BudgetConfig"];
|
||||
} | null;
|
||||
/** Role */
|
||||
role?: ("admin" | "user") | null;
|
||||
/**
|
||||
|
|
@ -45617,6 +45634,10 @@ export interface components {
|
|||
budget_duration?: string | null;
|
||||
/** Max Budget In Team */
|
||||
max_budget_in_team?: number | null;
|
||||
/** Model Max Budget */
|
||||
model_max_budget?: {
|
||||
[key: string]: components["schemas"]["BudgetConfig"];
|
||||
} | null;
|
||||
/** Rpm Limit */
|
||||
rpm_limit?: number | null;
|
||||
/** Team Id */
|
||||
|
|
@ -47559,6 +47580,10 @@ export interface components {
|
|||
team_member_budget_duration?: string | null;
|
||||
/** Team Member Key Duration */
|
||||
team_member_key_duration?: string | null;
|
||||
/** Team Member Model Max Budget */
|
||||
team_member_model_max_budget?: {
|
||||
[key: string]: components["schemas"]["BudgetConfig"];
|
||||
} | null;
|
||||
/** Team Member Rpm Limit */
|
||||
team_member_rpm_limit?: number | null;
|
||||
/** Team Member Tpm Limit */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue