From 7542471e172ec4ebb51102eaecc1e96e5ffab79e Mon Sep 17 00:00:00 2001 From: kerry Date: Fri, 13 Nov 2026 01:17:43 +0000 Subject: [PATCH] Revert "feat(spend): add daily customer usage and budget resets (#51493)" This reverts commit ccac47608f66efaf08fa4780b3ef8aded6da9ffd. --- 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 - .../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 - 27 files changed, 52 insertions(+), 1024 deletions(-) delete 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 d757cfed0e4..b860bbc9c8a 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 +from datetime import datetime, timezone 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, utc_date_from_start_time +from litellm.proxy.db.model_usage_rollup import 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 @@ -163,12 +163,7 @@ 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. - 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")}, - } + daily_row: Final = {**row, "organization_id": _fields(metadata.get("accounting_entities")).get("organization_id")} return tuple(_daily_rows(daily_row, shared_key, counters)) @@ -178,21 +173,10 @@ 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)), - *((column, _text(row.get(column))) for column in table.parent_id_columns), - *daily_key, - ), + ((table.entity_id_column, _text(identity)), *shared_key), counters, ( ("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))), @@ -322,7 +306,18 @@ 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") - return utc_date_from_start_time(start) if isinstance(start, (datetime, str)) else None + 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() + ) def _guardrail_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: diff --git a/litellm/constants.py b/litellm/constants.py index 9efeefacacd..7b83f6e2d6c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -496,7 +496,6 @@ 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 b88a5d75b6e..2c9d64d50dd 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -312,7 +312,6 @@ _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 073d31f2e8f..59861fcf462 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,17 +5852,6 @@ "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)", @@ -6063,17 +6052,6 @@ "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)", @@ -6274,17 +6252,6 @@ "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 deleted file mode 100644 index a99e2606d7b..00000000000 --- a/litellm/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_spend/migration.sql +++ /dev/null @@ -1,53 +0,0 @@ -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 95ce61ff7b8..bac27daa5de 100644 --- a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma @@ -1150,49 +1150,6 @@ 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 f749ee18cac..06d77e751c7 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.115" +version = "0.4.114" 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.115" +version = "0.4.114" version_files = [ "pyproject.toml:^version", "../../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d4171887c65..a7555bb89f1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5402,12 +5402,6 @@ 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 c1edf3ae0c4..ae91cbb1f2a 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -17,7 +17,6 @@ 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 @@ -40,8 +39,6 @@ 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, @@ -58,7 +55,6 @@ 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,12 +73,9 @@ from litellm.redis.stores.spend_counters import ( user_counter, ) from litellm.repositories.organization_repository import OrganizationRepository -from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable, TableActions +from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( - CustomerMemberRepository, - CustomerOrganizationRepository, - CustomerTeamRepository, EndUserRepository, ModelAccessGroupBudgetRepository, TagRepository, @@ -90,11 +83,9 @@ 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, ) @@ -103,18 +94,6 @@ 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") @@ -616,128 +595,10 @@ 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 fdd0be4abcc..38b132e7e10 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, TypeAlias +from typing import Final, Literal from pydantic import TypeAdapter -DailySpendEntity: TypeAlias = Literal["user", "team", "org", "tag", "end_user", "agent", "customer_member"] +DailySpendEntity = Literal["user", "team", "org", "tag", "end_user", "agent"] SqlValue = str | int | float | None @@ -32,16 +32,10 @@ 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"), @@ -53,7 +47,6 @@ 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", @@ -66,11 +59,6 @@ _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 () @@ -139,10 +127,7 @@ 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, *table.parent_id_columns, *_KEY_COLUMNS) - ) + return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -191,7 +176,6 @@ def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", table.entity_id_column, - *table.parent_id_columns, *_KEY_COLUMNS, "model_group", *_COUNTER_COLUMNS, @@ -230,7 +214,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, *table.parent_id_columns, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_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 98817dac3f1..8068d570dda 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -45,7 +45,6 @@ from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, - DailyCustomerMemberSpendTransaction, DailyEndUserSpendTransaction, DailyOrganizationSpendTransaction, DailyTagSpendTransaction, @@ -62,7 +61,6 @@ 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, ) @@ -81,7 +79,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, utc_date_from_start_time +from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction 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 @@ -448,7 +446,6 @@ 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]] @@ -771,7 +768,6 @@ 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() @@ -1461,13 +1457,6 @@ 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( @@ -1929,7 +1918,6 @@ 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, ) @@ -1951,7 +1939,6 @@ 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 = { @@ -1985,13 +1972,6 @@ 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 @@ -2079,15 +2059,6 @@ 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( @@ -2124,7 +2095,7 @@ class DBSpendUpdateWriter: async def _flush_daily_spend_queue( self, queue: DailySpendUpdateQueue, - entity_type: Literal["user", "team", "org", "tag", "end_user", "agent", "customer_member"], + entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"], commit: _DailySpendCommit[_DailySpendTransactionT], n_retry_times: int, prisma_client: PrismaClient, @@ -2247,15 +2218,6 @@ 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 = ( @@ -3020,18 +2982,6 @@ 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( @@ -3055,9 +3005,8 @@ class DBSpendUpdateWriter: | dict[str, DailyTagSpendTransaction] | dict[str, DailyOrganizationSpendTransaction] | dict[str, DailyEndUserSpendTransaction] - | dict[str, DailyAgentSpendTransaction] - | dict[str, DailyCustomerMemberSpendTransaction], - entity_type: Literal["user", "team", "org", "tag", "end_user", "agent", "customer_member"], + | dict[str, DailyAgentSpendTransaction], + entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"], entity_id_field: str, ) -> None: """ @@ -3277,24 +3226,6 @@ 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, @@ -3341,7 +3272,7 @@ class DBSpendUpdateWriter: self, payload: dict | SpendLogsPayload, prisma_client: PrismaClient, - type: Literal["user", "team", "org", "request_tags", "end_user", "agent", "customer_member"] = "user", + type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", ) -> BaseDailySpendTransaction | None: entity: Final = "tag" if type == "request_tags" else type identity_payload: Final = ( @@ -3571,38 +3502,6 @@ 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 146860d2a60..be58351b1d9 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -21,7 +21,6 @@ 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, @@ -37,7 +36,6 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( BaseDailySpendTransaction, DailyAgentSpendTransaction, - DailyCustomerMemberSpendTransaction, DailyEndUserSpendTransaction, DailyOrganizationSpendTransaction, DailyTagSpendTransaction, @@ -252,7 +250,6 @@ 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, ): """ @@ -322,9 +319,6 @@ 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 @@ -366,11 +360,6 @@ 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, @@ -416,7 +405,6 @@ 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, @@ -424,7 +412,6 @@ 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 @@ -445,7 +432,6 @@ 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, @@ -453,7 +439,6 @@ 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: """ @@ -464,7 +449,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, Mapping[str, float] | None]]] = [ + entity_entries: Final[list[tuple[Litellm_EntityType, dict[str, float] | None]]] = [ ( Litellm_EntityType.USER, db_spend_update_transactions.get("user_list_transactions"), @@ -534,7 +519,7 @@ class RedisUpdateBuffer: ) ) - daily_pairs: Final[list[tuple[Mapping[str, BaseDailySpendTransaction] | None, DailySpendUpdateQueue]]] = [ + daily_pairs: Final[list[tuple[dict[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), @@ -543,11 +528,10 @@ 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(dict(daily_txns)) + await daily_queue.update_queue.put(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) @@ -561,7 +545,6 @@ 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: @@ -584,7 +567,6 @@ 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 @@ -807,12 +789,11 @@ class RedisUpdateBuffer: dict[str, DailyEndUserSpendTransaction] | None, dict[str, DailyAgentSpendTransaction] | None, tuple[WindowSpendTransaction, ...] | None, - dict[str, DailyCustomerMemberSpendTransaction] | None, ]: """ - Drains the main 8 Redis buffer queues in a single pipeline round-trip. + Drains the main 7 Redis buffer queues in a single pipeline round-trip. - Returns an 8-tuple of parsed results in this order: + Returns a 7-tuple of parsed results in this order: 0: DBSpendUpdateTransactions 1: daily user spend 2: daily team spend @@ -820,10 +801,9 @@ 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, None + return None, None, None, None, None, None, None lpop_list: Final[list[RedisPipelineLpopOperation]] = [ RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), @@ -851,16 +831,12 @@ 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) < 8: + while len(raw_results) < 7: raw_results.append(None) # Slot 0: DBSpendUpdateTransactions @@ -872,7 +848,7 @@ class RedisUpdateBuffer: # Slots 1-5: daily spend categories daily_results: Final[list[dict[str, BaseDailySpendTransaction] | None]] = [] - for slot in (1, 2, 3, 4, 5, 7): + for slot in range(1, 6): slot_result = raw_results[slot] if slot_result is None: daily_results.append(None) @@ -897,9 +873,6 @@ 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 a17232394a8..6b8f5bfa70e 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, timezone +from datetime import datetime from itertools import groupby from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol @@ -97,17 +97,6 @@ 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 8d22af18208..e0be2e1e5ef 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -74,15 +74,6 @@ 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 95ce61ff7b8..bac27daa5de 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1150,49 +1150,6 @@ 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 2c80eca2749..63cfcbc7b45 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -141,15 +141,6 @@ 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 b1fc5f04b02..7833f58384d 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -81,14 +81,6 @@ 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 afe8ceb1faf..950b98e7b64 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, Literal, Protocol, TypeAlias, runtime_checkable +from typing import Final, Protocol, runtime_checkable from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch @@ -165,70 +165,6 @@ 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 95ce61ff7b8..bac27daa5de 100644 --- a/litellm/schema.prisma +++ b/litellm/schema.prisma @@ -1150,49 +1150,6 @@ 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 5a329772b83..e86044c9303 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -40,7 +40,6 @@ 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/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 2b68b71f573..1a3fb41ba9c 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -1547,253 +1547,3 @@ 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 62e4938f489..8c9b543471e 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -688,75 +688,3 @@ 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 5d88efe69df..f0ce610be01 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -12,14 +12,15 @@ 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, @@ -55,7 +56,6 @@ 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,9 +160,6 @@ 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] = [] @@ -200,7 +197,6 @@ 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}) @@ -4281,9 +4277,6 @@ 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 7879adaed93..65da5630f33 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,7 +133,6 @@ 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) @@ -201,7 +200,6 @@ 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 @@ -241,7 +239,6 @@ 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() @@ -295,7 +292,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) == 8 + assert len(result) == 7 ( db_spend, daily_user, @@ -304,7 +301,6 @@ 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. @@ -379,7 +375,6 @@ 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() @@ -457,7 +452,6 @@ 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 @@ -490,7 +484,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, None) + assert result == (None, None, None, None, None, None, None) @pytest.mark.asyncio @@ -728,7 +722,6 @@ 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"] @@ -820,7 +813,6 @@ 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 9557b0c1c75..1b5dd90205c 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -209,23 +209,3 @@ 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 9a026a5fbd7..7ee84e0ffb1 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, None) + return_value=(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, None) + return_value=(db_spend, daily_user, 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, None) + return_value=(db_spend, daily_user, 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, None) + return_value=(db_spend, 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, None) + return_value=(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, None) + return_value=(None, None, None, None, None, None, window_transactions) ) 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, None) + return_value=(None, None, None, None, None, None, window_transactions) ) 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, None) + return_value=(db_spend_transactions, 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, None) + return_value=(None, None, None, None, None, None, window_transactions) ) 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 bdd40389ff1..60ca45761f8 100644 --- a/tests/unit/proxy/db/test_model_access_group_spend.py +++ b/tests/unit/proxy/db/test_model_access_group_spend.py @@ -420,7 +420,6 @@ 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, @@ -428,7 +427,6 @@ 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, )