mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
feat: multiple concurrent budget windows per API key and team (#24883)
* feat(proxy): add BudgetLimitEntry type and wire budget_limits into key/team models * feat(schema): add budget_limits Json column to VerificationToken and TeamTable * feat(migrations): add migration for budget_limits column on keys and teams * feat(keys): initialize budget_limits windows with reset_at on key create/update * feat(teams): initialize budget_limits windows with reset_at on team create/update * feat(auth): add _virtual_key_multi_budget_check and _team_multi_budget_check * feat(auth): call multi-budget checks from common_checks for keys and teams * feat(proxy): increment per-window Redis spend counters after each request * feat(budget): reset individual budget windows on schedule via reset_budget_job * feat(ui): add hourly option to BudgetDurationDropdown * feat(ui): add budget_limits field to KeyResponse type * feat(ui): add Budget Windows editor to key edit view * feat(ui): add Budget Windows editor to create key form * fix(proxy): strip budget_limits=None before Prisma upsert to fix login 500 Prisma rejects nullable JSON fields (Json? without @default) when passed as Python None — it needs the field omitted entirely so the DB stores NULL via the column's nullable constraint. This was breaking /v2/login because the UI session key creation path hit the upsert with budget_limits=None. * ui(key-edit): use antd InputNumber+Button for budget windows, add reset hints * ui(create-key): use antd InputNumber+Button for budget windows, add reset hints * docs(users): add multiple budget windows section with API + dashboard walkthrough * fix: BudgetExceededError returns HTTP 429 instead of 400 - Add status_code=429 to BudgetExceededError class - auth_exception_handler hardcoded code=400 → code=429 * fix: no-op else branch in multi-budget auth checks causes KeyError - BudgetLimitEntry objects must be coerced via model_dump() not left as-is - Move _virtual_key_multi_budget_check into common_checks (was asymmetric with _team_multi_budget_check which already lived there) * fix: len() on JSON string returns char count not window count Guard with isinstance check + json.loads() before iterating per-window Redis counters in increment_spend_counters * fix: silent except:pass hides Redis reset failures in reset_budget_windows Log Redis counter reset failures as warnings so they are observable * test: add unit tests for multi-budget window enforcement 5 tests covering: no budget_limits passes, under budget passes, over hourly window raises 429, over monthly window raises 429, BudgetLimitEntry objects coerced without KeyError * fix: key per-window counters stable across reorders (duration key, not index) * fix: team+key per-window spend increments use duration key, not index * fix: budget window reset uses duration key; log failures instead of swallowing * refactor: extract BudgetWindowsEditor to shared component * refactor: key_edit_view imports BudgetWindowsEditor from shared component * refactor: create_key_button imports BudgetWindowsEditor from shared component --------- Co-authored-by: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
This commit is contained in:
ishaan-berri
GitHub
parent
fc885af994
commit
1e6203b735
20 changed files with 806 additions and 156 deletions
|
|
@ -333,6 +333,67 @@ curl 'http://0.0.0.0:4000/key/generate' \
|
|||
}'
|
||||
```
|
||||
|
||||
#### **Set multiple budget windows on a key**
|
||||
|
||||
Apply multiple concurrent budget limits at different time scales on the same key — for example, cap a key at **$10/day** AND **$100/month**.
|
||||
|
||||
**When is this useful?**
|
||||
|
||||
A single `budget_duration` window can't prevent a bad day from burning your entire month. Multiple budget windows let you:
|
||||
|
||||
- Block a runaway usage spike within the day while still allowing normal monthly spend.
|
||||
- Give Claude Code rollouts a daily guardrail (`24h`) and a monthly ceiling (`30d`) so a single heavy session doesn't exhaust the whole month.
|
||||
- Layer fine-grained hourly limits for bursty workloads on top of a weekly cap.
|
||||
|
||||
:::info
|
||||
|
||||
See [User Budget docs](https://docs.litellm.ai/docs/proxy/users) for more on how budgets work across keys, teams, and users.
|
||||
|
||||
:::
|
||||
|
||||
**Via API**
|
||||
|
||||
Pass `budget_limits` as a list of `{budget_duration, max_budget}` objects:
|
||||
|
||||
```bash
|
||||
curl 'http://0.0.0.0:4000/key/generate' \
|
||||
--header 'Authorization: Bearer <your-master-key>' \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data-raw '{
|
||||
"budget_limits": [
|
||||
{"budget_duration": "24h", "max_budget": 10},
|
||||
{"budget_duration": "30d", "max_budget": 100}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Each window is tracked independently and resets on its own schedule:
|
||||
|
||||
| `budget_duration` | Resets |
|
||||
|---|---|
|
||||
| `1h` | Every hour |
|
||||
| `24h` | Daily at midnight UTC |
|
||||
| `7d` | Every Sunday at midnight UTC |
|
||||
| `30d` | 1st of every month at midnight UTC |
|
||||
|
||||
**Via Dashboard**
|
||||
|
||||
Open **Virtual Keys → Create Key → Optional Settings → Budget Windows**.
|
||||
|
||||
|
||||
|
||||
Click **+ Add Budget Window** to add a row, choose the period from the dropdown, and enter the spend cap.
|
||||
|
||||
|
||||
|
||||
Add a second row for a different time period (e.g. monthly $100 on top of a daily $10).
|
||||
|
||||
|
||||
|
||||
Each window shows the reset schedule below the input so it's always clear when spend resets.
|
||||
|
||||
|
||||
|
||||
|
||||
### ✨ Virtual Key (Model Specific)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable: add budget_limits column to LiteLLM_VerificationToken
|
||||
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB;
|
||||
|
||||
-- AlterTable: add budget_limits column to LiteLLM_TeamTable
|
||||
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB;
|
||||
|
|
@ -140,6 +140,7 @@ model LiteLLM_TeamTable {
|
|||
team_member_permissions String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
policies String[] @default([])
|
||||
budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}]
|
||||
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
|
||||
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
|
||||
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
|
||||
|
|
@ -386,6 +387,7 @@ model LiteLLM_VerificationToken {
|
|||
rotation_interval String? // How often to rotate (e.g., "30d", "90d")
|
||||
last_rotation_at DateTime? // When this key was last rotated
|
||||
key_rotation_at DateTime? // When this key should next be rotated
|
||||
budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}]
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
|
||||
litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id])
|
||||
|
|
|
|||
|
|
@ -281,7 +281,7 @@ class Timeout(openai.APITimeoutError): # type: ignore
|
|||
return _message
|
||||
|
||||
|
||||
class PermissionDeniedError(openai.PermissionDeniedError): # type:ignore
|
||||
class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore
|
||||
def __init__(
|
||||
self,
|
||||
message,
|
||||
|
|
@ -847,6 +847,7 @@ class BudgetExceededError(Exception):
|
|||
):
|
||||
self.current_cost = current_cost
|
||||
self.max_budget = max_budget
|
||||
self.status_code = 429
|
||||
message = (
|
||||
message
|
||||
or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}"
|
||||
|
|
|
|||
|
|
@ -862,6 +862,14 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
|
|||
models: Optional[List[str]] = None
|
||||
|
||||
|
||||
class BudgetLimitEntry(LiteLLMPydanticObjectBase):
|
||||
"""A single budget window with its own limit and independent reset schedule."""
|
||||
|
||||
budget_duration: str # e.g. "24h", "7d", "30d"
|
||||
max_budget: float # max spend in USD for this window
|
||||
reset_at: Optional[datetime] = None # populated at creation/reset time
|
||||
|
||||
|
||||
class GenerateRequestBase(LiteLLMPydanticObjectBase):
|
||||
"""
|
||||
Overlapping schema between key and user generate/update requests
|
||||
|
|
@ -881,12 +889,15 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
|
|||
rpm_limit: Optional[int] = None
|
||||
|
||||
budget_duration: Optional[str] = None
|
||||
budget_limits: Optional[List[BudgetLimitEntry]] = (
|
||||
None # multiple concurrent budget windows
|
||||
)
|
||||
allowed_cache_controls: Optional[list] = []
|
||||
config: Optional[dict] = {}
|
||||
permissions: Optional[dict] = {}
|
||||
model_max_budget: Optional[
|
||||
dict
|
||||
] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
|
||||
model_max_budget: Optional[dict] = (
|
||||
{}
|
||||
) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
model_rpm_limit: Optional[dict] = None
|
||||
|
|
@ -985,6 +996,7 @@ class GenerateKeyResponse(KeyRequestBase):
|
|||
"permissions",
|
||||
"model_max_budget",
|
||||
"router_settings",
|
||||
"budget_limits",
|
||||
]
|
||||
for field in dict_fields:
|
||||
value = values.get(field)
|
||||
|
|
@ -1028,9 +1040,9 @@ class RegenerateKeyRequest(GenerateKeyRequest):
|
|||
spend: Optional[float] = None
|
||||
metadata: Optional[dict] = None
|
||||
new_master_key: Optional[str] = None
|
||||
grace_period: Optional[
|
||||
str
|
||||
] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke
|
||||
grace_period: Optional[str] = (
|
||||
None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke
|
||||
)
|
||||
|
||||
|
||||
class ResetSpendRequest(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -1540,12 +1552,12 @@ class NewCustomerRequest(BudgetNewRequest):
|
|||
blocked: bool = False # allow/disallow requests for this end-user
|
||||
budget_id: Optional[str] = None # give either a budget_id or max_budget
|
||||
spend: Optional[float] = None
|
||||
allowed_model_region: Optional[
|
||||
AllowedModelRegion
|
||||
] = None # require all user requests to use models in this specific region
|
||||
default_model: Optional[
|
||||
str
|
||||
] = None # if no equivalent model in allowed region - default all requests to this model
|
||||
allowed_model_region: Optional[AllowedModelRegion] = (
|
||||
None # require all user requests to use models in this specific region
|
||||
)
|
||||
default_model: Optional[str] = (
|
||||
None # if no equivalent model in allowed region - default all requests to this model
|
||||
)
|
||||
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
|
||||
|
||||
@model_validator(mode="before")
|
||||
|
|
@ -1568,12 +1580,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase):
|
|||
blocked: bool = False # allow/disallow requests for this end-user
|
||||
max_budget: Optional[float] = None
|
||||
budget_id: Optional[str] = None # give either a budget_id or max_budget
|
||||
allowed_model_region: Optional[
|
||||
AllowedModelRegion
|
||||
] = None # require all user requests to use models in this specific region
|
||||
default_model: Optional[
|
||||
str
|
||||
] = None # if no equivalent model in allowed region - default all requests to this model
|
||||
allowed_model_region: Optional[AllowedModelRegion] = (
|
||||
None # require all user requests to use models in this specific region
|
||||
)
|
||||
default_model: Optional[str] = (
|
||||
None # if no equivalent model in allowed region - default all requests to this model
|
||||
)
|
||||
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
|
||||
|
||||
|
||||
|
|
@ -1638,6 +1650,9 @@ class TeamBase(LiteLLMPydanticObjectBase):
|
|||
max_budget: Optional[float] = None
|
||||
soft_budget: Optional[float] = None
|
||||
budget_duration: Optional[str] = None
|
||||
budget_limits: Optional[List[BudgetLimitEntry]] = (
|
||||
None # multiple concurrent budget windows
|
||||
)
|
||||
|
||||
models: list = []
|
||||
blocked: bool = False
|
||||
|
|
@ -1663,15 +1678,15 @@ class NewTeamRequest(TeamBase):
|
|||
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
|
||||
|
||||
model_tpm_limit: Optional[Dict[str, int]] = None
|
||||
team_member_budget: Optional[
|
||||
float
|
||||
] = None # allow user to set a budget for all team members
|
||||
team_member_rpm_limit: Optional[
|
||||
int
|
||||
] = None # allow user to set RPM limit for all team members
|
||||
team_member_tpm_limit: Optional[
|
||||
int
|
||||
] = None # allow user to set TPM limit for all team members
|
||||
team_member_budget: Optional[float] = (
|
||||
None # allow user to set a budget for all team members
|
||||
)
|
||||
team_member_rpm_limit: Optional[int] = (
|
||||
None # allow user to set RPM limit for all team members
|
||||
)
|
||||
team_member_tpm_limit: Optional[int] = (
|
||||
None # allow user to set TPM limit for all team members
|
||||
)
|
||||
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
|
||||
team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo"
|
||||
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
|
||||
|
|
@ -1736,6 +1751,9 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
|
|||
enforced_file_expires_after: Optional[dict] = None
|
||||
router_settings: Optional[dict] = None
|
||||
access_group_ids: Optional[List[str]] = None
|
||||
budget_limits: Optional[List[BudgetLimitEntry]] = (
|
||||
None # multiple concurrent budget windows
|
||||
)
|
||||
|
||||
|
||||
class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -1768,9 +1786,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
class AddTeamCallback(LiteLLMPydanticObjectBase):
|
||||
callback_name: str
|
||||
callback_type: Optional[
|
||||
Literal["success", "failure", "success_and_failure"]
|
||||
] = "success_and_failure"
|
||||
callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = (
|
||||
"success_and_failure"
|
||||
)
|
||||
callback_vars: Dict[str, str]
|
||||
|
||||
@model_validator(mode="before")
|
||||
|
|
@ -1879,6 +1897,7 @@ class LiteLLM_TeamTable(TeamBase):
|
|||
"model_max_budget",
|
||||
"model_aliases",
|
||||
"router_settings",
|
||||
"budget_limits",
|
||||
]
|
||||
|
||||
if isinstance(values, BaseModel):
|
||||
|
|
@ -2110,9 +2129,9 @@ class ConfigList(LiteLLMPydanticObjectBase):
|
|||
stored_in_db: Optional[bool]
|
||||
field_default_value: Any
|
||||
premium_field: bool = False
|
||||
nested_fields: Optional[
|
||||
List[FieldDetail]
|
||||
] = None # For nested dictionary or Pydantic fields
|
||||
nested_fields: Optional[List[FieldDetail]] = (
|
||||
None # For nested dictionary or Pydantic fields
|
||||
)
|
||||
|
||||
|
||||
class UserHeaderMapping(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -2361,6 +2380,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
|
|||
last_rotation_at: Optional[datetime] = None # When this key was last rotated
|
||||
key_rotation_at: Optional[datetime] = None # When this key should next be rotated
|
||||
router_settings: Optional[dict] = None
|
||||
budget_limits: Optional[List[dict]] = None # multiple concurrent budget windows
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
|
||||
|
|
@ -2470,9 +2490,9 @@ class UserAPIKeyAuth(
|
|||
user_max_budget: Optional[float] = None
|
||||
request_route: Optional[str] = None
|
||||
user: Optional[Any] = None # Expanded user object when expand=user is used
|
||||
created_by_user: Optional[
|
||||
Any
|
||||
] = None # Expanded created_by user when expand=user is used
|
||||
created_by_user: Optional[Any] = (
|
||||
None # Expanded created_by user when expand=user is used
|
||||
)
|
||||
end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
|
||||
# Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery
|
||||
# and forwarded into outbound tokens by guardrails such as MCPJWTSigner.
|
||||
|
|
@ -2611,9 +2631,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase):
|
|||
budget_id: Optional[str] = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
user: Optional[
|
||||
Any
|
||||
] = None # You might want to replace 'Any' with a more specific type if available
|
||||
user: Optional[Any] = (
|
||||
None # You might want to replace 'Any' with a more specific type if available
|
||||
)
|
||||
litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
|
||||
user_email: Optional[str] = None
|
||||
|
||||
|
|
@ -3764,9 +3784,9 @@ class TeamModelDeleteRequest(BaseModel):
|
|||
# Organization Member Requests
|
||||
class OrganizationMemberAddRequest(OrgMemberAddRequest):
|
||||
organization_id: str
|
||||
max_budget_in_organization: Optional[
|
||||
float
|
||||
] = None # Users max budget within the organization
|
||||
max_budget_in_organization: Optional[float] = (
|
||||
None # Users max budget within the organization
|
||||
)
|
||||
|
||||
|
||||
class OrganizationMemberDeleteRequest(MemberDeleteRequest):
|
||||
|
|
@ -4017,9 +4037,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase):
|
|||
Maps provider names to their budget configs.
|
||||
"""
|
||||
|
||||
providers: Dict[
|
||||
str, ProviderBudgetResponseObject
|
||||
] = {} # Dictionary mapping provider names to their budget configurations
|
||||
providers: Dict[str, ProviderBudgetResponseObject] = (
|
||||
{}
|
||||
) # Dictionary mapping provider names to their budget configurations
|
||||
|
||||
|
||||
class ProxyStateVariables(TypedDict):
|
||||
|
|
@ -4181,9 +4201,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
enforce_rbac: bool = False
|
||||
roles_jwt_field: Optional[str] = None # v2 on role mappings
|
||||
role_mappings: Optional[List[RoleMapping]] = None
|
||||
object_id_jwt_field: Optional[
|
||||
str
|
||||
] = None # can be either user / team, inferred from the role mapping
|
||||
object_id_jwt_field: Optional[str] = (
|
||||
None # can be either user / team, inferred from the role mapping
|
||||
)
|
||||
scope_mappings: Optional[List[ScopeMapping]] = None
|
||||
enforce_scope_based_access: bool = False
|
||||
enforce_team_based_model_access: bool = False
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ Run checks for:
|
|||
2. If user is in budget
|
||||
3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
|
|
@ -456,9 +457,9 @@ async def common_checks( # noqa: PLR0915
|
|||
model=_model,
|
||||
team_object=team_object,
|
||||
llm_router=llm_router,
|
||||
team_model_aliases=valid_token.team_model_aliases
|
||||
if valid_token
|
||||
else None,
|
||||
team_model_aliases=(
|
||||
valid_token.team_model_aliases if valid_token else None
|
||||
),
|
||||
):
|
||||
raise ProxyException(
|
||||
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
|
||||
|
|
@ -518,6 +519,17 @@ async def common_checks( # noqa: PLR0915
|
|||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
# 3.1. Multi-window budget check for team
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.team_multi_budget_check"):
|
||||
await _team_multi_budget_check(team_object=team_object)
|
||||
|
||||
# 3.2. Multi-window budget check for key
|
||||
with tracer.trace(
|
||||
"litellm.proxy.auth.common_checks.virtual_key_multi_budget_check"
|
||||
):
|
||||
if valid_token is not None:
|
||||
await _virtual_key_multi_budget_check(valid_token=valid_token)
|
||||
|
||||
# 3.0.5. If team is over soft budget (alert only, doesn't block)
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"):
|
||||
await _team_soft_budget_check(
|
||||
|
|
@ -2966,6 +2978,43 @@ async def _virtual_key_max_budget_check(
|
|||
)
|
||||
|
||||
|
||||
async def _virtual_key_multi_budget_check(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
):
|
||||
"""
|
||||
Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded.
|
||||
|
||||
Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}.
|
||||
Using budget_duration (not list index) keeps counters stable when windows are reordered
|
||||
or removed during a key update.
|
||||
|
||||
Note: counters are not seeded from DB on Redis cold-start. After a Redis flush,
|
||||
per-window spend resets to zero within the current window period. This is an acceptable
|
||||
trade-off: the DB stores reset_at timestamps but not per-window accumulated spend.
|
||||
"""
|
||||
if not valid_token.budget_limits:
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
for window in valid_token.budget_limits:
|
||||
w: dict = window if isinstance(window, dict) else window.model_dump()
|
||||
counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}"
|
||||
window_spend = await get_current_spend(
|
||||
counter_key=counter_key,
|
||||
fallback_spend=0.0,
|
||||
)
|
||||
if window_spend >= w["max_budget"]:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=window_spend,
|
||||
max_budget=w["max_budget"],
|
||||
message=(
|
||||
f"ExceededBudget: Key over {w['budget_duration']} budget. "
|
||||
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _virtual_key_soft_budget_check(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
|
|
@ -3155,6 +3204,39 @@ async def _team_max_budget_check(
|
|||
)
|
||||
|
||||
|
||||
async def _team_multi_budget_check(
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
):
|
||||
"""
|
||||
Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded.
|
||||
|
||||
Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}.
|
||||
Using budget_duration (not list index) keeps counters stable when windows are reordered
|
||||
or removed during a team update.
|
||||
"""
|
||||
if team_object is None or not team_object.budget_limits:
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
for window in team_object.budget_limits:
|
||||
w: dict = window if isinstance(window, dict) else window.model_dump()
|
||||
counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}"
|
||||
window_spend = await get_current_spend(
|
||||
counter_key=counter_key,
|
||||
fallback_spend=0.0,
|
||||
)
|
||||
if window_spend >= w["max_budget"]:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=window_spend,
|
||||
max_budget=w["max_budget"],
|
||||
message=(
|
||||
f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. "
|
||||
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _team_soft_budget_check(
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
valid_token: Optional[UserAPIKeyAuth],
|
||||
|
|
|
|||
|
|
@ -106,7 +106,7 @@ class UserAPIKeyAuthExceptionHandler:
|
|||
message=e.message,
|
||||
type=ProxyErrorTypes.budget_exceeded,
|
||||
param=None,
|
||||
code=400,
|
||||
code=429,
|
||||
)
|
||||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
|
|
|
|||
|
|
@ -984,9 +984,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
route=route,
|
||||
)
|
||||
if _end_user_object is not None:
|
||||
end_user_params[
|
||||
"allowed_model_region"
|
||||
] = _end_user_object.allowed_model_region
|
||||
end_user_params["allowed_model_region"] = (
|
||||
_end_user_object.allowed_model_region
|
||||
)
|
||||
if _end_user_object.litellm_budget_table is not None:
|
||||
_apply_budget_limits_to_end_user_params(
|
||||
end_user_params=end_user_params,
|
||||
|
|
@ -1540,9 +1540,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
|
||||
if _end_user_object is not None:
|
||||
valid_token_dict.update(end_user_params)
|
||||
valid_token_dict[
|
||||
"end_user_object_permission"
|
||||
] = _end_user_object.object_permission
|
||||
valid_token_dict["end_user_object_permission"] = (
|
||||
_end_user_object.object_permission
|
||||
)
|
||||
|
||||
# check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions
|
||||
# sso/login, ui/login, /key functions and /user functions
|
||||
|
|
|
|||
|
|
@ -48,6 +48,9 @@ class ResetBudgetJob:
|
|||
### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ###
|
||||
await self.reset_budget_for_litellm_budget_table()
|
||||
|
||||
### RESET MULTI-WINDOW BUDGETS ###
|
||||
await self.reset_budget_windows()
|
||||
|
||||
async def reset_budget_for_litellm_team_members(
|
||||
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
|
||||
):
|
||||
|
|
@ -549,6 +552,115 @@ class ResetBudgetJob:
|
|||
)
|
||||
verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e)
|
||||
|
||||
async def reset_budget_windows(self) -> None:
|
||||
"""
|
||||
For keys and teams with budget_limits, reset any individual windows where
|
||||
reset_at <= now. Only the expired windows are reset; other windows are untouched.
|
||||
"""
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
now = datetime.utcnow()
|
||||
|
||||
# --- Keys ---
|
||||
try:
|
||||
all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"budget_limits": {"not": None}} # type: ignore[arg-type]
|
||||
)
|
||||
for key in all_keys:
|
||||
raw = key.budget_limits # type: ignore[attr-defined]
|
||||
if not raw:
|
||||
continue
|
||||
windows: list = raw if isinstance(raw, list) else json.loads(raw)
|
||||
changed = False
|
||||
for window in windows:
|
||||
reset_at_str = window.get("reset_at")
|
||||
if not reset_at_str:
|
||||
continue
|
||||
reset_at = datetime.fromisoformat(
|
||||
reset_at_str.replace("Z", "+00:00")
|
||||
).replace(tzinfo=None)
|
||||
if reset_at <= now:
|
||||
# Reset this window's counter (keyed by duration, not index)
|
||||
counter_key = (
|
||||
f"spend:key:{key.token}:window:{window['budget_duration']}"
|
||||
)
|
||||
spend_counter_cache.in_memory_cache.set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset Redis counter %s: %s",
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
window["reset_at"] = get_budget_reset_time(
|
||||
budget_duration=window["budget_duration"]
|
||||
).isoformat()
|
||||
changed = True
|
||||
if changed:
|
||||
await self.prisma_client.db.litellm_verificationtoken.update(
|
||||
where={"token": key.token},
|
||||
data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type]
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to reset budget windows for keys: %s", e
|
||||
)
|
||||
|
||||
# --- Teams ---
|
||||
try:
|
||||
all_teams = await self.prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"budget_limits": {"not": None}} # type: ignore[arg-type]
|
||||
)
|
||||
for team in all_teams:
|
||||
raw = team.budget_limits # type: ignore[attr-defined]
|
||||
if not raw:
|
||||
continue
|
||||
windows = raw if isinstance(raw, list) else json.loads(raw)
|
||||
changed = False
|
||||
for window in windows:
|
||||
reset_at_str = window.get("reset_at")
|
||||
if not reset_at_str:
|
||||
continue
|
||||
reset_at = datetime.fromisoformat(
|
||||
reset_at_str.replace("Z", "+00:00")
|
||||
).replace(tzinfo=None)
|
||||
if reset_at <= now:
|
||||
counter_key = f"spend:team:{team.team_id}:window:{window['budget_duration']}"
|
||||
spend_counter_cache.in_memory_cache.set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(
|
||||
key=counter_key, value=0.0
|
||||
)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset Redis counter %s: %s",
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
window["reset_at"] = get_budget_reset_time(
|
||||
budget_duration=window["budget_duration"]
|
||||
).isoformat()
|
||||
changed = True
|
||||
if changed:
|
||||
await self.prisma_client.db.litellm_teamtable.update(
|
||||
where={"team_id": team.team_id},
|
||||
data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type]
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to reset budget windows for teams: %s", e
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _reset_budget_common(
|
||||
item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken],
|
||||
|
|
@ -570,14 +682,14 @@ class ResetBudgetJob:
|
|||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
counter_key = None
|
||||
if item_type == "key" and hasattr(item, "token") and item.token is not None:
|
||||
counter_key = f"spend:key:{item.token}"
|
||||
if item_type == "key" and hasattr(item, "token") and item.token is not None: # type: ignore[union-attr]
|
||||
counter_key = f"spend:key:{item.token}" # type: ignore[union-attr]
|
||||
elif (
|
||||
item_type == "team"
|
||||
and hasattr(item, "team_id")
|
||||
and item.team_id is not None
|
||||
and item.team_id is not None # type: ignore[union-attr]
|
||||
):
|
||||
counter_key = f"spend:team:{item.team_id}"
|
||||
counter_key = f"spend:team:{item.team_id}" # type: ignore[union-attr]
|
||||
|
||||
if counter_key is not None:
|
||||
# Always reset in-memory (local fallback)
|
||||
|
|
|
|||
|
|
@ -732,9 +732,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
|||
request_type="key", **data_json, table_name="key"
|
||||
)
|
||||
|
||||
response[
|
||||
"soft_budget"
|
||||
] = data.soft_budget # include the user-input soft budget in the response
|
||||
response["soft_budget"] = (
|
||||
data.soft_budget
|
||||
) # include the user-input soft budget in the response
|
||||
|
||||
response = GenerateKeyResponse(**response)
|
||||
|
||||
|
|
@ -1559,6 +1559,19 @@ async def prepare_key_update_data(
|
|||
non_default_values["budget_reset_at"] = key_reset_at
|
||||
non_default_values["budget_duration"] = budget_duration
|
||||
|
||||
if "budget_limits" in non_default_values and non_default_values["budget_limits"]:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
raw_windows = non_default_values["budget_limits"]
|
||||
initialized_windows = []
|
||||
for window in raw_windows:
|
||||
w = window if isinstance(window, dict) else window.model_dump()
|
||||
w["reset_at"] = get_budget_reset_time(
|
||||
budget_duration=w["budget_duration"]
|
||||
).isoformat()
|
||||
initialized_windows.append(w)
|
||||
non_default_values["budget_limits"] = json.dumps(initialized_windows)
|
||||
|
||||
if "object_permission" in non_default_values:
|
||||
non_default_values = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
|
|
@ -2785,6 +2798,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
|
|||
rotation_interval: Optional[str] = None,
|
||||
router_settings: Optional[dict] = None,
|
||||
access_group_ids: Optional[list] = None,
|
||||
budget_limits: Optional[list] = None, # multiple concurrent budget windows
|
||||
):
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
|
||||
|
|
@ -2816,6 +2830,18 @@ async def generate_key_helper_fn( # noqa: PLR0915
|
|||
else:
|
||||
reset_at = get_budget_reset_time(budget_duration=budget_duration)
|
||||
|
||||
# Initialize reset_at for each budget window
|
||||
budget_limits_json: Optional[str] = None
|
||||
if budget_limits:
|
||||
initialized_windows = []
|
||||
for window in budget_limits:
|
||||
w = dict(window) if not isinstance(window, dict) else {**window}
|
||||
w["reset_at"] = get_budget_reset_time(
|
||||
budget_duration=w["budget_duration"]
|
||||
).isoformat()
|
||||
initialized_windows.append(w)
|
||||
budget_limits_json = json.dumps(initialized_windows)
|
||||
|
||||
aliases_json = json.dumps(aliases)
|
||||
config_json = json.dumps(config)
|
||||
permissions_json = json.dumps(permissions)
|
||||
|
|
@ -2898,6 +2924,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
|
|||
"organization_id": organization_id,
|
||||
"budget_id": budget_id,
|
||||
"blocked": blocked,
|
||||
"budget_limits": budget_limits_json,
|
||||
"created_by": created_by,
|
||||
"updated_by": updated_by,
|
||||
"allowed_routes": allowed_routes or [],
|
||||
|
|
@ -3175,10 +3202,10 @@ async def delete_verification_tokens(
|
|||
try:
|
||||
if prisma_client:
|
||||
tokens = [_hash_token_if_needed(token=key) for key in tokens]
|
||||
_keys_being_deleted: List[
|
||||
LiteLLM_VerificationToken
|
||||
] = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"token": {"in": tokens}}
|
||||
_keys_being_deleted: List[LiteLLM_VerificationToken] = (
|
||||
await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"token": {"in": tokens}}
|
||||
)
|
||||
)
|
||||
|
||||
if len(_keys_being_deleted) == 0:
|
||||
|
|
@ -3378,9 +3405,9 @@ async def _rotate_master_key( # noqa: PLR0915
|
|||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
try:
|
||||
models: Optional[
|
||||
List
|
||||
] = await prisma_client.db.litellm_proxymodeltable.find_many()
|
||||
models: Optional[List] = (
|
||||
await prisma_client.db.litellm_proxymodeltable.find_many()
|
||||
)
|
||||
except Exception:
|
||||
models = None
|
||||
# 2. process model table
|
||||
|
|
@ -4020,11 +4047,11 @@ async def validate_key_list_check(
|
|||
param="user_id",
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
)
|
||||
complete_user_info_db_obj: Optional[
|
||||
BaseModel
|
||||
] = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_api_key_dict.user_id},
|
||||
include={"organization_memberships": True},
|
||||
complete_user_info_db_obj: Optional[BaseModel] = (
|
||||
await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_api_key_dict.user_id},
|
||||
include={"organization_memberships": True},
|
||||
)
|
||||
)
|
||||
|
||||
if complete_user_info_db_obj is None:
|
||||
|
|
@ -4107,10 +4134,10 @@ async def _fetch_user_team_objects(
|
|||
if complete_user_info is None or not complete_user_info.teams:
|
||||
return []
|
||||
|
||||
teams: Optional[
|
||||
List[BaseModel]
|
||||
] = await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"team_id": {"in": complete_user_info.teams}}
|
||||
teams: Optional[List[BaseModel]] = (
|
||||
await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"team_id": {"in": complete_user_info.teams}}
|
||||
)
|
||||
)
|
||||
if teams is None:
|
||||
return []
|
||||
|
|
|
|||
|
|
@ -978,6 +978,19 @@ async def new_team( # noqa: PLR0915
|
|||
budget_duration=complete_team_data.budget_duration,
|
||||
)
|
||||
|
||||
# If budget_limits is set, initialize reset_at for each window
|
||||
if complete_team_data.budget_limits:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
initialized_windows = []
|
||||
for window in complete_team_data.budget_limits:
|
||||
w = window if isinstance(window, dict) else window.model_dump()
|
||||
w["reset_at"] = get_budget_reset_time(
|
||||
budget_duration=w["budget_duration"]
|
||||
).isoformat()
|
||||
initialized_windows.append(w)
|
||||
complete_team_data.budget_limits = initialized_windows
|
||||
|
||||
## Add Team Member Budget Table
|
||||
members_with_roles: List[Member] = []
|
||||
if complete_team_data.members_with_roles is not None:
|
||||
|
|
@ -1545,12 +1558,12 @@ async def update_team( # noqa: PLR0915
|
|||
updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"])
|
||||
|
||||
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
|
||||
team_row: Optional[
|
||||
LiteLLM_TeamTable
|
||||
] = await prisma_client.db.litellm_teamtable.update(
|
||||
where={"team_id": data.team_id},
|
||||
data=updated_kv,
|
||||
include={"litellm_model_table": True}, # type: ignore
|
||||
team_row: Optional[LiteLLM_TeamTable] = (
|
||||
await prisma_client.db.litellm_teamtable.update(
|
||||
where={"team_id": data.team_id},
|
||||
data=updated_kv,
|
||||
include={"litellm_model_table": True}, # type: ignore
|
||||
)
|
||||
)
|
||||
|
||||
if team_row is None or team_row.team_id is None:
|
||||
|
|
@ -1593,6 +1606,18 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None:
|
|||
reset_at = get_budget_reset_time(budget_duration=data.budget_duration)
|
||||
updated_kv["budget_reset_at"] = reset_at
|
||||
|
||||
if data.budget_limits is not None and len(data.budget_limits) > 0:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
initialized_windows = []
|
||||
for window in data.budget_limits:
|
||||
w = window if isinstance(window, dict) else window.model_dump()
|
||||
w["reset_at"] = get_budget_reset_time(
|
||||
budget_duration=w["budget_duration"]
|
||||
).isoformat()
|
||||
initialized_windows.append(w)
|
||||
updated_kv["budget_limits"] = json.dumps(initialized_windows)
|
||||
|
||||
|
||||
async def handle_update_object_permission(
|
||||
data_json: dict, existing_team_row: LiteLLM_TeamTable
|
||||
|
|
@ -2297,13 +2322,13 @@ async def team_member_delete(
|
|||
)
|
||||
|
||||
# Fetch keys before deletion to persist them
|
||||
keys_to_delete: List[
|
||||
LiteLLM_VerificationToken
|
||||
] = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={
|
||||
"user_id": {"in": list(user_ids_to_delete)},
|
||||
"team_id": data.team_id,
|
||||
}
|
||||
keys_to_delete: List[LiteLLM_VerificationToken] = (
|
||||
await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={
|
||||
"user_id": {"in": list(user_ids_to_delete)},
|
||||
"team_id": data.team_id,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
if keys_to_delete:
|
||||
|
|
@ -2687,10 +2712,10 @@ async def delete_team(
|
|||
team_rows: List[LiteLLM_TeamTable] = []
|
||||
for team_id in data.team_ids:
|
||||
try:
|
||||
team_row_base: Optional[
|
||||
BaseModel
|
||||
] = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
team_row_base: Optional[BaseModel] = (
|
||||
await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
)
|
||||
if team_row_base is None:
|
||||
raise Exception
|
||||
|
|
@ -2749,10 +2774,10 @@ async def delete_team(
|
|||
_persist_deleted_verification_tokens,
|
||||
)
|
||||
|
||||
keys_to_delete: List[
|
||||
LiteLLM_VerificationToken
|
||||
] = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"team_id": {"in": data.team_ids}}
|
||||
keys_to_delete: List[LiteLLM_VerificationToken] = (
|
||||
await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"team_id": {"in": data.team_ids}}
|
||||
)
|
||||
)
|
||||
|
||||
if keys_to_delete:
|
||||
|
|
@ -2972,11 +2997,11 @@ async def team_info(
|
|||
)
|
||||
|
||||
try:
|
||||
team_info: Optional[
|
||||
BaseModel
|
||||
] = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id},
|
||||
include={"object_permission": True},
|
||||
team_info: Optional[BaseModel] = (
|
||||
await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id},
|
||||
include={"object_permission": True},
|
||||
)
|
||||
)
|
||||
if team_info is None:
|
||||
raise Exception
|
||||
|
|
@ -3732,9 +3757,7 @@ async def list_team(
|
|||
except Exception as e:
|
||||
team_exception = """Invalid team object for team_id: {}. team_object={}.
|
||||
Error: {}
|
||||
""".format(
|
||||
team.team_id, team.model_dump(), str(e)
|
||||
)
|
||||
""".format(team.team_id, team.model_dump(), str(e))
|
||||
verbose_proxy_logger.exception(team_exception)
|
||||
continue
|
||||
# Sort the responses by team_alias
|
||||
|
|
|
|||
|
|
@ -641,9 +641,9 @@ except ImportError:
|
|||
server_root_path = get_server_root_path()
|
||||
_license_check = LicenseCheck()
|
||||
premium_user: bool = _license_check.is_premium()
|
||||
premium_user_data: Optional[
|
||||
"EnterpriseLicenseData"
|
||||
] = _license_check.airgapped_license_data
|
||||
premium_user_data: Optional["EnterpriseLicenseData"] = (
|
||||
_license_check.airgapped_license_data
|
||||
)
|
||||
global_max_parallel_request_retries_env: Optional[str] = os.getenv(
|
||||
"LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
|
||||
)
|
||||
|
|
@ -1537,9 +1537,9 @@ master_key: Optional[str] = None
|
|||
config_agents: Optional[List[AgentConfig]] = None
|
||||
otel_logging = False
|
||||
prisma_client: Optional[PrismaClient] = None
|
||||
shared_aiohttp_session: Optional[
|
||||
"ClientSession"
|
||||
] = None # Global shared session for connection reuse
|
||||
shared_aiohttp_session: Optional["ClientSession"] = (
|
||||
None # Global shared session for connection reuse
|
||||
)
|
||||
user_api_key_cache = DualCache(
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
)
|
||||
|
|
@ -1550,13 +1550,13 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
|
|||
dual_cache=user_api_key_cache
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
|
||||
redis_usage_cache: Optional[
|
||||
RedisCache
|
||||
] = None # redis cache used for tracking spend, tpm/rpm limits
|
||||
redis_usage_cache: Optional[RedisCache] = (
|
||||
None # redis cache used for tracking spend, tpm/rpm limits
|
||||
)
|
||||
polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False
|
||||
native_background_mode: List[
|
||||
str
|
||||
] = [] # Models that should use native provider background mode instead of polling
|
||||
native_background_mode: List[str] = (
|
||||
[]
|
||||
) # Models that should use native provider background mode instead of polling
|
||||
polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache
|
||||
user_custom_auth = None
|
||||
user_custom_key_generate = None
|
||||
|
|
@ -1783,6 +1783,26 @@ async def increment_spend_counters(
|
|||
increment=response_cost,
|
||||
)
|
||||
|
||||
# Increment per-window budget counters for multi-budget keys
|
||||
key_obj = await user_api_key_cache.async_get_cache(key=hashed_token)
|
||||
if key_obj is not None:
|
||||
key_budget_limits = getattr(key_obj, "budget_limits", None) or (
|
||||
key_obj.get("budget_limits") if isinstance(key_obj, dict) else None
|
||||
)
|
||||
if isinstance(key_budget_limits, str):
|
||||
key_budget_limits = json.loads(key_budget_limits)
|
||||
if isinstance(key_budget_limits, list):
|
||||
for window in key_budget_limits:
|
||||
duration = (
|
||||
window["budget_duration"]
|
||||
if isinstance(window, dict)
|
||||
else window.budget_duration
|
||||
)
|
||||
await spend_counter_cache.async_increment_cache(
|
||||
key=f"spend:key:{hashed_token}:window:{duration}",
|
||||
value=response_cost,
|
||||
)
|
||||
|
||||
if team_id is not None:
|
||||
await _init_and_increment_spend_counter(
|
||||
counter_key=f"spend:team:{team_id}",
|
||||
|
|
@ -1790,6 +1810,26 @@ async def increment_spend_counters(
|
|||
increment=response_cost,
|
||||
)
|
||||
|
||||
# Increment per-window budget counters for multi-budget teams
|
||||
team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{team_id}")
|
||||
if team_obj is not None:
|
||||
team_budget_limits = getattr(team_obj, "budget_limits", None) or (
|
||||
team_obj.get("budget_limits") if isinstance(team_obj, dict) else None
|
||||
)
|
||||
if isinstance(team_budget_limits, str):
|
||||
team_budget_limits = json.loads(team_budget_limits)
|
||||
if isinstance(team_budget_limits, list):
|
||||
for window in team_budget_limits:
|
||||
duration = (
|
||||
window["budget_duration"]
|
||||
if isinstance(window, dict)
|
||||
else window.budget_duration
|
||||
)
|
||||
await spend_counter_cache.async_increment_cache(
|
||||
key=f"spend:team:{team_id}:window:{duration}",
|
||||
value=response_cost,
|
||||
)
|
||||
|
||||
if user_id is not None and team_id is not None:
|
||||
await _init_and_increment_spend_counter(
|
||||
counter_key=f"spend:team_member:{user_id}:{team_id}",
|
||||
|
|
@ -2040,9 +2080,9 @@ async def update_cache( # noqa: PLR0915
|
|||
_id = "team_id:{}".format(team_id)
|
||||
try:
|
||||
# Fetch the existing cost for the given user
|
||||
existing_spend_obj: Optional[
|
||||
LiteLLM_TeamTable
|
||||
] = await user_api_key_cache.async_get_cache(key=_id)
|
||||
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
|
||||
await user_api_key_cache.async_get_cache(key=_id)
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
# do nothing if team not in api key cache
|
||||
return
|
||||
|
|
@ -2163,11 +2203,9 @@ def run_ollama_serve():
|
|||
with open(os.devnull, "w") as devnull:
|
||||
subprocess.Popen(command, stdout=devnull, stderr=devnull)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"""
|
||||
verbose_proxy_logger.debug(f"""
|
||||
LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve`
|
||||
"""
|
||||
)
|
||||
""")
|
||||
|
||||
|
||||
def _get_process_rss_mb() -> Optional[float]:
|
||||
|
|
@ -5215,10 +5253,10 @@ class ProxyConfig:
|
|||
)
|
||||
|
||||
try:
|
||||
guardrails_in_db: List[
|
||||
Guardrail
|
||||
] = await GuardrailRegistry.get_all_guardrails_from_db(
|
||||
prisma_client=prisma_client
|
||||
guardrails_in_db: List[Guardrail] = (
|
||||
await GuardrailRegistry.get_all_guardrails_from_db(
|
||||
prisma_client=prisma_client
|
||||
)
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"guardrails from the DB %s", str(guardrails_in_db)
|
||||
|
|
@ -5600,9 +5638,9 @@ async def initialize( # noqa: PLR0915
|
|||
user_api_base = api_base
|
||||
dynamic_config[user_model]["api_base"] = api_base
|
||||
if api_version:
|
||||
os.environ[
|
||||
"AZURE_API_VERSION"
|
||||
] = api_version # set this for azure - litellm can read this from the env
|
||||
os.environ["AZURE_API_VERSION"] = (
|
||||
api_version # set this for azure - litellm can read this from the env
|
||||
)
|
||||
if max_tokens: # model-specific param
|
||||
dynamic_config[user_model]["max_tokens"] = max_tokens
|
||||
if temperature: # model-specific param
|
||||
|
|
@ -5944,9 +5982,9 @@ class ProxyStartupEvent:
|
|||
"""
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
_use_redis_transaction_buffer: Optional[
|
||||
Union[bool, str]
|
||||
] = general_settings.get("use_redis_transaction_buffer", False)
|
||||
_use_redis_transaction_buffer: Optional[Union[bool, str]] = (
|
||||
general_settings.get("use_redis_transaction_buffer", False)
|
||||
)
|
||||
if isinstance(_use_redis_transaction_buffer, str):
|
||||
_use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer)
|
||||
|
||||
|
|
@ -12627,9 +12665,9 @@ async def get_config_list(
|
|||
hasattr(sub_field_info, "description")
|
||||
and sub_field_info.description is not None
|
||||
):
|
||||
nested_fields[
|
||||
idx
|
||||
].field_description = sub_field_info.description
|
||||
nested_fields[idx].field_description = (
|
||||
sub_field_info.description
|
||||
)
|
||||
idx += 1
|
||||
|
||||
_stored_in_db = None
|
||||
|
|
|
|||
|
|
@ -140,6 +140,7 @@ model LiteLLM_TeamTable {
|
|||
team_member_permissions String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
policies String[] @default([])
|
||||
budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}]
|
||||
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
|
||||
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
|
||||
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
|
||||
|
|
@ -387,6 +388,7 @@ model LiteLLM_VerificationToken {
|
|||
rotation_interval String? // How often to rotate (e.g., "30d", "90d")
|
||||
last_rotation_at DateTime? // When this key was last rotated
|
||||
key_rotation_at DateTime? // When this key should next be rotated
|
||||
budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}]
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
|
||||
litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id])
|
||||
|
|
|
|||
|
|
@ -1898,9 +1898,9 @@ class ProxyLogging:
|
|||
normalized_call_type = CallTypes.aembedding.value
|
||||
if normalized_call_type is not None:
|
||||
litellm_logging_obj.call_type = normalized_call_type
|
||||
litellm_logging_obj.model_call_details[
|
||||
"call_type"
|
||||
] = normalized_call_type
|
||||
litellm_logging_obj.model_call_details["call_type"] = (
|
||||
normalized_call_type
|
||||
)
|
||||
# Pass-through endpoints are logged via the callback loop's
|
||||
# async_post_call_failure_hook — skip pre_call and failure handlers.
|
||||
if litellm_logging_obj.call_type == CallTypes.pass_through.value:
|
||||
|
|
@ -2524,8 +2524,7 @@ class PrismaClient:
|
|||
required_view = "LiteLLM_VerificationTokenView"
|
||||
expected_views_str = ", ".join(f"'{view}'" for view in expected_views)
|
||||
pg_schema = os.getenv("DATABASE_SCHEMA", "public")
|
||||
ret = await self.db.query_raw(
|
||||
f"""
|
||||
ret = await self.db.query_raw(f"""
|
||||
WITH existing_views AS (
|
||||
SELECT viewname
|
||||
FROM pg_views
|
||||
|
|
@ -2537,8 +2536,7 @@ class PrismaClient:
|
|||
(SELECT COUNT(*) FROM existing_views) AS view_count,
|
||||
ARRAY_AGG(viewname) AS view_names
|
||||
FROM existing_views
|
||||
"""
|
||||
)
|
||||
""")
|
||||
expected_total_views = len(expected_views)
|
||||
if ret[0]["view_count"] == expected_total_views:
|
||||
verbose_proxy_logger.info("All necessary views exist!")
|
||||
|
|
@ -2547,8 +2545,7 @@ class PrismaClient:
|
|||
## check if required view exists ##
|
||||
if ret[0]["view_names"] and required_view not in ret[0]["view_names"]:
|
||||
await self.health_check() # make sure we can connect to db
|
||||
await self.db.execute_raw(
|
||||
"""
|
||||
await self.db.execute_raw("""
|
||||
CREATE VIEW "LiteLLM_VerificationTokenView" AS
|
||||
SELECT
|
||||
v.*,
|
||||
|
|
@ -2558,8 +2555,7 @@ class PrismaClient:
|
|||
t.rpm_limit AS team_rpm_limit
|
||||
FROM "LiteLLM_VerificationToken" v
|
||||
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
|
||||
"""
|
||||
)
|
||||
""")
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"LiteLLM_VerificationTokenView Created in DB!"
|
||||
|
|
@ -3142,6 +3138,10 @@ class PrismaClient:
|
|||
hashed_token = self.hash_token(token=token)
|
||||
db_data = self.jsonify_object(data=data)
|
||||
db_data["token"] = hashed_token
|
||||
# Prisma rejects nullable JSON fields set to None (no default).
|
||||
# Strip them so the DB stores NULL via the column's nullable constraint.
|
||||
if db_data.get("budget_limits") is None:
|
||||
db_data.pop("budget_limits", None)
|
||||
print_verbose(
|
||||
"PrismaClient: Before upsert into litellm_verificationtoken"
|
||||
)
|
||||
|
|
|
|||
137
tests/test_litellm/proxy/auth/test_multi_budget_windows.py
Normal file
137
tests/test_litellm/proxy/auth/test_multi_budget_windows.py
Normal file
|
|
@ -0,0 +1,137 @@
|
|||
"""
|
||||
Unit tests for multi-budget-window enforcement on API keys.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import _virtual_key_multi_budget_check
|
||||
|
||||
|
||||
def _make_valid_token(**kwargs) -> UserAPIKeyAuth:
|
||||
defaults = dict(
|
||||
token="sk-test-token",
|
||||
key_name="test",
|
||||
spend=0.0,
|
||||
max_budget=None,
|
||||
budget_limits=[],
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return UserAPIKeyAuth(**defaults)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_budget_limits_passes():
|
||||
"""Keys with empty budget_limits should pass without raising."""
|
||||
token = _make_valid_token(budget_limits=[])
|
||||
# Should not raise
|
||||
await _virtual_key_multi_budget_check(valid_token=token)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_under_budget_passes():
|
||||
"""Key with spend under all windows should pass."""
|
||||
token = _make_valid_token(
|
||||
budget_limits=[
|
||||
{"budget_duration": "24h", "max_budget": 10.0, "reset_at": None},
|
||||
{"budget_duration": "30d", "max_budget": 100.0, "reset_at": None},
|
||||
]
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_current_spend",
|
||||
new_callable=AsyncMock,
|
||||
return_value=1.0, # well under both windows
|
||||
):
|
||||
await _virtual_key_multi_budget_check(valid_token=token)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_over_first_window_raises():
|
||||
"""Key exceeding the first (daily) window should raise BudgetExceededError."""
|
||||
token = _make_valid_token(
|
||||
budget_limits=[
|
||||
{"budget_duration": "24h", "max_budget": 5.0, "reset_at": None},
|
||||
{"budget_duration": "30d", "max_budget": 100.0, "reset_at": None},
|
||||
]
|
||||
)
|
||||
|
||||
spend_by_window = [6.0, 6.0] # over daily, under monthly
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def fake_get_spend(counter_key, fallback_spend):
|
||||
nonlocal call_count
|
||||
val = spend_by_window[call_count]
|
||||
call_count += 1
|
||||
return val
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _virtual_key_multi_budget_check(valid_token=token)
|
||||
|
||||
err = exc_info.value
|
||||
assert err.status_code == 429
|
||||
assert "24h" in str(err)
|
||||
assert "Key over" in str(err)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_over_second_window_raises():
|
||||
"""Key exceeding only the monthly window should raise BudgetExceededError referencing 30d."""
|
||||
token = _make_valid_token(
|
||||
budget_limits=[
|
||||
{"budget_duration": "24h", "max_budget": 50.0, "reset_at": None},
|
||||
{"budget_duration": "30d", "max_budget": 5.0, "reset_at": None},
|
||||
]
|
||||
)
|
||||
|
||||
spend_by_window = [1.0, 10.0] # under daily, over monthly
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def fake_get_spend(counter_key, fallback_spend):
|
||||
nonlocal call_count
|
||||
val = spend_by_window[call_count]
|
||||
call_count += 1
|
||||
return val
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _virtual_key_multi_budget_check(valid_token=token)
|
||||
|
||||
err = exc_info.value
|
||||
assert err.status_code == 429
|
||||
assert "30d" in str(err)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_limit_entry_objects_coerced():
|
||||
"""BudgetLimitEntry Pydantic objects (not dicts) must be handled without KeyError.
|
||||
|
||||
While budget_limits is normally serialized as List[dict], the auth check must
|
||||
tolerate BudgetLimitEntry objects in case they arrive without prior serialization.
|
||||
"""
|
||||
from litellm.proxy._types import BudgetLimitEntry
|
||||
|
||||
token = _make_valid_token(budget_limits=[])
|
||||
# Bypass Pydantic validation to simulate BudgetLimitEntry objects reaching the check
|
||||
object.__setattr__(
|
||||
token,
|
||||
"budget_limits",
|
||||
[BudgetLimitEntry(budget_duration="24h", max_budget=10.0)],
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.get_current_spend",
|
||||
new_callable=AsyncMock,
|
||||
return_value=1.0,
|
||||
):
|
||||
# Should not raise TypeError / KeyError — model_dump() coerces the object
|
||||
await _virtual_key_multi_budget_check(valid_token=token)
|
||||
|
|
@ -25,6 +25,7 @@ const BudgetDurationDropdown: React.FC<BudgetDurationDropdownProps> = ({
|
|||
placeholder="n/a"
|
||||
allowClear
|
||||
>
|
||||
<Option value="1h">hourly</Option>
|
||||
<Option value="24h">daily</Option>
|
||||
<Option value="7d">weekly</Option>
|
||||
<Option value="30d">monthly</Option>
|
||||
|
|
@ -36,6 +37,7 @@ export const getBudgetDurationLabel = (value: string | null | undefined): string
|
|||
if (!value) return "Not set";
|
||||
|
||||
const budgetDurationMap: Record<string, string> = {
|
||||
"1h": "hourly",
|
||||
"24h": "daily",
|
||||
"7d": "weekly",
|
||||
"30d": "monthly",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,84 @@
|
|||
import { Button, InputNumber, Select } from "antd";
|
||||
import React from "react";
|
||||
|
||||
export interface BudgetWindowEntry {
|
||||
budget_duration: string;
|
||||
max_budget: number | null;
|
||||
}
|
||||
|
||||
export const BUDGET_WINDOW_OPTIONS = [
|
||||
{ value: "1h", label: "Hourly", resetHint: "Resets every hour" },
|
||||
{ value: "24h", label: "Daily", resetHint: "Resets daily at midnight UTC" },
|
||||
{ value: "7d", label: "Weekly", resetHint: "Resets every Sunday at midnight UTC" },
|
||||
{ value: "30d", label: "Monthly", resetHint: "Resets on the 1st of every month at midnight UTC" },
|
||||
];
|
||||
|
||||
interface BudgetWindowsEditorProps {
|
||||
value: BudgetWindowEntry[];
|
||||
onChange: (v: BudgetWindowEntry[]) => void;
|
||||
}
|
||||
|
||||
export function BudgetWindowsEditor({ value, onChange }: BudgetWindowsEditorProps) {
|
||||
const addWindow = () => {
|
||||
onChange([...value, { budget_duration: "24h", max_budget: null }]);
|
||||
};
|
||||
|
||||
const removeWindow = (idx: number) => {
|
||||
onChange(value.filter((_, i) => i !== idx));
|
||||
};
|
||||
|
||||
const updateWindow = (idx: number, field: keyof BudgetWindowEntry, fieldValue: string | number | null) => {
|
||||
const updated = value.map((w, i) => (i === idx ? { ...w, [field]: fieldValue } : w));
|
||||
onChange(updated);
|
||||
};
|
||||
|
||||
return (
|
||||
<div>
|
||||
{value.map((window, idx) => {
|
||||
const hint = BUDGET_WINDOW_OPTIONS.find((o) => o.value === window.budget_duration)?.resetHint;
|
||||
return (
|
||||
<div key={idx} style={{ marginBottom: 12 }}>
|
||||
<div style={{ display: "flex", gap: 8, alignItems: "center" }}>
|
||||
<Select
|
||||
value={window.budget_duration}
|
||||
onChange={(v) => updateWindow(idx, "budget_duration", v)}
|
||||
style={{ width: 130 }}
|
||||
options={BUDGET_WINDOW_OPTIONS.map((o) => ({ value: o.value, label: o.label }))}
|
||||
/>
|
||||
<InputNumber
|
||||
step={0.01}
|
||||
min={0}
|
||||
precision={2}
|
||||
value={window.max_budget ?? undefined}
|
||||
onChange={(v) => updateWindow(idx, "max_budget", v ?? null)}
|
||||
placeholder="Max spend ($)"
|
||||
style={{ width: 160 }}
|
||||
prefix="$"
|
||||
/>
|
||||
<Button
|
||||
type="text"
|
||||
danger
|
||||
size="small"
|
||||
onClick={() => removeWindow(idx)}
|
||||
style={{ padding: "0 4px" }}
|
||||
>
|
||||
✕
|
||||
</Button>
|
||||
</div>
|
||||
{hint && (
|
||||
<div style={{ fontSize: 11, color: "#888", marginTop: 3, marginLeft: 2 }}>
|
||||
↻ {hint}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
<Button
|
||||
size="small"
|
||||
onClick={(e) => { e.preventDefault(); addWindow(); }}
|
||||
>
|
||||
+ Add Budget Window
|
||||
</Button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -91,6 +91,7 @@ export interface KeyResponse {
|
|||
agent_access_groups?: string[];
|
||||
};
|
||||
access_group_ids?: string[];
|
||||
budget_limits?: Array<{ budget_duration: string; max_budget: number; reset_at?: string }>;
|
||||
auto_rotate?: boolean;
|
||||
rotation_interval?: string;
|
||||
last_rotation_at?: string;
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
|
|||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
|
||||
import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
|
||||
import { Button as Button2, Form, Input, InputNumber, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
|
||||
import debounce from "lodash/debounce";
|
||||
import React, { useCallback, useEffect, useState } from "react";
|
||||
import { rolesWithWriteAccess } from "../../utils/roles";
|
||||
|
|
@ -28,6 +28,7 @@ import TeamDropdown from "../common_components/team_dropdown";
|
|||
import OrganizationDropdown from "../common_components/OrganizationDropdown";
|
||||
import ProjectDropdown from "../common_components/ProjectDropdown";
|
||||
import { CreateUserButton } from "../CreateUserButton";
|
||||
import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
|
||||
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
|
||||
import { Team } from "../key_team_helpers/key_list";
|
||||
import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
|
||||
|
|
@ -153,6 +154,7 @@ export const fetchUserModels = async (
|
|||
}
|
||||
};
|
||||
|
||||
|
||||
/**
|
||||
* ─────────────────────────────────────────────────────────────────────────
|
||||
* @deprecated
|
||||
|
|
@ -201,6 +203,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
const [autoRotationEnabled, setAutoRotationEnabled] = useState<boolean>(false);
|
||||
const [rotationInterval, setRotationInterval] = useState<string>("30d");
|
||||
const [routerSettings, setRouterSettings] = useState<RouterSettingsAccordionValue | null>(null);
|
||||
const [budgetLimits, setBudgetLimits] = useState<BudgetWindowEntry[]>([]);
|
||||
const [routerSettingsKey, setRouterSettingsKey] = useState<number>(0);
|
||||
const [agentsList, setAgentsList] = useState<{ agent_id: string; agent_name: string }[]>([]);
|
||||
const [selectedAgentId, setSelectedAgentId] = useState<string | null>(null);
|
||||
|
|
@ -218,6 +221,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
setSelectedAgentId(null);
|
||||
setSelectedOrganizationId(null);
|
||||
setSelectedProjectId(null);
|
||||
setBudgetLimits([]);
|
||||
};
|
||||
|
||||
const handleCancel = () => {
|
||||
|
|
@ -236,6 +240,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
setSelectedAgentId(null);
|
||||
setSelectedOrganizationId(null);
|
||||
setSelectedProjectId(null);
|
||||
setBudgetLimits([]);
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
|
|
@ -519,6 +524,12 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
}
|
||||
}
|
||||
|
||||
// Add multi-window budget limits (filter out incomplete entries)
|
||||
const validWindows = budgetLimits.filter((w) => w.budget_duration && w.max_budget !== null && w.max_budget !== undefined);
|
||||
if (validWindows.length > 0) {
|
||||
formValues.budget_limits = validWindows;
|
||||
}
|
||||
|
||||
let response;
|
||||
if (keyOwner === "service_account") {
|
||||
response = await keyCreateServiceAccountCall(accessToken, formValues);
|
||||
|
|
@ -540,6 +551,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
setSoftBudget(response["soft_budget"]);
|
||||
NotificationsManager.success("Virtual Key Created");
|
||||
form.resetFields();
|
||||
setBudgetLimits([]);
|
||||
localStorage.removeItem("userData" + userID);
|
||||
} catch (error) {
|
||||
console.log("error in create key:", error);
|
||||
|
|
@ -1045,6 +1057,22 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
>
|
||||
<BudgetDurationDropdown onChange={(value) => form.setFieldValue("budget_duration", value)} />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
className="mt-4"
|
||||
label={
|
||||
<span>
|
||||
Budget Windows{" "}
|
||||
<Tooltip title="Set multiple independent budget windows (e.g., hourly $10 AND monthly $200). Each window tracks spend separately and resets on its own schedule.">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
>
|
||||
<BudgetWindowsEditor
|
||||
value={budgetLimits}
|
||||
onChange={setBudgetLimits}
|
||||
/>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
className="mt-4"
|
||||
label={
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"
|
|||
import PolicySelector from "@/components/policies/PolicySelector";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { TextInput, Button as TremorButton } from "@tremor/react";
|
||||
import { Form, Input, Select, Switch, Tooltip } from "antd";
|
||||
import { Button as AntButton, Form, Input, InputNumber, Select, Switch, Tooltip } from "antd";
|
||||
import { useEffect, useState } from "react";
|
||||
import { rolesWithWriteAccess } from "../../utils/roles";
|
||||
import AgentSelector from "../agent_management/AgentSelector";
|
||||
|
|
@ -16,6 +16,7 @@ import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSel
|
|||
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
|
||||
import OrganizationDropdown from "../common_components/OrganizationDropdown";
|
||||
import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils";
|
||||
import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
|
||||
import { KeyResponse } from "../key_team_helpers/key_list";
|
||||
import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
|
||||
import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
|
||||
|
|
@ -77,6 +78,7 @@ const getKeyTypeFromRoutes = (allowedRoutes: string[] | null | undefined): strin
|
|||
return "default";
|
||||
};
|
||||
|
||||
|
||||
export function KeyEditView({
|
||||
keyData,
|
||||
onCancel,
|
||||
|
|
@ -103,6 +105,9 @@ export function KeyEditView({
|
|||
const [rotationInterval, setRotationInterval] = useState<string>(keyData.rotation_interval || "");
|
||||
const [neverExpire, setNeverExpire] = useState<boolean>(!keyData.expires);
|
||||
const [isKeySaving, setIsKeySaving] = useState(false);
|
||||
const [budgetLimits, setBudgetLimits] = useState<BudgetWindowEntry[]>(
|
||||
Array.isArray(keyData.budget_limits) ? keyData.budget_limits : []
|
||||
);
|
||||
const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations();
|
||||
const { data: projects } = useProjects();
|
||||
const { data: uiSettingsData } = useUISettings();
|
||||
|
|
@ -274,6 +279,10 @@ export function KeyEditView({
|
|||
values.duration = null;
|
||||
}
|
||||
|
||||
// Include multi-window budget limits (filter out incomplete entries)
|
||||
const validWindows = budgetLimits.filter((w) => w.budget_duration && w.max_budget !== null && w.max_budget !== undefined);
|
||||
values.budget_limits = validWindows.length > 0 ? validWindows : undefined;
|
||||
|
||||
await onSubmit(values);
|
||||
} finally {
|
||||
setIsKeySaving(false);
|
||||
|
|
@ -423,6 +432,22 @@ export function KeyEditView({
|
|||
</Select>
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Budget Windows{" "}
|
||||
<Tooltip title="Set multiple independent budget windows (e.g., hourly $10 AND monthly $200). Each window tracks spend separately and resets on its own schedule.">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
>
|
||||
<BudgetWindowsEditor
|
||||
value={budgetLimits}
|
||||
onChange={setBudgetLimits}
|
||||
/>
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item label="TPM Limit" name="tpm_limit">
|
||||
<NumericalInput min={0} />
|
||||
</Form.Item>
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue