mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #41620 from BerriAI/litellm_team_member_temp_budget_increase
feat(proxy): temporary budget increase for team members
This commit is contained in:
Yassin Kortam
GitHub
commit
84ae0805ba
29 changed files with 998 additions and 35 deletions
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "temp_budget_increase" DOUBLE PRECISION;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "temp_budget_expiry" TIMESTAMP(3);
|
||||
|
|
@ -22,6 +22,8 @@ model LiteLLM_BudgetTable {
|
|||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models
|
||||
temp_budget_increase Float?
|
||||
temp_budget_expiry DateTime?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
|
|
|
|||
|
|
@ -5,7 +5,8 @@ Canonical definition for ``litellm_budgettable``. Re-exported from
|
|||
``litellm.proxy._types`` for backwards compatibility.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict
|
||||
|
||||
|
|
@ -30,9 +31,26 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
|||
model_max_budget: dict | None = None
|
||||
budget_duration: str | None = None
|
||||
allowed_models: list[str] | None = None # per-member model scope; empty = inherit team models
|
||||
temp_budget_increase: float | None = None
|
||||
temp_budget_expiry: datetime | None = None
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
def active_temp_budget_increase(self, now: datetime) -> float:
|
||||
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
|
||||
return 0.0
|
||||
expiry: Final = (
|
||||
self.temp_budget_expiry.replace(tzinfo=timezone.utc)
|
||||
if self.temp_budget_expiry.tzinfo is None
|
||||
else self.temp_budget_expiry
|
||||
)
|
||||
return 0.0 if expiry <= now else self.temp_budget_increase
|
||||
|
||||
def effective_max_budget(self, now: datetime) -> float | None:
|
||||
if self.max_budget is None:
|
||||
return None
|
||||
return self.max_budget + self.active_temp_budget_increase(now)
|
||||
|
||||
|
||||
class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
|
||||
"""LiteLLM_BudgetTable + server-managed fields returned on API responses."""
|
||||
|
|
|
|||
|
|
@ -4412,6 +4412,23 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
|
|||
default=None,
|
||||
description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.",
|
||||
)
|
||||
temp_budget_increase: float | None = Field(
|
||||
default=None,
|
||||
ge=0,
|
||||
allow_inf_nan=False,
|
||||
description="Temporary additive budget increase for this team member, active until temp_budget_expiry",
|
||||
)
|
||||
temp_budget_expiry: datetime | None = Field(
|
||||
default=None,
|
||||
description="UTC expiry for temp_budget_increase",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_temp_budget(self) -> "TeamMemberUpdateRequest":
|
||||
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
|
||||
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
|
||||
raise ValueError("temp_budget_increase and temp_budget_expiry must be set together")
|
||||
return self
|
||||
|
||||
|
||||
class TeamMemberUpdateResponse(MemberUpdateResponse):
|
||||
|
|
@ -4421,6 +4438,8 @@ class TeamMemberUpdateResponse(MemberUpdateResponse):
|
|||
rpm_limit: int | None = None
|
||||
budget_duration: str | None = None
|
||||
allowed_models: list[str] | None = None
|
||||
temp_budget_increase: float | None = None
|
||||
temp_budget_expiry: datetime | None = None
|
||||
|
||||
|
||||
class TeamModelAddRequest(BaseModel):
|
||||
|
|
|
|||
|
|
@ -5395,12 +5395,10 @@ async def _check_team_member_budget(
|
|||
# Per-member override wins; otherwise fall back to the team-level
|
||||
# default configured via team.metadata["team_member_budget_id"].
|
||||
team_member_budget: float | None = None
|
||||
if (
|
||||
loaded_membership is not None
|
||||
and loaded_membership.litellm_budget_table is not None
|
||||
and loaded_membership.litellm_budget_table.max_budget is not None
|
||||
):
|
||||
team_member_budget = loaded_membership.litellm_budget_table.max_budget
|
||||
member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None
|
||||
now: Final = get_utc_datetime()
|
||||
if member_budget_row is not None and member_budget_row.max_budget is not None:
|
||||
team_member_budget = member_budget_row.effective_max_budget(now=now)
|
||||
else:
|
||||
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
|
||||
if isinstance(default_budget_id, str):
|
||||
|
|
@ -5416,7 +5414,9 @@ async def _check_team_member_budget(
|
|||
and default_budget.max_budget is not None
|
||||
and default_budget.max_budget > 0
|
||||
):
|
||||
team_member_budget = default_budget.max_budget
|
||||
team_member_budget = default_budget.max_budget + (
|
||||
member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0
|
||||
)
|
||||
|
||||
if team_member_budget is not None:
|
||||
team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
|
||||
|
|
|
|||
|
|
@ -2250,7 +2250,9 @@ async def _user_api_key_auth_builder(
|
|||
)
|
||||
|
||||
if team_member_info is not None and team_member_info.litellm_budget_table is not None:
|
||||
team_member_budget: Final = team_member_info.litellm_budget_table.max_budget
|
||||
team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget(
|
||||
now=datetime.now(timezone.utc),
|
||||
)
|
||||
if team_member_budget is not None and team_member_budget > 0:
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
|
|
|||
|
|
@ -476,8 +476,12 @@ _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: Final = (
|
|||
"model_max_budget",
|
||||
"budget_duration",
|
||||
"allowed_models",
|
||||
"temp_budget_increase",
|
||||
"temp_budget_expiry",
|
||||
)
|
||||
|
||||
_TEMP_BUDGET_FIELDS: Final = frozenset({"temp_budget_increase", "temp_budget_expiry"})
|
||||
|
||||
|
||||
MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType(
|
||||
{
|
||||
|
|
@ -486,6 +490,8 @@ MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType(
|
|||
"rpm_limit": "rpm_limit",
|
||||
"budget_duration": "budget_duration",
|
||||
"allowed_models": "allowed_models",
|
||||
"temp_budget_increase": "temp_budget_increase",
|
||||
"temp_budget_expiry": "temp_budget_expiry",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -548,6 +554,8 @@ async def _upsert_budget_and_membership(
|
|||
``shared_budget_ids`` extends that protection to any other row more than one
|
||||
membership points at, which a caller patching several members at once has
|
||||
already counted; a row listed there is cloned rather than written in place.
|
||||
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:
|
||||
return
|
||||
|
|
@ -562,6 +570,7 @@ async def _upsert_budget_and_membership(
|
|||
is_shared_default: Final = existing_budget_id is not None and (
|
||||
existing_budget_id == team_default_budget_id or existing_budget_id in (shared_budget_ids or frozenset())
|
||||
)
|
||||
temp_only: Final = frozenset(write_data) <= _TEMP_BUDGET_FIELDS
|
||||
|
||||
async def _disconnect():
|
||||
await tx.litellm_teammembership.update(
|
||||
|
|
@ -583,7 +592,9 @@ async def _upsert_budget_and_membership(
|
|||
return
|
||||
|
||||
source_row: Final = (
|
||||
await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) if is_shared_default else None
|
||||
await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id})
|
||||
if is_shared_default and not temp_only
|
||||
else None
|
||||
)
|
||||
source: Final[Mapping[str, Any]] = source_row.model_dump() if source_row is not None else MappingProxyType({})
|
||||
|
||||
|
|
@ -604,7 +615,7 @@ async def _upsert_budget_and_membership(
|
|||
create_data.pop("budget_reset_at", None)
|
||||
|
||||
if not _has_meaningful_budget_limit(create_data):
|
||||
if existing_budget_id is not None:
|
||||
if existing_budget_id is not None and not temp_only:
|
||||
await _disconnect()
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -306,6 +306,9 @@ async def _verify_org_access(
|
|||
_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_BUDGET_SETTABLE_FIELDS: Final = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"}
|
||||
_ORG_COLUMN_FIELDS: Final = frozenset({"organization_alias", "models"})
|
||||
_ORG_METADATA_FIELDS: Final = tuple(
|
||||
field for field in LiteLLM_ManagementEndpoint_MetadataFields if field not in _BUDGET_SETTABLE_FIELDS
|
||||
)
|
||||
|
||||
|
||||
def build_budget_write_data(budget_updates: Mapping[str, object], updated_by: str) -> Mapping[str, object]:
|
||||
|
|
@ -391,6 +394,8 @@ async def new_organization(
|
|||
- model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - organization-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.
|
||||
- allowed_models: Optional[List[str]] - List of models the organization is allowed to access. If not set, defaults to the models field.
|
||||
- temp_budget_increase: *Optional[float]* - Stored on the org budget row but only enforced for team member budgets today.
|
||||
- temp_budget_expiry: *Optional[str]* - Stored on the org budget row but only enforced for team member budgets today.
|
||||
Case 1: Create new org **without** a budget_id
|
||||
|
||||
```bash
|
||||
|
|
@ -527,7 +532,7 @@ async def new_organization(
|
|||
organization_payload["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
organization_row: Final = LiteLLM_OrganizationTable.model_validate(organization_payload)
|
||||
|
||||
for field in LiteLLM_ManagementEndpoint_MetadataFields:
|
||||
for field in _ORG_METADATA_FIELDS:
|
||||
if getattr(data, field, None) is not None:
|
||||
_set_object_metadata_field(
|
||||
object_data=organization_row,
|
||||
|
|
|
|||
|
|
@ -3848,6 +3848,8 @@ async def team_member_update(
|
|||
rpm_limit=data.rpm_limit,
|
||||
budget_duration=data.budget_duration,
|
||||
allowed_models=data.allowed_models,
|
||||
temp_budget_increase=data.temp_budget_increase,
|
||||
temp_budget_expiry=data.temp_budget_expiry,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ model LiteLLM_BudgetTable {
|
|||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models
|
||||
temp_budget_increase Float?
|
||||
temp_budget_expiry DateTime?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
|
|
|
|||
|
|
@ -688,16 +688,22 @@ async def _get_team_member_budget_counter(
|
|||
elif isinstance(cached_team_membership, dict):
|
||||
team_membership = LiteLLM_TeamMembership(**cached_team_membership)
|
||||
|
||||
member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
team_member_budget: float | None = None
|
||||
if team_membership is not None and team_membership.litellm_budget_table is not None:
|
||||
team_member_budget = team_membership.litellm_budget_table.max_budget
|
||||
if member_budget_row is not None and member_budget_row.max_budget is not None:
|
||||
team_member_budget = member_budget_row.effective_max_budget(now=now)
|
||||
else:
|
||||
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
|
||||
if isinstance(default_budget_id, str):
|
||||
default_budget: Final = await user_api_key_cache.async_get_cache(
|
||||
key=f"team_member_default_budget:{default_budget_id}",
|
||||
)
|
||||
team_member_budget = _to_float(_get_value(default_budget, "max_budget"))
|
||||
default_cap: Final = _to_float(_get_value(default_budget, "max_budget"))
|
||||
if default_cap is not None and default_cap > 0:
|
||||
team_member_budget = default_cap + (
|
||||
member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0
|
||||
)
|
||||
|
||||
if team_member_budget is None or team_member_budget <= 0:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ model LiteLLM_BudgetTable {
|
|||
budget_duration String?
|
||||
budget_reset_at DateTime?
|
||||
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models
|
||||
temp_budget_increase Float?
|
||||
temp_budget_expiry DateTime?
|
||||
created_at DateTime @default(now()) @map("created_at")
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
Tests for backend domain models.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
|
@ -71,6 +71,34 @@ class TestBudget:
|
|||
assert budget.max_budget is None
|
||||
assert budget.allowed_models is None
|
||||
|
||||
def test_effective_max_budget_applies_unexpired_increase(self):
|
||||
budget = LiteLLM_BudgetTable(
|
||||
max_budget=100.0,
|
||||
temp_budget_increase=50.0,
|
||||
temp_budget_expiry=datetime(2100, 1, 1),
|
||||
)
|
||||
assert budget.effective_max_budget(now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 150.0
|
||||
|
||||
def test_effective_max_budget_ignores_expired_increase(self):
|
||||
expiry = datetime(2020, 1, 1, tzinfo=timezone.utc)
|
||||
budget = LiteLLM_BudgetTable(max_budget=100.0, temp_budget_increase=50.0, temp_budget_expiry=expiry)
|
||||
assert budget.effective_max_budget(now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 100.0
|
||||
assert budget.effective_max_budget(now=expiry) == 100.0
|
||||
|
||||
def test_effective_max_budget_without_increase(self):
|
||||
now = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
assert LiteLLM_BudgetTable(max_budget=100.0).effective_max_budget(now=now) == 100.0
|
||||
assert LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0).effective_max_budget(now=now) is None
|
||||
|
||||
def test_active_temp_budget_increase_is_independent_of_max_budget(self):
|
||||
now = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
bare = LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0, temp_budget_expiry=datetime(2100, 1, 1))
|
||||
assert bare.active_temp_budget_increase(now=now) == 50.0
|
||||
assert bare.effective_max_budget(now=now) is None
|
||||
expired = LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0, temp_budget_expiry=now)
|
||||
assert expired.active_temp_budget_increase(now=now) == 0.0
|
||||
assert LiteLLM_BudgetTable(max_budget=None).active_temp_budget_increase(now=now) == 0.0
|
||||
|
||||
|
||||
class TestCredentials:
|
||||
def test_credentials_creation(self):
|
||||
|
|
|
|||
|
|
@ -8686,3 +8686,161 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None:
|
|||
def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() -> None:
|
||||
assert request_skips_budget_checks(route="/v1/models", model=None, llm_router=None) is True
|
||||
assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_check_temp_budget_increase_extends_cap():
|
||||
"""Spend above max_budget but below max_budget + active temp increase
|
||||
must not raise; once the increase expires the same spend must raise."""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
team_object = LiteLLM_TeamTable(team_id="test-team", metadata={})
|
||||
user_object = LiteLLM_UserTable(user_id="test-user")
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-token",
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
)
|
||||
|
||||
team_membership = LiteLLM_TeamMembership(
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
spend=0.0,
|
||||
budget_id="budget-1",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(
|
||||
max_budget=100.0,
|
||||
temp_budget_increase=100.0,
|
||||
temp_budget_expiry=datetime.now(timezone.utc) + timedelta(hours=1),
|
||||
),
|
||||
)
|
||||
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
if counter_key == "spend:team_member:test-user:test-team":
|
||||
return 150.0
|
||||
return fallback_spend
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter
|
||||
patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch
|
||||
"litellm.proxy.auth.auth_checks.get_team_membership",
|
||||
new_callable=AsyncMock,
|
||||
return_value=team_membership,
|
||||
),
|
||||
):
|
||||
await _check_team_member_budget(
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
valid_token=valid_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=DualCache(),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
expired_membership = LiteLLM_TeamMembership(
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
spend=0.0,
|
||||
budget_id="budget-1",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(
|
||||
max_budget=100.0,
|
||||
temp_budget_increase=100.0,
|
||||
temp_budget_expiry=datetime.now(timezone.utc) - timedelta(hours=1),
|
||||
),
|
||||
)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter
|
||||
patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch
|
||||
"litellm.proxy.auth.auth_checks.get_team_membership",
|
||||
new_callable=AsyncMock,
|
||||
return_value=expired_membership,
|
||||
),
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_team_member_budget(
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
valid_token=valid_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=DualCache(),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
assert exc_info.value.current_cost == 150.0
|
||||
assert exc_info.value.max_budget == 100.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"default_cap, expiry_offset, spend, expected_cap",
|
||||
[
|
||||
(0.4, timedelta(hours=1), 1.0, None),
|
||||
(0.4, timedelta(hours=-1), 1.0, 0.4),
|
||||
(0.0, timedelta(hours=1), 1.0, None),
|
||||
],
|
||||
)
|
||||
async def test_team_member_budget_check_adds_temp_increase_to_live_team_default(
|
||||
default_cap: float, expiry_offset: timedelta, spend: float, expected_cap: float | None
|
||||
):
|
||||
"""A member row that carries only the temporary pair inherits the team default
|
||||
cap live: the increase is added to it while active, the default alone applies
|
||||
once it expires, and a zero default stays uncapped."""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
cache = DualCache()
|
||||
await cache.async_set_cache(
|
||||
key="team_member_default_budget:default-budget-1",
|
||||
value=LiteLLM_BudgetTable(budget_id="default-budget-1", max_budget=default_cap),
|
||||
)
|
||||
team_object = LiteLLM_TeamTable(team_id="test-team", metadata={"team_member_budget_id": "default-budget-1"})
|
||||
valid_token = UserAPIKeyAuth(token="test-token", user_id="test-user", team_id="test-team")
|
||||
team_membership = LiteLLM_TeamMembership(
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
spend=spend,
|
||||
budget_id="budget-1",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(
|
||||
max_budget=None,
|
||||
temp_budget_increase=1.0,
|
||||
temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset,
|
||||
),
|
||||
)
|
||||
|
||||
async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
return fallback_spend
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter
|
||||
patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch
|
||||
"litellm.proxy.auth.auth_checks.get_team_membership",
|
||||
new_callable=AsyncMock,
|
||||
return_value=team_membership,
|
||||
),
|
||||
):
|
||||
if expected_cap is None:
|
||||
await _check_team_member_budget(
|
||||
team_object=team_object,
|
||||
user_object=LiteLLM_UserTable(user_id="test-user"),
|
||||
valid_token=valid_token,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
|
||||
)
|
||||
return
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_team_member_budget(
|
||||
team_object=team_object,
|
||||
user_object=LiteLLM_UserTable(user_id="test-user"),
|
||||
valid_token=valid_token,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
|
||||
)
|
||||
assert exc_info.value.max_budget == expected_cap
|
||||
|
|
|
|||
|
|
@ -7754,6 +7754,112 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
|
|||
assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"expiry_offset, expect_blocked",
|
||||
[
|
||||
(timedelta(days=1), False),
|
||||
(timedelta(days=-1), True),
|
||||
],
|
||||
)
|
||||
async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset, expect_blocked):
|
||||
"""A member over their permanent cap is admitted while a temp_budget_increase is unexpired
|
||||
and blocked again once it expires, on the cached-key auth path."""
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
api_key = "sk-team-member-temp-budget"
|
||||
hashed_token = hash_token(api_key)
|
||||
team_id = "team-temp-budget"
|
||||
user_id = "user-temp-budget"
|
||||
team_member_spend = 2.5
|
||||
|
||||
user_api_key_cache = DualCache()
|
||||
await _cache_key_object(
|
||||
hashed_token=hashed_token,
|
||||
user_api_key_obj=UserAPIKeyAuth(
|
||||
token=hashed_token,
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
team_member_spend=team_member_spend,
|
||||
),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=f"team_id:{team_id}",
|
||||
value=LiteLLM_TeamTableCachedObj(team_id=team_id),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_id,
|
||||
value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
spend=team_member_spend,
|
||||
budget_id="budget-temp",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(
|
||||
max_budget=2.0,
|
||||
temp_budget_increase=1.0,
|
||||
temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/messages"
|
||||
mock_request.method = "POST"
|
||||
mock_request.headers = {"authorization": f"Bearer {api_key}"}
|
||||
mock_request.query_params = {}
|
||||
mock_request.state = SimpleNamespace()
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.budget_alerts = AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def _auth():
|
||||
return await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
|
||||
"litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True}
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state
|
||||
patch( # test-quality-ok: seed the cached key, team and membership without a DB
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache
|
||||
),
|
||||
patch( # test-quality-ok: module-global proxy state
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
|
||||
),
|
||||
patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares
|
||||
"litellm.proxy.proxy_server.get_current_spend",
|
||||
new=AsyncMock(return_value=team_member_spend),
|
||||
),
|
||||
):
|
||||
if not expect_blocked:
|
||||
result = await _auth()
|
||||
assert result.team_member_spend == team_member_spend
|
||||
return
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _auth()
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
|
||||
assert "Max budget: 2.0" in exc_info.value.message
|
||||
|
||||
|
||||
async def _proxy_exception_for_key(
|
||||
api_key: str,
|
||||
general_settings: dict[str, bool],
|
||||
|
|
|
|||
|
|
@ -57,6 +57,12 @@ def assert_future_reset_time(value):
|
|||
assert value > datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def stored_budget_row(mock_tx):
|
||||
"""The budget row the create call persists, minus the audit columns."""
|
||||
data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
|
||||
return {k: v for k, v in data.items() if k not in ("created_by", "updated_by")}
|
||||
|
||||
|
||||
# TEST: an empty patch (caller sent no budget fields) leaves everything alone.
|
||||
# This is the merge-patch contract: absent != clear. Updating only a member's
|
||||
# role must not silently wipe their budget.
|
||||
|
|
@ -211,6 +217,130 @@ async def test_create_seeds_reset_at_and_links(mock_tx, fake_user):
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_from_temp_budget_pair_only(mock_tx, fake_user):
|
||||
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
|
||||
await _upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-new",
|
||||
user_id="user-new",
|
||||
existing_budget_id=None,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"temp_budget_increase": 5.0, "temp_budget_expiry": expiry},
|
||||
)
|
||||
|
||||
mock_tx.litellm_budgettable.create.assert_awaited_once()
|
||||
assert stored_budget_row(mock_tx) == {"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()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_from_temp_pair_never_snapshots_team_default(mock_tx, fake_user):
|
||||
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10)
|
||||
)
|
||||
await _upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-default",
|
||||
user_id="user-unlinked",
|
||||
existing_budget_id=None,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry},
|
||||
team_default_budget_id="team-default-budget-1",
|
||||
)
|
||||
|
||||
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
|
||||
assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temp_pair_on_shared_default_member_creates_bare_row(mock_tx, fake_user):
|
||||
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10)
|
||||
)
|
||||
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={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry},
|
||||
team_default_budget_id="team-default-budget-1",
|
||||
)
|
||||
|
||||
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}
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clearing_temp_pair_on_shared_default_member_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={"temp_budget_increase": None, "temp_budget_expiry": None},
|
||||
team_default_budget_id="team-default-budget-1",
|
||||
)
|
||||
|
||||
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_temp_pair_with_permanent_field_still_clones_shared_default(mock_tx, fake_user):
|
||||
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10)
|
||||
)
|
||||
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={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry, "tpm_limit": 500},
|
||||
team_default_budget_id="team-default-budget-1",
|
||||
)
|
||||
|
||||
mock_tx.litellm_budgettable.find_unique.assert_awaited_once_with(where={"budget_id": "team-default-budget-1"})
|
||||
data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
|
||||
assert data["max_budget"] == 0.4
|
||||
assert data["rpm_limit"] == 10
|
||||
assert data["tpm_limit"] == 500
|
||||
assert data["temp_budget_increase"] == 1.0
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_from_plain_patch_does_not_snapshot_team_default(mock_tx, fake_user):
|
||||
mock_tx.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10)
|
||||
)
|
||||
await _upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-default",
|
||||
user_id="user-unlinked",
|
||||
existing_budget_id=None,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_patch={"tpm_limit": 500},
|
||||
team_default_budget_id="team-default-budget-1",
|
||||
)
|
||||
|
||||
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
|
||||
assert stored_budget_row(mock_tx) == {"tpm_limit": 500}
|
||||
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
# TEST: clone-on-write when the membership still points at the team's shared
|
||||
# default budget. Editing this member must fork a private budget instead of
|
||||
# mutating the shared row, and cloning a duration must seed a fresh reset time.
|
||||
|
|
|
|||
|
|
@ -806,6 +806,8 @@ _EXPECTED_CUSTOMER = {
|
|||
"model_max_budget": None,
|
||||
"budget_duration": "30d",
|
||||
"allowed_models": [],
|
||||
"temp_budget_increase": None,
|
||||
"temp_budget_expiry": None,
|
||||
"budget_reset_at": "2024-02-01T00:00:00",
|
||||
"created_at": "2024-01-01T00:00:00",
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1352,6 +1352,49 @@ async def test_new_organization_rejects_shared_alias_tool_permission_key():
|
|||
prisma_client.db.litellm_objectpermissiontable.create.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_organization_temp_budget_fields_go_to_budget_row_not_metadata(monkeypatch):
|
||||
"""temp_budget_increase/expiry are budget columns and also key-metadata field names, so
|
||||
/organization/new must write them to the budget row and keep the datetime out of the org
|
||||
metadata JSON (a datetime there broke JSON serialization and 500'd the request)."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, NewOrganizationRequest, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.organization_endpoints import new_organization
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
expiry = datetime(2099, 1, 1, tzinfo=timezone.utc)
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.jsonify_object = MagicMock(side_effect=lambda data: PrismaClient.jsonify_object(prisma_client, data))
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1"))
|
||||
prisma_client.db.litellm_organizationtable.create = AsyncMock(return_value={"organization_id": "org-1"})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False)
|
||||
|
||||
response = await new_organization(
|
||||
data=NewOrganizationRequest(
|
||||
organization_alias="org",
|
||||
max_budget=10,
|
||||
temp_budget_increase=5,
|
||||
temp_budget_expiry=expiry,
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert response == {"organization_id": "org-1"}
|
||||
budget_write = prisma_client.db.litellm_budgettable.create.await_args.kwargs["data"]
|
||||
assert (budget_write["max_budget"], budget_write["temp_budget_increase"], budget_write["temp_budget_expiry"]) == (
|
||||
10,
|
||||
5,
|
||||
expiry,
|
||||
)
|
||||
org_write = prisma_client.db.litellm_organizationtable.create.await_args.kwargs["data"]
|
||||
assert org_write["budget_id"] == "budget-1"
|
||||
assert json.loads(org_write.get("metadata", "{}")) == {}
|
||||
|
||||
|
||||
def test_v2_update_organization_is_in_openapi_schema():
|
||||
"""PATCH /v2/organization/{organization_id} is documented in the generated OpenAPI spec."""
|
||||
from fastapi import FastAPI
|
||||
|
|
|
|||
|
|
@ -15571,3 +15571,37 @@ async def test_team_info_reports_what_the_caller_may_edit(caller, org_admin, ena
|
|||
)
|
||||
|
||||
assert response["team_info"].caller_edit_access.model_dump(mode="json") == expected
|
||||
|
||||
|
||||
def test_member_budget_patch_maps_temp_budget_fields() -> None:
|
||||
from litellm.proxy.management_endpoints.common_utils import member_budget_patch
|
||||
|
||||
expiry: Final = datetime(2030, 1, 1, tzinfo=timezone.utc)
|
||||
request: Final = TeamMemberUpdateRequest(
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
temp_budget_increase=50.0,
|
||||
temp_budget_expiry=expiry,
|
||||
)
|
||||
assert member_budget_patch(request) == {
|
||||
"temp_budget_increase": 50.0,
|
||||
"temp_budget_expiry": expiry,
|
||||
}
|
||||
|
||||
|
||||
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)
|
||||
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_expiry="2030-01-01T00:00:00Z")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("increase", "message"),
|
||||
[(-1.0, "greater than or equal to 0"), (float("inf"), "finite number")],
|
||||
)
|
||||
def test_team_member_update_request_rejects_unusable_temp_budget_increase(increase: float, message: str) -> None:
|
||||
with pytest.raises(ValidationError, match=message):
|
||||
TeamMemberUpdateRequest(
|
||||
team_id="team-1", user_id="user-1", temp_budget_increase=increase, temp_budget_expiry="2030-01-01T00:00:00Z"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import math
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -9,10 +10,20 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.caching import DualCache
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
_get_team_member_budget_counter,
|
||||
count_request_input_tokens,
|
||||
estimate_request_max_cost,
|
||||
reserve_budget_for_request,
|
||||
|
|
@ -445,3 +456,93 @@ async def test_models_without_a_rust_tokenizer_stay_in_python(
|
|||
assert factory.calls == []
|
||||
assert dict(counts) == dict(python_counts)
|
||||
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"expiry_offset, expected_max_budget",
|
||||
[
|
||||
(timedelta(days=1), 3.0),
|
||||
(timedelta(days=-1), 2.0),
|
||||
],
|
||||
)
|
||||
async def test_team_member_reservation_counter_honours_temp_budget_increase(
|
||||
expiry_offset: timedelta, expected_max_budget: float
|
||||
) -> None:
|
||||
user_id: Final = "member-temp"
|
||||
team_id: Final = "team-temp"
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
spend=0.5,
|
||||
budget_id="budget-temp",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(
|
||||
max_budget=2.0,
|
||||
temp_budget_increase=1.0,
|
||||
temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
counter: Final = await _get_team_member_budget_counter(
|
||||
valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id),
|
||||
team_object=LiteLLM_TeamTable(team_id=team_id),
|
||||
user_object=LiteLLM_UserTable(user_id=user_id),
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
|
||||
assert counter is not None
|
||||
assert counter.max_budget == expected_max_budget
|
||||
assert counter.fallback_spend == 0.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"default_cap, expiry_offset, expected_max_budget",
|
||||
[
|
||||
(2.0, timedelta(days=1), 3.0),
|
||||
(2.0, timedelta(days=-1), 2.0),
|
||||
(0.0, timedelta(days=1), None),
|
||||
],
|
||||
)
|
||||
async def test_team_member_reservation_counter_adds_temp_increase_to_live_team_default(
|
||||
default_cap: float, expiry_offset: timedelta, expected_max_budget: float | None
|
||||
) -> None:
|
||||
user_id: Final = "member-bare"
|
||||
team_id: Final = "team-bare"
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key="team_member_default_budget:default-bare",
|
||||
value=LiteLLM_BudgetTable(budget_id="default-bare", max_budget=default_cap),
|
||||
)
|
||||
await cache.async_set_cache(
|
||||
key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
spend=0.5,
|
||||
budget_id="budget-bare",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(
|
||||
max_budget=None,
|
||||
temp_budget_increase=1.0,
|
||||
temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
counter: Final = await _get_team_member_budget_counter(
|
||||
valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id),
|
||||
team_object=LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": "default-bare"}),
|
||||
user_object=LiteLLM_UserTable(user_id=user_id),
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
|
||||
if expected_max_budget is None:
|
||||
assert counter is None
|
||||
return
|
||||
assert counter is not None
|
||||
assert counter.max_budget == expected_max_budget
|
||||
assert counter.fallback_spend == 0.5
|
||||
|
|
|
|||
|
|
@ -2821,6 +2821,8 @@ export interface Member {
|
|||
rpm_limit?: number | null;
|
||||
budget_duration?: string | null;
|
||||
allowed_models?: string[] | null;
|
||||
temp_budget_increase?: number | null;
|
||||
temp_budget_expiry?: string | null;
|
||||
}
|
||||
|
||||
export const teamMemberAddCall = async (accessToken: string, teamId: string, formValues: Member) => {
|
||||
|
|
@ -2955,6 +2957,12 @@ export const teamMemberUpdateCall = async (
|
|||
if (formValues.allowed_models !== undefined) {
|
||||
requestBody.allowed_models = formValues.allowed_models;
|
||||
}
|
||||
if ("temp_budget_increase" in formValues) {
|
||||
requestBody.temp_budget_increase = orNull(formValues.temp_budget_increase);
|
||||
}
|
||||
if ("temp_budget_expiry" in formValues) {
|
||||
requestBody.temp_budget_expiry = orNull(formValues.temp_budget_expiry);
|
||||
}
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
|
|
|
|||
|
|
@ -28,6 +28,27 @@ const additionalFields = [
|
|||
|
||||
const teamMemberConfig = { title: "Edit Member", showEmail: true, showUserId: true, roleOptions, additionalFields };
|
||||
|
||||
const tempBudgetConfig = {
|
||||
...teamMemberConfig,
|
||||
additionalFields: [
|
||||
...additionalFields,
|
||||
{ name: "temp_budget_increase", label: "Temporary Budget Increase (USD)", type: "numerical" as const, step: 0.01 },
|
||||
{ name: "temp_budget_expiry", label: "Temporary Budget Expiry (UTC)", type: "utc-datetime" as const },
|
||||
],
|
||||
};
|
||||
|
||||
const TEMP_BUDGET_PAIR_MESSAGE = "Set both a temporary budget increase and its expiry, or neither";
|
||||
|
||||
const cappedMember = { user_id: "u1", user_email: "a@b.com", role: "user", max_budget_in_team: 10 };
|
||||
|
||||
const tempBudgetMember = {
|
||||
user_id: "u1",
|
||||
user_email: "a@b.com",
|
||||
role: "user",
|
||||
temp_budget_increase: 25,
|
||||
temp_budget_expiry: "2030-01-02T03:04:00Z",
|
||||
};
|
||||
|
||||
const orgMemberConfig = { title: "Edit Member", showEmail: true, showUserId: true, roleOptions };
|
||||
|
||||
type Member = Record<string, unknown>;
|
||||
|
|
@ -242,6 +263,50 @@ describe("EditMembership submit payload", () => {
|
|||
await waitFor(() => expect(onSubmit).not.toHaveBeenCalled());
|
||||
});
|
||||
|
||||
it("submits a typed temporary increase with its expiry as a UTC ISO timestamp", async () => {
|
||||
renderEdit(tempBudgetConfig, cappedMember);
|
||||
|
||||
fireEvent.change(screen.getByLabelText("Temporary Budget Increase (USD)"), { target: { value: "25" } });
|
||||
fireEvent.change(screen.getByLabelText("Temporary Budget Expiry (UTC)"), { target: { value: "2030-01-02T03:04" } });
|
||||
|
||||
save();
|
||||
|
||||
await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce());
|
||||
expect(submitted().max_budget_in_team).toBe(10);
|
||||
expect(submitted().temp_budget_increase).toBe("25");
|
||||
expect(submitted().temp_budget_expiry).toBe("2030-01-02T03:04:00.000Z");
|
||||
});
|
||||
|
||||
it("seeds a stored temporary budget into the controls and clears both to null when the operator blanks them", async () => {
|
||||
renderEdit(tempBudgetConfig, tempBudgetMember);
|
||||
|
||||
expect(screen.getByLabelText("Temporary Budget Increase (USD)")).toHaveValue(25);
|
||||
expect(screen.getByLabelText("Temporary Budget Expiry (UTC)")).toHaveValue("2030-01-02T03:04");
|
||||
|
||||
fireEvent.change(screen.getByLabelText("Temporary Budget Increase (USD)"), { target: { value: "" } });
|
||||
fireEvent.change(screen.getByLabelText("Temporary Budget Expiry (UTC)"), { target: { value: "" } });
|
||||
|
||||
save();
|
||||
|
||||
await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce());
|
||||
expect(submitted().temp_budget_increase).toBeNull();
|
||||
expect(submitted().temp_budget_expiry).toBeNull();
|
||||
});
|
||||
|
||||
it.each([
|
||||
["Temporary Budget Increase (USD)", "25"],
|
||||
["Temporary Budget Expiry (UTC)", "2030-01-02T03:04"],
|
||||
])("blocks submission when only %s is set", async (label, value) => {
|
||||
renderEdit(tempBudgetConfig, { user_id: "u1", user_email: "a@b.com", role: "user" });
|
||||
|
||||
fireEvent.change(screen.getByLabelText(label), { target: { value } });
|
||||
|
||||
save();
|
||||
|
||||
expect(await screen.findByText(TEMP_BUDGET_PAIR_MESSAGE)).toBeInTheDocument();
|
||||
expect(onSubmit).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("clears the fields once the submit handler resolves", async () => {
|
||||
renderEdit(orgMemberConfig, { user_id: "u1", user_email: "a@b.com", role: "user" });
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import React, { useEffect, useMemo, useState } from "react";
|
||||
import dayjs from "dayjs";
|
||||
import utc from "dayjs/plugin/utc";
|
||||
import { z } from "zod/v4";
|
||||
import NumericalInput from "../shared/numerical_input";
|
||||
import BudgetDurationDropdown from "../common_components/budget_duration_dropdown";
|
||||
|
|
@ -9,17 +11,22 @@ import { Button } from "@/components/ui/button";
|
|||
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 { useZodForm } from "@/lib/forms/useZodForm";
|
||||
import {
|
||||
buildMemberFormData,
|
||||
buildMemberFormValues,
|
||||
emptyMemberFormValues,
|
||||
TEMP_BUDGET_PAIR_MESSAGE,
|
||||
tempBudgetPairError,
|
||||
type MemberAdditionalField,
|
||||
type MemberFieldsConfig,
|
||||
type MemberFormValues,
|
||||
} from "./memberFormValues";
|
||||
import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
|
||||
dayjs.extend(utc);
|
||||
|
||||
interface BaseMember {
|
||||
user_email?: string;
|
||||
user_id?: string;
|
||||
|
|
@ -53,7 +60,10 @@ const buildMemberSchema = (config: ModalConfig): z.ZodType<MemberFormValues, Mem
|
|||
...Object.fromEntries((config.additionalFields ?? []).map((field) => [field.name, memberFieldSchema])),
|
||||
};
|
||||
|
||||
return z.object(shape);
|
||||
return z.object(shape).superRefine((values, ctx) => {
|
||||
const path = tempBudgetPairError(values);
|
||||
if (path !== null) ctx.addIssue({ code: "custom", path: [path], message: TEMP_BUDGET_PAIR_MESSAGE });
|
||||
});
|
||||
};
|
||||
|
||||
const MemberModal = <T extends BaseMember>({
|
||||
|
|
@ -160,6 +170,16 @@ const MemberModal = <T extends BaseMember>({
|
|||
onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)}
|
||||
/>
|
||||
);
|
||||
case "utc-datetime":
|
||||
return (
|
||||
<UtcDateTimeInput
|
||||
{...rest}
|
||||
id={id}
|
||||
ref={ref}
|
||||
value={typeof value === "string" && value !== "" ? dayjs.utc(value) : null}
|
||||
onChange={(next) => onChange(next === null ? null : next.toISOString())}
|
||||
/>
|
||||
);
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -264,6 +264,8 @@ export interface TeamMembership {
|
|||
budget_duration: string | null;
|
||||
budget_reset_at: string | null;
|
||||
allowed_models?: string[] | null;
|
||||
temp_budget_increase?: number | null;
|
||||
temp_budget_expiry?: string | null;
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -799,6 +801,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
rpm_limit: values.rpm_limit,
|
||||
budget_duration: values.budget_duration,
|
||||
allowed_models: values.allowed_models,
|
||||
temp_budget_increase: values.temp_budget_increase,
|
||||
temp_budget_expiry: values.temp_budget_expiry,
|
||||
};
|
||||
toast.dismiss(); // Remove all existing toasts
|
||||
|
||||
|
|
@ -2306,6 +2310,33 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
),
|
||||
type: "budget-duration" as const,
|
||||
},
|
||||
{
|
||||
name: "temp_budget_increase",
|
||||
label: (
|
||||
<span>
|
||||
Temporary Budget Increase (USD){" "}
|
||||
<SimpleTooltip content="Extra USD added on top of the team member budget until the expiry below. The permanent budget is left unchanged and the increase stops applying at expiry.">
|
||||
<Info className="ml-1 inline size-3.5 align-text-bottom" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
),
|
||||
type: "numerical" as const,
|
||||
step: 0.01,
|
||||
min: 0,
|
||||
placeholder: "Extra budget for this member until the expiry",
|
||||
},
|
||||
{
|
||||
name: "temp_budget_expiry",
|
||||
label: (
|
||||
<span>
|
||||
Temporary Budget Expiry (UTC){" "}
|
||||
<SimpleTooltip content="When the temporary budget increase stops applying. Required whenever a temporary budget increase is set.">
|
||||
<Info className="ml-1 inline size-3.5 align-text-bottom" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
),
|
||||
type: "utc-datetime" as const,
|
||||
},
|
||||
{
|
||||
name: "tpm_limit",
|
||||
label: (
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event";
|
|||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { renderWithProviders } from "../../../tests/test-utils";
|
||||
import { TeamData } from "./TeamInfo";
|
||||
import TeamMembersComponent from "./TeamMemberTab";
|
||||
import TeamMembersComponent, { seedMemberBudgetFields } from "./TeamMemberTab";
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
|
||||
useUISettings: vi.fn(),
|
||||
|
|
@ -380,6 +380,35 @@ describe("TeamMembersComponent", () => {
|
|||
expect(mockSetSelectedEditMember).toHaveBeenCalledWith(expect.objectContaining(zeroLimitsMember));
|
||||
});
|
||||
|
||||
it("seeds the edit payload with the stored temporary budget increase and expiry, keeping a 0 increase as 0", () => {
|
||||
const budget = {
|
||||
...createMockTeamData().team_memberships[0].litellm_budget_table,
|
||||
temp_budget_increase: 0,
|
||||
temp_budget_expiry: "2030-01-02T03:04:00Z",
|
||||
};
|
||||
|
||||
const seeded = {
|
||||
user_id: "user1@test.com",
|
||||
role: "member",
|
||||
max_budget_in_team: 1000,
|
||||
tpm_limit: 10000,
|
||||
rpm_limit: 100,
|
||||
budget_duration: null,
|
||||
allowed_models: [],
|
||||
temp_budget_increase: 0,
|
||||
temp_budget_expiry: "2030-01-02T03:04:00Z",
|
||||
};
|
||||
expect(seedMemberBudgetFields({ user_id: "user1@test.com", role: "member" }, budget)).toStrictEqual(seeded);
|
||||
});
|
||||
|
||||
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,
|
||||
temp_budget_increase: null,
|
||||
temp_budget_expiry: null,
|
||||
});
|
||||
});
|
||||
|
||||
it("should call setIsAddMemberModalVisible when Add Member button is clicked", async () => {
|
||||
const user = userEvent.setup();
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,21 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
|
|||
import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles";
|
||||
import { CircleHelp } from "lucide-react";
|
||||
import type { ComponentProps } from "react";
|
||||
import { TeamData } from "./TeamInfo";
|
||||
import { TeamData, TeamMembership } from "./TeamInfo";
|
||||
|
||||
export const seedMemberBudgetFields = (
|
||||
record: Member,
|
||||
budget: TeamMembership["litellm_budget_table"] | undefined,
|
||||
): Member => ({
|
||||
...record,
|
||||
max_budget_in_team: budget?.max_budget ?? null,
|
||||
tpm_limit: budget?.tpm_limit ?? null,
|
||||
rpm_limit: budget?.rpm_limit ?? null,
|
||||
budget_duration: budget?.budget_duration || null,
|
||||
allowed_models: budget?.allowed_models || [],
|
||||
temp_budget_increase: budget?.temp_budget_increase ?? null,
|
||||
temp_budget_expiry: budget?.temp_budget_expiry ?? null,
|
||||
});
|
||||
|
||||
interface TeamMemberTabProps {
|
||||
teamData: TeamData;
|
||||
|
|
@ -192,15 +206,7 @@ export default function TeamMemberTab({
|
|||
canEdit={canEditTeam}
|
||||
onEdit={(record) => {
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === record.user_id);
|
||||
const enhancedMember = {
|
||||
...record,
|
||||
max_budget_in_team: membership?.litellm_budget_table?.max_budget ?? null,
|
||||
tpm_limit: membership?.litellm_budget_table?.tpm_limit ?? null,
|
||||
rpm_limit: membership?.litellm_budget_table?.rpm_limit ?? null,
|
||||
budget_duration: membership?.litellm_budget_table?.budget_duration || null,
|
||||
allowed_models: membership?.litellm_budget_table?.allowed_models || [],
|
||||
};
|
||||
setSelectedEditMember(enhancedMember);
|
||||
setSelectedEditMember(seedMemberBudgetFields(record, membership?.litellm_budget_table));
|
||||
setIsEditMemberModalVisible(true);
|
||||
}}
|
||||
onDelete={handleMemberDelete}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import {
|
|||
buildMemberFormValues,
|
||||
emptyMemberFormValues,
|
||||
memberFieldNames,
|
||||
tempBudgetPairError,
|
||||
type MemberFieldsConfig,
|
||||
} from "./memberFormValues";
|
||||
|
||||
|
|
@ -25,6 +26,16 @@ const teamConfig: MemberFieldsConfig = {
|
|||
],
|
||||
};
|
||||
|
||||
const tempBudgetConfig: MemberFieldsConfig = {
|
||||
roleOptions,
|
||||
showUserId: true,
|
||||
additionalFields: [
|
||||
{ name: "max_budget_in_team", label: "Budget", type: "numerical" },
|
||||
{ name: "temp_budget_increase", label: "Temp Increase", type: "numerical" },
|
||||
{ name: "temp_budget_expiry", label: "Temp Expiry", type: "utc-datetime" },
|
||||
],
|
||||
};
|
||||
|
||||
const orgConfig: MemberFieldsConfig = { roleOptions, showEmail: true, showUserId: true };
|
||||
|
||||
describe("memberFieldNames", () => {
|
||||
|
|
@ -117,6 +128,30 @@ describe("buildMemberFormValues", () => {
|
|||
).toStrictEqual(unlimitedMember);
|
||||
});
|
||||
|
||||
it("seeds a stored temporary budget increase and its expiry, keeping a 0 increase as 0", () => {
|
||||
const tempBudgetMember = {
|
||||
user_id: "u1",
|
||||
role: "user",
|
||||
max_budget_in_team: 10,
|
||||
temp_budget_increase: 0,
|
||||
temp_budget_expiry: "2030-01-01T00:00:00Z",
|
||||
};
|
||||
expect(buildMemberFormValues("edit", tempBudgetMember, tempBudgetConfig)).toStrictEqual(tempBudgetMember);
|
||||
});
|
||||
|
||||
it("collapses a missing temporary budget increase and expiry to null", () => {
|
||||
const noTempBudget = {
|
||||
user_id: "u1",
|
||||
role: "user",
|
||||
max_budget_in_team: null,
|
||||
temp_budget_increase: null,
|
||||
temp_budget_expiry: null,
|
||||
};
|
||||
expect(buildMemberFormValues("edit", { user_id: "u1", role: "user" }, tempBudgetConfig)).toStrictEqual(
|
||||
noTempBudget,
|
||||
);
|
||||
});
|
||||
|
||||
it("falls back to the configured default role when the member has none", () => {
|
||||
expect(buildMemberFormValues("edit", { user_id: "u1", role: "" }, { ...orgConfig, defaultRole: "user" }).role).toBe(
|
||||
"user",
|
||||
|
|
@ -157,6 +192,17 @@ describe("emptyMemberFormValues", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("clears a utc-datetime field to null", () => {
|
||||
const cleared = {
|
||||
user_id: "",
|
||||
role: "",
|
||||
max_budget_in_team: null,
|
||||
temp_budget_increase: null,
|
||||
temp_budget_expiry: null,
|
||||
};
|
||||
expect(emptyMemberFormValues(tempBudgetConfig)).toStrictEqual(cleared);
|
||||
});
|
||||
|
||||
it("clears numeric, duration and multi-select fields to values their controls accept", () => {
|
||||
expect(
|
||||
emptyMemberFormValues({
|
||||
|
|
@ -199,9 +245,12 @@ describe("buildMemberFormData", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it.each(["max_budget_in_team", "tpm_limit", "rpm_limit"])("turns a blank %s into null", (key) => {
|
||||
expect(buildMemberFormData({ [key]: " " })[key]).toBeNull();
|
||||
});
|
||||
it.each(["max_budget_in_team", "tpm_limit", "rpm_limit", "temp_budget_increase"])(
|
||||
"turns a blank %s into null",
|
||||
(key) => {
|
||||
expect(buildMemberFormData({ [key]: " " })[key]).toBeNull();
|
||||
},
|
||||
);
|
||||
|
||||
it.each(["user_email", "user_id", "budget_duration"])("leaves a blank %s as an empty string", (key) => {
|
||||
expect(buildMemberFormData({ [key]: " " })[key]).toBe("");
|
||||
|
|
@ -226,3 +275,27 @@ describe("buildMemberFormData", () => {
|
|||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("tempBudgetPairError", () => {
|
||||
it.each([
|
||||
[{ temp_budget_increase: 50, temp_budget_expiry: "2030-01-01T00:00:00.000Z" }],
|
||||
[{ temp_budget_increase: "0", temp_budget_expiry: "2030-01-01T00:00:00.000Z" }],
|
||||
[{ temp_budget_increase: 0, temp_budget_expiry: "2030-01-01T00:00:00.000Z" }],
|
||||
[{ temp_budget_increase: null, temp_budget_expiry: null }],
|
||||
[{ temp_budget_increase: "", temp_budget_expiry: null }],
|
||||
[{}],
|
||||
])("accepts %j", (values) => {
|
||||
expect(tempBudgetPairError(values)).toBeNull();
|
||||
});
|
||||
|
||||
it("points at the missing increase when only the expiry is set", () => {
|
||||
expect(tempBudgetPairError({ temp_budget_increase: "", temp_budget_expiry: "2030-01-01T00:00:00.000Z" })).toBe(
|
||||
"temp_budget_increase",
|
||||
);
|
||||
});
|
||||
|
||||
it("points at the missing expiry when only the increase is set", () => {
|
||||
expect(tempBudgetPairError({ temp_budget_increase: 25, temp_budget_expiry: null })).toBe("temp_budget_expiry");
|
||||
expect(tempBudgetPairError({ temp_budget_increase: 0, temp_budget_expiry: undefined })).toBe("temp_budget_expiry");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ export type MemberFieldValue = string | number | null | undefined | string[];
|
|||
|
||||
export type MemberFormValues = Record<string, MemberFieldValue>;
|
||||
|
||||
export type MemberFieldType = "input" | "select" | "numerical" | "multi-select" | "budget-duration";
|
||||
export type MemberFieldType = "input" | "select" | "numerical" | "multi-select" | "budget-duration" | "utc-datetime";
|
||||
|
||||
export interface MemberAdditionalField {
|
||||
name: string;
|
||||
|
|
@ -22,7 +22,23 @@ export interface MemberFieldsConfig {
|
|||
additionalFields?: Array<MemberAdditionalField>;
|
||||
}
|
||||
|
||||
const NULLABLE_NUMERIC_FIELDS: ReadonlySet<string> = new Set(["max_budget_in_team", "tpm_limit", "rpm_limit"]);
|
||||
const NULLABLE_NUMERIC_FIELDS: ReadonlySet<string> = new Set([
|
||||
"max_budget_in_team",
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
"temp_budget_increase",
|
||||
]);
|
||||
|
||||
export const TEMP_BUDGET_PAIR_MESSAGE = "Set both a temporary budget increase and its expiry, or neither";
|
||||
|
||||
const isUnset = (value: MemberFieldValue): boolean => value === null || value === undefined || 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);
|
||||
if (increaseUnset === expiryUnset) return null;
|
||||
return increaseUnset ? "temp_budget_increase" : "temp_budget_expiry";
|
||||
};
|
||||
|
||||
export const memberFieldNames = (config: MemberFieldsConfig): string[] => [
|
||||
...(config.showEmail ? ["user_email"] : []),
|
||||
|
|
@ -48,6 +64,8 @@ export const buildMemberFormValues = (
|
|||
rpm_limit: initialData.rpm_limit ?? null,
|
||||
budget_duration: initialData.budget_duration || null,
|
||||
allowed_models: initialData.allowed_models || [],
|
||||
temp_budget_increase: initialData.temp_budget_increase ?? null,
|
||||
temp_budget_expiry: initialData.temp_budget_expiry || null,
|
||||
};
|
||||
|
||||
return pickFieldNames(config, seeded);
|
||||
|
|
@ -62,6 +80,7 @@ const emptyValueForType = (type: MemberFieldType | undefined): MemberFieldValue
|
|||
return [];
|
||||
case "numerical":
|
||||
case "budget-duration":
|
||||
case "utc-datetime":
|
||||
return null;
|
||||
default:
|
||||
return "";
|
||||
|
|
|
|||
36
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
36
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -10743,6 +10743,8 @@ export interface paths {
|
|||
* - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
|
||||
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - organization-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.
|
||||
* - allowed_models: Optional[List[str]] - List of models the organization is allowed to access. If not set, defaults to the models field.
|
||||
* - temp_budget_increase: *Optional[float]* - Stored on the org budget row but only enforced for team member budgets today.
|
||||
* - temp_budget_expiry: *Optional[str]* - Stored on the org budget row but only enforced for team member budgets today.
|
||||
* Case 1: Create new org **without** a budget_id
|
||||
*
|
||||
* ```bash
|
||||
|
|
@ -29617,6 +29619,10 @@ export interface components {
|
|||
rpm_limit?: number | null;
|
||||
/** Soft Budget */
|
||||
soft_budget?: number | null;
|
||||
/** Temp Budget Expiry */
|
||||
temp_budget_expiry?: string | null;
|
||||
/** Temp Budget Increase */
|
||||
temp_budget_increase?: number | null;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
@ -29652,6 +29658,10 @@ export interface components {
|
|||
rpm_limit?: number | null;
|
||||
/** Soft Budget */
|
||||
soft_budget?: number | null;
|
||||
/** Temp Budget Expiry */
|
||||
temp_budget_expiry?: string | null;
|
||||
/** Temp Budget Increase */
|
||||
temp_budget_increase?: number | null;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
@ -33465,6 +33475,10 @@ export interface components {
|
|||
rpm_limit?: number | null;
|
||||
/** Soft Budget */
|
||||
soft_budget?: number | null;
|
||||
/** Temp Budget Expiry */
|
||||
temp_budget_expiry?: string | null;
|
||||
/** Temp Budget Increase */
|
||||
temp_budget_increase?: number | null;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
@ -33588,6 +33602,10 @@ export interface components {
|
|||
tags?: string[] | null;
|
||||
/** Team Id */
|
||||
team_id: string;
|
||||
/** Temp Budget Expiry */
|
||||
temp_budget_expiry?: string | null;
|
||||
/** Temp Budget Increase */
|
||||
temp_budget_increase?: number | null;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
@ -38371,6 +38389,16 @@ export interface components {
|
|||
rpm_limit?: number | null;
|
||||
/** Team Id */
|
||||
team_id: string;
|
||||
/**
|
||||
* Temp Budget Expiry
|
||||
* @description UTC expiry for temp_budget_increase
|
||||
*/
|
||||
temp_budget_expiry?: string | null;
|
||||
/**
|
||||
* Temp Budget Increase
|
||||
* @description Temporary additive budget increase for this team member, active until temp_budget_expiry
|
||||
*/
|
||||
temp_budget_increase?: number | null;
|
||||
/**
|
||||
* Tpm Limit
|
||||
* @description Tokens per minute limit for this team member
|
||||
|
|
@ -38393,6 +38421,10 @@ export interface components {
|
|||
rpm_limit?: number | null;
|
||||
/** Team Id */
|
||||
team_id: string;
|
||||
/** Temp Budget Expiry */
|
||||
temp_budget_expiry?: string | null;
|
||||
/** Temp Budget Increase */
|
||||
temp_budget_increase?: number | null;
|
||||
/** Tpm Limit */
|
||||
tpm_limit?: number | null;
|
||||
/** User Email */
|
||||
|
|
@ -39561,6 +39593,10 @@ export interface components {
|
|||
tags?: string[] | null;
|
||||
/** Team Id */
|
||||
team_id?: string | null;
|
||||
/** Temp Budget Expiry */
|
||||
temp_budget_expiry?: string | null;
|
||||
/** Temp Budget Increase */
|
||||
temp_budget_increase?: number | null;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue