From ccac47608f66efaf08fa4780b3ef8aded6da9ffd Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 13 Nov 2026 09:01:18 +0800 Subject: [PATCH] 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 --- litellm/accounting/projection.py | 37 +-- litellm/constants.py | 1 + litellm/integrations/otel/model/spans.py | 1 + .../grafana_dashboard.json | 33 +++ .../migration.sql | 53 ++++ .../litellm_proxy_extras/schema.prisma | 43 +++ litellm/proxy-extras/pyproject.toml | 4 +- litellm/proxy/_types.py | 6 + .../proxy/common_utils/reset_budget_job.py | 141 +++++++++- litellm/proxy/db/daily_spend_bulk_upsert.py | 24 +- litellm/proxy/db/db_spend_update_writer.py | 111 +++++++- .../redis_update_buffer.py | 43 ++- litellm/proxy/db/model_usage_rollup.py | 13 +- litellm/proxy/db/prisma_client.py | 9 + litellm/proxy/schema.prisma | 43 +++ litellm/repositories/prisma_protocols.py | 9 + litellm/repositories/table_repositories.py | 8 + litellm/repositories/unit_of_work.py | 66 ++++- litellm/schema.prisma | 43 +++ litellm/types/services.py | 1 + pyproject.toml | 2 +- .../database/test_spendlogs_ledger_parity.py | 250 ++++++++++++++++++ tests/unit/accounting/test_projection.py | 72 +++++ .../common_utils/test_reset_budget_job.py | 13 +- .../test_redis_update_buffer.py | 12 +- .../proxy/db/test_daily_spend_bulk_upsert.py | 20 ++ .../proxy/db/test_db_spend_update_writer.py | 18 +- .../proxy/db/test_model_access_group_spend.py | 2 + uv.lock | 2 +- 29 files changed, 1026 insertions(+), 54 deletions(-) create mode 100644 litellm/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_spend/migration.sql diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py index 57cc128a66e..b1d7d9e7357 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 @@ -16,7 +16,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import GuardrailCostByUnitEntry 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 @@ -162,7 +162,12 @@ 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")) + daily_row: Final = { + **row, + "organization_id": entities.get("organization_id"), + **{key: entities.get(key) for key in ("customer_member_id", "customer_team_id", "customer_org_id")}, + } return tuple(_daily_rows(daily_row, shared_key, counters)) @@ -172,10 +177,21 @@ 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 + ) 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, + ), counters, ( ("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))), @@ -262,18 +278,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/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/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index d1ed9f7b62b..5a3fbc93a33 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -311,6 +311,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/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_spend/migration.sql b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_spend/migration.sql new file mode 100644 index 00000000000..a99e2606d7b --- /dev/null +++ b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_spend/migration.sql @@ -0,0 +1,53 @@ +SET LOCAL lock_timeout = '10s'; +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend" ( + "id" TEXT NOT NULL, + "customer_member_id" TEXT NOT NULL, + "customer_team_id" TEXT NOT NULL, + "customer_org_id" TEXT NOT NULL, + "date" TEXT NOT NULL, + "api_key" TEXT NOT NULL, + "model" TEXT, + "model_group" TEXT, + "custom_llm_provider" TEXT, + "mcp_namespaced_tool_name" TEXT, + "endpoint" TEXT, + "prompt_tokens" BIGINT NOT NULL DEFAULT 0, + "completion_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0, + "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0, + "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "api_requests" BIGINT NOT NULL DEFAULT 0, + "successful_requests" BIGINT NOT NULL DEFAULT 0, + "failed_requests" BIGINT NOT NULL DEFAULT 0, + "total_response_time_ms" BIGINT NOT NULL DEFAULT 0, + "timed_requests" BIGINT NOT NULL DEFAULT 0, + "ptu_flat_cost" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyCustomerMemberSpend_pkey" PRIMARY KEY ("id") +); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_date_idx" ON "LiteLLM_DailyCustomerMemberSpend"("date"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_customer_member_id_date_idx" ON "LiteLLM_DailyCustomerMemberSpend"("customer_member_id", "date"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_customer_team_id_date_idx" ON "LiteLLM_DailyCustomerMemberSpend"("customer_team_id", "date"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_customer_org_id_date_idx" ON "LiteLLM_DailyCustomerMemberSpend"("customer_org_id", "date"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_api_key_idx" ON "LiteLLM_DailyCustomerMemberSpend"("api_key"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_model_idx" ON "LiteLLM_DailyCustomerMemberSpend"("model"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyCustomerMemberSpend"("mcp_namespaced_tool_name"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_endpoint_idx" ON "LiteLLM_DailyCustomerMemberSpend"("endpoint"); + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_identity_key" ON "LiteLLM_DailyCustomerMemberSpend"("customer_member_id", "customer_team_id", "customer_org_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma index 2c0246bb94e..9fd5b3238a2 100644 --- a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma @@ -1111,6 +1111,49 @@ model LiteLLM_DailyAgentSpend { } // Track daily team spend metrics per model and key +model LiteLLM_DailyCustomerMemberSpend { + id String @id @default(uuid()) + customer_member_id String + customer_team_id String + customer_org_id String + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([customer_member_id, customer_team_id, customer_org_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyCustomerMemberSpend_identity_key") + @@index([date]) + @@index([customer_member_id, date]) + @@index([customer_team_id, date]) + @@index([customer_org_id, date]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) + @@index([endpoint]) +} + model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) team_id String? diff --git a/litellm/proxy-extras/pyproject.toml b/litellm/proxy-extras/pyproject.toml index 06d77e751c7..f749ee18cac 100644 --- a/litellm/proxy-extras/pyproject.toml +++ b/litellm/proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.114" +version = "0.4.115" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -30,7 +30,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.114" +version = "0.4.115" version_files = [ "pyproject.toml:^version", "../../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fbbb41a870c..23016445c50 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5398,6 +5398,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 diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index ae91cbb1f2a..c1edf3ae0c4 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 @@ -39,6 +40,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, @@ -55,6 +58,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 ( @@ -73,9 +77,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, @@ -83,9 +90,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, ) @@ -94,6 +103,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_template: BudgetRecord | None = None + budget_override_expires_at: datetime | None = None + + _RowT = TypeVar("_RowT") @@ -595,10 +616,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_duration, + "OR": [ + {"budget_override_expires_at": None}, + {"budget_override_expires_at": {"gt": now}}, + ], + }, + {"budget_template": {"is": {"budget_duration": {"not": None}}}}, + ) + 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_template": 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_template, row.budget, 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..fdd0be4abcc 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -12,11 +12,11 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from itertools import groupby 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 +32,16 @@ class DailySpendTable: name: str entity_id_column: str carries_request_id: bool = False + parent_id_columns: tuple[str, ...] = () 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"), + ), "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 +53,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 +66,11 @@ _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 all( + isinstance(payload.get(column), str) and payload.get(column) + for column in ("customer_member_id", "customer_team_id", "customer_org_id") + ): + return () key: Final = _ENTITY_INPUT_KEYS[entity] if key not in payload: return () @@ -127,7 +139,10 @@ def _as_float(value: object) -> float: 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 (table.entity_id_column, *table.parent_id_columns, *_KEY_COLUMNS) + ) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -176,6 +191,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 +230,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((table.entity_id_column, *table.parent_id_columns, *_KEY_COLUMNS))}) 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 8068d570dda..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 @@ -446,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]] @@ -768,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() @@ -1457,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( @@ -1918,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, ) @@ -1939,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 = { @@ -1972,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 @@ -2059,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( @@ -2095,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, @@ -2218,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 = ( @@ -2982,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( @@ -3005,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: """ @@ -3226,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, @@ -3272,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 = ( @@ -3502,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 be58351b1d9..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, @@ -250,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, ): """ @@ -319,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 @@ -360,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, @@ -405,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, @@ -412,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 @@ -432,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, @@ -439,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: """ @@ -449,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"), @@ -519,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), @@ -528,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) @@ -545,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: @@ -567,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 @@ -789,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 @@ -801,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), @@ -831,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 @@ -848,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) @@ -873,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( 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 e0be2e1e5ef..8d22af18208 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/schema.prisma b/litellm/proxy/schema.prisma index 2c0246bb94e..9fd5b3238a2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1111,6 +1111,49 @@ model LiteLLM_DailyAgentSpend { } // Track daily team spend metrics per model and key +model LiteLLM_DailyCustomerMemberSpend { + id String @id @default(uuid()) + customer_member_id String + customer_team_id String + customer_org_id String + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([customer_member_id, customer_team_id, customer_org_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyCustomerMemberSpend_identity_key") + @@index([date]) + @@index([customer_member_id, date]) + @@index([customer_team_id, date]) + @@index([customer_org_id, date]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) + @@index([endpoint]) +} + model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) team_id String? 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/schema.prisma b/litellm/schema.prisma index 2c0246bb94e..9fd5b3238a2 100644 --- a/litellm/schema.prisma +++ b/litellm/schema.prisma @@ -1111,6 +1111,49 @@ model LiteLLM_DailyAgentSpend { } // Track daily team spend metrics per model and key +model LiteLLM_DailyCustomerMemberSpend { + id String @id @default(uuid()) + customer_member_id String + customer_team_id String + customer_org_id String + date String + api_key String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([customer_member_id, customer_team_id, customer_org_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint], map: "LiteLLM_DailyCustomerMemberSpend_identity_key") + @@index([date]) + @@index([customer_member_id, date]) + @@index([customer_team_id, date]) + @@index([customer_org_id, date]) + @@index([api_key]) + @@index([model]) + @@index([mcp_namespaced_tool_name]) + @@index([endpoint]) +} + model LiteLLM_DailyTeamSpend { id String @id @default(uuid()) team_id String? 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/pyproject.toml b/pyproject.toml index 39c760ef943..ccfbd5d7d9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,7 @@ proxy = [ "mcp>=2.3.0,<3", "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3", - "litellm-proxy-extras==0.4.114", + "litellm-proxy-extras==0.4.115", "litellm-enterprise==0.1.83", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 1a3fb41ba9c..2b68b71f573 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_template_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_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_projection.py b/tests/unit/accounting/test_projection.py index ecbb17a4f7e..edf0fc4110d 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -588,3 +588,75 @@ 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/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index f0ce610be01..5d88efe69df 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, @@ -56,6 +55,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 @@ -160,6 +160,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] = [] @@ -197,6 +200,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}) @@ -4277,6 +4281,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..9557b0c1c75 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -209,3 +209,23 @@ 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 ("customer_member_id", "customer_team_id", "customer_org_id",' in sql + assert all(identity in params for identity in ("member", "team-before", "org-before", "team-after", "org-after")) 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 7ee84e0ffb1..9a026a5fbd7 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -2607,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 @@ -2688,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 @@ -2743,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 @@ -2791,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 @@ -3424,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() @@ -3454,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() @@ -3712,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 @@ -3747,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 @@ -3830,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 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/uv.lock b/uv.lock index 01516d54939..6dc41480bf5 100644 --- a/uv.lock +++ b/uv.lock @@ -5581,7 +5581,7 @@ source = { editable = "litellm/packaging/litellm-proxy" } [[package]] name = "litellm-proxy-extras" -version = "0.4.114" +version = "0.4.115" source = { editable = "litellm/proxy-extras" } dependencies = [ { name = "psycopg" },