mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
Revert "feat(spend): add daily customer usage and budget resets (#51493)"
This reverts commit ccac47608f.
This commit is contained in:
kerry
parent
c9ea4e20db
commit
7542471e17
27 changed files with 52 additions and 1024 deletions
|
|
@ -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, ...]:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -312,7 +312,6 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
|
|||
"LiteLLM_DailyEndUserSpend",
|
||||
"LiteLLM_DailyAgentSpend",
|
||||
"LiteLLM_DailyTeamSpend",
|
||||
"LiteLLM_DailyCustomerMemberSpend",
|
||||
"LiteLLM_DailyTagSpend",
|
||||
"LiteLLM_ProxyWorkerHeartbeat",
|
||||
"LiteLLM_CronJob",
|
||||
|
|
|
|||
|
|
@ -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_<service>_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_<service>_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_<service>_failed_requests)",
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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')"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
||||
|
|
|
|||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue