scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

feat(proxy): enforce customer org, team and member budgets and limits on the request path (#51489)

Browse source 
* feat(proxy): enforce customer org, team and member budgets and limits on the request path

Resolve x-litellm-customer-member-id at auth, walk member -> team -> org, enforce max_budget,
allowed_models, tpm/rpm and per-model limits at every level in the v3 limiter, honor the hierarchy
allowlists in alternate-model steering, and settle spend to the member, team and org rows.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(auth): customer hierarchy survives virtual key auth without a legacy end user

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(auth): keys pinned to a customer organization may only act for that organization's members

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(proxy): keep the customer hierarchy spend maps within the line limit

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(auth): check the key's customer org scope after the key is resolved

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): declare the hierarchy cache key family and keep type counts flat

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(auth): resolve the customer hierarchy in the shared authorization gate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): drop a member's cached hierarchy through the declared cache target

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(auth): keep the hierarchy reset within the line limit

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(auth): type the requested models as a sequence

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(spend): add daily customer usage and budget resets (#51493)

* feat(spend): add daily customer usage and budget resets

* chore(spend): document customer reset query bounds

* fix(spend): chart customer queue metrics and update recovery test

* fix(proxy): read stamped customer ids from either metadata bucket, alias-aware model check, broadcast member evictions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* Revert "feat(spend): add daily customer usage and budget resets (#51493)"

This reverts commit ccac47608f.

* fix(proxy): evict cached member hierarchies when team, org or template budgets change

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): declare the auth cache dependency alias as a TypeAlias

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* Reapply "feat(spend): add daily customer usage and budget resets (#51493)"

This reverts commit 7542471e17.

* fix(spend): reconcile the daily customer spend table with main

Keep main's LiteLLM_DailyCustomerMemberSpend shape, make model_group nullable and add the identity unique index the bulk upsert arbitrates on in a new migration, drop the duplicate model and migration, and feed both the member-header hierarchy and the key's customer usage attribution through the shared daily spend projection so neither path double-writes.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(pagerduty): tolerate failure metadata without customer hierarchy ids

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(accounting): key daily customer spend rows on main's hashed identity

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(pagerduty): read customer ids from failure metadata through a typed lookup

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: moe-berri <moe@berri.ai>
This commit is contained in:
Kerry Lu 2026-11-12 20:43:44 -08:00 • committed by GitHub
parent f64f741d08
commit d30b2b9b26
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
53 changed files with 2895 additions and 590 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
litellm/accounting/application.py
View file

@ -1,5 +1,3 @@
from hashlib import sha256
from json import dumps
from typing import TYPE_CHECKING, Final
from litellm.accounting.delivery_projection import entity_transactions, project_deliveries
@ -32,7 +30,6 @@ async def apply_deliveries(
await writer.apply_accounting_entities(
prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions)
)
await _apply_customer_usage(transaction, contributions)
for table in sorted(DAILY_SPEND_TABLES.values(), key=lambda item: item.name):
rows = tuple(
{**dict(item.key), **dict(item.values), **dict(item.increments)}
@ -63,45 +60,6 @@ async def apply_deliveries(
return contributions
async def _apply_customer_usage(transaction: LedgerTransaction, contributions: tuple[Contribution, ...]) -> None:
for item in contributions:
if item.table == "LiteLLM_DailyCustomerMemberSpend":
await _upsert_customer_daily(transaction, item)
if item.table == "LiteLLM_CustomerOrganizationTable":
await transaction.litellm_customerorganizationtable.update_many(
where=dict(item.key), data={"spend": {"increment": dict(item.increments)["spend"]}}
)
if item.table == "LiteLLM_CustomerTeamTable":
await transaction.litellm_customerteamtable.update_many(
where=dict(item.key), data={"spend": {"increment": dict(item.increments)["spend"]}}
)
if item.table == "LiteLLM_CustomerMemberTable":
await transaction.litellm_customermembertable.update_many(
where=dict(item.key), data={"spend": {"increment": dict(item.increments)["spend"]}}
)
async def _upsert_customer_daily(transaction: LedgerTransaction, contribution: Contribution) -> None:
row_id: Final = _customer_daily_row_id(contribution.key)
await transaction.litellm_dailycustomermemberspend.upsert(
where={"id": row_id},
data={
"create": {
"id": row_id,
**dict(contribution.key),
**dict(contribution.values),
**dict(contribution.increments),
},
"update": {name: {"increment": value} for name, value in contribution.increments},
},
)
def _customer_daily_row_id(key: tuple[tuple[str, str], ...]) -> str:
encoded: Final = dumps(sorted(key), separators=(",", ":")).encode("utf-8")
return sha256(encoded).hexdigest()
def _window(item: Contribution) -> WindowSpendTransaction:
key: Final = dict(item.key)
values: Final = dict(item.values)

3
litellm/accounting/delivery_projection.py
View file

@ -56,6 +56,9 @@ def entity_transactions(contributions: tuple[Contribution, ...]) -> "DBSpendUpda
"model_access_group_list_transactions": dict(
_entity_totals(contributions, "LiteLLM_ModelAccessGroupBudgetTable")
),
"customer_org_list_transactions": dict(_entity_totals(contributions, "LiteLLM_CustomerOrganizationTable")),
"customer_team_list_transactions": dict(_entity_totals(contributions, "LiteLLM_CustomerTeamTable")),
"customer_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_CustomerMemberTable")),
}

112
litellm/accounting/projection.py
View file

@ -1,6 +1,6 @@
from collections.abc import Iterator, Mapping
from dataclasses import dataclass, replace
from datetime import datetime, timezone
from datetime import datetime
from itertools import chain
from types import MappingProxyType
from typing import Final
@ -17,7 +17,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import GuardrailCos
from litellm.proxy.common_utils.customer_usage import CUSTOMER_USAGE_METADATA_KEY, CustomerUsageAttribution
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, daily_spend_entity_ids
from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks
from litellm.proxy.db.model_usage_rollup import date_from_start_time
from litellm.proxy.db.model_usage_rollup import date_from_start_time, utc_date_from_start_time
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
@ -71,9 +71,32 @@ def _spend(table: str, key: tuple[tuple[str, str], ...], spend: float) -> Contri
return Contribution(table, key, (("spend", spend),))
_CUSTOMER_ID_KEYS: Final = ("customer_org_id", "customer_team_id", "customer_member_id")
def _customer_ids(metadata: Mapping[str, object]) -> Mapping[str, str | None]:
"""The customer hierarchy a request is charged to: the member the request named, else the key's attribution."""
entities: Final = _fields(metadata.get("accounting_entities"))
if isinstance(entities.get("customer_org_id"), str):
return {key: value if isinstance(value := entities.get(key), str) else None for key in _CUSTOMER_ID_KEYS}
serialized: Final = metadata.get(CUSTOMER_USAGE_METADATA_KEY)
if not isinstance(serialized, str):
return dict.fromkeys(_CUSTOMER_ID_KEYS)
try:
attribution: Final = CustomerUsageAttribution.model_validate_json(serialized)
except ValidationError:
return dict.fromkeys(_CUSTOMER_ID_KEYS)
return {
"customer_org_id": attribution.customer_org_id,
"customer_team_id": attribution.customer_team_id,
"customer_member_id": attribution.customer_member_id,
}
def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:
spend: Final = _number(row.get("spend"))
entities: Final = _fields(metadata.get("accounting_entities"))
customer: Final = _customer_ids(metadata)
identities: Final = (
("LiteLLM_VerificationToken", "token", entities.get("api_key")),
("LiteLLM_UserTable", "user_id", entities.get("user")),
@ -81,6 +104,9 @@ def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, obje
("LiteLLM_TeamTable", "team_id", entities.get("team_id")),
("LiteLLM_OrganizationTable", "organization_id", entities.get("organization_id")),
("LiteLLM_EndUserTable", "user_id", entities.get("end_user")),
("LiteLLM_CustomerOrganizationTable", "customer_org_id", customer["customer_org_id"]),
("LiteLLM_CustomerTeamTable", "customer_team_id", customer["customer_team_id"]),
("LiteLLM_CustomerMemberTable", "customer_member_id", customer["customer_member_id"]),
("LiteLLM_AgentsTable", "agent_id", entities.get("agent_id")),
("LiteLLM_ProjectTable", "project_id", entities.get("project_id")),
)
@ -160,7 +186,14 @@ def _daily_contributions(row: Mapping[str, object], metadata: Mapping[str, objec
)
# The legacy daily organization writer uses the charged org argument, not
# SpendLogs' normalized empty-string identity for a request with no org.
daily_row: Final = {**row, "organization_id": _fields(metadata.get("accounting_entities")).get("organization_id")}
entities: Final = _fields(metadata.get("accounting_entities"))
customer: Final = _customer_ids(metadata)
daily_row: Final = {
**row,
"organization_id": entities.get("organization_id"),
"customer_org_id": customer["customer_org_id"],
**{key: customer[key] or "" for key in ("customer_team_id", "customer_member_id")},
}
return tuple(_daily_rows(daily_row, shared_key, counters))
@ -170,13 +203,26 @@ def _daily_rows(
counters: tuple[tuple[str, int | float], ...],
) -> Iterator[Contribution]:
for entity, table in DAILY_SPEND_TABLES.items():
customer_date = _utc_date(row) if entity == "customer_member" else None
if entity == "customer_member" and customer_date is None:
continue
daily_key = tuple(
(column, customer_date if column == "date" and customer_date is not None else value)
for column, value in shared_key
)
model_group = (("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))),)
for identity in daily_spend_entity_ids(row, entity):
yield Contribution(
table.name,
((table.entity_id_column, _text(identity)), *shared_key),
(
(table.entity_id_column, _text(identity)),
*((column, _text(row.get(column))) for column in table.parent_id_columns),
*daily_key,
*((("model_group", _text(row.get("model_group"))),) if table.hashed_identity else ()),
),
counters,
(
("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))),
*(() if table.hashed_identity else model_group),
*((("request_id", _text(row.get("request_id"))),) if table.carries_request_id else ()),
),
)
@ -203,53 +249,10 @@ def project(row: Mapping[str, object]) -> tuple[Contribution, ...]:
*_window_contributions(row, metadata),
*_model_contributions(row, metadata),
*_guardrail_contributions(row, metadata),
*_customer_usage_contributions(row, metadata),
)
return tuple(_capped(contribution) for contribution in contributions if _indexable(contribution))
def _customer_usage_contributions(
row: Mapping[str, object], metadata: Mapping[str, object]
) -> tuple[Contribution, ...]:
serialized: Final = metadata.get(CUSTOMER_USAGE_METADATA_KEY)
date: Final = _date(row)
if not isinstance(serialized, str) or date is None:
return ()
try:
attribution: Final = CustomerUsageAttribution.model_validate_json(serialized)
except ValidationError:
return ()
counters: Final = _daily_counters(row, metadata)
identities: Final = (
("LiteLLM_CustomerOrganizationTable", "customer_org_id", attribution.customer_org_id),
("LiteLLM_CustomerTeamTable", "customer_team_id", attribution.customer_team_id),
("LiteLLM_CustomerMemberTable", "customer_member_id", attribution.customer_member_id),
)
return (
Contribution(
"LiteLLM_DailyCustomerMemberSpend",
(
("date", date),
("customer_org_id", attribution.customer_org_id),
("customer_team_id", attribution.customer_team_id or ""),
("customer_member_id", attribution.customer_member_id or ""),
("model", _text(row.get("model"))),
("model_group", _text(row.get("model_group"))),
("custom_llm_provider", _text(row.get("custom_llm_provider"))),
("api_key", _text(row.get("api_key"))),
("mcp_namespaced_tool_name", _text(row.get("mcp_namespaced_tool_name"))),
("endpoint", _text(ROUTE_ENDPOINT_MAPPING.get(_text(row.get("call_type"))))),
),
counters,
),
*(
_spend(table, ((column, identity),), float(_number(row.get("spend"))))
for table, column, identity in identities
if identity is not None
),
)
def _indexable(contribution: Contribution) -> bool:
texts: Final = (value for _, value in (*contribution.key, *contribution.values) if value is not None)
return all(len(text.encode("utf-8", "surrogatepass")) <= _MAX_INDEXED_TEXT_BYTES for text in texts)
@ -303,18 +306,7 @@ def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, objec
def _utc_date(row: Mapping[str, object]) -> str | None:
start: Final = row.get("startTime")
try:
parsed: Final = datetime.fromisoformat(start.replace("Z", "+00:00")) if isinstance(start, str) else start
except ValueError:
return None
if not isinstance(parsed, datetime):
return None
return (
(parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed)
.astimezone(timezone.utc)
.date()
.isoformat()
)
return utc_date_from_start_time(start) if isinstance(start, (datetime, str)) else None
def _guardrail_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:

3
litellm/accounting/snapshot.py
View file

@ -52,6 +52,9 @@ _ACCOUNTING_FIELDS: Final = frozenset(
"team_id",
"organization_id",
"end_user",
"customer_member_id",
"customer_team_id",
"customer_org_id",
"agent_id",
"request_tags",
"status",

1
litellm/constants.py
View file

@ -496,6 +496,7 @@ REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_upda
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer"
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer"
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer"
REDIS_DAILY_CUSTOMER_MEMBER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_customer_member_spend_update_buffer"
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
REDIS_TOOL_DISCOVERY_BUFFER_KEY: Final = "litellm_tool_discovery_buffer"

660
litellm/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py
View file

@ -1,326 +1,334 @@
"""
PagerDuty Alerting Integration
Handles two types of alerts:
- High LLM API Failure Rate. Configure X fails in Y seconds to trigger an alert.
- High Number of Hanging LLM Requests. Configure X hangs in Y seconds to trigger an alert.
Note: This is a Free feature on the regular litellm docker image.
However, this is under the enterprise license
"""
import asyncio
import os
from datetime import datetime, timedelta, timezone
from typing import List, Optional, Union
from litellm._logging import verbose_logger
from litellm.caching import DualCache
from litellm.exceptions import MissingMessageRoleError
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.integrations.pagerduty import (
AlertingConfig,
PagerDutyInternalEvent,
PagerDutyPayload,
PagerDutyRequestBody,
)
from litellm.types.utils import (
CallTypesLiteral,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
PAGERDUTY_DEFAULT_FAILURE_THRESHOLD = 60
PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS = 60
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS = 60
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS = 600
class PagerDutyAlerting(SlackAlerting):
"""
Tracks failed requests and hanging requests separately.
If threshold is crossed for either type, triggers a PagerDuty alert.
"""
def __init__(
self, alerting_args: Optional[Union[AlertingConfig, dict]] = None, **kwargs
):
super().__init__()
_api_key = os.getenv("PAGERDUTY_API_KEY")
if not _api_key:
raise ValueError("PAGERDUTY_API_KEY is not set")
self.api_key: str = _api_key
alerting_args = alerting_args or {}
self.pagerduty_alerting_args: AlertingConfig = AlertingConfig(
failure_threshold=alerting_args.get(
"failure_threshold", PAGERDUTY_DEFAULT_FAILURE_THRESHOLD
),
failure_threshold_window_seconds=alerting_args.get(
"failure_threshold_window_seconds",
PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS,
),
hanging_threshold_seconds=alerting_args.get(
"hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS
),
hanging_threshold_window_seconds=alerting_args.get(
"hanging_threshold_window_seconds",
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS,
),
)
# Separate storage for failures vs. hangs
self._failure_events: List[PagerDutyInternalEvent] = []
self._hanging_events: List[PagerDutyInternalEvent] = []
# ------------------ MAIN LOGIC ------------------ #
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""
Record a failure event. Only send an alert to PagerDuty if the
configured *failure* threshold is exceeded in the specified window.
"""
if isinstance(kwargs.get("exception"), MissingMessageRoleError):
return
now = datetime.now(timezone.utc)
standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object"
)
if not standard_logging_payload:
raise ValueError(
"standard_logging_object is required for PagerDutyAlerting"
)
# Extract error details
error_info: Optional[StandardLoggingPayloadErrorInformation] = (
standard_logging_payload.get("error_information") or {}
)
_meta = standard_logging_payload.get("metadata") or {}
self._failure_events.append(
PagerDutyInternalEvent(
failure_event_type="failed_response",
timestamp=now,
error_class=error_info.get("error_class"),
error_code=error_info.get("error_code"),
error_llm_provider=error_info.get("llm_provider"),
user_api_key_hash=_meta.get("user_api_key_hash"),
user_api_key_alias=_meta.get("user_api_key_alias"),
user_api_key_spend=_meta.get("user_api_key_spend"),
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
user_api_key_user_spend=_meta.get("user_api_key_user_spend"),
user_api_key_user_max_budget=_meta.get("user_api_key_user_max_budget"),
user_api_key_team_spend=_meta.get("user_api_key_team_spend"),
user_api_key_team_max_budget=_meta.get("user_api_key_team_max_budget"),
user_api_key_org_id=_meta.get("user_api_key_org_id"),
user_api_key_org_alias=_meta.get("user_api_key_org_alias"),
user_api_key_team_id=_meta.get("user_api_key_team_id"),
user_api_key_project_id=_meta.get("user_api_key_project_id"),
user_api_key_project_alias=_meta.get("user_api_key_project_alias"),
user_api_key_user_id=_meta.get("user_api_key_user_id"),
user_api_key_team_alias=_meta.get("user_api_key_team_alias"),
user_api_key_end_user_id=_meta.get("user_api_key_end_user_id"),
user_api_key_user_email=_meta.get("user_api_key_user_email"),
user_api_key_request_route=_meta.get("user_api_key_request_route"),
user_api_key_auth_metadata=_meta.get("user_api_key_auth_metadata"),
)
)
# Prune + Possibly alert
window_seconds = self.pagerduty_alerting_args.get(
"failure_threshold_window_seconds", 60
)
threshold = self.pagerduty_alerting_args.get("failure_threshold", 1)
# If threshold is crossed, send PD alert for failures
await self._send_alert_if_thresholds_crossed(
events=self._failure_events,
window_seconds=window_seconds,
threshold=threshold,
alert_prefix="High LLM API Failure Rate",
)
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
"""
Example of detecting hanging requests by waiting a given threshold.
If the request didn't finish by then, we treat it as 'hanging'.
"""
verbose_logger.info("Inside Proxy Logging Pre-call hook!")
asyncio.create_task(
self.hanging_response_handler(
request_data=data, user_api_key_dict=user_api_key_dict
)
)
return None
async def hanging_response_handler(
self, request_data: Optional[dict], user_api_key_dict: UserAPIKeyAuth
):
"""
Checks if request completed by the time 'hanging_threshold_seconds' elapses.
If not, we classify it as a hanging request.
"""
verbose_logger.debug(
f"Inside Hanging Response Handler!..sleeping for {self.pagerduty_alerting_args.get('hanging_threshold_seconds', PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS)} seconds"
)
await asyncio.sleep(
self.pagerduty_alerting_args.get(
"hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS
)
)
if await self._request_is_completed(request_data=request_data):
return # It's not hanging if completed
# Otherwise, record it as hanging
self._hanging_events.append(
PagerDutyInternalEvent(
failure_event_type="hanging_response",
timestamp=datetime.now(timezone.utc),
error_class="HangingRequest",
error_code="HangingRequest",
error_llm_provider="HangingRequest",
user_api_key_hash=user_api_key_dict.api_key,
user_api_key_alias=user_api_key_dict.key_alias,
user_api_key_spend=user_api_key_dict.spend,
user_api_key_max_budget=user_api_key_dict.max_budget,
user_api_key_budget_reset_at=(
user_api_key_dict.budget_reset_at.isoformat()
if user_api_key_dict.budget_reset_at
else None
),
user_api_key_user_spend=user_api_key_dict.user_spend,
user_api_key_user_max_budget=user_api_key_dict.user_max_budget,
user_api_key_team_spend=user_api_key_dict.team_spend,
user_api_key_team_max_budget=user_api_key_dict.team_max_budget,
user_api_key_org_id=user_api_key_dict.org_id,
user_api_key_org_alias=user_api_key_dict.organization_alias,
user_api_key_team_id=user_api_key_dict.team_id,
user_api_key_project_id=user_api_key_dict.project_id,
user_api_key_project_alias=user_api_key_dict.project_alias,
user_api_key_user_id=user_api_key_dict.user_id,
user_api_key_team_alias=user_api_key_dict.team_alias,
user_api_key_end_user_id=user_api_key_dict.end_user_id,
user_api_key_user_email=user_api_key_dict.user_email,
user_api_key_request_route=user_api_key_dict.request_route,
user_api_key_auth_metadata=user_api_key_dict.metadata,
)
)
# Prune + Possibly alert
window_seconds = self.pagerduty_alerting_args.get(
"hanging_threshold_window_seconds",
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS,
)
threshold: int = self.pagerduty_alerting_args.get(
"hanging_threshold_fails", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS
)
# If threshold is crossed, send PD alert for hangs
await self._send_alert_if_thresholds_crossed(
events=self._hanging_events,
window_seconds=window_seconds,
threshold=threshold,
alert_prefix="High Number of Hanging LLM Requests",
)
# ------------------ HELPERS ------------------ #
async def _send_alert_if_thresholds_crossed(
self,
events: List[PagerDutyInternalEvent],
window_seconds: int,
threshold: int,
alert_prefix: str,
):
"""
1. Prune old events
2. If threshold is reached, build alert, send to PagerDuty
3. Clear those events
"""
cutoff = datetime.now(timezone.utc) - timedelta(seconds=window_seconds)
pruned = [e for e in events if e.get("timestamp", datetime.min) > cutoff]
# Update the reference list
events.clear()
events.extend(pruned)
# Check threshold
verbose_logger.debug(
f"Have {len(events)} events in the last {window_seconds} seconds. Threshold is {threshold}"
)
if len(events) >= threshold:
# Build short summary of last N events
error_summaries = self._build_error_summaries(events, max_errors=5)
alert_message = (
f"{alert_prefix}: {len(events)} in the last {window_seconds} seconds."
)
custom_details = {"recent_errors": error_summaries}
await self.send_alert_to_pagerduty(
alert_message=alert_message,
custom_details=custom_details,
)
# Clear them after sending an alert, so we don't spam
events.clear()
def _build_error_summaries(
self, events: List[PagerDutyInternalEvent], max_errors: int = 5
) -> List[PagerDutyInternalEvent]:
"""
Build short text summaries for the last `max_errors`.
Example: "ValueError (code: 500, provider: openai)"
"""
recent = events[-max_errors:]
summaries = []
for fe in recent:
# If any of these is None, show "N/A" to avoid messing up the summary string
fe.pop("timestamp")
summaries.append(fe)
return summaries
async def send_alert_to_pagerduty(self, alert_message: str, custom_details: dict):
"""
Send [critical] Alert to PagerDuty
https://developer.pagerduty.com/api-reference/YXBpOjI3NDgyNjU-pager-duty-v2-events-api
"""
try:
verbose_logger.debug(f"Sending alert to PagerDuty: {alert_message}")
async_client: AsyncHTTPHandler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
payload: PagerDutyRequestBody = PagerDutyRequestBody(
payload=PagerDutyPayload(
summary=alert_message,
severity="critical",
source="LiteLLM Alert",
component="LiteLLM",
custom_details=custom_details,
),
routing_key=self.api_key,
event_action="trigger",
)
return await async_client.post(
url="https://events.pagerduty.com/v2/enqueue",
json=dict(payload),
headers={"Content-Type": "application/json"},
)
except Exception as e:
verbose_logger.exception(f"Error sending alert to PagerDuty: {e}")
"""
PagerDuty Alerting Integration
Handles two types of alerts:
- High LLM API Failure Rate. Configure X fails in Y seconds to trigger an alert.
- High Number of Hanging LLM Requests. Configure X hangs in Y seconds to trigger an alert.
Note: This is a Free feature on the regular litellm docker image.
However, this is under the enterprise license
"""
import asyncio
import os
from datetime import datetime, timedelta, timezone
from typing import Final, Literal
from litellm._logging import verbose_logger
from litellm.caching import DualCache
from litellm.exceptions import MissingMessageRoleError
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.integrations.pagerduty import (
AlertingConfig,
PagerDutyInternalEvent,
PagerDutyPayload,
PagerDutyRequestBody,
)
from litellm.types.utils import (
CallTypesLiteral,
StandardLoggingMetadata,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
PAGERDUTY_DEFAULT_FAILURE_THRESHOLD = 60
PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS = 60
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS = 60
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS = 600
def _metadata_of(payload: StandardLoggingPayload) -> StandardLoggingMetadata:
return payload["metadata"]
_CustomerIdKey = Literal[
"user_api_key_customer_member_id", "user_api_key_customer_team_id", "user_api_key_customer_org_id"
]
def _customer_id(metadata: StandardLoggingMetadata, key: _CustomerIdKey) -> str | None:
return metadata[key] if key in metadata else None
class PagerDutyAlerting(SlackAlerting):
"""
Tracks failed requests and hanging requests separately.
If threshold is crossed for either type, triggers a PagerDuty alert.
"""
def __init__(self, alerting_args: AlertingConfig | dict | None = None, **kwargs):
super().__init__()
_api_key = os.getenv("PAGERDUTY_API_KEY")
if not _api_key:
raise ValueError("PAGERDUTY_API_KEY is not set")
self.api_key: str = _api_key
alerting_args = alerting_args or {}
self.pagerduty_alerting_args: AlertingConfig = AlertingConfig(
failure_threshold=alerting_args.get("failure_threshold", PAGERDUTY_DEFAULT_FAILURE_THRESHOLD),
failure_threshold_window_seconds=alerting_args.get(
"failure_threshold_window_seconds",
PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS,
),
hanging_threshold_seconds=alerting_args.get(
"hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS
),
hanging_threshold_window_seconds=alerting_args.get(
"hanging_threshold_window_seconds",
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS,
),
)
# Separate storage for failures vs. hangs
self._failure_events: list[PagerDutyInternalEvent] = []
self._hanging_events: list[PagerDutyInternalEvent] = []
# ------------------ MAIN LOGIC ------------------ #
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""
Record a failure event. Only send an alert to PagerDuty if the
configured *failure* threshold is exceeded in the specified window.
"""
if isinstance(kwargs.get("exception"), MissingMessageRoleError):
return
now = datetime.now(timezone.utc)
standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object")
if not standard_logging_payload:
raise ValueError("standard_logging_object is required for PagerDutyAlerting")
# Extract error details
error_info: StandardLoggingPayloadErrorInformation | None = (
standard_logging_payload.get("error_information") or {}
)
_meta: Final = _metadata_of(standard_logging_payload)
self._failure_events.append(
PagerDutyInternalEvent(
failure_event_type="failed_response",
timestamp=now,
error_class=error_info.get("error_class"),
error_code=error_info.get("error_code"),
error_llm_provider=error_info.get("llm_provider"),
user_api_key_hash=_meta.get("user_api_key_hash"),
user_api_key_alias=_meta.get("user_api_key_alias"),
user_api_key_spend=_meta.get("user_api_key_spend"),
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
user_api_key_user_spend=_meta.get("user_api_key_user_spend"),
user_api_key_user_max_budget=_meta.get("user_api_key_user_max_budget"),
user_api_key_team_spend=_meta.get("user_api_key_team_spend"),
user_api_key_team_max_budget=_meta.get("user_api_key_team_max_budget"),
user_api_key_org_id=_meta.get("user_api_key_org_id"),
user_api_key_org_alias=_meta.get("user_api_key_org_alias"),
user_api_key_team_id=_meta.get("user_api_key_team_id"),
user_api_key_project_id=_meta.get("user_api_key_project_id"),
user_api_key_project_alias=_meta.get("user_api_key_project_alias"),
user_api_key_user_id=_meta.get("user_api_key_user_id"),
user_api_key_team_alias=_meta.get("user_api_key_team_alias"),
user_api_key_end_user_id=_meta.get("user_api_key_end_user_id"),
user_api_key_customer_member_id=_customer_id(_meta, "user_api_key_customer_member_id"),
user_api_key_customer_team_id=_customer_id(_meta, "user_api_key_customer_team_id"),
user_api_key_customer_org_id=_customer_id(_meta, "user_api_key_customer_org_id"),
user_api_key_user_email=_meta.get("user_api_key_user_email"),
user_api_key_request_route=_meta.get("user_api_key_request_route"),
user_api_key_auth_metadata=_meta.get("user_api_key_auth_metadata"),
)
)
# Prune + Possibly alert
window_seconds = self.pagerduty_alerting_args.get("failure_threshold_window_seconds", 60)
threshold = self.pagerduty_alerting_args.get("failure_threshold", 1)
# If threshold is crossed, send PD alert for failures
await self._send_alert_if_thresholds_crossed(
events=self._failure_events,
window_seconds=window_seconds,
threshold=threshold,
alert_prefix="High LLM API Failure Rate",
)
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
) -> Exception | str | dict | None:
"""
Example of detecting hanging requests by waiting a given threshold.
If the request didn't finish by then, we treat it as 'hanging'.
"""
verbose_logger.info("Inside Proxy Logging Pre-call hook!")
asyncio.create_task(self.hanging_response_handler(request_data=data, user_api_key_dict=user_api_key_dict))
return None
async def hanging_response_handler(self, request_data: dict | None, user_api_key_dict: UserAPIKeyAuth):
"""
Checks if request completed by the time 'hanging_threshold_seconds' elapses.
If not, we classify it as a hanging request.
"""
verbose_logger.debug(
f"Inside Hanging Response Handler!..sleeping for {self.pagerduty_alerting_args.get('hanging_threshold_seconds', PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS)} seconds"
)
await asyncio.sleep(
self.pagerduty_alerting_args.get("hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS)
)
if await self._request_is_completed(request_data=request_data):
return # It's not hanging if completed
# Otherwise, record it as hanging
self._hanging_events.append(
PagerDutyInternalEvent(
failure_event_type="hanging_response",
timestamp=datetime.now(timezone.utc),
error_class="HangingRequest",
error_code="HangingRequest",
error_llm_provider="HangingRequest",
user_api_key_hash=user_api_key_dict.api_key,
user_api_key_alias=user_api_key_dict.key_alias,
user_api_key_spend=user_api_key_dict.spend,
user_api_key_max_budget=user_api_key_dict.max_budget,
user_api_key_budget_reset_at=(
user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None
),
user_api_key_user_spend=user_api_key_dict.user_spend,
user_api_key_user_max_budget=user_api_key_dict.user_max_budget,
user_api_key_team_spend=user_api_key_dict.team_spend,
user_api_key_team_max_budget=user_api_key_dict.team_max_budget,
user_api_key_org_id=user_api_key_dict.org_id,
user_api_key_org_alias=user_api_key_dict.organization_alias,
user_api_key_team_id=user_api_key_dict.team_id,
user_api_key_project_id=user_api_key_dict.project_id,
user_api_key_project_alias=user_api_key_dict.project_alias,
user_api_key_user_id=user_api_key_dict.user_id,
user_api_key_team_alias=user_api_key_dict.team_alias,
user_api_key_end_user_id=user_api_key_dict.end_user_id,
user_api_key_customer_member_id=(
user_api_key_dict.customer_hierarchy.customer_member_id
if user_api_key_dict.customer_hierarchy
else None
),
user_api_key_customer_team_id=(
user_api_key_dict.customer_hierarchy.customer_team_id
if user_api_key_dict.customer_hierarchy
else None
),
user_api_key_customer_org_id=(
user_api_key_dict.customer_hierarchy.customer_org_id
if user_api_key_dict.customer_hierarchy
else None
),
user_api_key_user_email=user_api_key_dict.user_email,
user_api_key_request_route=user_api_key_dict.request_route,
user_api_key_auth_metadata=user_api_key_dict.metadata,
)
)
# Prune + Possibly alert
window_seconds = self.pagerduty_alerting_args.get(
"hanging_threshold_window_seconds",
PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS,
)
threshold: int = self.pagerduty_alerting_args.get(
"hanging_threshold_fails", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS
)
# If threshold is crossed, send PD alert for hangs
await self._send_alert_if_thresholds_crossed(
events=self._hanging_events,
window_seconds=window_seconds,
threshold=threshold,
alert_prefix="High Number of Hanging LLM Requests",
)
# ------------------ HELPERS ------------------ #
async def _send_alert_if_thresholds_crossed(
self,
events: list[PagerDutyInternalEvent],
window_seconds: int,
threshold: int,
alert_prefix: str,
):
"""
1. Prune old events
2. If threshold is reached, build alert, send to PagerDuty
3. Clear those events
"""
cutoff = datetime.now(timezone.utc) - timedelta(seconds=window_seconds)
pruned = [e for e in events if e.get("timestamp", datetime.min) > cutoff]
# Update the reference list
events.clear()
events.extend(pruned)
# Check threshold
verbose_logger.debug(
f"Have {len(events)} events in the last {window_seconds} seconds. Threshold is {threshold}"
)
if len(events) >= threshold:
# Build short summary of last N events
error_summaries = self._build_error_summaries(events, max_errors=5)
alert_message = f"{alert_prefix}: {len(events)} in the last {window_seconds} seconds."
custom_details = {"recent_errors": error_summaries}
await self.send_alert_to_pagerduty(
alert_message=alert_message,
custom_details=custom_details,
)
# Clear them after sending an alert, so we don't spam
events.clear()
def _build_error_summaries(
self, events: list[PagerDutyInternalEvent], max_errors: int = 5
) -> list[PagerDutyInternalEvent]:
"""
Build short text summaries for the last `max_errors`.
Example: "ValueError (code: 500, provider: openai)"
"""
recent = events[-max_errors:]
summaries = []
for fe in recent:
# If any of these is None, show "N/A" to avoid messing up the summary string
fe.pop("timestamp")
summaries.append(fe)
return summaries
async def send_alert_to_pagerduty(self, alert_message: str, custom_details: dict):
"""
Send [critical] Alert to PagerDuty
https://developer.pagerduty.com/api-reference/YXBpOjI3NDgyNjU-pager-duty-v2-events-api
"""
try:
verbose_logger.debug(f"Sending alert to PagerDuty: {alert_message}")
async_client: AsyncHTTPHandler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
payload: PagerDutyRequestBody = PagerDutyRequestBody(
payload=PagerDutyPayload(
summary=alert_message,
severity="critical",
source="LiteLLM Alert",
component="LiteLLM",
custom_details=custom_details,
),
routing_key=self.api_key,
event_action="trigger",
)
return await async_client.post(
url="https://events.pagerduty.com/v2/enqueue",
json=dict(payload),
headers={"Content-Type": "application/json"},
)
except Exception as e:
verbose_logger.exception(f"Error sending alert to PagerDuty: {e}")

1
litellm/integrations/otel/model/spans.py
View file

@ -312,6 +312,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_DailyEndUserSpend",
"LiteLLM_DailyAgentSpend",
"LiteLLM_DailyTeamSpend",
"LiteLLM_DailyCustomerMemberSpend",
"LiteLLM_DailyTagSpend",
"LiteLLM_ProxyWorkerHeartbeat",
"LiteLLM_CronJob",

33
litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json
View file

@ -5852,6 +5852,17 @@
"legendFormat": "self",
"range": true,
"refId": "M"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_daily_customer_member_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
"legendFormat": "redis_daily_customer_member_spend_update_queue",
"range": true,
"refId": "N"
}
],
"title": "Service latency p95 (litellm_<service>_latency)",
@ -6052,6 +6063,17 @@
"legendFormat": "self",
"range": true,
"refId": "N"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_redis_daily_customer_member_spend_update_queue_total_requests_total[$__rate_interval]))",
"legendFormat": "redis_daily_customer_member_spend_update_queue",
"range": true,
"refId": "O"
}
],
"title": "Service request rate (litellm_<service>_total_requests)",
@ -6252,6 +6274,17 @@
"legendFormat": "self / {{error_class}}",
"range": true,
"refId": "N"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_redis_daily_customer_member_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
"legendFormat": "redis_daily_customer_member_spend_update_queue / {{error_class}}",
"range": true,
"refId": "O"
}
],
"title": "Service failure rate (litellm_<service>_failed_requests)",

6
litellm/litellm_core_utils/litellm_logging.py
View file

@ -7274,6 +7274,9 @@ class StandardLoggingPayloadSetup:
user_api_key_team_alias=None,
user_api_key_user_email=None,
user_api_key_end_user_id=None,
user_api_key_customer_member_id=None,
user_api_key_customer_team_id=None,
user_api_key_customer_org_id=None,
user_api_key_request_route=None,
spend_logs_metadata=None,
requester_ip_address=None,
@ -8418,6 +8421,9 @@ def get_standard_logging_metadata(
user_agent=None,
requester_metadata=None,
user_api_key_end_user_id=None,
user_api_key_customer_member_id=None,
user_api_key_customer_team_id=None,
user_api_key_customer_org_id=None,
prompt_management_metadata=None,
applied_guardrails=None,
mcp_tool_call_metadata=None,

9
litellm/proxy/_types.py
View file

@ -5403,6 +5403,12 @@ class DailyAgentSpendTransaction(BaseDailySpendTransaction):
agent_id: str
class DailyCustomerMemberSpendTransaction(BaseDailySpendTransaction):
customer_member_id: ReadOnly[str]
customer_team_id: ReadOnly[str]
customer_org_id: ReadOnly[str]
class DBSpendUpdateTransactions(TypedDict):
"""
Internal Data Structure for buffering spend updates in Redis or in memory before committing them to the database
@ -5419,6 +5425,9 @@ class DBSpendUpdateTransactions(TypedDict):
tag_list_transactions: dict[str, float] | None
agent_list_transactions: dict[str, float] | None
model_access_group_list_transactions: ReadOnly[dict[str, float] | None]
customer_org_list_transactions: ReadOnly[dict[str, float] | None] # mutable-ok: aggregated like siblings
customer_team_list_transactions: ReadOnly[dict[str, float] | None] # mutable-ok: aggregated like siblings
customer_member_list_transactions: ReadOnly[dict[str, float] | None] # mutable-ok: aggregated like siblings
class SpendUpdateQueueItem(TypedDict, total=False):

3
litellm/proxy/auth/auth_checks.py
View file

@ -5620,12 +5620,15 @@ async def can_key_call_resolved_model(
valid_token: UserAPIKeyAuth,
llm_router: RegistryReader | None,
) -> None:
from litellm.proxy.auth.customer_hierarchy_auth import check_customer_hierarchy_model_access
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if valid_token.customer_hierarchy is not None:
check_customer_hierarchy_model_access(valid_token.customer_hierarchy, model, llm_router)
skip_key_model_check: Final = valid_token.config or (
isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models
)

252
litellm/proxy/auth/customer_hierarchy_auth.py Normal file
View file

@ -0,0 +1,252 @@
"""Request-path resolution of the customer organization, team and member a request runs under.
A request names its customer member with the ``x-litellm-customer-member-id`` header. The member row
carries its team and that team's organization, so one lookup yields every level whose budget and rate
limits apply to the request.
"""
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final
from fastapi import status
from pydantic import BaseModel, ConfigDict, TypeAdapter
import litellm
from litellm._internal_context import with_service_target
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.auth.auth_checks import can_object_call_model
from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl
from litellm.proxy.management_helpers.customer_hierarchy import (
PrismaRecord,
effective_member_budget,
related_budget,
utc_now,
)
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth, CustomerScopeKey
from litellm.types.proxy.management_endpoints.customer_hierarchy import Budget
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
from litellm.router_api import RegistryReader
CUSTOMER_MEMBER_ID_HEADER: Final = "x-litellm-customer-member-id"
_MEMBER_INCLUDE: Final = {
"budget": True,
"budget_override": True,
"parent_team": {"include": {"budget": True, "parent_org": {"include": {"budget": True}}}},
}
_EXPIRY: Final[TypeAdapter[datetime | None]] = TypeAdapter(datetime | None)
class _OrganizationRow(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
customer_org_id: str
spend: float = 0.0
budget: Mapping[str, object] | None = None
class _TeamRow(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
customer_team_id: str
spend: float = 0.0
budget: Mapping[str, object] | None = None
parent_org: _OrganizationRow
class _MemberRow(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
customer_member_id: str
spend: float = 0.0
budget: Mapping[str, object] | None = None
budget_override: Mapping[str, object] | None = None
budget_override_expires_at: datetime | None = None
parent_team: _TeamRow
def customer_member_cache_key(customer_member_id: str) -> str:
return f"customer_member_id:{customer_member_id}"
def customer_member_id_from_headers(headers: Mapping[str, object] | None) -> str | None:
if headers is None:
return None
values: Final = tuple(
value.strip()
for name, value in headers.items()
if name.lower() == CUSTOMER_MEMBER_ID_HEADER and isinstance(value, str) and value.strip()
)
return values[0] if values else None
def _scope(key: CustomerScopeKey, entity_id: str, spend: float, budget: Budget | None) -> CustomerScopeAuth:
if budget is None:
return CustomerScopeAuth(key=key, entity_id=entity_id, spend=spend)
return CustomerScopeAuth(
key=key,
entity_id=entity_id,
spend=spend,
max_budget=budget.max_budget,
tpm_limit=budget.tpm_limit,
rpm_limit=budget.rpm_limit,
allowed_models=budget.allowed_models or None,
model_tpm_limit=budget.model_tpm_limit,
model_rpm_limit=budget.model_rpm_limit,
)
def customer_hierarchy_from_member_row(row: Mapping[str, object], now: datetime) -> CustomerHierarchyAuth:
"""The org, team and member scopes of one member row loaded with ``_MEMBER_INCLUDE``. The member's limits are
its effective budget: the selected organization template under any unexpired override."""
member: Final = _MemberRow.model_validate(row)
team: Final = member.parent_team
organization: Final = team.parent_org
template: Final = related_budget(member.budget, now)
override: Final = related_budget(member.budget_override, now, template.max_budget if template is not None else None)
return CustomerHierarchyAuth(
customer_member_id=member.customer_member_id,
customer_team_id=team.customer_team_id,
customer_org_id=organization.customer_org_id,
scopes=(
_scope(
"customer_org",
organization.customer_org_id,
organization.spend,
related_budget(organization.budget, now),
),
_scope("customer_team", team.customer_team_id, team.spend, related_budget(team.budget, now)),
_scope(
"customer_member",
member.customer_member_id,
member.spend,
effective_member_budget(template, override, member.budget_override_expires_at, now),
),
),
)
def _member_table(prisma_client: "PrismaClient") -> TableActions[PrismaRecord]:
return prisma_client.db.litellm_customermembertable # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportAny] # prisma tables are generated at build time
@with_service_target(AUTH_OBJECTS_TARGET)
async def get_customer_hierarchy(
customer_member_id: str,
prisma_client: "PrismaClient",
user_api_key_cache: UserApiKeyCache,
now: datetime | None = None,
) -> CustomerHierarchyAuth | None:
cache_key: Final = customer_member_cache_key(customer_member_id)
cached: Final = await user_api_key_cache.async_get_cache(key=cache_key, model_type=CustomerHierarchyAuth)
if cached is not None:
return cached
row: Final = await _member_table(prisma_client).find_unique(
where={"customer_member_id": customer_member_id}, include=_MEMBER_INCLUDE
)
if row is None:
return None
loaded_at: Final = now if now is not None else utc_now()
row_values: Final = row.model_dump()
hierarchy: Final = customer_hierarchy_from_member_row(row_values, loaded_at)
await user_api_key_cache.async_set_cache(
key=cache_key,
value=hierarchy,
ttl=_cache_ttl(row_values, loaded_at, get_management_object_ttl(user_api_key_cache)),
)
return hierarchy
def _cache_ttl(row: Mapping[str, object], now: datetime, default_ttl: float) -> float:
"""A cached hierarchy may not outlive the member's budget override, or an expired override would keep
granting its limits until the management-object TTL ran out."""
expires_at: Final = _MemberRow.model_validate(row).budget_override_expires_at
if expires_at is None or expires_at <= now:
return default_ttl
return max(1.0, min(default_ttl, (expires_at - now).total_seconds()))
def check_customer_hierarchy_budget(hierarchy: CustomerHierarchyAuth) -> None:
"""Raises ``BudgetExceededError`` for the first level, org first, whose spend has reached its ``max_budget``."""
for scope in hierarchy.scopes:
if scope.max_budget is not None and scope.spend >= scope.max_budget:
raise litellm.BudgetExceededError(
current_cost=scope.spend,
max_budget=scope.max_budget,
message=(
f"ExceededBudget: {scope.key}={scope.entity_id} over budget. "
f"Spend={scope.spend}, Budget={scope.max_budget}"
),
entity_type=scope.key,
entity_id=scope.entity_id,
)
def check_customer_hierarchy_model_access(
hierarchy: CustomerHierarchyAuth, model: str, llm_router: "RegistryReader | None"
) -> None:
"""Every level with an ``allowed_models`` list must permit ``model``; a level without one permits every model."""
for scope in hierarchy.scopes:
if scope.allowed_models:
can_object_call_model(
model=model, llm_router=llm_router, models=list(scope.allowed_models), object_type="customer"
)
async def resolve_customer_hierarchy(
customer_member_id: str,
prisma_client: "PrismaClient | None",
user_api_key_cache: UserApiKeyCache,
llm_router: "RegistryReader | None",
request_model: str | Sequence[str] | None,
route: str,
key_metadata: Mapping[str, object] | None = None,
) -> CustomerHierarchyAuth:
"""Load the hierarchy the header names and admit the request against it. An unknown member is refused
rather than treated as unrestricted: the caller asked for a budgeted identity that does not exist.
``request_model`` is the model the route will actually serve, which may differ from the body."""
hierarchy: Final = (
await get_customer_hierarchy(customer_member_id, prisma_client, user_api_key_cache)
if prisma_client is not None
else None
)
if hierarchy is None:
raise ProxyException(
message=f"Authentication Error - customer member not found: {customer_member_id}",
type=ProxyErrorTypes.auth_error,
param=CUSTOMER_MEMBER_ID_HEADER,
code=status.HTTP_401_UNAUTHORIZED,
)
check_customer_hierarchy_key_scope(hierarchy, key_metadata or {})
if not RouteChecks.is_info_route(route):
check_customer_hierarchy_budget(hierarchy)
requested: Final = (request_model,) if isinstance(request_model, str) else tuple(request_model or ())
for model in requested:
if model:
check_customer_hierarchy_model_access(hierarchy, model, llm_router)
return hierarchy
def check_customer_hierarchy_key_scope(hierarchy: CustomerHierarchyAuth, key_metadata: Mapping[str, object]) -> None:
"""A key whose metadata pins ``customer_org_id`` may only act for members of that organization."""
key_customer_org_id: Final = key_metadata.get("customer_org_id")
if key_customer_org_id is None or key_customer_org_id == hierarchy.customer_org_id:
return
raise ProxyException(
message="API key is not allowed to act for this customer member's organization",
type=ProxyErrorTypes.auth_error,
param=CUSTOMER_MEMBER_ID_HEADER,
code=status.HTTP_403_FORBIDDEN,
)
@with_service_target(AUTH_OBJECTS_TARGET)
async def forget_customer_hierarchies(customer_member_ids: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None:
"""Drop the cached hierarchies of members whose own, team, org or template budget changed so auth reloads them."""
await evict_and_broadcast(tuple(customer_member_cache_key(m) for m in customer_member_ids), user_api_key_cache)

26
litellm/proxy/auth/user_api_key_auth.py
View file

@ -130,6 +130,10 @@ from litellm.proxy.auth.auth_utils import (
request_dispatched_to_provider_pass_through,
route_in_additonal_public_routes,
)
from litellm.proxy.auth.customer_hierarchy_auth import (
customer_member_id_from_headers,
resolve_customer_hierarchy,
)
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context
from litellm.proxy.auth.oauth2_check import Oauth2Handler
@ -301,6 +305,7 @@ _CLAUDE_MODEL_ROUTES: Final = frozenset(
)
_CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized"
_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object])
_KEY_METADATA_ADAPTER: Final = TypeAdapter(dict[str, object])
async def _normalize_claude_model(
@ -4535,9 +4540,30 @@ async def _authorize_authenticated_request(
llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
llm_router,
prisma_client,
user_api_key_cache,
user_model,
)
customer_member_id: Final = customer_member_id_from_headers(
safe_get_request_headers(request) # pyright: ignore[reportUnknownArgumentType] # request headers are untyped
)
user_api_key_auth_obj.customer_hierarchy = None # rebind-ok: request state, never inherited from a cached key
if customer_member_id is not None:
with tracer.trace("litellm.proxy.auth.resolve_customer_hierarchy"): # pyright: ignore[reportUnknownMemberType] # tracer is the untyped proxy tracer
user_api_key_auth_obj.customer_hierarchy = await resolve_customer_hierarchy( # rebind-ok: the gate updates the request-owned auth object in place
customer_member_id=customer_member_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
llm_router=llm_router,
request_model=_get_model_from_request_context(
request_data, route, request, llm_router, team_id=user_api_key_auth_obj.team_id
),
route=route,
key_metadata=_KEY_METADATA_ADAPTER.validate_python(
user_api_key_auth_obj.metadata or {} # pyright: ignore[reportUnknownMemberType] # untyped key metadata, validated here
),
)
grants: Final = await access_group_ids_for_identity(user_api_key_auth_obj)
user_api_key_auth_obj.apply_access_group_rule_grants(grants)
key_check_required: Final = (

3
litellm/proxy/common_utils/proxy_rate_limit_error.py
View file

@ -56,6 +56,9 @@ PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS: Final = frozenset(
"model_per_project",
"model_per_project_itpm",
"model_per_project_otpm",
"model_per_customer_org",
"model_per_customer_team",
"model_per_customer_member",
}
)
DISABLE_FALLBACKS_ON_PER_MODEL_RATE_LIMITS_SETTING: Final = "disable_fallbacks_on_per_model_rate_limits"

141
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -17,6 +17,7 @@ from typing import (
cast, # noqa: TID251 # get_data returns validated rows selected by table_name
)
from pydantic import AliasChoices, BaseModel, ConfigDict, Field
from typing_extensions import assert_never
import litellm
@ -43,6 +44,8 @@ from litellm.proxy._types import (
LiteLLM_UserTable,
LiteLLM_VerificationToken,
)
from litellm.proxy.auth.customer_hierarchy_auth import customer_member_cache_key
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import broadcast_evictions
from litellm.proxy.common_utils.timezone_utils import (
BudgetResetSettings,
compute_budget_reset_at,
@ -59,6 +62,7 @@ from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
from litellm.proxy.db.db_span import db_span, db_spanned
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.management_helpers.customer_hierarchy import BudgetRecord, PrismaRecord, effective_member_budget
from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import (
@ -77,9 +81,12 @@ from litellm.redis.stores.spend_counters import (
user_counter,
)
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable
from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable, TableActions
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
CustomerMemberRepository,
CustomerOrganizationRepository,
CustomerTeamRepository,
EndUserRepository,
ModelAccessGroupBudgetRepository,
TagRepository,
@ -87,9 +94,11 @@ from litellm.repositories.table_repositories import (
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.unit_of_work import (
CustomerSpendEntity,
LinkedSpendResetWrites,
SpendResetWrites,
budget_cascade_unit_of_work,
customer_spend_reset_unit_of_work,
enduser_spend_reset_unit_of_work,
spend_reset_unit_of_work,
)
@ -98,6 +107,18 @@ from litellm.repositories.verification_token_repository import (
)
from litellm.types.services import ServiceTypes
class _CustomerBudgetRow(BaseModel):
model_config = ConfigDict(frozen=True)
identity: str = Field(validation_alias=AliasChoices("customer_member_id", "customer_team_id", "customer_org_id"))
spend: float
budget_reset_at: datetime | None
budget: BudgetRecord | None = None
budget_override: BudgetRecord | None = None
budget_override_expires_at: datetime | None = None
_RowT = TypeVar("_RowT")
@ -599,10 +620,128 @@ class ResetBudgetJob:
await self.reset_budget_for_litellm_teams()
await self.reset_budget_for_litellm_budget_table()
await self.reset_budget_windows()
await self.reset_budget_for_customer_hierarchy()
finally:
if lease is _Lease.LEADER and self.pod_lock_manager is not None:
await self.pod_lock_manager.release_lock(cronjob_id=RESET_BUDGET_JOB_NAME)
async def reset_budget_for_customer_hierarchy(self) -> None:
tables: Final[tuple[tuple[CustomerSpendEntity, TableActions[PrismaRecord]], ...]] = (
("org", CustomerOrganizationRepository(self.prisma_client, use_writer=True).table),
("team", CustomerTeamRepository(self.prisma_client, use_writer=True).table),
("member", CustomerMemberRepository(self.prisma_client, use_writer=True).table),
)
for entity, table in tables:
await self._reset_customer_budget_table(entity, table)
async def _reset_customer_budget_table(
self, entity: CustomerSpendEntity, table: TableActions[PrismaRecord]
) -> None:
now: Final = datetime.now(timezone.utc)
id_field: Final = f"customer_{entity}_id"
budget_duration: Final = {"budget": {"is": {"budget_duration": {"not": None}}}}
duration_sources: Final = (
(
{
"budget_override": {"is": {"budget_duration": {"not": None}}},
"OR": [
{"budget_override_expires_at": None},
{"budget_override_expires_at": {"gt": now}},
],
},
budget_duration,
)
if entity == "member"
else (budget_duration,)
)
cursor: str | None = None # rebind-ok: advances a bounded keyset walk
for _ in range(RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN):
where = {
**({id_field: {"gt": cursor}} if cursor is not None else {}),
"OR": [
{"budget_reset_at": {"lte": now}},
{"budget_reset_at": None, "OR": list(duration_sources)},
],
}
records = await self._with_db_retry(
lambda: table.find_many(
where=where,
include={"budget": True, **({"budget_override": True} if entity == "member" else {})},
order={id_field: "asc"},
take=RESET_BUDGET_JOB_BATCH_SIZE,
),
reason="reset_budget_read_customer_failure",
)
if not records:
return
rows = tuple(_CustomerBudgetRow.model_validate(record.model_dump()) for record in records)
await self._with_db_write_retry(
lambda: self._write_customer_budget_resets(entity, rows, now),
reason="reset_budget_write_customer_failure",
)
await self._invalidate_customer_budget_caches(entity, tuple(row.identity for row in rows))
if len(records) < RESET_BUDGET_JOB_BATCH_SIZE:
return
cursor = rows[-1].identity
async def _invalidate_customer_budget_caches(self, entity: CustomerSpendEntity, identities: Sequence[str]) -> None:
if entity == "member":
keys: Final = tuple(customer_member_cache_key(identity) for identity in identities)
await self._invalidate_cache("user_api_key_cache", keys)
await broadcast_evictions(keys)
return
table: Final = CustomerMemberRepository(self.prisma_client, use_writer=True).table
parents: Final = (
{"parent_team_id": {"in": list(identities)}} # bounded-ok: one RESET_BUDGET_JOB_BATCH_SIZE page
if entity == "team"
else {"parent_team": {"is": {"parent_org_id": {"in": list(identities)}}}} # bounded-ok: one reset page
)
cursor: str | None = None # rebind-ok: keyset walk must invalidate every descendant after commit
while True:
records = await self._with_db_retry(
lambda: table.find_many(
where={**parents, **({"customer_member_id": {"gt": cursor}} if cursor is not None else {})},
order={"customer_member_id": "asc"},
take=RESET_BUDGET_JOB_BATCH_SIZE,
),
reason="reset_budget_read_customer_cache_failure",
)
if not records:
return
rows = tuple(_CustomerBudgetRow.model_validate(record.model_dump()) for record in records)
await self._invalidate_customer_budget_caches("member", tuple(row.identity for row in rows))
if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE:
return
cursor = rows[-1].identity
async def _write_customer_budget_resets(
self, entity: CustomerSpendEntity, rows: Sequence[_CustomerBudgetRow], now: datetime
) -> None:
async with customer_spend_reset_unit_of_work(self._new_batch, entity) as writes:
for row in rows:
budget = effective_member_budget(row.budget, row.budget_override, row.budget_override_expires_at, now)
next_reset = (
compute_budget_reset_at(budget.budget_duration, self.reset_settings)
if budget.budget_duration is not None
else None
)
# A new duration starts a window without erasing already accrued spend.
resetting = row.budget_reset_at is not None and next_reset is not None
carried = (
_carried_spend(row.spend, _rollover_cap(budget.max_budget))
if resetting and _rollover_enabled()
else 0.0
if resetting
else row.spend
)
writes.queue_spend_reset(
identity=row.identity,
budget_reset_at=next_reset,
spend_decrement=row.spend - carried,
expected_reset_at=row.budget_reset_at,
spend_after_reset=carried,
)
async def _with_db_retry(self, operation: Callable[[], Awaitable[_RowT]], *, reason: str) -> _RowT:
"""Reconnect and retry once on a transport error, so a dropped connection
costs one retry instead of the whole tick.

51
litellm/proxy/db/daily_spend_bulk_upsert.py
View file

@ -10,13 +10,15 @@ the window in which those locks are held.
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from hashlib import sha256
from itertools import groupby
from json import dumps
from types import MappingProxyType
from typing import Final, Literal
from typing import Final, Literal, TypeAlias
from pydantic import TypeAdapter
DailySpendEntity = Literal["user", "team", "org", "tag", "end_user", "agent"]
DailySpendEntity: TypeAlias = Literal["user", "team", "org", "tag", "end_user", "agent", "customer_member"]
SqlValue = str | int | float | None
@ -32,10 +34,21 @@ class DailySpendTable:
name: str
entity_id_column: str
carries_request_id: bool = False
parent_id_columns: tuple[str, ...] = ()
# The customer table arbitrates on a bounded hash of its identity (model group
# included) stored in ``id``, so three long customer ids never exceed an index entry
# and rows already written under that identity keep accumulating.
hashed_identity: bool = False
DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType(
{
"customer_member": DailySpendTable(
name="LiteLLM_DailyCustomerMemberSpend",
entity_id_column="customer_member_id",
parent_id_columns=("customer_team_id", "customer_org_id"),
hashed_identity=True,
),
"user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"),
"team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"),
"org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"),
@ -47,6 +60,7 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP
_ENTITY_INPUT_KEYS: Final[Mapping[DailySpendEntity, str]] = MappingProxyType(
{
"customer_member": "customer_member_id",
"user": "user",
"team": "team_id",
"org": "organization_id",
@ -59,6 +73,12 @@ _TAGS: Final = TypeAdapter(tuple[object, ...])
def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEntity) -> tuple[str | None, ...]:
if entity == "customer_member" and not (
isinstance(payload.get("customer_org_id"), str)
and payload.get("customer_org_id")
and all(isinstance(payload.get(column), str) for column in ("customer_member_id", "customer_team_id"))
):
return ()
key: Final = _ENTITY_INPUT_KEYS[entity]
if key not in payload:
return ()
@ -125,9 +145,24 @@ def _as_float(value: object) -> float:
return float(value) if isinstance(value, (int, float)) else 0.0
def _identity_columns(table: DailySpendTable) -> tuple[str, ...]:
return (
table.entity_id_column,
*table.parent_id_columns,
*_KEY_COLUMNS,
*(("model_group",) if table.hashed_identity else ()),
)
def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]:
"""The tuple the database arbitrates the upsert on, normalized free of NULLs."""
return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS))
return tuple(_as_text(transaction.get(column)) for column in _identity_columns(table))
def hashed_row_id(table: DailySpendTable, key: tuple[str, ...]) -> str:
"""The ``id`` of a hashed-identity row, identical for every writer of the same identity."""
pairs: Final = sorted(zip(_identity_columns(table), key, strict=True))
return sha256(dumps(pairs, separators=(",", ":")).encode("utf-8")).hexdigest()
def _merge(group: Sequence[SpendRow]) -> SpendRow:
@ -162,10 +197,11 @@ def _row_params(
transaction: SpendRow,
) -> tuple[SqlValue, ...]:
request_id: Final = transaction.get("request_id")
model_group: Final = transaction.get("model_group")
return (
str(uuid.uuid4()),
*key,
None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")),
hashed_row_id(table, key) if table.hashed_identity else str(uuid.uuid4()),
*key[: len(key) - 1 if table.hashed_identity else len(key)],
_as_text(model_group) if table.hashed_identity or model_group is not None else None,
*(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS),
*(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS),
*((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()),
@ -176,6 +212,7 @@ def _insert_columns(table: DailySpendTable) -> tuple[str, ...]:
return (
"id",
table.entity_id_column,
*table.parent_id_columns,
*_KEY_COLUMNS,
"model_group",
*_COUNTER_COLUMNS,
@ -214,7 +251,7 @@ def build_bulk_upsert(
sql: Final = (
f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n'
f"VALUES {rows}\n"
f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n"
f"ON CONFLICT ({_quoted(('id',) if table.hashed_identity else _identity_columns(table))}) DO UPDATE SET\n"
f" {increments}{request_id_update},\n"
f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')"
)

265
litellm/proxy/db/db_spend_update_writer.py
View file

@ -45,6 +45,7 @@ from litellm.proxy._types import (
DB_RETRY_SAFE_ERROR_TYPES,
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
DailyCustomerMemberSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTagSpendTransaction,
@ -61,6 +62,7 @@ from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, p
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
build_bulk_upsert,
conflict_key,
daily_spend_entity_ids,
merge_by_conflict_key,
)
@ -79,7 +81,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
WindowSpendUpdateQueue,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_failure_is_requeue_safe
from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction
from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction, utc_date_from_start_time
from litellm.proxy.db.request_error_tracking import request_error_accumulator
from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
@ -96,6 +98,7 @@ from litellm.proxy.spend_tracking.session_spend_total import SessionSpendTotals
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.repositories.prisma_protocols import BatchTable
from litellm.responses.streaming_iterator import spawn_logging_task
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyIds, customer_hierarchy_ids_from_request
from litellm.types.utils import CallTypes
from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES
@ -363,6 +366,9 @@ _SpendTableName = Literal[
"tag_list_transactions",
"model_access_group_list_transactions",
"agent_list_transactions",
"customer_org_list_transactions",
"customer_team_list_transactions",
"customer_member_list_transactions",
]
_SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = (
"org_list_transactions",
@ -376,6 +382,36 @@ _SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = (
"tag_list_transactions",
"model_access_group_list_transactions",
"agent_list_transactions",
"customer_org_list_transactions",
"customer_team_list_transactions",
"customer_member_list_transactions",
)
@dataclasses.dataclass(frozen=True, slots=True)
class _CustomerHierarchyLevel:
transactions: _SpendTableName
table: str
prisma_table: str
id_column: str
_CUSTOMER_HIERARCHY_LEVELS: Final = (
_CustomerHierarchyLevel(
"customer_org_list_transactions",
"LiteLLM_CustomerOrganizationTable",
"litellm_customerorganizationtable",
"customer_org_id",
),
_CustomerHierarchyLevel(
"customer_team_list_transactions", "LiteLLM_CustomerTeamTable", "litellm_customerteamtable", "customer_team_id"
),
_CustomerHierarchyLevel(
"customer_member_list_transactions",
"LiteLLM_CustomerMemberTable",
"litellm_customermembertable",
"customer_member_id",
),
)
@ -412,6 +448,7 @@ class _SpendRestore(TypedDict, total=False):
daily_org_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]]
daily_end_user_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]]
daily_agent_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]]
daily_customer_member_spend_update_transactions: ReadOnly[Mapping[str, BaseDailySpendTransaction]]
window_spend_update_transactions: ReadOnly[Sequence[WindowSpendTransaction]]
@ -734,6 +771,7 @@ class DBSpendUpdateWriter:
self.daily_team_spend_update_queue = DailySpendUpdateQueue()
self.daily_end_user_spend_update_queue = DailySpendUpdateQueue()
self.daily_agent_spend_update_queue = DailySpendUpdateQueue()
self.daily_customer_member_spend_update_queue = DailySpendUpdateQueue()
self.daily_org_spend_update_queue = DailySpendUpdateQueue()
self.daily_tag_spend_update_queue = DailySpendUpdateQueue()
self.window_spend_update_queue = WindowSpendUpdateQueue()
@ -774,6 +812,7 @@ class DBSpendUpdateWriter:
)
use_ledger: Final = proxy_settings.get("use_spend_logs_ledger", True) is not False
request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {})
customer_hierarchy: Final = customer_hierarchy_ids_from_request(request_fields)
if kwargs is not None:
kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker
try:
@ -834,6 +873,7 @@ class DBSpendUpdateWriter:
"team_id": team_id,
"organization_id": org_id,
"project_id": project_id,
**(customer_hierarchy.entities() if customer_hierarchy is not None else {}),
"agent_id": payload["agent_id"],
"proxy_budget_name": litellm_proxy_budget_name if litellm.max_budget > 0 else None,
},
@ -891,6 +931,7 @@ class DBSpendUpdateWriter:
litellm_proxy_budget_name=litellm_proxy_budget_name,
payload=payload,
request_model_access_groups=get_request_model_access_groups(request_fields),
customer_hierarchy=customer_hierarchy,
),
task_name="spend rollups",
)
@ -1357,6 +1398,7 @@ class DBSpendUpdateWriter:
payload: SpendLogsPayload,
request_model_access_groups: Sequence[str] = (),
project_id: str | None = None,
customer_hierarchy: CustomerHierarchyIds | None = None,
):
"""
Runs every spend-update helper for one request. Each helper goes through
@ -1400,6 +1442,13 @@ class DBSpendUpdateWriter:
self._update_project_db(response_cost=response_cost, project_id=project_id, prisma_client=prisma_client),
payload=payload_copy,
)
await self._run_spend_update_step(
"_update_customer_hierarchy_db",
self._update_customer_hierarchy_db(
response_cost=response_cost, customer_hierarchy=customer_hierarchy, prisma_client=prisma_client
),
payload=payload_copy,
)
await self._run_spend_update_step(
"_update_tag_db",
self._update_tag_db(response_cost=response_cost, request_tags=request_tags, prisma_client=prisma_client),
@ -1412,6 +1461,13 @@ class DBSpendUpdateWriter:
prisma_client=prisma_client,
router=get_llm_router(),
)
await self._run_spend_update_step(
"add_spend_log_transaction_to_daily_customer_member_transaction",
self.add_spend_log_transaction_to_daily_customer_member_transaction(
payload_copy, customer_hierarchy, prisma_client
),
payload=payload_copy,
)
await self._run_spend_update_step(
"_update_agent_db",
self._update_agent_db(
@ -1589,6 +1645,19 @@ class DBSpendUpdateWriter:
)
raise e
async def _update_customer_hierarchy_db(
self,
response_cost: float | None,
customer_hierarchy: CustomerHierarchyIds | None,
prisma_client: PrismaClient | None,
) -> None:
if customer_hierarchy is None or prisma_client is None:
return
for entity_type, entity_id in customer_hierarchy.spend_entities():
await self.spend_update_queue.add_update(
update=SpendUpdateQueueItem(entity_type=entity_type, entity_id=entity_id, response_cost=response_cost)
)
async def _update_org_db(
self,
response_cost: float | None,
@ -1860,6 +1929,7 @@ class DBSpendUpdateWriter:
daily_org_spend_update_queue=self.daily_org_spend_update_queue,
daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue,
daily_agent_spend_update_queue=self.daily_agent_spend_update_queue,
daily_customer_member_spend_update_queue=self.daily_customer_member_spend_update_queue,
window_spend_update_queue=self.window_spend_update_queue,
)
@ -1881,6 +1951,7 @@ class DBSpendUpdateWriter:
daily_end_user_spend_update_transactions,
daily_agent_spend_update_transactions,
window_spend_update_transactions,
daily_customer_member_spend_update_transactions,
) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
uncommitted = {
@ -1914,6 +1985,13 @@ class DBSpendUpdateWriter:
if daily_agent_spend_update_transactions is None
else {"daily_agent_spend_update_transactions": daily_agent_spend_update_transactions}
),
"daily_customer_member_spend_update_transactions": (
{}
if daily_customer_member_spend_update_transactions is None
else {
"daily_customer_member_spend_update_transactions": daily_customer_member_spend_update_transactions
}
),
"window_spend_update_transactions": (
{}
if window_spend_update_transactions is None
@ -1925,7 +2003,7 @@ class DBSpendUpdateWriter:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, "
"projects=%d, tags=%d, agents=%d, model_access_groups=%d",
"projects=%d, tags=%d, agents=%d, model_access_groups=%d, customer_members=%d",
len(db_spend_update_transactions.get("key_list_transactions") or ()),
len(db_spend_update_transactions.get("user_list_transactions") or ()),
len(db_spend_update_transactions.get("team_list_transactions") or ()),
@ -1937,6 +2015,7 @@ class DBSpendUpdateWriter:
len(db_spend_update_transactions.get("tag_list_transactions") or ()),
len(db_spend_update_transactions.get("agent_list_transactions") or ()),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or ()),
len(db_spend_update_transactions.get("customer_member_list_transactions") or ()),
)
try:
await self._commit_spend_updates_to_db(
@ -2000,6 +2079,15 @@ class DBSpendUpdateWriter:
daily_spend_transactions=daily_agent_spend_update_transactions,
)
uncommitted.pop("daily_agent_spend_update_transactions", None)
if daily_customer_member_spend_update_transactions is not None:
await DBSpendUpdateWriter.update_daily_customer_member_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_customer_member_spend_update_transactions,
)
uncommitted.pop("daily_customer_member_spend_update_transactions", None)
if window_spend_update_transactions is not None:
try:
await DBSpendUpdateWriter._commit_window_spend_updates(
@ -2036,7 +2124,7 @@ class DBSpendUpdateWriter:
async def _flush_daily_spend_queue(
self,
queue: DailySpendUpdateQueue,
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"],
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent", "customer_member"],
commit: _DailySpendCommit[_DailySpendTransactionT],
n_retry_times: int,
prisma_client: PrismaClient,
@ -2159,6 +2247,15 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
await self._flush_daily_spend_queue(
queue=self.daily_customer_member_spend_update_queue,
entity_type="customer_member",
commit=DBSpendUpdateWriter.update_daily_customer_member_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
################## Budget Window Spend Update Transactions ##################
# Aggregate all in memory budget window spend transactions and commit to db
window_spend_update_transactions: Final = (
@ -2699,6 +2796,99 @@ class DBSpendUpdateWriter:
if on_table_committed is not None:
on_table_committed("agent_list_transactions")
await self._commit_customer_hierarchy_spend(
db_spend_update_transactions,
prisma_client=prisma_client,
n_retry_times=n_retry_times,
proxy_logging_obj=proxy_logging_obj,
outer_transaction=outer_transaction,
on_table_committed=on_table_committed,
)
@staticmethod
async def _commit_customer_hierarchy_spend(
db_spend_update_transactions: DBSpendUpdateTransactions,
*,
prisma_client: PrismaClient,
n_retry_times: int,
proxy_logging_obj: ProxyLogging,
outer_transaction: _SpendTransaction | None,
on_table_committed: Callable[[_SpendTableName], None] | None,
) -> None:
"""Increment the customer org, team and member rows a request ran under, one table per commit so a
partial failure re-sends only the tables still owed."""
for level in _CUSTOMER_HIERARCHY_LEVELS:
await DBSpendUpdateWriter._commit_customer_hierarchy_level(
level,
db_spend_update_transactions.get(level.transactions),
prisma_client=prisma_client,
n_retry_times=n_retry_times,
proxy_logging_obj=proxy_logging_obj,
outer_transaction=outer_transaction,
)
if on_table_committed is not None:
on_table_committed(level.transactions)
@staticmethod
async def _commit_customer_hierarchy_level(
level: _CustomerHierarchyLevel,
transactions: Mapping[str, float] | None,
*,
prisma_client: PrismaClient,
n_retry_times: int,
proxy_logging_obj: ProxyLogging,
outer_transaction: _SpendTransaction | None,
) -> None:
if not transactions:
return
for attempt in range(n_retry_times + 1):
if await DBSpendUpdateWriter._commit_customer_hierarchy_level_once(
level,
transactions,
attempt=attempt,
prisma_client=prisma_client,
n_retry_times=n_retry_times,
proxy_logging_obj=proxy_logging_obj,
outer_transaction=outer_transaction,
):
return
@staticmethod
async def _commit_customer_hierarchy_level_once(
level: _CustomerHierarchyLevel,
transactions: Mapping[str, float],
*,
attempt: int,
prisma_client: PrismaClient,
n_retry_times: int,
proxy_logging_obj: ProxyLogging,
outer_transaction: _SpendTransaction | None,
) -> bool:
"""One attempt at the level's increments; False means the attempt failed and was reported for retry."""
start_time: Final = time.time()
try:
async with (
_spend_update_tx(prisma_client, level.table, outer_transaction=outer_transaction) as tx,
tx.batch_() as batcher,
):
for entity_id, response_cost in sorted(transactions.items()):
getattr(batcher, level.prisma_table).update_many(
where={level.id_column: entity_id},
data={"spend": {"increment": response_cost}},
)
except Exception as e:
if outer_transaction is not None:
raise
await DBSpendUpdateWriter.handle_spend_update_failure(
e=e,
attempt=attempt,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
return False
return True
@staticmethod
@with_service_target(AUTH_OBJECTS_TARGET)
async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None:
@ -2830,6 +3020,18 @@ class DBSpendUpdateWriter:
) -> None:
...
@overload
@staticmethod
async def _update_daily_spend(
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
daily_spend_transactions: dict[str, DailyCustomerMemberSpendTransaction], # mutable-ok: shared writer removes committed batches for safe requeue
entity_type: Literal["customer_member"],
entity_id_field: str,
) -> None:
...
@overload
@staticmethod
async def _update_daily_spend(
@ -2853,8 +3055,9 @@ class DBSpendUpdateWriter:
| dict[str, DailyTagSpendTransaction]
| dict[str, DailyOrganizationSpendTransaction]
| dict[str, DailyEndUserSpendTransaction]
| dict[str, DailyAgentSpendTransaction],
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"],
| dict[str, DailyAgentSpendTransaction]
| dict[str, DailyCustomerMemberSpendTransaction],
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent", "customer_member"],
entity_id_field: str,
) -> None:
"""
@ -3074,6 +3277,24 @@ class DBSpendUpdateWriter:
entity_id_field="agent_id",
)
@staticmethod
async def update_daily_customer_member_spend(
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
daily_spend_transactions: dict[ # mutable-ok: shared writer removes committed batches for safe requeue
str, DailyCustomerMemberSpendTransaction
],
) -> None:
await DBSpendUpdateWriter._update_daily_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_transactions,
entity_type="customer_member",
entity_id_field="customer_member_id",
)
@staticmethod
async def update_daily_tag_spend(
n_retry_times: int,
@ -3120,7 +3341,7 @@ class DBSpendUpdateWriter:
self,
payload: dict | SpendLogsPayload,
prisma_client: PrismaClient,
type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user",
type: Literal["user", "team", "org", "request_tags", "end_user", "agent", "customer_member"] = "user",
) -> BaseDailySpendTransaction | None:
entity: Final = "tag" if type == "request_tags" else type
identity_payload: Final = (
@ -3350,6 +3571,38 @@ class DBSpendUpdateWriter:
daily_transaction: Final = DailyEndUserSpendTransaction(end_user_id=end_user_id, **base_daily_transaction)
await self.daily_end_user_spend_update_queue.add_update(update={daily_transaction_key: daily_transaction})
async def add_spend_log_transaction_to_daily_customer_member_transaction(
self,
payload: SpendLogsPayload,
customer_hierarchy: CustomerHierarchyIds | None,
prisma_client: PrismaClient | None = None,
) -> None:
if prisma_client is None or customer_hierarchy is None:
return
member_id: Final = customer_hierarchy.customer_member_id
team_id: Final = customer_hierarchy.customer_team_id
org_id: Final = customer_hierarchy.customer_org_id
base: Final = await self._common_add_spend_log_transaction_to_daily_transaction( # pyright: ignore[reportUnknownMemberType] # existing helper accepts legacy untyped payloads
{**payload, "customer_member_id": member_id, "customer_team_id": team_id, "customer_org_id": org_id},
prisma_client,
"customer_member",
)
if base is None:
return
date: Final = utc_date_from_start_time(payload["startTime"])
if date is None:
return
transaction: Final[DailyCustomerMemberSpendTransaction] = {
**base,
"date": date,
"customer_member_id": member_id,
"customer_team_id": team_id,
"customer_org_id": org_id,
}
table: Final = DAILY_SPEND_TABLES["customer_member"]
key: Final = json.dumps(conflict_key(table, transaction))
await self.daily_customer_member_spend_update_queue.add_update({key: transaction})
async def add_spend_log_transaction_to_daily_agent_transaction(
self,
payload: SpendLogsPayload,

70
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -21,6 +21,7 @@ from litellm.caching import RedisCache
from litellm.constants import (
MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_CUSTOMER_MEMBER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
@ -36,6 +37,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
DailyCustomerMemberSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTagSpendTransaction,
@ -89,6 +91,9 @@ _SpendTransactionField: TypeAlias = Literal[
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
"customer_org_list_transactions",
"customer_team_list_transactions",
"customer_member_list_transactions",
]
_SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
@ -103,6 +108,9 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
"customer_org_list_transactions",
"customer_team_list_transactions",
"customer_member_list_transactions",
)
_ValueT = TypeVar("_ValueT")
@ -244,6 +252,7 @@ class RedisUpdateBuffer:
daily_org_spend_update_queue: DailySpendUpdateQueue,
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
daily_agent_spend_update_queue: DailySpendUpdateQueue,
daily_customer_member_spend_update_queue: DailySpendUpdateQueue,
window_spend_update_queue: WindowSpendUpdateQueue | None = None,
):
"""
@ -313,6 +322,9 @@ class RedisUpdateBuffer:
daily_agent_spend_update_transactions: Final = (
await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
daily_customer_member_spend_update_transactions: Final = (
await daily_customer_member_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
window_spend_update_transactions: Final = (
await window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions()
if window_spend_update_queue is not None
@ -354,6 +366,11 @@ class RedisUpdateBuffer:
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
),
(
daily_customer_member_spend_update_transactions,
REDIS_DAILY_CUSTOMER_MEMBER_SPEND_UPDATE_BUFFER_KEY,
ServiceTypes.REDIS_DAILY_CUSTOMER_MEMBER_SPEND_UPDATE_QUEUE,
),
(
tuple(map(to_wire_payload, window_spend_update_transactions)),
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
@ -399,6 +416,7 @@ class RedisUpdateBuffer:
daily_org_spend_update_transactions=daily_org_spend_update_transactions,
daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions,
daily_agent_spend_update_transactions=daily_agent_spend_update_transactions,
daily_customer_member_spend_update_transactions=daily_customer_member_spend_update_transactions,
window_spend_update_transactions=window_spend_update_transactions,
spend_update_queue=spend_update_queue,
daily_spend_update_queue=daily_spend_update_queue,
@ -406,6 +424,7 @@ class RedisUpdateBuffer:
daily_org_spend_update_queue=daily_org_spend_update_queue,
daily_end_user_spend_update_queue=daily_end_user_spend_update_queue,
daily_agent_spend_update_queue=daily_agent_spend_update_queue,
daily_customer_member_spend_update_queue=daily_customer_member_spend_update_queue,
window_spend_update_queue=window_spend_update_queue,
)
return
@ -426,6 +445,7 @@ class RedisUpdateBuffer:
daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None,
daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None,
daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None,
daily_customer_member_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None,
window_spend_update_transactions: tuple[WindowSpendTransaction, ...] | None,
spend_update_queue: SpendUpdateQueue,
daily_spend_update_queue: DailySpendUpdateQueue,
@ -433,6 +453,7 @@ class RedisUpdateBuffer:
daily_org_spend_update_queue: DailySpendUpdateQueue,
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
daily_agent_spend_update_queue: DailySpendUpdateQueue,
daily_customer_member_spend_update_queue: DailySpendUpdateQueue,
window_spend_update_queue: WindowSpendUpdateQueue | None,
) -> None:
"""
@ -443,7 +464,7 @@ class RedisUpdateBuffer:
because the source queues were already drained before the rpush.
"""
if db_spend_update_transactions is not None:
entity_entries: Final[list[tuple[Litellm_EntityType, dict[str, float] | None]]] = [
entity_entries: Final[list[tuple[Litellm_EntityType, Mapping[str, float] | None]]] = [
(
Litellm_EntityType.USER,
db_spend_update_transactions.get("user_list_transactions"),
@ -488,6 +509,18 @@ class RedisUpdateBuffer:
Litellm_EntityType.MODEL_ACCESS_GROUP,
db_spend_update_transactions.get("model_access_group_list_transactions"),
),
(
Litellm_EntityType.CUSTOMER_ORG,
db_spend_update_transactions.get("customer_org_list_transactions"),
),
(
Litellm_EntityType.CUSTOMER_TEAM,
db_spend_update_transactions.get("customer_team_list_transactions"),
),
(
Litellm_EntityType.CUSTOMER_MEMBER,
db_spend_update_transactions.get("customer_member_list_transactions"),
),
]
for entity_type, entities in entity_entries:
if not entities:
@ -501,7 +534,7 @@ class RedisUpdateBuffer:
)
)
daily_pairs: Final[list[tuple[dict[str, BaseDailySpendTransaction] | None, DailySpendUpdateQueue]]] = [
daily_pairs: Final[list[tuple[Mapping[str, BaseDailySpendTransaction] | None, DailySpendUpdateQueue]]] = [
(daily_spend_update_transactions, daily_spend_update_queue),
(daily_team_spend_update_transactions, daily_team_spend_update_queue),
(daily_org_spend_update_transactions, daily_org_spend_update_queue),
@ -510,10 +543,11 @@ class RedisUpdateBuffer:
daily_end_user_spend_update_queue,
),
(daily_agent_spend_update_transactions, daily_agent_spend_update_queue),
(daily_customer_member_spend_update_transactions, daily_customer_member_spend_update_queue),
]
for daily_txns, daily_queue in daily_pairs:
if daily_txns:
await daily_queue.update_queue.put(daily_txns)
await daily_queue.update_queue.put(dict(daily_txns))
if window_spend_update_transactions and window_spend_update_queue is not None:
await window_spend_update_queue.update_queue.put(window_spend_update_transactions)
@ -527,6 +561,7 @@ class RedisUpdateBuffer:
daily_org_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
daily_customer_member_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None,
) -> None:
@ -549,6 +584,7 @@ class RedisUpdateBuffer:
(daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY),
(daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY),
(daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY),
(daily_customer_member_spend_update_transactions, REDIS_DAILY_CUSTOMER_MEMBER_SPEND_UPDATE_BUFFER_KEY),
(daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY),
(
None
@ -771,11 +807,12 @@ class RedisUpdateBuffer:
dict[str, DailyEndUserSpendTransaction] | None,
dict[str, DailyAgentSpendTransaction] | None,
tuple[WindowSpendTransaction, ...] | None,
dict[str, DailyCustomerMemberSpendTransaction] | None,
]:
"""
Drains the main 7 Redis buffer queues in a single pipeline round-trip.
Drains the main 8 Redis buffer queues in a single pipeline round-trip.
Returns a 7-tuple of parsed results in this order:
Returns an 8-tuple of parsed results in this order:
0: DBSpendUpdateTransactions
1: daily user spend
2: daily team spend
@ -783,9 +820,10 @@ class RedisUpdateBuffer:
4: daily end-user spend
5: daily agent spend
6: budget window spend
7: daily customer member spend
"""
if self.redis_cache is None:
return None, None, None, None, None, None, None
return None, None, None, None, None, None, None, None
lpop_list: Final[list[RedisPipelineLpopOperation]] = [
RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
@ -813,12 +851,16 @@ class RedisUpdateBuffer:
key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
RedisPipelineLpopOperation(
key=REDIS_DAILY_CUSTOMER_MEMBER_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
]
raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
# Pad with None if pipeline returned fewer results than expected
while len(raw_results) < 7:
while len(raw_results) < 8:
raw_results.append(None)
# Slot 0: DBSpendUpdateTransactions
@ -830,7 +872,7 @@ class RedisUpdateBuffer:
# Slots 1-5: daily spend categories
daily_results: Final[list[dict[str, BaseDailySpendTransaction] | None]] = []
for slot in range(1, 6):
for slot in (1, 2, 3, 4, 5, 7):
slot_result = raw_results[slot]
if slot_result is None:
daily_results.append(None)
@ -855,6 +897,9 @@ class RedisUpdateBuffer:
cast(dict[str, DailyEndUserSpendTransaction] | None, daily_results[3]),
cast(dict[str, DailyAgentSpendTransaction] | None, daily_results[4]),
window_spend,
cast( # cast-ok: dedicated customer queue contains the three validated identity fields
dict[str, DailyCustomerMemberSpendTransaction] | None, daily_results[5]
),
)
async def store_in_memory_daily_tag_spend_updates_in_redis(
@ -1049,6 +1094,15 @@ class RedisUpdateBuffer:
model_access_group_list_transactions=_merged_entity_transactions(
list_of_transactions, "model_access_group_list_transactions"
),
customer_org_list_transactions=_merged_entity_transactions(
list_of_transactions, "customer_org_list_transactions"
),
customer_team_list_transactions=_merged_entity_transactions(
list_of_transactions, "customer_team_list_transactions"
),
customer_member_list_transactions=_merged_entity_transactions(
list_of_transactions, "customer_member_list_transactions"
),
)
async def _emit_new_item_added_to_redis_buffer_event(

12
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -145,6 +145,9 @@ class SpendUpdateQueue(BaseUpdateQueue):
tag_list_transactions={},
agent_list_transactions={},
model_access_group_list_transactions={},
customer_org_list_transactions={},
customer_team_list_transactions={},
customer_member_list_transactions={},
)
# Map entity types to their corresponding transaction dictionary keys
@ -160,6 +163,9 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions",
Litellm_EntityType.CUSTOMER_ORG: "customer_org_list_transactions",
Litellm_EntityType.CUSTOMER_TEAM: "customer_team_list_transactions",
Litellm_EntityType.CUSTOMER_MEMBER: "customer_member_list_transactions",
}
for update in updates:
@ -205,6 +211,12 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions["agent_list_transactions"]
elif dict_key == "model_access_group_list_transactions":
transactions_dict = db_spend_update_transactions["model_access_group_list_transactions"]
elif dict_key == "customer_org_list_transactions":
transactions_dict = db_spend_update_transactions["customer_org_list_transactions"]
elif dict_key == "customer_team_list_transactions":
transactions_dict = db_spend_update_transactions["customer_team_list_transactions"]
elif dict_key == "customer_member_list_transactions":
transactions_dict = db_spend_update_transactions["customer_member_list_transactions"]
else:
continue

13
litellm/proxy/db/model_usage_rollup.py
View file

@ -4,7 +4,7 @@ import asyncio
import random
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from datetime import datetime, timezone
from itertools import groupby
from types import TracebackType
from typing import TYPE_CHECKING, Final, Protocol
@ -97,6 +97,17 @@ def _internal_call_origin(metadata: str) -> object:
return decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)
def utc_date_from_start_time(start_time: datetime | str) -> str | None:
try:
parsed: Final = (
datetime.fromisoformat(start_time.replace("Z", "+00:00")) if isinstance(start_time, str) else start_time
)
except ValueError:
return None
aware: Final = parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed
return aware.astimezone(timezone.utc).date().isoformat()
def date_from_start_time(start_time: datetime | str) -> str | None:
if isinstance(start_time, datetime):
return start_time.date().isoformat()

9
litellm/proxy/db/prisma_client.py
View file

@ -74,6 +74,15 @@ if TYPE_CHECKING:
def delete_many(self, where: Mapping[str, object] | None = None) -> None: ...
class PrismaBatchSession(Protocol):
@property
def litellm_customermembertable(self) -> PrismaBatchActions: ...
@property
def litellm_customerteamtable(self) -> PrismaBatchActions: ...
@property
def litellm_customerorganizationtable(self) -> PrismaBatchActions: ...
@property
def litellm_agentstable(self) -> PrismaBatchActions: ...

3
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -126,6 +126,7 @@ async def alternate_grant_for(user_api_key_dict: UserAPIKeyAuth, model: str) ->
get_team_membership,
get_team_object,
)
from litellm.proxy.auth.customer_hierarchy_auth import check_customer_hierarchy_model_access
from litellm.proxy.auth.user_api_key_auth import (
_is_fallback_model_authorized, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract
_is_within_every_model_budget, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract
@ -189,6 +190,8 @@ async def alternate_grant_for(user_api_key_dict: UserAPIKeyAuth, model: str) ->
can_customer_access_model(
model=model, end_user_object=end_user_object, llm_router=llm_router, valid_token=user_api_key_dict
)
if user_api_key_dict.customer_hierarchy is not None:
check_customer_hierarchy_model_access(user_api_key_dict.customer_hierarchy, model, llm_router)
except Exception: # noqa: BLE001 # a failed team, member, project or customer lookup must not authorize an alternate
return None
if not await _is_fallback_model_authorized(

54
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -2693,6 +2693,50 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
},
)
def _customer_hierarchy_rate_limit_descriptors(
self, user_api_key_dict: UserAPIKeyAuth
) -> tuple[RateLimitDescriptor, ...]:
"""One descriptor per customer org, team and member level that carries a TPM or RPM limit."""
scopes: Final = user_api_key_dict.customer_hierarchy.scopes if user_api_key_dict.customer_hierarchy else ()
return tuple(
RateLimitDescriptor(
key=scope.key,
value=scope.entity_id,
rate_limit={
"requests_per_unit": scope.rpm_limit,
"tokens_per_unit": scope.tpm_limit,
"window_size": self.window_size,
},
)
for scope in scopes
if scope.rpm_limit is not None or scope.tpm_limit is not None
)
def _customer_hierarchy_model_rate_limit_descriptors(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None
) -> tuple[RateLimitDescriptor, ...]:
"""The ``model_per_customer_*`` descriptors for every level whose budget limits this model group."""
model: Final = self._rate_limited_model(requested_model)
if model is None or user_api_key_dict.customer_hierarchy is None:
return ()
limits: Final = tuple(
(scope, model.limit_in(scope.model_tpm_limit), model.limit_in(scope.model_rpm_limit))
for scope in user_api_key_dict.customer_hierarchy.scopes
)
return tuple(
RateLimitDescriptor(
key=f"model_per_{scope.key}",
value=f"{scope.entity_id}:{model.group}",
rate_limit={
"requests_per_unit": rpm_limit,
"tokens_per_unit": tpm_limit,
"window_size": self.window_size,
},
)
for scope, tpm_limit, rpm_limit in limits
if tpm_limit is not None or rpm_limit is not None
)
def _end_user_model_rate_limit_descriptor(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None
) -> RateLimitDescriptor | None:
@ -3201,6 +3245,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
)
descriptors.extend(self._customer_hierarchy_rate_limit_descriptors(user_api_key_dict))
# Model rate limits
requested_model: Final = data.get("model", None)
model_candidates: Final = expand_model_latest_alias(requested_model if isinstance(requested_model, str) else "")
@ -3212,6 +3258,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
if (end_user_model := self._end_user_model_rate_limit_descriptor(user_api_key_dict, candidate)) is not None:
descriptors.append(end_user_model)
descriptors.extend(self._customer_hierarchy_model_rate_limit_descriptors(user_api_key_dict, candidate))
# Per-request-tag rate limits scoped to this key
self._add_tag_per_key_rate_limit_descriptor(
@ -4792,12 +4839,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
targets.append(("end_user", user_api_key_end_user_id))
if user_api_key_organization_id:
targets.append(("organization", user_api_key_organization_id))
customer_scope_owners: Final = (
("customer_org", standard_logging_metadata.get("user_api_key_customer_org_id")),
("customer_team", standard_logging_metadata.get("user_api_key_customer_team_id")),
("customer_member", standard_logging_metadata.get("user_api_key_customer_member_id")),
)
targets.extend((scope, owner) for scope, owner in customer_scope_owners if owner)
model_scope_owners: Final = (
("model_per_key", user_api_key),
("model_per_team", user_api_key_team_id),
("model_per_end_user", f"{user_api_key_end_user_id}" if user_api_key_end_user_id else None),
("model_per_organization", user_api_key_organization_id),
("model_per_project", user_api_key_project_id),
*((f"model_per_{scope}", owner) for scope, owner in customer_scope_owners),
)
if model_group:
for group in dict.fromkeys(model_groups or (model_group,)):

4
litellm/proxy/litellm_pre_call_utils.py
View file

@ -2119,6 +2119,7 @@ class LiteLLMProxyRequestSetup:
) -> StandardLoggingUserAPIKeyMetadata:
stripped_metadata: Final = strip_callback_config(user_api_key_dict.metadata)
auth_metadata: Final = cast("dict[str, str] | None", stripped_metadata) # cast-ok: metadata is free-form JSON
customer: Final = user_api_key_dict.customer_hierarchy
user_api_key_logged_metadata: Final = StandardLoggingUserAPIKeyMetadata(
user_api_key_hash=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
user_api_key_alias=user_api_key_dict.key_alias,
@ -2136,6 +2137,9 @@ class LiteLLMProxyRequestSetup:
user_api_key_org_alias=user_api_key_dict.organization_alias,
user_api_key_team_alias=user_api_key_dict.team_alias,
user_api_key_end_user_id=user_api_key_dict.end_user_id,
user_api_key_customer_member_id=customer.customer_member_id if customer else None,
user_api_key_customer_team_id=customer.customer_team_id if customer else None,
user_api_key_customer_org_id=customer.customer_org_id if customer else None,
user_api_key_user_email=user_api_key_dict.user_email,
user_api_key_request_route=user_api_key_dict.request_route,
user_api_key_budget_reset_at=(

35
litellm/proxy/management_endpoints/management_v1/customers.py
View file

@ -1,11 +1,13 @@
from math import ceil
from typing import Annotated, Final, NoReturn, TypeVar
from typing import Annotated, Final, NoReturn, TypeAlias, TypeVar
from fastapi import APIRouter, Depends, Query, Request, Response
from typing_extensions import assert_never
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.customer_hierarchy_auth import forget_customer_hierarchies
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.list_api.common import (
PROBLEM_TYPE_BASE,
ManagementProblem,
@ -140,6 +142,15 @@ router: Final = APIRouter(
)
def get_user_api_key_cache() -> UserApiKeyCache:
from litellm.proxy.proxy_server import user_api_key_cache
return user_api_key_cache
AuthCache: TypeAlias = Annotated[UserApiKeyCache, Depends(get_user_api_key_cache)]
@router.post(
"/customer_organizations",
status_code=201,
@ -179,8 +190,11 @@ async def update_customer_organization(
data: OrganizationPatch,
actor: Annotated[str, Depends(customer_actor)],
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
cache: AuthCache,
) -> Organization:
return resolve(await hierarchy.update_organization(customer_org_id, data, actor))
organization: Final = resolve(await hierarchy.update_organization(customer_org_id, data, actor))
await forget_customer_hierarchies(resolve(await hierarchy.member_ids_in_organization(customer_org_id)), cache)
return organization
@router.delete(
@ -255,8 +269,11 @@ async def update_customer_team(
data: TeamPatch,
actor: Annotated[str, Depends(customer_actor)],
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
cache: AuthCache,
) -> Team:
return resolve(await hierarchy.update_team(customer_team_id, data, actor))
team: Final = resolve(await hierarchy.update_team(customer_team_id, data, actor))
await forget_customer_hierarchies(resolve(await hierarchy.member_ids_in_team(customer_team_id)), cache)
return team
@router.delete(
@ -331,8 +348,11 @@ async def update_customer_member(
data: MemberPatch,
actor: Annotated[str, Depends(customer_actor)],
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
cache: AuthCache,
) -> CustomerMember:
return resolve(await hierarchy.update_member(customer_member_id, data, actor))
member: Final = resolve(await hierarchy.update_member(customer_member_id, data, actor))
await forget_customer_hierarchies((customer_member_id,), cache)
return member
@router.delete(
@ -344,8 +364,10 @@ async def update_customer_member(
async def delete_customer_member(
customer_member_id: str,
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
cache: AuthCache,
) -> Response:
resolve(await hierarchy.delete_member(customer_member_id))
await forget_customer_hierarchies((customer_member_id,), cache)
return Response(status_code=204)
@ -396,8 +418,11 @@ async def update_customer_budget_template(
data: TemplatePatch,
actor: Annotated[str, Depends(customer_actor)],
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
cache: AuthCache,
) -> Organization:
return resolve(await hierarchy.update_template(customer_org_id, budget_id, data, actor))
organization: Final = resolve(await hierarchy.update_template(customer_org_id, budget_id, data, actor))
await forget_customer_hierarchies(resolve(await hierarchy.member_ids_using_template(budget_id)), cache)
return organization
@router.delete(

64
litellm/proxy/management_helpers/customer_hierarchy.py
View file

@ -102,6 +102,23 @@ def budget_values(budget: Budget) -> Mapping[str, object]:
return _VALUES.validate_python(budget.model_dump())
_JSON_BUDGET_COLUMNS: Final = frozenset({"model_tpm_limit", "model_rpm_limit"})
_MODEL_LIMIT_MAP: Final = TypeAdapter(dict[str, int])
def budget_row_values(values: Mapping[str, object]) -> Mapping[str, object]:
"""Prisma writes Json columns only through ``Json``; the per-model limit maps are the budget's Json columns."""
from prisma import Json
return {
key: Json(_MODEL_LIMIT_MAP.validate_python(value)) if key in _JSON_BUDGET_COLUMNS else value
for key, value in values.items()
if not (key in _JSON_BUDGET_COLUMNS and value is None)
}
def related_budget(value: object, now: datetime | None = None, inherited_maximum: float | None = None) -> Budget | None:
record: Final = _BUDGET.validate_python(value)
if record is None:
@ -224,6 +241,11 @@ def alias_patch(payload: OrganizationPatch | TeamPatch | MemberPatch) -> Mapping
_BUDGET_INCLUDE: Final = {"budget": True}
_MEMBER_INCLUDE: Final = {"budget": True, "budget_override": True}
_ID: Final = TypeAdapter(str)
def _member_ids(rows: Sequence[PrismaRecord]) -> tuple[str, ...]:
return tuple(_ID.validate_python(row.model_dump()["customer_member_id"]) for row in rows)
class CustomerHierarchyStore:
@ -288,11 +310,13 @@ class CustomerHierarchyStore:
budget_id: Final = str(uuid4())
await create_budget_row(
self.budgets,
{
**budget_values(budget),
"allowed_models": budget.allowed_models or (),
"budget_id": budget_id,
},
budget_row_values(
{
**budget_values(budget),
"allowed_models": budget.allowed_models or (),
"budget_id": budget_id,
}
),
actor,
)
return budget_id
@ -304,10 +328,12 @@ class CustomerHierarchyStore:
await update_budget_row(
self.budgets,
budget_id,
{
**patch,
**({"allowed_models": budget.allowed_models or ()} if "allowed_models" in patch else {}),
},
budget_row_values(
{
**patch,
**({"allowed_models": budget.allowed_models or ()} if "allowed_models" in patch else {}),
}
),
actor,
)
return budget_id
@ -604,6 +630,17 @@ class CustomerHierarchyStore:
)
return Page(tuple(member_response(row, self.now) for row in rows), count)
async def member_ids_in_team(self, team_id: str) -> tuple[str, ...]:
return _member_ids(await self.members.find_many(where={"parent_team_id": team_id}))
async def member_ids_in_organization(self, organization_id: str) -> tuple[str, ...]:
teams: Final = await self.teams.find_many(where={"parent_org_id": organization_id})
team_ids: Final = tuple(_ID.validate_python(team.model_dump()["customer_team_id"]) for team in teams)
return _member_ids(await find_many_in(self.members, "parent_team_id", team_ids))
async def member_ids_using_template(self, budget_id: str) -> tuple[str, ...]:
return _member_ids(await self.members.find_many(where={"budget_id": budget_id}))
async def delete_member(self, member_id: str) -> None | Failure:
member: Final = await self.get_member(member_id)
if isinstance(member, Failure):
@ -707,6 +744,15 @@ class CustomerHierarchy:
lambda store: store.delete_member(entity_id), operation_type=TransactionOperation.DELETE
)
async def member_ids_in_team(self, entity_id: str) -> tuple[str, ...] | Failure:
return await self.transaction(lambda store: store.member_ids_in_team(entity_id))
async def member_ids_in_organization(self, entity_id: str) -> tuple[str, ...] | Failure:
return await self.transaction(lambda store: store.member_ids_in_organization(entity_id))
async def member_ids_using_template(self, budget_id: str) -> tuple[str, ...] | Failure:
return await self.transaction(lambda store: store.member_ids_using_template(budget_id))
async def list_organizations(
self, page: int = 1, page_size: int = 50, search: str = ""
) -> Page[Organization] | Failure:

9
litellm/repositories/prisma_protocols.py
View file

@ -141,6 +141,15 @@ class BatchTable(Protocol):
class PrismaBatch(Protocol):
@property
def litellm_customermembertable(self) -> BatchTable: ...
@property
def litellm_customerteamtable(self) -> BatchTable: ...
@property
def litellm_customerorganizationtable(self) -> BatchTable: ...
@property
def litellm_verificationtoken(self) -> BatchTable: ...

8
litellm/repositories/table_repositories.py
View file

@ -81,6 +81,14 @@ class CustomerOrganizationRepository(PrismaTableRepository["prisma_models.LiteLL
table_name = "litellm_customerorganizationtable"
class CustomerMemberRepository(PrismaTableRepository["prisma_models.LiteLLM_CustomerMemberTable"]):
table_name = "litellm_customermembertable"
class CustomerTeamRepository(PrismaTableRepository["prisma_models.LiteLLM_CustomerTeamTable"]):
table_name = "litellm_customerteamtable"
class PolicyRepository(PrismaTableRepository["prisma_models.LiteLLM_PolicyTable"]):
table_name = "litellm_policytable"

66
litellm/repositories/unit_of_work.py
View file

@ -20,7 +20,7 @@ from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from math import isfinite, ulp
from typing import Final, Protocol, runtime_checkable
from typing import Final, Literal, Protocol, TypeAlias, runtime_checkable
from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch
@ -165,6 +165,70 @@ class TeamSpendResetWrites:
)
CustomerSpendEntity: TypeAlias = Literal["member", "team", "org"]
_CUSTOMER_RESET_QUERIES: Final = {
"member": 'UPDATE "LiteLLM_CustomerMemberTable"'
+ _SPEND_RESET_SET
+ " WHERE customer_member_id = $1"
+ _SPEND_RESET_WHERE,
"team": 'UPDATE "LiteLLM_CustomerTeamTable"'
+ _SPEND_RESET_SET
+ " WHERE customer_team_id = $1"
+ _SPEND_RESET_WHERE,
"org": 'UPDATE "LiteLLM_CustomerOrganizationTable"'
+ _SPEND_RESET_SET
+ " WHERE customer_org_id = $1"
+ _SPEND_RESET_WHERE,
}
@dataclass(frozen=True, slots=True)
class CustomerSpendResetWrites:
entity: CustomerSpendEntity
table: BatchTable
batch: PrismaBatch
def queue_spend_reset(
self,
identity: str,
budget_reset_at: datetime | None,
spend_decrement: float,
expected_reset_at: datetime | None,
spend_after_reset: float = 0.0,
) -> None:
if _queue_guarded_spend_reset(
self.batch,
_CUSTOMER_RESET_QUERIES[self.entity],
identity,
budget_reset_at,
spend_decrement,
expected_reset_at,
spend_after_reset,
):
return
self.table.update_many(
where={f"customer_{self.entity}_id": identity, "budget_reset_at": expected_reset_at},
data=_spend_reset_data(budget_reset_at, spend_decrement),
)
@asynccontextmanager
async def customer_spend_reset_unit_of_work(
new_batch: Callable[[], PrismaBatch], entity: CustomerSpendEntity
) -> AsyncGenerator[CustomerSpendResetWrites, None]:
batch: Final = new_batch()
table: Final = (
batch.litellm_customermembertable
if entity == "member"
else batch.litellm_customerteamtable
if entity == "team"
else batch.litellm_customerorganizationtable
)
yield CustomerSpendResetWrites(entity=entity, table=table, batch=batch)
await batch.commit()
@runtime_checkable
class SpendResetWrites(Protocol):
def queue_spend_zero(self, where: Mapping[str, object]) -> None: ...

3
litellm/types/integrations/slack_alerting.py
View file

@ -314,6 +314,9 @@ class Litellm_EntityType(enum.Enum):
TAG = "tag"
AGENT = "agent"
MODEL_ACCESS_GROUP = "model_access_group"
CUSTOMER_ORG = "customer_org"
CUSTOMER_TEAM = "customer_team"
CUSTOMER_MEMBER = "customer_member"
# global proxy level entity
PROXY = "proxy"

83
litellm/types/proxy/auth/customer_hierarchy.py Normal file
View file

@ -0,0 +1,83 @@
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.types.integrations.slack_alerting import Litellm_EntityType
CustomerScopeKey: TypeAlias = Literal["customer_org", "customer_team", "customer_member"]
class CustomerScopeAuth(BaseModel):
"""One level of the customer hierarchy a request runs under, with the limits its budget row grants."""
model_config = ConfigDict(frozen=True)
key: CustomerScopeKey
entity_id: str
spend: float = 0.0
max_budget: float | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
allowed_models: tuple[str, ...] | None = None
model_tpm_limit: Mapping[str, int] | None = None
model_rpm_limit: Mapping[str, int] | None = None
class CustomerHierarchyAuth(BaseModel):
model_config = ConfigDict(frozen=True)
customer_member_id: str
customer_team_id: str
customer_org_id: str
scopes: tuple[CustomerScopeAuth, ...]
@dataclass(frozen=True, slots=True)
class CustomerHierarchyIds:
"""The three rows a request's spend is charged to, as the request metadata recorded them."""
customer_member_id: str
customer_team_id: str
customer_org_id: str
def entities(self) -> Mapping[str, str]:
return {
"customer_member_id": self.customer_member_id,
"customer_team_id": self.customer_team_id,
"customer_org_id": self.customer_org_id,
}
def spend_entities(self) -> tuple[tuple[Litellm_EntityType, str], ...]:
return (
(Litellm_EntityType.CUSTOMER_ORG, self.customer_org_id),
(Litellm_EntityType.CUSTOMER_TEAM, self.customer_team_id),
(Litellm_EntityType.CUSTOMER_MEMBER, self.customer_member_id),
)
_FIELDS: Final = TypeAdapter(Mapping[str, object])
def _mapping(value: object) -> Mapping[str, object]:
try:
return _FIELDS.validate_python(value)
except ValidationError:
return {}
def _stamped_metadata(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
"""The proxy stamps auth ids into ``litellm_metadata`` on routes whose body owns ``metadata``, else into ``metadata``."""
buckets: Final = tuple(_mapping(litellm_params.get(name)) for name in ("litellm_metadata", "metadata"))
return next((bucket for bucket in buckets if "user_api_key_customer_member_id" in bucket), {})
def customer_hierarchy_ids_from_request(kwargs: Mapping[str, object]) -> CustomerHierarchyIds | None:
"""The ids auth stamped into the request's proxy metadata, or None when the request named no customer member."""
metadata: Final = _stamped_metadata(_mapping(kwargs.get("litellm_params")))
ids: Final = tuple(metadata.get(f"user_api_key_customer_{level}_id") for level in ("member", "team", "org"))
if not all(isinstance(value, str) and value for value in ids):
return None
member_id, team_id, org_id = ids
return CustomerHierarchyIds(str(member_id), str(team_id), str(org_id))

2
litellm/types/proxy/auth/user_api_key_auth.py
View file

@ -14,6 +14,7 @@ from litellm.models.team import Member
from litellm.models.verification_token import LiteLLM_VerificationToken
from litellm.types.agents import AgentCaller, AgentResponse
from litellm.types.proxy.agent_identity import ManagedAgentContext
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
from litellm.types.proxy.carried_budget_state import (
CommittedBudgetSpend,
@ -89,6 +90,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
end_user_model_max_budget: _UntypedDict | None = None
end_user_model_tpm_limit: Mapping[str, int] | None = None
end_user_model_rpm_limit: Mapping[str, int] | None = None
customer_hierarchy: CustomerHierarchyAuth | None = None
# Organization Params
organization_alias: str | None = None

5
litellm/types/proxy/management_endpoints/customer_hierarchy.py
View file

@ -1,7 +1,8 @@
from collections.abc import Mapping
from datetime import datetime
from typing import Final
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, NonNegativeInt, field_validator
from litellm.proxy.common_utils.timezone_utils import budget_duration_error
@ -15,6 +16,8 @@ class Budget(BaseModel):
rpm_limit: int | None = Field(default=None, ge=0)
allowed_models: tuple[str, ...] | None = None
budget_duration: str | None = None
model_tpm_limit: Mapping[str, NonNegativeInt] | None = None
model_rpm_limit: Mapping[str, NonNegativeInt] | None = None
@field_validator("budget_duration")
@classmethod

1
litellm/types/services.py
View file

@ -40,6 +40,7 @@ class ServiceTypes(str, enum.Enum):
REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE = "redis_daily_org_spend_update_queue"
REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE = "redis_daily_team_spend_update_queue"
REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE = "redis_daily_agent_spend_update_queue"
REDIS_DAILY_CUSTOMER_MEMBER_SPEND_UPDATE_QUEUE = "redis_daily_customer_member_spend_update_queue"
REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE = "redis_daily_tag_spend_update_queue"
# spend update queue - current spend of key, user, team
IN_MEMORY_SPEND_UPDATE_QUEUE = "in_memory_spend_update_queue"

3
litellm/types/utils.py
View file

@ -3145,6 +3145,9 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict):
user_api_key_user_email: str | None
user_api_key_team_alias: str | None
user_api_key_end_user_id: str | None
user_api_key_customer_member_id: ReadOnly[str | None]
user_api_key_customer_team_id: ReadOnly[str | None]
user_api_key_customer_org_id: ReadOnly[str | None]
user_api_key_request_route: str | None
user_api_key_auth_metadata: dict[str, str] | None

78
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -29446,6 +29446,14 @@ export interface components {
budget_duration?: string | null;
/** Max Budget */
max_budget?: number | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: number;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
@ -29604,6 +29612,14 @@ export interface components {
is_default: boolean;
/** Max Budget */
max_budget?: number | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: number;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Name */
name: string;
/** Rpm Limit */
@ -33128,6 +33144,17 @@ export interface components {
*/
type: "custom";
};
/** CustomerHierarchyAuth */
CustomerHierarchyAuth: {
/** Customer Member Id */
customer_member_id: string;
/** Customer Org Id */
customer_org_id: string;
/** Customer Team Id */
customer_team_id: string;
/** Scopes */
scopes: components["schemas"]["CustomerScopeAuth"][];
};
/** CustomerMember */
CustomerMember: {
/** Alias */
@ -33199,6 +33226,40 @@ export interface components {
/** User Id */
user_id: string;
};
/**
* CustomerScopeAuth
* @description One level of the customer hierarchy a request runs under, with the limits its budget row grants.
*/
CustomerScopeAuth: {
/** Allowed Models */
allowed_models?: string[] | null;
/** Entity Id */
entity_id: string;
/**
* Key
* @enum {string}
*/
key: "customer_org" | "customer_team" | "customer_member";
/** Max Budget */
max_budget?: number | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: number;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Rpm Limit */
rpm_limit?: number | null;
/**
* Spend
* @default 0
*/
spend: number;
/** Tpm Limit */
tpm_limit?: number | null;
};
/**
* CyberArkConfig
* @description Configuration for CyberArk Conjur secret manager integration.
@ -42355,6 +42416,14 @@ export interface components {
is_default: boolean;
/** Max Budget */
max_budget?: number | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: number;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Name */
name: string;
/** Rpm Limit */
@ -50003,6 +50072,14 @@ export interface components {
is_default?: boolean | null;
/** Max Budget */
max_budget?: number | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: number;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Name */
name?: string | null;
/** Rpm Limit */
@ -52526,6 +52603,7 @@ export interface components {
created_by?: string | null;
/** Created By User */
created_by_user?: unknown | null;
customer_hierarchy?: components["schemas"]["CustomerHierarchyAuth"] | null;
/** End User Id */
end_user_id?: string | null;
/** End User Max Budget */

250
tests/integration/database/test_spendlogs_ledger_parity.py
View file

@ -1547,3 +1547,253 @@ async def test_analytics_overflow_across_requests_still_applies_every_charge(
).fetchone() == (4, 3 * 2**40)
finally:
await client.db.disconnect()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"ledger,keep_logs,failing_table",
[
(True, True, None),
(True, False, None),
(False, True, None),
(False, False, None),
(True, True, "LiteLLM_DailyCustomerMemberSpend"),
],
)
async def test_customer_daily_usage_is_atomic_with_accounting_and_survives_budget_reset(
parity_database: str, monkeypatch: pytest.MonkeyPatch, ledger: bool, keep_logs: bool, failing_table: str | None
) -> None:
from litellm.accounting.worker import flush_accounting_deliveries
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
monkeypatch.setenv("DATABASE_URL", parity_database)
monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger})
monkeypatch.setattr(proxy_server, "disable_spend_logs", not keep_logs)
client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
client.spend_log_transactions = []
monkeypatch.setattr(proxy_server, "prisma_client", client)
writer = DBSpendUpdateWriter()
_seed_customer_hierarchy(parity_database)
await client.db.connect()
try:
ids = {"customer_member_id": "member", "customer_team_id": "customer-team", "customer_org_id": "customer-org"}
request_ids = {f"user_api_key_{key}": value for key, value in ids.items()}
assert await writer.update_database(
token="key",
user_id="user",
team_id="team",
org_id="org",
end_user_id="end-user",
response_cost=0.25,
start_time=_NOW,
end_time=_NOW + timedelta(seconds=1),
kwargs={
"model": "ledger-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "customer-request",
"litellm_params": {"metadata": {"status": "success", "user_api_key": "key", **request_ids}},
},
completion_response=ModelResponse(
id="customer-request", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
)
await drain_logging_tasks()
await SpendLogsRepository(client).insert_with_deliveries(tuple(client.spend_log_transactions))
if failing_table is not None:
before_failure = _snapshot(parity_database)
with psycopg.connect(parity_database) as connection:
connection.execute(
sql.SQL("ALTER TABLE {} ADD CONSTRAINT reject_customer_spend CHECK (spend = 0)").format(
sql.Identifier(failing_table)
)
)
await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj)
assert _snapshot(parity_database) == before_failure
with psycopg.connect(parity_database) as connection:
assert connection.execute(
'SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"'
).fetchall() == [("pending", 1)]
for table in (
"LiteLLM_CustomerMemberTable",
"LiteLLM_CustomerTeamTable",
"LiteLLM_CustomerOrganizationTable",
):
assert connection.execute(
sql.SQL("SELECT spend FROM {}").format(sql.Identifier(table))
).fetchall() == [(0.0,)]
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyCustomerMemberSpend"').fetchone() == (0,)
connection.execute(
sql.SQL("ALTER TABLE {} DROP CONSTRAINT reject_customer_spend").format(
sql.Identifier(failing_table)
)
)
connection.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP')
await writer.db_update_spend_transaction_handler(client, 0, proxy_server.proxy_logging_obj)
await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj)
await writer.db_update_spend_transaction_handler(client, 0, proxy_server.proxy_logging_obj)
tables = ("LiteLLM_CustomerMemberTable", "LiteLLM_CustomerTeamTable", "LiteLLM_CustomerOrganizationTable")
with psycopg.connect(parity_database) as connection:
for table in tables:
assert connection.execute(sql.SQL("SELECT spend FROM {}").format(sql.Identifier(table))).fetchall() == [
(0.25,)
]
daily_before = connection.execute(
'SELECT customer_member_id, customer_team_id, customer_org_id, date, spend, prompt_tokens, completion_tokens, api_requests FROM "LiteLLM_DailyCustomerMemberSpend"'
).fetchall()
assert daily_before == [("member", "customer-team", "customer-org", "2026-11-02", 0.25, 10, 5, 1)]
logs = connection.execute('SELECT metadata FROM "LiteLLM_SpendLogs"').fetchall()
assert len(logs) == int(keep_logs)
job = ResetBudgetJob(proxy_logging_obj=proxy_server.proxy_logging_obj, prisma_client=client)
await job.reset_budget_for_customer_hierarchy()
with psycopg.connect(parity_database) as connection:
for table in tables:
assert connection.execute(
sql.SQL("SELECT spend, budget_reset_at IS NOT NULL FROM {}").format(sql.Identifier(table))
).fetchall() == [(0.25, True)]
connection.execute(sql.SQL("UPDATE {} SET budget_reset_at=%s").format(sql.Identifier(table)), (_NOW,))
await job.reset_budget_for_customer_hierarchy()
await job.reset_budget_for_customer_hierarchy()
with psycopg.connect(parity_database) as connection:
for table in tables:
assert connection.execute(
sql.SQL("SELECT spend, budget_reset_at > %s FROM {}").format(sql.Identifier(table)), (_NOW,)
).fetchall() == [(0.0, True)]
assert (
connection.execute(
'SELECT customer_member_id, customer_team_id, customer_org_id, date, spend, prompt_tokens, completion_tokens, api_requests FROM "LiteLLM_DailyCustomerMemberSpend"'
).fetchall()
== daily_before
)
finally:
await client.db.disconnect()
def _seed_customer_hierarchy(parity_database: str) -> None:
with psycopg.connect(parity_database) as connection:
connection.execute("UPDATE \"LiteLLM_BudgetTable\" SET budget_duration='1d' WHERE budget_id='budget'")
connection.execute(
"INSERT INTO \"LiteLLM_CustomerOrganizationTable\" (customer_org_id, budget_id, created_by, updated_by) VALUES ('customer-org', 'budget', 'test', 'test')"
)
connection.execute(
"INSERT INTO \"LiteLLM_CustomerTeamTable\" (customer_team_id, parent_org_id, budget_id, created_by, updated_by) VALUES ('customer-team', 'customer-org', 'budget', 'test', 'test')"
)
connection.execute(
"INSERT INTO \"LiteLLM_CustomerMemberTable\" (customer_member_id, parent_team_id, budget_id, created_by, updated_by) VALUES ('member', 'customer-team', 'budget', 'test', 'test')"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"entity,table,id_column,identity",
[
("member", "LiteLLM_CustomerMemberTable", "customer_member_id", "member"),
("team", "LiteLLM_CustomerTeamTable", "customer_team_id", "customer-team"),
("org", "LiteLLM_CustomerOrganizationTable", "customer_org_id", "customer-org"),
],
)
async def test_customer_budget_reset_preserves_concurrent_charges_and_is_not_reapplied(
parity_database: str, monkeypatch: pytest.MonkeyPatch, entity: str, table: str, id_column: str, identity: str
) -> None:
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob, _CustomerBudgetRow
monkeypatch.setenv("DATABASE_URL", parity_database)
_seed_customer_hierarchy(parity_database)
client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
reset_at = _NOW.replace(tzinfo=None)
row = _CustomerBudgetRow.model_validate(
{
id_column: identity,
"spend": 5.0,
"budget_reset_at": reset_at,
"budget": {"max_budget": 100, "budget_duration": "1d"},
}
)
with psycopg.connect(parity_database) as connection:
connection.execute(sql.SQL("UPDATE {} SET spend=7, budget_reset_at=%s").format(sql.Identifier(table)), (reset_at,))
await client.db.connect()
try:
job = ResetBudgetJob(proxy_logging_obj=proxy_server.proxy_logging_obj, prisma_client=client)
await job._write_customer_budget_resets(entity, (row,), _NOW)
await job._write_customer_budget_resets(entity, (row,), _NOW)
with psycopg.connect(parity_database) as connection:
assert connection.execute(
sql.SQL("SELECT spend, budget_reset_at > %s FROM {}").format(sql.Identifier(table)), (reset_at,)
).fetchall() == [(2.0, True)]
finally:
await client.db.disconnect()
@pytest.mark.asyncio
@pytest.mark.parametrize("expired", [True, False])
async def test_customer_member_reset_uses_active_override_or_inherited_template(
parity_database: str, monkeypatch: pytest.MonkeyPatch, expired: bool
) -> None:
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
monkeypatch.setenv("DATABASE_URL", parity_database)
_seed_customer_hierarchy(parity_database)
now = datetime.now(timezone.utc)
with psycopg.connect(parity_database) as connection:
connection.execute("UPDATE \"LiteLLM_BudgetTable\" SET budget_duration='3d' WHERE budget_id='budget'")
with psycopg.connect(parity_database) as connection:
connection.execute(
"INSERT INTO \"LiteLLM_BudgetTable\" (budget_id, budget_duration, created_by, updated_by) VALUES ('override', '1d', 'test', 'test')"
)
connection.execute(
'UPDATE "LiteLLM_CustomerMemberTable" SET spend=5, budget_override_id=%s, budget_reset_at=%s, budget_override_expires_at=%s',
("override", _NOW, now + timedelta(days=-1 if expired else 1)),
)
client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
await client.db.connect()
try:
await ResetBudgetJob(
proxy_logging_obj=proxy_server.proxy_logging_obj, prisma_client=client
).reset_budget_for_customer_hierarchy()
with psycopg.connect(parity_database) as connection:
spend, reset_at = connection.execute(
'SELECT spend, budget_reset_at FROM "LiteLLM_CustomerMemberTable"'
).fetchone()
assert spend == 0.0
remaining = reset_at.replace(tzinfo=timezone.utc) - now
assert timedelta(days=2 if expired else 0) < remaining <= timedelta(days=3 if expired else 1)
finally:
await client.db.disconnect()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"entity,table",
[
("org", "LiteLLM_CustomerOrganizationTable"),
("team", "LiteLLM_CustomerTeamTable"),
("member", "LiteLLM_CustomerMemberTable"),
],
)
async def test_customer_budget_reset_reloads_cached_hierarchy_spend(
parity_database: str, monkeypatch: pytest.MonkeyPatch, entity: str, table: str
) -> None:
from litellm.proxy.auth.customer_hierarchy_auth import get_customer_hierarchy
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
monkeypatch.setenv("DATABASE_URL", parity_database)
monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache())
_seed_customer_hierarchy(parity_database)
with psycopg.connect(parity_database) as connection:
connection.execute(
sql.SQL("UPDATE {} SET spend=5, budget_reset_at=%s").format(sql.Identifier(table)), (_NOW,)
)
client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
await client.db.connect()
try:
before = await get_customer_hierarchy("member", client, proxy_server.user_api_key_cache)
assert before is not None
assert next(scope.spend for scope in before.scopes if scope.key == f"customer_{entity}") == 5.0
await ResetBudgetJob(
proxy_logging_obj=proxy_server.proxy_logging_obj, prisma_client=client
).reset_budget_for_customer_hierarchy()
after = await get_customer_hierarchy("member", client, proxy_server.user_api_key_cache)
assert after is not None
assert tuple(scope.spend for scope in after.scopes) == (0.0, 0.0, 0.0)
finally:
await client.db.disconnect()

105
tests/unit/accounting/test_application.py
View file

@ -1,62 +1,20 @@
from collections.abc import Mapping
from dataclasses import dataclass, replace
from dataclasses import replace
from datetime import datetime, timezone
from typing import Final, cast
from typing import Final
import pytest
from pydantic import TypeAdapter
from litellm.accounting.application import _apply_customer_usage
from litellm.accounting.delivery_projection import project_deliveries
from litellm.accounting.projection import Contribution
from litellm.proxy.common_utils.customer_usage import CUSTOMER_USAGE_METADATA_KEY, CustomerUsageAttribution
from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, merge_by_conflict_key
from litellm.repositories.accounting_delivery_repository import DeliveryProjection
_FIELDS: Final = TypeAdapter(Mapping[str, object])
_TEXT: Final = TypeAdapter(str)
_UPDATES: Final = TypeAdapter(Mapping[str, Mapping[str, float]])
_IDENTIFIER: Final = "\U0001f680" * 256
_ATTRIBUTION: Final = CustomerUsageAttribution(
customer_org_id=_IDENTIFIER, customer_team_id=_IDENTIFIER, customer_member_id=_IDENTIFIER
)
@dataclass(frozen=True, slots=True)
class Bucket:
values: Mapping[str, object]
counters: Mapping[str, float]
class Buckets:
def __init__(self) -> None:
self.rows: Mapping[str, Bucket] = {}
async def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None:
row_id: Final = _TEXT.validate_python(where["id"])
create: Final = _FIELDS.validate_python(data["create"])
previous: Final = self.rows.get(row_id)
increments: Final = (
{name: fields["increment"] for name, fields in _UPDATES.validate_python(data["update"]).items()}
if previous is not None
else {name: float(value) for name, value in create.items() if isinstance(value, (int, float))}
)
self.rows = {
**self.rows,
row_id: Bucket(
values=previous.values if previous is not None else create,
counters={
name: (previous.counters.get(name, 0) if previous is not None else 0) + value
for name, value in increments.items()
},
),
}
@dataclass(frozen=True, slots=True)
class Transaction:
litellm_dailycustomermemberspend: Buckets
def contribution(alias: str, provider: str, spend: float) -> Contribution:
delivery: Final = DeliveryProjection(
{
@ -75,29 +33,38 @@ def contribution(alias: str, provider: str, spend: float) -> Contribution:
return next(item for item in project_deliveries((delivery,)) if item.table == "LiteLLM_DailyCustomerMemberSpend")
@pytest.mark.asyncio
async def test_long_customer_ids_roll_up_without_truncation_and_keep_alias_provider_buckets_distinct() -> None:
buckets: Final = Buckets()
transaction: Final = cast(LedgerTransaction, Transaction(buckets)) # cast-ok: in-memory Prisma upsert boundary
def _number(value: object) -> float:
return float(value) if isinstance(value, (int, float)) else 0.0
def _row(item: Contribution) -> Mapping[str, object]:
return {**dict(item.key), **dict(item.values), **dict(item.increments)}
def test_long_customer_ids_roll_up_without_truncation_and_aliases_and_providers_keep_distinct_rows() -> None:
first: Final = contribution("alias-a", "provider-a", 1)
second: Final = replace(contribution("alias-a", "provider-a", 2), key=tuple(reversed(first.key)))
await _apply_customer_usage(transaction, (first,))
await _apply_customer_usage(transaction, (second, contribution("alias-b", "provider-a", 4)))
await _apply_customer_usage(transaction, (contribution("alias-a", "provider-b", 8),))
assert len(buckets.rows) == 3
assert all(len(row_id.encode("utf-8")) == 64 for row_id in buckets.rows)
assert all(
row.values["customer_org_id"]
== row.values["customer_team_id"]
== row.values["customer_member_id"]
== _IDENTIFIER
for row in buckets.rows.values()
merged: Final = merge_by_conflict_key(
DAILY_SPEND_TABLES["customer_member"],
tuple(
_row(item)
for item in (
first,
second,
contribution("alias-b", "provider-a", 4),
contribution("alias-a", "provider-b", 8),
)
),
)
assert {
(row.values["model_group"], row.values["custom_llm_provider"]): row.counters["spend"]
for row in buckets.rows.values()
} == {("alias-a", "provider-a"): 3, ("alias-b", "provider-a"): 4, ("alias-a", "provider-b"): 8}
assert sum(row.counters["api_requests"] for row in buckets.rows.values()) == 4
assert sum(row.counters["prompt_tokens"] + row.counters["completion_tokens"] for row in buckets.rows.values()) == 60
rows: Final = tuple(row for _, row in merged)
assert len(rows) == 3
assert all(
row["customer_org_id"] == row["customer_team_id"] == row["customer_member_id"] == _IDENTIFIER for row in rows
)
assert {(row["model_group"], row["custom_llm_provider"]): _number(row["spend"]) for row in rows} == {
("alias-a", "provider-a"): 3,
("alias-b", "provider-a"): 4,
("alias-a", "provider-b"): 8,
}
assert sum(_number(row["api_requests"]) for row in rows) == 4
assert sum(_number(row["prompt_tokens"]) + _number(row["completion_tokens"]) for row in rows) == 60

83
tests/unit/accounting/test_projection.py
View file

@ -23,7 +23,7 @@ _NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc)
_FIELDS: Final = TypeAdapter(dict[str, object])
def test_customer_model_aliases_and_providers_keep_separate_daily_totals() -> None:
def test_customer_daily_totals_stay_separate_per_model_alias_and_provider() -> None:
stamp: Final = CustomerUsageAttribution(customer_org_id="org-1", customer_member_id="member-1").model_dump_json()
base: Final = {
"startTime": _NOW,
@ -84,11 +84,11 @@ def test_customer_usage_survives_logging_and_accounting_with_correct_request_cou
"customer_team_id": attribution.customer_team_id,
"customer_member_id": attribution.customer_member_id,
"model": payload["model"],
"model_group": payload["model_group"],
"custom_llm_provider": payload["custom_llm_provider"],
"api_key": payload["api_key"],
"mcp_namespaced_tool_name": "",
"endpoint": "/chat/completions",
"model_group": payload["model_group"],
}
assert dict(usage.increments) == {
"spend": 0.75,
@ -293,6 +293,9 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest
"user_api_key_org_id": "org",
"user_api_key_project_id": "project",
"user_api_key_end_user_id": "end-user",
"user_api_key_customer_member_id": "member",
"user_api_key_customer_team_id": "customer-team",
"user_api_key_customer_org_id": "customer-org",
"tags": ["tag", "tag"],
},
"user": "end-user",
@ -322,6 +325,9 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest
"team_id": "team",
"organization_id": "org",
"project_id": "project",
"customer_member_id": "member",
"customer_team_id": "customer-team",
"customer_org_id": "customer-org",
"agent_id": "agent",
},
model_access_groups=(),
@ -343,6 +349,9 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest
("LiteLLM_EndUserTable", "end_user", "end-user"),
("LiteLLM_AgentsTable", "agent", "agent"),
("LiteLLM_ProjectTable", "project", "project"),
("LiteLLM_CustomerMemberTable", "customer_member", "member"),
("LiteLLM_CustomerTeamTable", "customer_team", "customer-team"),
("LiteLLM_CustomerOrganizationTable", "customer_org", "customer-org"),
):
assert legacy[(entity, identity)] == next(
dict(item.increments)["spend"] for item in projected if item.table == table
@ -679,3 +688,73 @@ def test_caller_supplied_savings_beyond_json_range_are_clamped_and_never_drop_th
assert [
dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"
] == [0.5]
@pytest.mark.asyncio
@pytest.mark.parametrize("status,cost", [("success", 0.25), ("failure", 0.0), ("failure", 0.25)])
async def test_customer_daily_usage_matches_ledger_and_legacy_with_utc_date(status: str, cost: float) -> None:
identities: Final = {
"customer_member_id": "member",
"customer_team_id": "customer-team",
"customer_org_id": "customer-org",
}
start: Final = datetime.fromisoformat("2026-11-02T23:30:00-05:00")
payload: Final[SpendLogsPayload] = {
**get_logging_payload(
kwargs={
"model": "projection-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "customer-request",
"litellm_params": {"metadata": {"user_api_key": "key", "status": status, **identities}},
},
response_obj=ModelResponse(
id="customer-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=start,
end_time=start,
),
"spend": cost,
"startTime": start,
}
accounting: Final = build_accounting_payload(
payload,
entities={"api_key": "key", "team_id": "internal-team", **identities},
model_access_groups=(),
windows=(),
router=lambda: None,
)
contributions: Final = project(_FIELDS.validate_json(json.dumps(accounting, default=str)))
daily: Final = next(item for item in contributions if item.table == "LiteLLM_DailyCustomerMemberSpend")
assert {key: dict(daily.key)[key] for key in identities} == identities
assert dict(daily.key)["date"] == "2026-11-03"
writer: Final = DBSpendUpdateWriter()
client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj)
from litellm.types.proxy.auth.customer_hierarchy import customer_hierarchy_ids_from_request
hierarchy: Final = customer_hierarchy_ids_from_request(
{"litellm_params": {"metadata": {f"user_api_key_{key}": value for key, value in identities.items()}}}
)
await writer.add_spend_log_transaction_to_daily_customer_member_transaction(payload, hierarchy, client)
rows: Final = (
await writer.daily_customer_member_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
(legacy,) = rows.values()
assert {key: legacy[key] or "" for key, _ in daily.key} == dict(daily.key)
assert {key: legacy[key] for key, _ in daily.increments} == dict(daily.increments)
assert legacy["api_requests"] == 1
assert legacy["successful_requests"] == int(status == "success")
assert legacy["failed_requests"] == int(status == "failure")
def test_receipts_without_customer_identity_do_not_create_customer_daily_rows() -> None:
row: Final = {
"startTime": _NOW,
"api_key": "key",
"model": "model",
"custom_llm_provider": "openai",
"model_group": "model",
"spend": 0.25,
"metadata": {"accounting_entities": {"api_key": "key"}},
}
assert not tuple(item for item in project(row) if item.table == "LiteLLM_DailyCustomerMemberSpend")

30
tests/unit/accounting/test_worker.py
View file

@ -1,4 +1,5 @@
import json
import re
from collections.abc import AsyncIterator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass
@ -134,6 +135,15 @@ class _CustomerLedger:
return ()
async def execute(self, query: str, *args: object) -> int:
if query.startswith('INSERT INTO "LiteLLM_DailyCustomerMemberSpend"'):
columns: Final = re.findall(r'"([a-z_]+)"', query.split("VALUES", 1)[0])[:-1]
values: Final = dict(zip(columns, args, strict=True))
self.daily = {
name: self.daily.get(name, 0) + value
for name, value in values.items()
if isinstance(value, (int, float))
}
return 1
if "SET state = 'applied'" not in query:
return 0
if self.fail_acknowledgement:
@ -142,11 +152,6 @@ class _CustomerLedger:
self.pending = False
return 1
def upsert_daily(self, *, where: object, data: object) -> None:
values: Final = _FIELDS.validate_python(_FIELDS.validate_python(data)["create"])
counters: Final = {name: value for name, value in values.items() if isinstance(value, (int, float))}
self.daily = {name: self.daily.get(name, 0) + value for name, value in counters.items()}
def update_entity(self, entity: str, *, where: object, data: object) -> int:
spend: Final = _INCREMENTS.validate_python(data)["spend"]["increment"]
self.entities = {**self.entities, entity: self.entities.get(entity, 0.0) + spend}
@ -158,6 +163,9 @@ _CUSTOMER_PAYLOAD: Final = json.dumps(
"request_id": "customer-completed-request",
"startTime": datetime(2026, 1, 1, tzinfo=timezone.utc).isoformat(),
"api_key": "hashed-key",
"model": "usage-model",
"model_group": "usage-model",
"custom_llm_provider": "openai",
"spend": 0.75,
"prompt_tokens": 10,
"completion_tokens": 5,
@ -179,13 +187,11 @@ async def test_customer_daily_and_lifetime_totals_commit_once_after_delivery_rep
transaction: Final = MagicMock()
transaction.query_raw = AsyncMock(side_effect=ledger.query)
transaction.execute_raw = AsyncMock(side_effect=ledger.execute)
transaction.litellm_dailycustomermemberspend.upsert = AsyncMock(side_effect=ledger.upsert_daily)
transaction.litellm_customerorganizationtable.update_many = AsyncMock(
side_effect=partial(ledger.update_entity, "organization")
)
transaction.litellm_customerteamtable.update_many = AsyncMock(side_effect=partial(ledger.update_entity, "team"))
transaction.litellm_customermembertable.update_many = AsyncMock(side_effect=partial(ledger.update_entity, "member"))
transaction.batch_.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
batcher: Final = MagicMock()
batcher.litellm_customerorganizationtable.update_many.side_effect = partial(ledger.update_entity, "organization")
batcher.litellm_customerteamtable.update_many.side_effect = partial(ledger.update_entity, "team")
batcher.litellm_customermembertable.update_many.side_effect = partial(ledger.update_entity, "member")
transaction.batch_.return_value.__aenter__ = AsyncMock(return_value=batcher)
transaction.batch_.return_value.__aexit__ = AsyncMock(return_value=False)
@asynccontextmanager

259
tests/unit/proxy/auth/test_customer_hierarchy_auth.py Normal file
View file

@ -0,0 +1,259 @@
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
import litellm
from litellm.caching.caching import DualCache
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.customer_hierarchy_auth import (
CUSTOMER_MEMBER_ID_HEADER,
check_customer_hierarchy_budget,
check_customer_hierarchy_key_scope,
check_customer_hierarchy_model_access,
customer_hierarchy_from_member_row,
customer_member_id_from_headers,
get_customer_hierarchy,
resolve_customer_hierarchy,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
def _row(
*,
template: Mapping[str, object] | None = None,
override: Mapping[str, object] | None = None,
override_expires_at: datetime | None = None,
team_budget: Mapping[str, object] | None = None,
org_budget: Mapping[str, object] | None = None,
spend: tuple[float, float, float] = (0.0, 0.0, 0.0),
) -> Mapping[str, object]:
org_spend, team_spend, member_spend = spend
return {
"customer_member_id": "member-1",
"parent_team_id": "team-1",
"spend": member_spend,
"budget": template,
"budget_override": override,
"budget_override_expires_at": override_expires_at,
"parent_team": {
"customer_team_id": "team-1",
"parent_org_id": "org-1",
"spend": team_spend,
"budget": team_budget,
"parent_org": {"customer_org_id": "org-1", "spend": org_spend, "budget": org_budget},
},
}
def _scopes(hierarchy: CustomerHierarchyAuth) -> Mapping[str, CustomerScopeAuth]:
return {scope.key: scope for scope in hierarchy.scopes}
def test_member_row_yields_org_team_and_member_scopes_with_the_effective_member_budget() -> None:
hierarchy: Final = customer_hierarchy_from_member_row(
_row(
template={"max_budget": 10.0, "tpm_limit": 1000, "model_tpm_limit": {"gpt-4o": 400}},
override={"max_budget": 25.0, "allowed_models": ["gpt-4o"]},
override_expires_at=NOW + timedelta(days=1),
team_budget={"rpm_limit": 50, "model_rpm_limit": {"gpt-4o": 5}},
org_budget={"max_budget": 100.0},
spend=(1.0, 2.0, 3.0),
),
NOW,
)
scopes: Final = _scopes(hierarchy)
assert (hierarchy.customer_org_id, hierarchy.customer_team_id, hierarchy.customer_member_id) == (
"org-1",
"team-1",
"member-1",
)
assert tuple(scope.key for scope in hierarchy.scopes) == ("customer_org", "customer_team", "customer_member")
assert (scopes["customer_org"].max_budget, scopes["customer_org"].spend) == (100.0, 1.0)
assert (scopes["customer_team"].rpm_limit, scopes["customer_team"].model_rpm_limit) == (50, {"gpt-4o": 5})
assert scopes["customer_team"].max_budget is None
member: Final = scopes["customer_member"]
assert (member.max_budget, member.tpm_limit, member.allowed_models, member.spend) == (
25.0,
1000,
("gpt-4o",),
3.0,
)
assert member.model_tpm_limit == {"gpt-4o": 400}
def test_an_expired_override_leaves_the_member_on_its_template() -> None:
hierarchy: Final = customer_hierarchy_from_member_row(
_row(
template={"max_budget": 10.0},
override={"max_budget": 25.0},
override_expires_at=NOW - timedelta(seconds=1),
),
NOW,
)
assert _scopes(hierarchy)["customer_member"].max_budget == 10.0
def test_levels_without_a_budget_row_carry_no_limits() -> None:
hierarchy: Final = customer_hierarchy_from_member_row(_row(template={"rpm_limit": 1}), NOW)
org: Final = _scopes(hierarchy)["customer_org"]
assert (org.max_budget, org.tpm_limit, org.rpm_limit, org.allowed_models, org.model_tpm_limit) == (
None,
None,
None,
None,
None,
)
def test_budget_check_fails_the_outermost_exhausted_level_first() -> None:
hierarchy: Final = customer_hierarchy_from_member_row(
_row(
template={"max_budget": 1.0},
team_budget={"max_budget": 5.0},
org_budget={"max_budget": 50.0},
spend=(50.0, 5.0, 0.5),
),
NOW,
)
with pytest.raises(litellm.BudgetExceededError) as exc:
check_customer_hierarchy_budget(hierarchy)
assert (exc.value.entity_type, exc.value.entity_id) == ("customer_org", "org-1")
within: Final = customer_hierarchy_from_member_row(
_row(template={"max_budget": 1.0}, team_budget={"max_budget": 5.0}, spend=(0.0, 4.99, 0.99)), NOW
)
check_customer_hierarchy_budget(within)
def test_model_access_requires_every_level_with_an_allowlist_to_permit_the_model() -> None:
hierarchy: Final = customer_hierarchy_from_member_row(
_row(template={"allowed_models": ["luna", "opus"]}, team_budget={"allowed_models": ["luna"]}), NOW
)
check_customer_hierarchy_model_access(hierarchy, "luna", llm_router=None)
with pytest.raises(ProxyException) as exc:
check_customer_hierarchy_model_access(hierarchy, "opus", llm_router=None)
assert exc.value.code == "403"
unrestricted: Final = customer_hierarchy_from_member_row(_row(template={"max_budget": 1.0}), NOW)
check_customer_hierarchy_model_access(unrestricted, "anything", llm_router=None)
@pytest.mark.parametrize(
"headers, expected",
[
({"X-LiteLLM-Customer-Member-Id": " member-9 "}, "member-9"),
({CUSTOMER_MEMBER_ID_HEADER: ""}, None),
({"x-litellm-customer-id": "end-user"}, None),
(None, None),
],
)
def test_member_header_is_read_case_insensitively_and_blank_values_are_ignored(
headers: Mapping[str, object] | None, expected: str | None
) -> None:
assert customer_member_id_from_headers(headers) == expected
class _MemberTable:
def __init__(self, rows: Mapping[str, Mapping[str, object]]) -> None:
self.rows = rows
self.lookups = 0
async def find_unique(self, where: Mapping[str, object], include: Mapping[str, object] | None = None):
self.lookups += 1
row: Final = self.rows.get(str(where["customer_member_id"]))
return _Record(row) if row is not None else None
class _Record:
def __init__(self, values: Mapping[str, object]) -> None:
self.values = values
def model_dump(self) -> Mapping[str, object]:
return self.values
class _Db:
def __init__(self, table: _MemberTable) -> None:
self.litellm_customermembertable = table
class _Prisma:
def __init__(self, table: _MemberTable) -> None:
self.db = _Db(table)
@pytest.mark.asyncio
async def test_hierarchy_is_loaded_once_and_then_served_from_the_cache() -> None:
table: Final = _MemberTable({"member-1": _row(template={"max_budget": 10.0})})
cache: Final = UserApiKeyCache(DualCache())
prisma: Final = _Prisma(table)
first: Final = await get_customer_hierarchy("member-1", prisma, cache, now=NOW) # pyright: ignore[reportArgumentType] # test double
second: Final = await get_customer_hierarchy("member-1", prisma, cache, now=NOW) # pyright: ignore[reportArgumentType] # test double
assert first is not None and second == first
assert table.lookups == 1
assert await get_customer_hierarchy("missing", prisma, cache, now=NOW) is None # pyright: ignore[reportArgumentType] # test double
@pytest.mark.asyncio
async def test_resolution_refuses_an_unknown_member_and_a_proxy_without_a_database() -> None:
cache: Final = UserApiKeyCache(DualCache())
with pytest.raises(ProxyException) as missing:
await resolve_customer_hierarchy(
customer_member_id="ghost",
prisma_client=_Prisma(_MemberTable({})), # pyright: ignore[reportArgumentType] # test double
user_api_key_cache=cache,
llm_router=None,
request_model="gpt-4o",
route="/chat/completions",
)
assert missing.value.code == "401"
with pytest.raises(ProxyException):
await resolve_customer_hierarchy(
customer_member_id="member-1",
prisma_client=None,
user_api_key_cache=cache,
llm_router=None,
request_model="gpt-4o",
route="/chat/completions",
)
@pytest.mark.asyncio
async def test_resolution_enforces_budget_on_model_routes_but_not_on_info_routes() -> None:
cache: Final = UserApiKeyCache(DualCache())
prisma: Final = _Prisma(_MemberTable({"member-1": _row(template={"max_budget": 1.0}, spend=(0.0, 0.0, 1.0))}))
with pytest.raises(litellm.BudgetExceededError):
await resolve_customer_hierarchy(
customer_member_id="member-1",
prisma_client=prisma, # pyright: ignore[reportArgumentType] # test double
user_api_key_cache=cache,
llm_router=None,
request_model="gpt-4o",
route="/chat/completions",
)
hierarchy: Final = await resolve_customer_hierarchy(
customer_member_id="member-1",
prisma_client=prisma, # pyright: ignore[reportArgumentType] # test double
user_api_key_cache=cache,
llm_router=None,
request_model=None,
route="/key/info",
)
assert hierarchy.customer_member_id == "member-1"
def test_a_key_pinned_to_another_organization_cannot_act_for_the_member() -> None:
hierarchy: Final = customer_hierarchy_from_member_row(_row(template={"max_budget": 10.0}), now=NOW)
with pytest.raises(ProxyException) as denied:
check_customer_hierarchy_key_scope(hierarchy, {"customer_org_id": "some-other-org"})
assert denied.value.code == "403"
check_customer_hierarchy_key_scope(hierarchy, {"customer_org_id": "org-1"})
check_customer_hierarchy_key_scope(hierarchy, {})

138
tests/unit/proxy/auth/test_user_api_key_auth.py
View file

@ -2647,3 +2647,141 @@ async def test_non_object_metadata_is_served_like_unset_mappings_whatever_the_ba
)
assert mapped == unset
assert mapped[1] == "app-deployment"
@pytest.mark.asyncio
async def test_customer_hierarchy_survives_virtual_key_auth_without_legacy_end_user(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy.auth import customer_hierarchy_auth as hierarchy_module
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
server: Final = litellm.proxy.proxy_server
key: Final = "sk-hierarchy-member"
stored: Final = UserAPIKeyAuth(token=server.hash_token(key))
hierarchy: Final = CustomerHierarchyAuth(
customer_member_id="m1",
customer_team_id="t1",
customer_org_id="o1",
scopes=(CustomerScopeAuth(key="customer_member", entity_id="m1", rpm_limit=1),),
)
loader: Final = AsyncMock(return_value=hierarchy)
monkeypatch.setattr(hierarchy_module, "get_customer_hierarchy", loader)
monkeypatch.setattr(server, "master_key", MASTER_KEY)
monkeypatch.setattr(server, "general_settings", {})
monkeypatch.setattr(server, "user_api_key_cache", UserApiKeyCache())
monkeypatch.setattr(server, "prisma_client", MagicMock(get_data=AsyncMock(return_value=stored)))
monkeypatch.setattr(server, "llm_router", None)
monkeypatch.setattr(server, "user_custom_auth", None)
request: Final = StarletteRequest(
{
"type": "http",
"method": "POST",
"path": "/chat/completions",
"headers": [(b"x-litellm-customer-member-id", b"m1")],
},
receive=AsyncMock(
return_value={
"type": "http.request",
"body": json.dumps({"model": "gpt-4o-mini", "messages": []}).encode(),
"more_body": False,
}
),
)
identity: Final = await user_api_key_auth(request=request, api_key=f"Bearer {key}")
assert identity.customer_hierarchy == hierarchy
assert identity.end_user_id is None
assert loader.await_args is not None and loader.await_args.args[0] == "m1"
@pytest.mark.asyncio
@pytest.mark.parametrize("key_org,allowed", (("o1", True), ("o2", False), (None, True)))
async def test_a_key_pinned_to_a_customer_org_only_serves_that_orgs_members(
monkeypatch: pytest.MonkeyPatch, key_org: str | None, allowed: bool
) -> None:
from litellm.proxy.auth import customer_hierarchy_auth as hierarchy_module
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
server: Final = litellm.proxy.proxy_server
key: Final = "sk-hierarchy-scoped"
metadata: Final = {"customer_org_id": key_org} if key_org is not None else {}
stored: Final = UserAPIKeyAuth(token=server.hash_token(key), metadata=metadata)
hierarchy: Final = CustomerHierarchyAuth(
customer_member_id="m1",
customer_team_id="t1",
customer_org_id="o1",
scopes=(CustomerScopeAuth(key="customer_member", entity_id="m1"),),
)
monkeypatch.setattr(hierarchy_module, "get_customer_hierarchy", AsyncMock(return_value=hierarchy))
monkeypatch.setattr(server, "master_key", MASTER_KEY)
monkeypatch.setattr(server, "general_settings", {})
monkeypatch.setattr(server, "user_api_key_cache", UserApiKeyCache())
monkeypatch.setattr(server, "prisma_client", MagicMock(get_data=AsyncMock(return_value=stored)))
monkeypatch.setattr(server, "llm_router", None)
monkeypatch.setattr(server, "user_custom_auth", None)
request: Final = StarletteRequest(
{
"type": "http",
"method": "POST",
"path": "/chat/completions",
"headers": [(b"x-litellm-customer-member-id", b"m1")],
},
receive=AsyncMock(
return_value={
"type": "http.request",
"body": json.dumps({"model": "gpt-4o-mini", "messages": []}).encode(),
"more_body": False,
}
),
)
if allowed:
identity: Final = await user_api_key_auth(request=request, api_key=f"Bearer {key}")
assert identity.customer_hierarchy == hierarchy
else:
with pytest.raises(ProxyException) as denied:
await user_api_key_auth(request=request, api_key=f"Bearer {key}")
assert denied.value.code == "403"
@pytest.mark.asyncio
async def test_customer_model_access_is_checked_against_the_model_the_route_serves(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy.auth import customer_hierarchy_auth as hierarchy_module
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
server: Final = litellm.proxy.proxy_server
key: Final = "sk-hierarchy-route-model"
stored: Final = UserAPIKeyAuth(token=server.hash_token(key))
hierarchy: Final = CustomerHierarchyAuth(
customer_member_id="m1",
customer_team_id="t1",
customer_org_id="o1",
scopes=(CustomerScopeAuth(key="customer_member", entity_id="m1", allowed_models=("allowed",)),),
)
monkeypatch.setattr(hierarchy_module, "get_customer_hierarchy", AsyncMock(return_value=hierarchy))
monkeypatch.setattr(server, "master_key", MASTER_KEY)
monkeypatch.setattr(server, "general_settings", {})
monkeypatch.setattr(server, "user_api_key_cache", UserApiKeyCache())
monkeypatch.setattr(server, "prisma_client", MagicMock(get_data=AsyncMock(return_value=stored)))
monkeypatch.setattr(server, "llm_router", None)
monkeypatch.setattr(server, "user_custom_auth", None)
request: Final = StarletteRequest(
{
"type": "http",
"method": "POST",
"path": "/openai/deployments/blocked/chat/completions",
"headers": [(b"x-litellm-customer-member-id", b"m1")],
},
receive=AsyncMock(
return_value={
"type": "http.request",
"body": json.dumps({"messages": []}).encode(),
"more_body": False,
}
),
)
with pytest.raises(ProxyException) as denied:
await user_api_key_auth(request=request, api_key=f"Bearer {key}")
assert denied.value.code == "403"
assert "blocked" in denied.value.message

13
tests/unit/proxy/common_utils/test_reset_budget_job.py
View file

@ -12,15 +12,14 @@ import httpx
import prisma
import pytest
from litellm.proxy._types import LiteLLM_BudgetTableFull, LiteLLM_VerificationToken
from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module
from litellm.constants import (
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
RESET_BUDGET_JOB_BATCH_SIZE,
RESET_BUDGET_JOB_LOCK_TTL_SECONDS,
RESET_BUDGET_JOB_NAME,
)
from litellm.proxy._types import LiteLLM_BudgetTableFull, LiteLLM_VerificationToken
from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module
from litellm.proxy.common_utils.reset_budget_job import (
ResetBudgetJob,
_RowReset,
@ -116,6 +115,7 @@ class MockTable:
where: Dict[str, Any],
order: Optional[Dict[str, str]] = None,
take: Optional[int] = None,
include: Optional[Dict[str, bool]] = None,
) -> List[Any]:
"""Replays canned rows, honouring the keyset cursor + ``take`` a paged
caller relies on: without that a paged walk never advances and the
@ -220,6 +220,9 @@ class MockDB:
self.litellm_tagtable = MockTable()
self.litellm_modelaccessgroupbudgettable = MockTable()
self.litellm_projecttable = MockTable()
self.litellm_customermembertable = MockTable()
self.litellm_customerteamtable = MockTable()
self.litellm_customerorganizationtable = MockTable()
self.batch_calls: List[Dict[str, Any]] = []
self.batchers: List[MockBatcher] = []
@ -257,6 +260,7 @@ class MockPrismaClient:
}
self.get_data_calls: List[Dict[str, Any]] = []
self.db = MockDB()
self.writer_db = self.db
async def get_data(self, table_name, query_type, **kwargs):
self.get_data_calls.append({"table_name": table_name, "query_type": query_type, **kwargs})
@ -4337,6 +4341,9 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c
**{t["team_id"]: t["spend"] for t in [team1, team2]},
}
_wire_cascade_reads_for_test(prisma_client, endusers=[enduser1])
prisma_client.writer_db.litellm_customermembertable.find_many = AsyncMock(return_value=[])
prisma_client.writer_db.litellm_customerteamtable.find_many = AsyncMock(return_value=[])
prisma_client.writer_db.litellm_customerorganizationtable.find_many = AsyncMock(return_value=[])
proxy_logging_obj: Final = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()

12
tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py
View file

@ -133,6 +133,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
daily_org_spend_update_queue=daily_org_queue,
daily_end_user_spend_update_queue=daily_end_user_queue,
daily_agent_spend_update_queue=daily_agent_queue,
daily_customer_member_spend_update_queue=DailySpendUpdateQueue(),
)
# Should be called exactly once (pipeline)
@ -200,6 +201,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_upd
daily_org_spend_update_queue=daily_org_queue,
daily_end_user_spend_update_queue=daily_end_user_queue,
daily_agent_spend_update_queue=daily_agent_queue,
daily_customer_member_spend_update_queue=DailySpendUpdateQueue(),
)
# After restore, the main spend queue should hold one item per
@ -239,6 +241,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_updat
daily_org_spend_update_queue=empty_daily_queue,
daily_end_user_spend_update_queue=empty_daily_queue,
daily_agent_spend_update_queue=empty_daily_queue,
daily_customer_member_spend_update_queue=DailySpendUpdateQueue(),
)
mock_redis_cache.async_rpush_pipeline.assert_not_called()
@ -292,7 +295,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert len(result) == 7
assert len(result) == 8
(
db_spend,
daily_user,
@ -301,6 +304,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
daily_end_user,
daily_agent,
window_spend,
customer_spend,
) = result
# Budget window spend from two pods is summed per window, not overwritten.
@ -375,6 +379,7 @@ async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_fail
daily_org_spend_update_queue=DailySpendUpdateQueue(),
daily_end_user_spend_update_queue=DailySpendUpdateQueue(),
daily_agent_spend_update_queue=DailySpendUpdateQueue(),
daily_customer_member_spend_update_queue=DailySpendUpdateQueue(),
)
restored_spend: Final = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
@ -452,6 +457,7 @@ async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_a
daily_user,
*_rest,
window_spend,
customer_spend,
) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert db_spend is not None
@ -484,7 +490,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
"""When redis_cache is None, should return all Nones"""
buffer = RedisUpdateBuffer(redis_cache=None)
result = await buffer.get_all_transactions_from_redis_buffer_pipeline()
assert result == (None, None, None, None, None, None, None)
assert result == (None, None, None, None, None, None, None, None)
@pytest.mark.asyncio
@ -722,6 +728,7 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up
daily_end_user_spend_update_queue=empty_daily_queue,
daily_agent_spend_update_queue=empty_daily_queue,
window_spend_update_queue=window_queue,
daily_customer_member_spend_update_queue=DailySpendUpdateQueue(),
)
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
@ -813,6 +820,7 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu
daily_end_user_spend_update_queue=empty_daily_queue,
daily_agent_spend_update_queue=empty_daily_queue,
window_spend_update_queue=window_queue,
daily_customer_member_spend_update_queue=DailySpendUpdateQueue(),
)
restored = await window_queue.flush_and_get_aggregated_window_spend_transactions()

47
tests/unit/proxy/db/test_daily_spend_bulk_upsert.py
View file

@ -127,9 +127,7 @@ def test_counters_increment_rather_than_overwrite(column):
def test_request_id_is_preserved_when_a_later_batch_carries_none():
sql, params = build_bulk_upsert(
TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(request_id=None),))
)
sql, params = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(request_id=None),)))
assert '"request_id" = COALESCE(EXCLUDED."request_id", "LiteLLM_DailyTagSpend"."request_id")' in sql
assert None in params
@ -209,3 +207,46 @@ async def test_writer_survives_a_transaction_whose_key_columns_are_null():
_, params = prisma_client.db.statements[0]
assert None not in params[:9]
assert transactions == {}
def test_customer_daily_rows_keep_historical_parent_attribution_when_member_id_is_reused():
table = DAILY_SPEND_TABLES["customer_member"]
first = {
**tag_txn(),
"customer_member_id": "member",
"customer_team_id": "team-before",
"customer_org_id": "org-before",
}
second = {**first, "customer_team_id": "team-after", "customer_org_id": "org-after", "spend": 0.75}
merged = merge_by_conflict_key(table, (first, first, second))
assert len(merged) == 2
assert {(row["customer_team_id"], row["customer_org_id"], row["spend"]) for _, row in merged} == {
("team-before", "org-before", 0.5),
("team-after", "org-after", 0.75),
}
sql, params = build_bulk_upsert(table, merged)
assert 'ON CONFLICT ("id")' in sql
assert all(identity in params for identity in ("member", "team-before", "org-before", "team-after", "org-after"))
row_ids = {params[index] for index in range(0, len(params), len(params) // 2)}
assert len(row_ids) == 2
assert row_ids == {params[0] for _, _ in (build_bulk_upsert(table, merged[:1]),)} | {
build_bulk_upsert(table, merged[1:])[1][0]
}
def test_customer_daily_aliases_of_one_model_keep_separate_rows_under_a_stable_id():
table = DAILY_SPEND_TABLES["customer_member"]
base = {**tag_txn(), "customer_member_id": "member", "customer_team_id": "team", "customer_org_id": "org"}
merged = merge_by_conflict_key(
table,
(
{**base, "model_group": "alias-a"},
{**base, "model_group": "alias-a", "spend": 0.25},
{**base, "model_group": "alias-b"},
),
)
assert sorted((row["model_group"], row["spend"]) for _, row in merged) == [("alias-a", 0.5), ("alias-b", 0.25)]
_, params = build_bulk_upsert(table, merged)
first_id, second_id = params[0], params[len(params) // 2]
assert first_id != second_id
assert build_bulk_upsert(table, merged[:1])[1][0] == first_id

157
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -21,11 +21,11 @@ from redis.exceptions import DataError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._service_logger import ServiceTypes
from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS
from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem
from litellm.proxy.db.db_spend_update_writer import (
_SPEND_TABLE_COMMIT_ORDER,
_TEAM_ADVISORY_LOCK_SQL,
_TEAM_MEMBER_SPEND_SQL,
DBSpendUpdateWriter,
@ -45,7 +45,6 @@ from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpush
from tests.unit.proxy.db.fake_prisma_engine import engine_call
@pytest.fixture(autouse=True)
def legacy_spend_writers(monkeypatch: pytest.MonkeyPatch) -> None:
"""These tests cover the legacy writers, still reachable as the rollback path with use_spend_logs_ledger off."""
@ -1075,6 +1074,88 @@ async def test_commit_spend_updates_to_db_increments_agent_spend():
assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_every_customer_hierarchy_table():
"""Spend a request ran under a customer member lands on the member, its team and its organization rows."""
db_writer = DBSpendUpdateWriter()
mock_batcher = MagicMock()
mock_transaction = AsyncMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
)
)
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
committed: list[str] = []
with patch("litellm.proxy.utils.raise_failed_update_spend_exception"):
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=MagicMock(),
db_spend_update_transactions={
**{name: {} for name in _SPEND_TABLE_COMMIT_ORDER},
"customer_org_list_transactions": {"org-1": 0.5},
"customer_team_list_transactions": {"team-1": 0.5},
"customer_member_list_transactions": {"member-1": 0.5},
},
on_table_committed=committed.append,
)
for table, column, entity_id in (
("litellm_customerorganizationtable", "customer_org_id", "org-1"),
("litellm_customerteamtable", "customer_team_id", "team-1"),
("litellm_customermembertable", "customer_member_id", "member-1"),
):
update_many = getattr(mock_batcher, table).update_many
update_many.assert_called_once_with(where={column: entity_id}, data={"spend": {"increment": 0.5}})
assert committed[-3:] == [
"customer_org_list_transactions",
"customer_team_list_transactions",
"customer_member_list_transactions",
]
@pytest.mark.asyncio
async def test_update_customer_hierarchy_db_enqueues_one_increment_per_level():
from litellm.types.proxy.auth.customer_hierarchy import customer_hierarchy_ids_from_request
db_writer = DBSpendUpdateWriter()
ids = customer_hierarchy_ids_from_request(
{
"litellm_params": {
"metadata": {
"user_api_key_customer_member_id": "member-1",
"user_api_key_customer_team_id": "team-1",
"user_api_key_customer_org_id": "org-1",
}
}
}
)
assert ids is not None
assert (
customer_hierarchy_ids_from_request({"litellm_params": {"metadata": {"user_api_key_customer_member_id": "m"}}})
is None
)
await db_writer._update_customer_hierarchy_db(response_cost=0.25, customer_hierarchy=ids, prisma_client=MagicMock())
queued = await db_writer.spend_update_queue.flush_all_updates_from_in_memory_queue()
assert {(item["entity_type"], item["entity_id"], item["response_cost"]) for item in queued} == {
(Litellm_EntityType.CUSTOMER_ORG, "org-1", 0.25),
(Litellm_EntityType.CUSTOMER_TEAM, "team-1", 0.25),
(Litellm_EntityType.CUSTOMER_MEMBER, "member-1", 0.25),
}
aggregated = db_writer.spend_update_queue.get_aggregated_db_spend_update_transactions(queued)
assert aggregated["customer_member_list_transactions"] == {"member-1": 0.25}
assert aggregated["customer_team_list_transactions"] == {"team-1": 0.25}
assert aggregated["customer_org_list_transactions"] == {"org-1": 0.25}
def _team_member_flush_fixtures() -> tuple[AsyncMock, MagicMock]:
"""A transaction and prisma client that record the raw statement the member spend flush runs."""
mock_transaction = AsyncMock()
@ -2526,7 +2607,7 @@ async def test_commit_spend_updates_uses_pipeline():
mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock()
# Return all-None tuple (no data to commit); the pipeline yields 6 slots
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(None, None, None, None, None, None, None)
return_value=(None, None, None, None, None, None, None, None)
)
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -2607,7 +2688,7 @@ async def test_commit_with_redis_requeues_all_on_db_failure():
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(db_spend, daily_user, None, None, None, None, None)
return_value=(db_spend, daily_user, None, None, None, None, None, None)
)
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -2662,7 +2743,7 @@ async def test_commit_with_redis_only_requeues_failed_category():
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(db_spend, daily_user, None, None, None, None, None)
return_value=(db_spend, daily_user, None, None, None, None, None, None)
)
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -2710,7 +2791,7 @@ async def test_commit_with_redis_no_requeue_on_success():
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(db_spend, None, None, None, None, None, None)
return_value=(db_spend, None, None, None, None, None, None, None)
)
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -3343,7 +3424,7 @@ async def test_window_spend_queue_is_handed_to_the_redis_buffer():
db_writer = DBSpendUpdateWriter()
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(None, None, None, None, None, None, None)
return_value=(None, None, None, None, None, None, None, None)
)
db_writer.redis_update_buffer = mock_redis_update_buffer
db_writer.pod_lock_manager = AsyncMock()
@ -3373,7 +3454,7 @@ async def test_window_spend_transactions_from_redis_are_committed_by_the_lock_wi
)
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(None, None, None, None, None, None, window_transactions)
return_value=(None, None, None, None, None, None, window_transactions, None)
)
db_writer.redis_update_buffer = mock_redis_update_buffer
db_writer.pod_lock_manager = AsyncMock()
@ -3631,7 +3712,7 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis():
)
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(None, None, None, None, None, None, window_transactions)
return_value=(None, None, None, None, None, None, window_transactions, None)
)
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -3666,7 +3747,7 @@ async def test_failed_per_entity_increment_from_redis_restores_only_what_may_sti
)
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(db_spend_transactions, None, None, None, None, None, None)
return_value=(db_spend_transactions, None, None, None, None, None, None, None)
)
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -3726,8 +3807,7 @@ async def test_failed_per_entity_increment_from_redis_restores_only_what_may_sti
restores = mock_redis_update_buffer.restore_transactions_to_redis.await_args_list
assert [
{name: rows for name, rows in call.kwargs["db_spend_update_transactions"].items() if rows}
for call in restores
{name: rows for name, rows in call.kwargs["db_spend_update_transactions"].items() if rows} for call in restores
] == ([{"key_list_transactions": {"key-1": 1.5}}] if safe_to_resend else [])
mock_batcher.litellm_teamtable.update_many.assert_called_once()
mock_batcher.litellm_usertable.update_many.assert_called_once()
@ -3750,7 +3830,7 @@ async def test_failed_window_spend_commit_from_redis_is_restored_only_when_safe_
)
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(None, None, None, None, None, None, window_transactions)
return_value=(None, None, None, None, None, None, window_transactions, None)
)
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
db_writer.redis_update_buffer = mock_redis_update_buffer
@ -3914,6 +3994,9 @@ async def test_commit_spend_updates_to_db_reports_each_completed_table():
call("tag_list_transactions"),
call("model_access_group_list_transactions"),
call("agent_list_transactions"),
call("customer_org_list_transactions"),
call("customer_team_list_transactions"),
call("customer_member_list_transactions"),
]
@ -5859,7 +5942,9 @@ async def test_completed_spend_survives_error_log_suppression(
@pytest.mark.parametrize("ledger_enabled", [False, True])
@pytest.mark.asyncio
async def test_request_error_counts_survive_the_ledger_path_switch(monkeypatch: pytest.MonkeyPatch, ledger_enabled: bool):
async def test_request_error_counts_survive_the_ledger_path_switch(
monkeypatch: pytest.MonkeyPatch, ledger_enabled: bool
):
from litellm.proxy import proxy_server
from litellm.proxy.db import db_spend_update_writer as writer_module
from litellm.proxy.db.request_error_tracking import RequestErrorAccumulator
@ -5874,20 +5959,40 @@ async def test_request_error_counts_survive_the_ledger_path_switch(monkeypatch:
monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
now = datetime(2026, 11, 2, tzinfo=timezone.utc)
await DBSpendUpdateWriter().update_database(
token="key", user_id="user", team_id=None, org_id=None, end_user_id=None,
token="key",
user_id="user",
team_id=None,
org_id=None,
end_user_id=None,
kwargs={
"model": "ledger-test", "call_type": "acompletion", "litellm_call_id": "failed-request",
"litellm_params": {"metadata": {
"user_api_key": "key", "user_api_key_user_id": "user", "status": "failure",
"error_information": {"error_code": "503"},
}},
"model": "ledger-test",
"call_type": "acompletion",
"litellm_call_id": "failed-request",
"litellm_params": {
"metadata": {
"user_api_key": "key",
"user_api_key_user_id": "user",
"status": "failure",
"error_information": {"error_code": "503"},
}
},
},
completion_response=None, start_time=now, end_time=now, response_cost=0,
completion_response=None,
start_time=now,
end_time=now,
response_cost=0,
)
await drain_logging_tasks()
assert accumulator.drain() == {RequestErrorKey(
date="2026-11-02", api_key=prisma.spend_log_transactions[0]["api_key"], team_id="", user_id="user", model_group="ledger-test", status_code=503,
): 1}
assert accumulator.drain() == {
RequestErrorKey(
date="2026-11-02",
api_key=prisma.spend_log_transactions[0]["api_key"],
team_id="",
user_id="user",
model_group="ledger-test",
status_code=503,
): 1
}
@pytest.mark.asyncio
@ -6053,9 +6158,7 @@ async def test_a_billed_row_in_a_caller_session_adds_its_spend_to_that_sessions_
pytest.param(_CALLER_SESSION, 0.0, id="free-row"),
],
)
async def test_a_row_without_a_caller_session_or_cost_adds_nothing(
litellm_params: Mapping[str, object], spend: float
):
async def test_a_row_without_a_caller_session_or_cost_adds_nothing(litellm_params: Mapping[str, object], spend: float):
writer, totals = _session_writer()
await _record_spend_row(writer, spend, litellm_params)

2
tests/unit/proxy/db/test_model_access_group_spend.py
View file

@ -420,6 +420,7 @@ async def test_redis_buffer_requeues_access_group_transactions_as_queue_items():
daily_org_spend_update_transactions=None,
daily_end_user_spend_update_transactions=None,
daily_agent_spend_update_transactions=None,
daily_customer_member_spend_update_transactions=None,
window_spend_update_transactions=None,
spend_update_queue=queue,
daily_spend_update_queue=daily_queue,
@ -427,6 +428,7 @@ async def test_redis_buffer_requeues_access_group_transactions_as_queue_items():
daily_org_spend_update_queue=daily_queue,
daily_end_user_spend_update_queue=daily_queue,
daily_agent_spend_update_queue=daily_queue,
daily_customer_member_spend_update_queue=daily_queue,
window_spend_update_queue=None,
)

32
tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -1915,6 +1915,38 @@ async def test_flexible_team_is_steered_to_its_alternate_once_the_preferred_mode
)
@pytest.mark.asyncio
@pytest.mark.parametrize("team_allowed_models, granted", [(("luna", "opus"), True), (("luna",), False)])
async def test_an_alternate_is_granted_only_when_every_customer_level_permits_it(
monkeypatch: pytest.MonkeyPatch, team_allowed_models: tuple[str, ...], granted: bool
):
from litellm.proxy import proxy_server
from litellm.proxy.auth import auth_checks, user_api_key_auth
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import alternate_grant_for
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
monkeypatch.setattr(proxy_server, "llm_router", _steering_router("luna", "opus"))
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(user_api_key_auth, "_is_fallback_model_authorized", AsyncMock(return_value=True))
monkeypatch.setattr(user_api_key_auth, "_is_within_every_model_budget", AsyncMock(return_value=True))
monkeypatch.setattr(auth_checks, "collect_matched_model_access_groups", AsyncMock(return_value=()))
monkeypatch.setattr(auth_checks, "_model_access_group_max_budget_check", AsyncMock(return_value=None))
customer: Final = UserAPIKeyAuth(
customer_hierarchy=CustomerHierarchyAuth(
customer_member_id="member-1",
customer_team_id="team-1",
customer_org_id="org-1",
scopes=(
CustomerScopeAuth(key="customer_org", entity_id="org-1"),
CustomerScopeAuth(key="customer_team", entity_id="team-1", allowed_models=team_allowed_models),
CustomerScopeAuth(key="customer_member", entity_id="member-1", allowed_models=("luna", "opus")),
),
)
)
assert (await alternate_grant_for(customer, "opus") is not None) is granted
@pytest.mark.asyncio
async def test_teams_without_a_permitted_idle_alternate_stay_on_the_preferred_model(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)

122
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -6,7 +6,7 @@ import asyncio
import contextvars
import logging
import time
from collections.abc import AsyncGenerator, Iterator, Sequence
from collections.abc import AsyncGenerator, Iterator, Mapping, Sequence
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from functools import partial
@ -33,6 +33,7 @@ from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_
from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired, total_deadline
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
@ -8618,6 +8619,111 @@ async def _rpm_request(
)
def _customer_hierarchy(
*,
org_rpm: int | None = None,
team_rpm: int | None = None,
member_rpm: int | None = None,
org_model_rpm: Mapping[str, int] | None = None,
team_model_rpm: Mapping[str, int] | None = None,
member_model_rpm: Mapping[str, int] | None = None,
) -> CustomerHierarchyAuth:
return CustomerHierarchyAuth(
customer_member_id="member-1",
customer_team_id="team-1",
customer_org_id="org-1",
scopes=(
CustomerScopeAuth(key="customer_org", entity_id="org-1", rpm_limit=org_rpm, model_rpm_limit=org_model_rpm),
CustomerScopeAuth(
key="customer_team", entity_id="team-1", rpm_limit=team_rpm, model_rpm_limit=team_model_rpm
),
CustomerScopeAuth(
key="customer_member", entity_id="member-1", rpm_limit=member_rpm, model_rpm_limit=member_model_rpm
),
),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"hierarchy, limited_scope",
[
(_customer_hierarchy(org_rpm=1, team_rpm=5, member_rpm=5), "customer_org"),
(_customer_hierarchy(org_rpm=5, team_rpm=1, member_rpm=5), "customer_team"),
(_customer_hierarchy(org_rpm=5, team_rpm=5, member_rpm=1), "customer_member"),
],
ids=["org", "team", "member"],
)
async def test_customer_hierarchy_rpm_limits_apply_at_every_level_and_name_the_exhausted_one(
hierarchy: CustomerHierarchyAuth, limited_scope: str
) -> None:
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-customer-hierarchy"), customer_hierarchy=hierarchy)
await _rpm_request(handler, cache, auth, "target")
with pytest.raises(HTTPException) as exc:
await _rpm_request(handler, cache, auth, "target")
assert exc.value.status_code == 429
assert exc.value.headers["rate_limited_by"] == limited_scope
assert f"{limited_scope}: {limited_scope.split('_')[-1]}-1" in str(exc.value.detail), exc.value.detail
@pytest.mark.asyncio
async def test_two_members_of_one_customer_team_share_the_team_bucket_but_not_each_others_member_bucket() -> None:
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
first: Final = UserAPIKeyAuth(
api_key=hash_token("sk-first"), customer_hierarchy=_customer_hierarchy(team_rpm=2, member_rpm=1)
)
second: Final = UserAPIKeyAuth(
api_key=hash_token("sk-second"),
customer_hierarchy=_customer_hierarchy(team_rpm=2, member_rpm=1).model_copy(
update={
"customer_member_id": "member-2",
"scopes": tuple(
scope.model_copy(update={"entity_id": "member-2"}) if scope.key == "customer_member" else scope
for scope in _customer_hierarchy(team_rpm=2, member_rpm=1).scopes
),
}
),
)
await _rpm_request(handler, cache, first, "target")
await _rpm_request(handler, cache, second, "target")
with pytest.raises(HTTPException) as exc:
await _rpm_request(handler, cache, second, "target")
assert exc.value.headers["rate_limited_by"] in {"customer_team", "customer_member"}
assert "team-1" in str(exc.value.detail) or "member-2" in str(exc.value.detail)
def test_tpm_settlement_targets_include_the_customer_levels_and_their_model_buckets() -> None:
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
targets: Final = handler._collect_tpm_scope_targets( # pyright: ignore[reportPrivateUsage] # settlement enumerates scopes here
standard_logging_metadata={
"user_api_key_hash": "hash",
"user_api_key_customer_org_id": "org-1",
"user_api_key_customer_team_id": "team-1",
"user_api_key_customer_member_id": "member-1",
},
kwargs={},
model_group="target",
)
assert {
("customer_org", "org-1"),
("customer_team", "team-1"),
("customer_member", "member-1"),
("model_per_customer_org", "org-1:target"),
("model_per_customer_team", "team-1:target"),
("model_per_customer_member", "member-1:target"),
} <= set(targets)
without: Final = handler._collect_tpm_scope_targets( # pyright: ignore[reportPrivateUsage] # settlement enumerates scopes here
standard_logging_metadata={"user_api_key_hash": "hash"}, kwargs={}, model_group="target"
)
assert not any(scope.startswith(("customer_", "model_per_customer_")) for scope, _ in without)
@pytest.mark.asyncio
@pytest.mark.parametrize("first_name, second_name", [("target", "alias"), ("alias", "target")])
async def test_model_group_alias_shares_deployment_default_rpm_bucket_with_its_target(
@ -8667,8 +8773,20 @@ async def test_model_group_alias_shares_deployment_default_rpm_bucket_with_its_t
({"org_id": "o", "organization_metadata": {"model_rpm_limit": {"target": 1}}}, "model_per_organization"),
({"project_id": "p", "project_metadata": {"model_rpm_limit": {"target": 1}}}, "model_per_project"),
({"end_user_id": "customer", "end_user_model_rpm_limit": {"target": 1}}, "model_per_end_user"),
({"customer_hierarchy": _customer_hierarchy(member_model_rpm={"target": 1})}, "model_per_customer_member"),
({"customer_hierarchy": _customer_hierarchy(team_model_rpm={"target": 1})}, "model_per_customer_team"),
({"customer_hierarchy": _customer_hierarchy(org_model_rpm={"target": 1})}, "model_per_customer_org"),
],
ids=[
"key_metadata",
"team_metadata",
"organization_metadata",
"project_metadata",
"end_user_budget",
"customer_member_budget",
"customer_team_budget",
"customer_org_budget",
],
ids=["key_metadata", "team_metadata", "organization_metadata", "project_metadata", "end_user_budget"],
)
async def test_model_group_alias_shares_metadata_model_rpm_bucket_with_its_target(
limits: dict[str, object], counter_scope: str, first_name: str, second_name: str

81
tests/unit/proxy/management_endpoints/management_v1/test_customers.py
View file

@ -1,3 +1,4 @@
import asyncio
from typing import TYPE_CHECKING, Final, cast
import pytest
@ -5,10 +6,12 @@ from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.customer_hierarchy_auth import customer_member_cache_key
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.list_api.common import ManagementProblem, problem_response
from litellm.proxy.management_endpoints.management_v1 import router
from litellm.proxy.management_endpoints.management_v1.customers import get_customer_hierarchy
from litellm.proxy.management_endpoints.management_v1.customers import get_customer_hierarchy, get_user_api_key_cache
from litellm.proxy.management_helpers.customer_hierarchy import CustomerHierarchy
from tests.unit.proxy.management_helpers.test_customer_hierarchy import NOW, Database
@ -16,7 +19,7 @@ if TYPE_CHECKING:
from prisma import Prisma
def client(role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN) -> TestClient:
def client(role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN, cache: UserApiKeyCache | None = None) -> TestClient:
app: Final = FastAPI()
app.include_router(router)
@ -30,6 +33,7 @@ def client(role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN) -> TestClient:
) # cast-ok: in-memory Prisma edge
app.dependency_overrides[user_api_key_auth] = lambda: auth
app.dependency_overrides[get_customer_hierarchy] = lambda: hierarchy
app.dependency_overrides[get_user_api_key_cache] = lambda: cache if cache is not None else UserApiKeyCache()
return TestClient(app)
@ -166,6 +170,8 @@ def test_template_edits_propagate_and_only_live_override_fields_win() -> None:
"tpm_limit": None,
"allowed_models": ["model-b"],
"budget_duration": None,
"model_tpm_limit": None,
"model_rpm_limit": None,
}
fallback: Final = http.get("/management/v1/customer_members/expired")
assert fallback.status_code == 200, fallback.text
@ -541,6 +547,72 @@ def test_template_routes_return_404_for_a_missing_organization(
assert response.status_code == 404, response.text
def test_per_model_rate_limits_round_trip_at_every_level_and_override_the_template_map() -> None:
http: Final = client()
family(http)
team: Final = http.patch(
"/management/v1/customer_teams/team",
json={"budget": {"model_tpm_limit": {"model-a": 5000}, "model_rpm_limit": {"model-a": 50}}},
)
assert team.status_code == 200, team.text
assert team.json()["budget"]["model_tpm_limit"] == {"model-a": 5000}
assert team.json()["budget"]["model_rpm_limit"] == {"model-a": 50}
template: Final = http.patch(
"/management/v1/customer_organizations/org/default_budgets/" + family_default_budget_id(http),
json={"model_tpm_limit": {"model-a": 1000}},
)
assert template.status_code == 200, template.text
created: Final = http.post(
"/management/v1/customer_members",
json={
"customer_member_id": "member",
"parent_team_id": "team",
"budget_override": {"model_rpm_limit": {"model-a": 1}},
},
)
assert created.status_code == 201, created.text
effective: Final = created.json()["effective_budget"]
assert effective["model_tpm_limit"] == {"model-a": 1000}
assert effective["model_rpm_limit"] == {"model-a": 1}
overridden: Final = http.patch(
"/management/v1/customer_members/member", json={"budget_override": {"model_tpm_limit": {"model-a": 10}}}
)
assert overridden.status_code == 200, overridden.text
assert overridden.json()["effective_budget"]["model_tpm_limit"] == {"model-a": 10}
assert http.get("/management/v1/customer_members/member").json()["effective_budget"]["model_rpm_limit"] == {
"model-a": 1
}
def test_budget_edits_at_every_level_drop_the_cached_member_hierarchy() -> None:
cache: Final = UserApiKeyCache()
http: Final = client(cache=cache)
family(http)
assert (
http.post(
"/management/v1/customer_members", json={"customer_member_id": "m", "parent_team_id": "team"}
).status_code
== 201
)
cache_key: Final = customer_member_cache_key("m")
edits: Final = (
("/management/v1/customer_members/m", {"budget_override": {"max_budget": 1}}),
("/management/v1/customer_teams/team", {"budget": {"max_budget": 2}}),
("/management/v1/customer_organizations/org", {"budget": {"max_budget": 3}}),
(
f"/management/v1/customer_organizations/org/default_budgets/{family_default_budget_id(http)}",
{"max_budget": 4},
),
)
for path, body in edits:
asyncio.run(cache.async_set_cache(key=cache_key, value="stale hierarchy"))
assert http.patch(path, json=body).status_code == 200, path
assert cache.get_cache(cache_key) is None, path
asyncio.run(cache.async_set_cache(key=cache_key, value="stale hierarchy"))
assert http.delete("/management/v1/customer_members/m").status_code == 204
assert cache.get_cache(cache_key) is None
def test_organization_list_counts_all_members_and_searches_before_paginating() -> None:
http: Final = client()
family(http)
@ -571,3 +643,8 @@ def test_organization_list_counts_all_members_and_searches_before_paginating() -
assert literal.json()["data"][0]["customer_org_id"] == "other"
assert literal.json()["data"][0]["member_count"] == 0
assert literal.json()["meta"]["total_count"] == 1
def family_default_budget_id(http: TestClient) -> str:
organization: Final = http.get("/management/v1/customer_organizations/org").json()
return str(organization["default_budget_id"])

2
tests/unit/proxy/management_helpers/test_customer_hierarchy.py
View file

@ -79,7 +79,7 @@ class Table:
"budget_names": {},
"created_at": NOW,
"updated_at": NOW,
**data,
**{key: value.data if isinstance(value, Json) else value for key, value in data.items()},
}
)
self.rows = {**self.rows, str(data[self.id_field]): row}

0
tests/unit/types/proxy/auth/__init__.py Normal file
View file

19
tests/unit/types/proxy/auth/test_customer_hierarchy.py Normal file
View file

@ -0,0 +1,19 @@
from typing import Final
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyIds, customer_hierarchy_ids_from_request
_IDS: Final = {
"user_api_key_customer_member_id": "m1",
"user_api_key_customer_team_id": "t1",
"user_api_key_customer_org_id": "o1",
}
def test_ids_are_read_from_whichever_metadata_bucket_the_proxy_stamped() -> None:
expected: Final = CustomerHierarchyIds("m1", "t1", "o1")
assert customer_hierarchy_ids_from_request({"litellm_params": {"metadata": _IDS}}) == expected
assert (
customer_hierarchy_ids_from_request({"litellm_params": {"metadata": {"note": "x"}, "litellm_metadata": _IDS}})
== expected
)
assert customer_hierarchy_ids_from_request({"litellm_params": {"metadata": {"note": "x"}}}) is None