From d30b2b9b26213c79aa4cc50f962730f55d2479c5 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 12 Nov 2026 20:43:44 -0800 Subject: [PATCH] feat(proxy): enforce customer org, team and member budgets and limits on the request path (#51489) * 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 ccac47608f66efaf08fa4780b3ef8aded6da9ffd. * 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 7542471e172ec4ebb51102eaecc1e96e5ffab79e. * 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 --- litellm/accounting/application.py | 42 -- litellm/accounting/delivery_projection.py | 3 + litellm/accounting/projection.py | 112 ++- litellm/accounting/snapshot.py | 3 + litellm/constants.py | 1 + .../pagerduty/pagerduty.py | 660 +++++++++--------- litellm/integrations/otel/model/spans.py | 1 + .../grafana_dashboard.json | 33 + litellm/litellm_core_utils/litellm_logging.py | 6 + litellm/proxy/_types.py | 9 + litellm/proxy/auth/auth_checks.py | 3 + litellm/proxy/auth/customer_hierarchy_auth.py | 252 +++++++ litellm/proxy/auth/user_api_key_auth.py | 26 + .../common_utils/proxy_rate_limit_error.py | 3 + .../proxy/common_utils/reset_budget_job.py | 141 +++- litellm/proxy/db/daily_spend_bulk_upsert.py | 51 +- litellm/proxy/db/db_spend_update_writer.py | 265 ++++++- .../redis_update_buffer.py | 70 +- .../spend_update_queue.py | 12 + litellm/proxy/db/model_usage_rollup.py | 13 +- litellm/proxy/db/prisma_client.py | 9 + .../proxy/hooks/dynamic_rate_limiter_v3.py | 3 + .../hooks/parallel_request_limiter_v3.py | 54 ++ litellm/proxy/litellm_pre_call_utils.py | 4 + .../management_v1/customers.py | 35 +- .../management_helpers/customer_hierarchy.py | 64 +- litellm/repositories/prisma_protocols.py | 9 + litellm/repositories/table_repositories.py | 8 + litellm/repositories/unit_of_work.py | 66 +- litellm/types/integrations/slack_alerting.py | 3 + .../types/proxy/auth/customer_hierarchy.py | 83 +++ litellm/types/proxy/auth/user_api_key_auth.py | 2 + .../customer_hierarchy.py | 5 +- litellm/types/services.py | 1 + litellm/types/utils.py | 3 + .../src/lib/http/schema.d.ts | 78 +++ .../database/test_spendlogs_ledger_parity.py | 250 +++++++ tests/unit/accounting/test_application.py | 105 +-- tests/unit/accounting/test_projection.py | 83 ++- tests/unit/accounting/test_worker.py | 30 +- .../auth/test_customer_hierarchy_auth.py | 259 +++++++ .../unit/proxy/auth/test_user_api_key_auth.py | 138 ++++ .../common_utils/test_reset_budget_job.py | 13 +- .../test_redis_update_buffer.py | 12 +- .../proxy/db/test_daily_spend_bulk_upsert.py | 47 +- .../proxy/db/test_db_spend_update_writer.py | 157 ++++- .../proxy/db/test_model_access_group_spend.py | 2 + .../hooks/test_dynamic_rate_limiter_v3.py | 32 + .../hooks/test_parallel_request_limiter_v3.py | 122 +++- .../management_v1/test_customers.py | 81 ++- .../test_customer_hierarchy.py | 2 +- tests/unit/types/proxy/auth/__init__.py | 0 .../proxy/auth/test_customer_hierarchy.py | 19 + 53 files changed, 2895 insertions(+), 590 deletions(-) create mode 100644 litellm/proxy/auth/customer_hierarchy_auth.py create mode 100644 litellm/types/proxy/auth/customer_hierarchy.py create mode 100644 tests/unit/proxy/auth/test_customer_hierarchy_auth.py create mode 100644 tests/unit/types/proxy/auth/__init__.py create mode 100644 tests/unit/types/proxy/auth/test_customer_hierarchy.py diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py index 58c060b81c6..c911f4e35ce 100644 --- a/litellm/accounting/application.py +++ b/litellm/accounting/application.py @@ -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) diff --git a/litellm/accounting/delivery_projection.py b/litellm/accounting/delivery_projection.py index da63542d5b8..56b25f0580a 100644 --- a/litellm/accounting/delivery_projection.py +++ b/litellm/accounting/delivery_projection.py @@ -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")), } diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py index 2c11372d7cb..4d7341336ce 100644 --- a/litellm/accounting/projection.py +++ b/litellm/accounting/projection.py @@ -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, ...]: diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py index b86f37c3f5e..86fbf94cb90 100644 --- a/litellm/accounting/snapshot.py +++ b/litellm/accounting/snapshot.py @@ -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", diff --git a/litellm/constants.py b/litellm/constants.py index 7b83f6e2d6c..9efeefacacd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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" diff --git a/litellm/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/litellm/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index a0c896b7e05..589e3e27c47 100644 --- a/litellm/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/litellm/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -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}") diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 2c9d64d50dd..b88a5d75b6e 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -312,6 +312,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_DailyEndUserSpend", "LiteLLM_DailyAgentSpend", "LiteLLM_DailyTeamSpend", + "LiteLLM_DailyCustomerMemberSpend", "LiteLLM_DailyTagSpend", "LiteLLM_ProxyWorkerHeartbeat", "LiteLLM_CronJob", diff --git a/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json b/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json index 59861fcf462..073d31f2e8f 100644 --- a/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json +++ b/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json @@ -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__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__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__failed_requests)", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1f2e6511be7..35821e856d7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 070f867a492..849ecb9a245 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fc818334f9c..7c52728b448 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 ) diff --git a/litellm/proxy/auth/customer_hierarchy_auth.py b/litellm/proxy/auth/customer_hierarchy_auth.py new file mode 100644 index 00000000000..51c6344465b --- /dev/null +++ b/litellm/proxy/auth/customer_hierarchy_auth.py @@ -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) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cad35b669c2..28d4d5efcd5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 = ( diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index 6f3ef4e03d4..32daf68e910 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -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" diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 92960bf9d21..1b55b2cfc3d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -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. diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index 38b132e7e10..784eec03dff 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -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')" ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 6eb7b6373f7..98817dac3f1 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 459f95972ce..146860d2a60 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -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( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index e88b9fc2bfb..cef57f5a92c 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -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 diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index 6b8f5bfa70e..a17232394a8 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -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() diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 1b9851cf082..f90a9f72bf4 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -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: ... diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 003099f8e3a..44fc44c337d 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -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( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index bf6703443c2..b8eeaeab217 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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,)): diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 5ba324e7154..5da902ee4ec 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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=( diff --git a/litellm/proxy/management_endpoints/management_v1/customers.py b/litellm/proxy/management_endpoints/management_v1/customers.py index e568d87d590..eb7e727fe87 100644 --- a/litellm/proxy/management_endpoints/management_v1/customers.py +++ b/litellm/proxy/management_endpoints/management_v1/customers.py @@ -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( diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py index 3ff6c7846e2..0bc5abb9e26 100644 --- a/litellm/proxy/management_helpers/customer_hierarchy.py +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -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: diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 63cfcbc7b45..2c80eca2749 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -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: ... diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 7833f58384d..b1fc5f04b02 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -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" diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 950b98e7b64..afe8ceb1faf 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -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: ... diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index 77322a4850a..5afb12db773 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -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" diff --git a/litellm/types/proxy/auth/customer_hierarchy.py b/litellm/types/proxy/auth/customer_hierarchy.py new file mode 100644 index 00000000000..c131df439f5 --- /dev/null +++ b/litellm/types/proxy/auth/customer_hierarchy.py @@ -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)) diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index a0dd405ac18..4bd562750a3 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -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 diff --git a/litellm/types/proxy/management_endpoints/customer_hierarchy.py b/litellm/types/proxy/management_endpoints/customer_hierarchy.py index db6cc101b2f..0e14525d4b5 100644 --- a/litellm/types/proxy/management_endpoints/customer_hierarchy.py +++ b/litellm/types/proxy/management_endpoints/customer_hierarchy.py @@ -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 diff --git a/litellm/types/services.py b/litellm/types/services.py index e86044c9303..5a329772b83 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -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" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d1f2f343f1c..7ae7bc92b01 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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 diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index e305581093a..d0ec4789a5a 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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 */ diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 1a3fb41ba9c..8dfdced7d1e 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -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() diff --git a/tests/unit/accounting/test_application.py b/tests/unit/accounting/test_application.py index d60893a147d..c338c625be2 100644 --- a/tests/unit/accounting/test_application.py +++ b/tests/unit/accounting/test_application.py @@ -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 diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index 75b3f76de26..2d719c685e6 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -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") diff --git a/tests/unit/accounting/test_worker.py b/tests/unit/accounting/test_worker.py index 372ae90b158..82b3533086c 100644 --- a/tests/unit/accounting/test_worker.py +++ b/tests/unit/accounting/test_worker.py @@ -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 diff --git a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py new file mode 100644 index 00000000000..84f0b86f219 --- /dev/null +++ b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py @@ -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, {}) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index d39c8cb8c85..bef2b08d2a9 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -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 diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 544b6fb2082..e024fd40a53 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -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() diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 65da5630f33..7879adaed93 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -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() diff --git a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py index 1b5dd90205c..f18f2a82601 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -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 diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 36aeacd8476..9a026a5fbd7 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -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) diff --git a/tests/unit/proxy/db/test_model_access_group_spend.py b/tests/unit/proxy/db/test_model_access_group_spend.py index 60ca45761f8..bdd40389ff1 100644 --- a/tests/unit/proxy/db/test_model_access_group_spend.py +++ b/tests/unit/proxy/db/test_model_access_group_spend.py @@ -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, ) diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index 78ed4cf682a..1c71e918cf5 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -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) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d6e1fcf07bb..21edfa00322 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py index f9763ad6855..72f0ba43a63 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -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"]) diff --git a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py index 849ed493ff4..1fba4ab4056 100644 --- a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py +++ b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py @@ -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} diff --git a/tests/unit/types/proxy/auth/__init__.py b/tests/unit/types/proxy/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/types/proxy/auth/test_customer_hierarchy.py b/tests/unit/types/proxy/auth/test_customer_hierarchy.py new file mode 100644 index 00000000000..a6d72cdc3b1 --- /dev/null +++ b/tests/unit/types/proxy/auth/test_customer_hierarchy.py @@ -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