From ca09ee0f22efec3e78b88e2d65019625a072b992 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 16:25:43 -0800 Subject: [PATCH] refactor(accounting): leave model-budget windows on the existing window writer Model-budget windows are written by the model budget limiter rather than derived from a request's spend, so they stay on main's window queue like live budget counters do. This drops window capture from the accounting payload, the deferred limiter registration and the job-window capture, and the ledger keeps owning key and team budget windows. Co-Authored-By: Claude Opus 5.5 --- litellm/litellm_core_utils/litellm_logging.py | 13 +-- litellm/proxy/db/db_spend_update_writer.py | 18 +--- .../hooks/model_budget_shared_windows.py | 36 +------ .../proxy/hooks/model_max_budget_limiter.py | 67 ++---------- litellm/proxy/proxy_server.py | 2 - litellm/proxy/spend_tracking/spend_event.py | 14 +-- .../database/test_spendlogs_ledger_parity.py | 101 +----------------- .../test_litellm_logging.py | 3 - .../proxy/db/test_db_spend_update_writer.py | 50 --------- .../hooks/test_model_budget_shared_windows.py | 29 ----- .../hooks/test_model_max_budget_limiter.py | 23 ---- .../unit/proxy/proxy_server/test_lifecycle.py | 53 ++------- .../proxy/spend_tracking/test_spend_event.py | 26 ----- 13 files changed, 25 insertions(+), 410 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e2fdaddaf15..1fcb06b80e8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -929,11 +929,7 @@ class Logging(LiteLLMLoggingBaseClass): new_call.start_time = start_time new_call.litellm_params = litellm_params new_call.model_call_details = { - **{ - key: value - for key, value in self.model_call_details.items() - if not key.startswith("_accounting_") and key != "accounting_model_windows" - }, + **{key: value for key, value in self.model_call_details.items() if not key.startswith("_accounting_")}, "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, @@ -954,11 +950,6 @@ class Logging(LiteLLMLoggingBaseClass): """Forget an earlier cache lookup's result, so a call that looks up again reports only its latest.""" self.caching_details = None - def _retain_accounting_model_windows(self, callback_details: Mapping[str, object]) -> None: - captured: Final = callback_details.get("accounting_model_windows") - if captured is not None: - self.model_call_details["accounting_model_windows"] = captured - def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) @@ -4459,8 +4450,6 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, ) - # Redaction gives each callback its own copy; retain captured charges for the later DB callback. - self._retain_accounting_model_windows(model_call_details) if callable(callback): # custom logger functions global customLogger if customLogger is None: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 64df398515b..80f09df6c86 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -700,7 +700,6 @@ class DBSpendUpdateWriter: False only for a batch or video poll whose cost row another poll already wrote, so the caller leaves the key, team, and user counters alone (LIT-7048). """ - from litellm.proxy.hooks.model_budget_shared_windows import captured_model_windows from litellm.proxy.proxy_server import ( accounting_window_keys, disable_spend_logs, @@ -715,8 +714,6 @@ class DBSpendUpdateWriter: ) use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {}) - model_windows: Final = captured_model_windows(request_fields) - windows_queued = False # rebind-ok: recovery must not enqueue a legacy window twice if kwargs is not None: kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker try: @@ -776,10 +773,7 @@ class DBSpendUpdateWriter: model_access_groups=debitable_model_access_groups( get_request_model_access_groups(request_fields), payload["model_id"], llm_router ), - windows=( - *(await accounting_window_keys(hashed_token, team_id) if billable_cost else ()), - *model_windows, - ), + windows=await accounting_window_keys(hashed_token, team_id) if billable_cost else (), router=get_llm_router, ) keep_spend_logs: Final = disable_spend_logs is False and not ( @@ -815,8 +809,6 @@ class DBSpendUpdateWriter: ) if not use_ledger: - windows_queued = True # rebind-ok: an attempted queue write must not be repeated after a later failure - await self._queue_model_windows_safely(model_windows) spawn_logging_task( asyncio.get_running_loop(), self._batch_database_updates( @@ -846,8 +838,6 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("Runs spend update on all tables") return True except Exception: - if not use_ledger and not windows_queued: - await self._queue_model_windows_safely(model_windows) spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " "may not have completed for this request. " @@ -861,12 +851,6 @@ class DBSpendUpdateWriter: ) return True - async def _queue_model_windows_safely(self, windows: tuple[WindowSpendTransaction, ...]) -> None: - try: - await asyncio.gather(*(self.window_spend_update_queue.add_update(window) for window in windows)) - except Exception: - verbose_proxy_logger.debug("Unable to queue shared model budget spend", exc_info=True) - async def _record_spend_log( self, payload: SpendLogsPayload, diff --git a/litellm/proxy/hooks/model_budget_shared_windows.py b/litellm/proxy/hooks/model_budget_shared_windows.py index b6dfbf39c7d..880dd0eb218 100644 --- a/litellm/proxy/hooks/model_budget_shared_windows.py +++ b/litellm/proxy/hooks/model_budget_shared_windows.py @@ -1,6 +1,4 @@ -from collections.abc import Awaitable, Callable, Generator, Mapping -from contextlib import contextmanager -from contextvars import ContextVar +from collections.abc import Awaitable, Callable from dataclasses import dataclass from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, Protocol, TypeAlias @@ -59,41 +57,9 @@ def _proxy_prisma_client() -> "PrismaClient | None": return prisma_client -@dataclass(frozen=True, slots=True) -class CapturedModelWindows: - transactions: tuple[WindowSpendTransaction, ...] - - -_MODEL_WINDOWS: Final[ContextVar[tuple[WindowSpendTransaction, ...] | None]] = ContextVar( - "model_budget_database_windows", default=None -) - - -def captured_model_windows(kwargs: Mapping[str, object] | None) -> tuple[WindowSpendTransaction, ...]: - value: Final = kwargs.get("accounting_model_windows") if kwargs is not None else None - return value.transactions if isinstance(value, CapturedModelWindows) else () - - -def capturing_model_windows() -> bool: - return _MODEL_WINDOWS.get() is not None - - -@contextmanager -def capture_model_windows(enabled: bool) -> Generator[Callable[[], CapturedModelWindows]]: - token: Final = _MODEL_WINDOWS.set(() if enabled else None) - try: - yield lambda: CapturedModelWindows(_MODEL_WINDOWS.get() or ()) - finally: - _MODEL_WINDOWS.reset(token) - - async def _queue_on_proxy_spend_writer(transaction: WindowSpendTransaction) -> None: from litellm.proxy.proxy_server import proxy_logging_obj - captured: Final = _MODEL_WINDOWS.get() - if captured is not None: - _MODEL_WINDOWS.set((*captured, transaction)) - return await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update(transaction) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index c3b382f78de..889fb9fb373 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -34,8 +34,6 @@ from litellm.proxy.hooks.model_budget_shared_windows import ( LocalTTLCache, SharedWindow, SharedWindows, - capture_model_windows, - capturing_model_windows, ) from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds from litellm.types.llms.openai import AllMessageValues @@ -847,13 +845,11 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): clock: Callable[[], float] = time.time, model_access_groups: ModelAccessGroups = proxy_model_access_groups, shared_windows: SharedWindows = PROXY_SHARED_WINDOWS, - defer_window_spend: bool = False, ): self.dual_cache = dual_cache self._clock = clock self._model_access_groups = model_access_groups self._shared_windows = shared_windows - self.defer_window_spend = defer_window_spend self.redis_increment_operation_queue = [] self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() @@ -1221,20 +1217,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return job: Final = job_to_charge_once(kwargs=kwargs, response_obj=response_obj, response_cost=response_cost) - from litellm.proxy.utils import ProxyUpdateSpend - - defer_windows: Final = self.defer_window_spend and not ProxyUpdateSpend.disable_spend_updates() - with capture_model_windows(defer_windows) as captured: - for entity_type, entity_id, resolved in resolved_budgets: - await self._charge_entity( - entity_type=entity_type, - entity_id=entity_id, - resolved=resolved, - response_cost=response_cost, - job=job, - ) - if defer_windows: - kwargs["accounting_model_windows"] = captured() # rebind-ok: forward windows to DB callback + for entity_type, entity_id, resolved in resolved_budgets: + await self._charge_entity( + entity_type=entity_type, + entity_id=entity_id, + resolved=resolved, + response_cost=response_cost, + job=job, + ) if self.dual_cache.redis_cache is not None: await self._push_in_memory_increments_to_redis() @@ -1266,14 +1256,6 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): job=job, ttl_seconds=duration_in_seconds(budget_duration), ): - if capturing_model_windows() and self._shares_windows_in_database(entity_type): - await self._capture_job_window( - entity_type=entity_type, - entity_id=entity_id, - resolved=resolved, - spend_key=spend_key, - response_cost=response_cost, - ) if job.repeat_poll_cost > 0: await self._charge_entity( entity_type=entity_type, @@ -1308,39 +1290,6 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): response_cost=response_cost, ) - async def _capture_job_window( - self, - entity_type: Litellm_EntityType, - entity_id: str | None, - resolved: ResolvedModelBudget, - spend_key: str, - response_cost: float, - ) -> None: - """The ledger keeps only the poll whose delivery claims the job, which need not be the poll that won this - limiter's charge, so every completed poll carries the job's window and the ledger applies it once.""" - budget_duration: Final = resolved.budget_config.budget_duration - if budget_duration is None: - return - start_time_key: Final = model_budget_start_time_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=resolved.budget_model, - budget_duration=budget_duration, - ) - if _local_windows(self.dual_cache).get_cache(key=_unshared_spend_key(start_time_key)) is not None: - return - now: Final = self._clock() - await self._shared_windows.add( - spend_key=spend_key, - budget_duration=budget_duration, - started_at=await self._get_or_set_budget_start_time( - start_time_key=start_time_key, - current_time=now, - ttl_seconds=budget_window_seconds(now, budget_duration), - ), - spend=response_cost, - ) - def _shares_windows_in_database(self, entity_type: Litellm_EntityType) -> bool: """End-user ids are whatever the client sends, so their windows stay on each replica rather than adding a database row per id.""" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6ce95b86692..d7e0dbefd39 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3313,8 +3313,6 @@ def cost_tracking(): from litellm.integrations.shadow_eval_logger import ShadowEvalLogger spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event) - model_max_budget_limiter.defer_window_spend = True - litellm.logging_callback_manager.add_litellm_async_success_callback(model_max_budget_limiter) db_logger: Final = ProxyDBLogger(spend_event_producer) litellm.logging_callback_manager.add_litellm_callback(db_logger) litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger) diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py index 947591e2be3..d054519cb9c 100644 --- a/litellm/proxy/spend_tracking/spend_event.py +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -24,9 +24,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction from litellm.proxy.db.spend_log_tool_index import response_tool_call_names -from litellm.proxy.hooks.model_budget_shared_windows import CapturedModelWindows, captured_model_windows from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import LiteLLMBatch, Usage @@ -155,7 +153,6 @@ class SpendCallbackKwargs(TypedDict): realtime_tool_calls: ReadOnly[Sequence[object] | None] tools: ReadOnly[Sequence[_FunctionTool] | None] complete_streaming_response: NotRequired[ReadOnly[CompactResponse | None]] - accounting_model_windows: NotRequired[ReadOnly[CapturedModelWindows]] _NO_LITELLM_PARAMS: Final[_LitellmParams] = {} @@ -189,7 +186,6 @@ class SpendEvent(LiteLLMBaseModel): realtime_tools: Sequence[object] | None realtime_tool_calls: Sequence[object] | None request_tool_names: tuple[str, ...] - accounting_model_windows: tuple[WindowSpendTransaction, ...] = () @dataclass(frozen=True, slots=True) @@ -384,15 +380,8 @@ def build_spend_event( realtime_tools=kwargs.get("realtime_tools"), realtime_tool_calls=kwargs.get("realtime_tool_calls"), request_tool_names=_request_tool_names(kwargs), - accounting_model_windows=captured_model_windows(raw_kwargs), - ) - return ( - event.model_dump_json( - fallback=_json_fallback, - exclude={"accounting_model_windows"} if not event.accounting_model_windows else None, - ).encode() - + b"\n" ) + return event.model_dump_json(fallback=_json_fallback).encode() + b"\n" def decode_spend_event(line: bytes) -> SpendEvent | SpendEventDecodeError: @@ -407,7 +396,6 @@ def spend_event_callback_args(event: SpendEvent) -> SpendEventCallbackArgs: tools: Final = tuple(_function_tool(name) for name in event.request_tool_names) kwargs: Final[SpendCallbackKwargs] = { "litellm_call_id": event.litellm_call_id, - "accounting_model_windows": CapturedModelWindows(event.accounting_model_windows), "call_type": event.call_type, "model": event.model, "custom_llm_provider": event.custom_llm_provider, diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 88cf159de59..e1894aa66bf 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -33,8 +33,6 @@ from litellm.repositories.accounting_delivery_repository import AccountingDelive from litellm.repositories.table_repositories import SpendLogsRepository from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.types.utils import ModelResponse, Usage -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.litellm_logging import Logging from tests.integration._support.database import scratch_database _ROOT: Final = Path(__file__).resolve().parents[3] @@ -65,15 +63,6 @@ _DAILY_TABLES: Final = ( ) -class _WindowConsumer(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.received = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.received = kwargs - - @pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), ("200", 251, 200, False), (None, 10001, 10000, False), (None, 3001, 3001, True)]) @pytest.mark.asyncio async def test_worker_applies_the_configured_backlog_in_one_transaction( @@ -114,89 +103,6 @@ async def test_worker_applies_the_configured_backlog_in_one_transaction( await client.db.disconnect() -class _FailingWindowQueue(asyncio.Queue): - def __init__(self, after_enqueue: bool): - super().__init__() - self.after_enqueue = after_enqueue - - async def put(self, item): - if self.after_enqueue: - await super().put(item) - raise RuntimeError("Window queue unavailable") - - -@pytest.mark.parametrize("ledger,payload_failure,queue_failure", [ - (False, False, None), (False, True, None), (True, False, None), (True, True, None), - (False, False, "before"), (False, False, "after"), -]) -@pytest.mark.parametrize("redaction", ["excluded_fields", "messages"]) -@pytest.mark.asyncio -async def test_model_window_charge_survives_redacted_callback_copies( - parity_database: str, monkeypatch: pytest.MonkeyPatch, ledger: bool, redaction: str, payload_failure: bool, queue_failure: str | None -): - from litellm.accounting.worker import flush_accounting_deliveries - from litellm.proxy.hooks.model_budget_shared_windows import captured_model_windows - - monkeypatch.setenv("DATABASE_URL", parity_database) - monkeypatch.setattr(litellm, "_async_success_callback", []) - monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["messages"] if redaction == "excluded_fields" else []) - monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger}) - monkeypatch.setattr(proxy_server, "disable_spend_logs", False) - client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) - client.spend_log_transactions = [] - client.model_usage_transactions = [] - monkeypatch.setattr(proxy_server, "prisma_client", client) - writer = DBSpendUpdateWriter() - if queue_failure is not None: - writer.window_spend_update_queue.update_queue = _FailingWindowQueue(queue_failure == "after") - consumer = _WindowConsumer() - limiter = PROXY_VirtualKeyModelMaxBudgetLimiter( - dual_cache=DualCache(), clock=lambda: _NOW.timestamp(), - shared_windows=DatabaseSharedWindows(prisma_client_source=lambda: client), defer_window_spend=True, - ) - limiter.turn_off_message_logging = redaction == "messages" - consumer.turn_off_message_logging = redaction == "messages" - logging = Logging(model="ledger-test", messages=[], stream=False, call_type="acompletion", - start_time=_NOW, litellm_call_id="redacted-window", function_id="redacted-window") - logging.update_environment_variables( - optional_params={}, custom_llm_provider="openai", - litellm_params={"metadata": {"user_api_key": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", "user_api_key_model_max_budget": { - "ledger-test": {"max_budget": 100, "budget_duration": "1d"} - }}}, - ) - logging.model_call_details["response_cost"] = 0.25 - logging.dynamic_async_success_callbacks = [limiter, consumer] - response = ModelResponse(id="redacted-window", model="ledger-test", usage=Usage(prompt_tokens=10, completion_tokens=5)) - await client.db.connect() - try: - await logging.async_success_handler(response, _NOW, _NOW, standard_logging_object={ - "model": "ledger-test", "model_group": "ledger-test", "response_cost": 0.25, - "metadata": {"user_api_key_hash": "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"}, - "model_map_information": {}, "hidden_params": {}, "request_tags": [], - }) - assert consumer.received is not None - windows = captured_model_windows(consumer.received) - assert len(windows) == 1 - assert windows[0]["spend"] == 0.25 - if payload_failure: - consumer.received["standard_logging_object"].pop("model_map_information") - assert await writer.update_database( - token="2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683", user_id="user", end_user_id=None, team_id=None, org_id=None, - kwargs=consumer.received, completion_response=response, start_time=_NOW, end_time=_NOW, response_cost=0.25, - ) - await drain_logging_tasks() - await SpendLogsRepository(client).insert_with_deliveries(tuple(client.spend_log_transactions)) - if ledger: - await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) - else: - await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) - with psycopg.connect(parity_database) as connection: - assert connection.execute('SELECT spend FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type=\'model_max_budget\'').fetchall() == [(0.0 if (payload_failure and ledger) or queue_failure == "before" else 0.25,)] - assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0.0 if payload_failure else 0.25,) - finally: - await client.db.disconnect() - - @pytest.mark.parametrize("ledger", [False, True]) @pytest.mark.parametrize("exclude_messages", [False, True]) @pytest.mark.asyncio @@ -247,8 +153,7 @@ async def test_registered_proxy_callbacks_charge_model_windows_before_persisting await SpendLogsRepository(client).insert_with_deliveries(tuple(client.spend_log_transactions)) if ledger: await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) - else: - await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) with psycopg.connect(parity_database) as connection: assert connection.execute('SELECT spend FROM "LiteLLM_BudgetWindowSpend" WHERE entity_type=\'model_max_budget\'').fetchall() == [(0.25,)] assert connection.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [(0.25,)] @@ -386,7 +291,7 @@ async def _assert_legacy_database_parity( await client.db.connect() results = [] try: - for path in ("legacy", "deferred_legacy", "ledger"): + for path in ("legacy", "ledger"): _reset(parity_database) if reset_kind is not None: _prime_reset(parity_database, reset_kind) @@ -466,7 +371,6 @@ async def _assert_legacy_database_parity( dual_cache=proxy_server.spend_counter_cache, clock=lambda: _NOW.timestamp(), shared_windows=DatabaseSharedWindows(prisma_client_source=lambda: client), - defer_window_spend=path != "legacy", ) if status == "success": await limiter.async_log_success_event(request_kwargs, None, _NOW, _NOW) @@ -570,6 +474,7 @@ async def _assert_legacy_database_parity( connection.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') result = await AccountingDeliveryRepository(client.writer_db).apply_pending(apply) assert result.applied == 1, result + await writer._commit_spend_updates_to_db_without_redis_buffer(client, 0, proxy_server.proxy_logging_obj) if correction_order == "after": await _publish_savings_correction(client, stored_rows[0]) results.append(_snapshot(parity_database)) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 82d465b8c64..a014d1d352d 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14197,13 +14197,10 @@ def test_copy_for_new_call_logs_a_turn_under_its_own_call_id_against_the_connect "metadata": {"user_api_key": "hashed-key", "user_api_key_budget_reservation": reservation}, } connection.model_call_details["litellm_params"] = connection.litellm_params - connection.model_call_details["accounting_model_windows"] = object() turn_start: Final = datetime.datetime(2026, 1, 1, 12, 0, 0) turn: Final = connection.copy_for_new_call(litellm_call_id="turn-1", start_time=turn_start) - assert "accounting_model_windows" not in turn.model_call_details - assert (turn.litellm_call_id, turn.start_time) == ("turn-1", turn_start) assert (turn.model_call_details["litellm_call_id"], turn.model_call_details["start_time"]) == ("turn-1", turn_start) assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1" 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 64109c06b37..d90c7f6d21d 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -44,56 +44,6 @@ from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpush from tests.unit.proxy.db.fake_prisma_engine import engine_call -class _UnavailableWindowQueue(asyncio.Queue): - async def put(self, item): - raise RuntimeError("Window queue unavailable") - - -@pytest.mark.asyncio -async def test_window_recovery_failure_does_not_fail_the_request(monkeypatch): - from litellm.proxy import proxy_server - from litellm.proxy.hooks.model_budget_shared_windows import CapturedModelWindows - - monkeypatch.setattr(proxy_server, "general_settings", {}) - writer = DBSpendUpdateWriter() - writer.window_spend_update_queue.update_queue = _UnavailableWindowQueue() - now = datetime(2026, 11, 2, tzinfo=timezone.utc) - window = build_window_spend_transaction( - entity_type="model_max_budget", entity_id="key:model:1d", window_duration="1d", - window_start=now, spend=0.25, - ) - assert await writer.update_database( - token="key", user_id=None, end_user_id=None, team_id=None, org_id=None, - kwargs={"accounting_model_windows": CapturedModelWindows((window,)), "standard_logging_object": {}}, - completion_response=None, start_time=now, end_time=now, response_cost=0.25, - ) is True - assert writer.window_spend_update_queue.update_queue.empty() - - -@pytest.mark.parametrize(("use_ledger", "expected"), [(False, (0.25,)), (True, ())]) -@pytest.mark.asyncio -async def test_failed_spend_write_queues_model_windows_only_for_legacy_requests(monkeypatch, use_ledger, expected): - from litellm.proxy import proxy_server - from litellm.proxy.hooks.model_budget_shared_windows import CapturedModelWindows - - monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": use_ledger}) - monkeypatch.setattr(proxy_server, "accounting_window_keys", AsyncMock(return_value=())) - writer = DBSpendUpdateWriter() - monkeypatch.setattr(writer, "_record_spend_log", AsyncMock(side_effect=RuntimeError("spend log write failed"))) - now = datetime(2026, 11, 2, tzinfo=timezone.utc) - window = build_window_spend_transaction( - entity_type="model_max_budget", entity_id="key:model:1d", window_duration="1d", window_start=now, spend=0.25, - ) - assert await writer.update_database( - token="key", user_id=None, end_user_id=None, team_id=None, org_id=None, - kwargs={"accounting_model_windows": CapturedModelWindows((window,)), "standard_logging_object": {}}, - completion_response=None, start_time=now, end_time=now, response_cost=0.25, - ) is True - charges = await writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() - assert tuple(charge["spend"] for charge in charges) == expected - - - @pytest.mark.asyncio async def test_daily_spend_tracking_with_disabled_spend_logs(): """ diff --git a/tests/unit/proxy/hooks/test_model_budget_shared_windows.py b/tests/unit/proxy/hooks/test_model_budget_shared_windows.py index ba55012de19..7002f602fca 100644 --- a/tests/unit/proxy/hooks/test_model_budget_shared_windows.py +++ b/tests/unit/proxy/hooks/test_model_budget_shared_windows.py @@ -197,32 +197,3 @@ async def test_open_reports_no_shared_window_while_the_database_is_unreachable() windows: Final = _database_windows(_Table(error=RuntimeError("database is down")), _Sink()) assert await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp()) is None - - -@pytest.mark.asyncio -async def test_captured_database_windows_are_isolated_between_requests(monkeypatch: pytest.MonkeyPatch): - import asyncio - - from litellm.proxy import proxy_server - from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter - from litellm.proxy.hooks.model_budget_shared_windows import capture_model_windows - from litellm.proxy.utils import PrismaClient - - client = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_server.proxy_logging_obj) - writer = DBSpendUpdateWriter() - monkeypatch.setattr(proxy_server, "prisma_client", client) - monkeypatch.setattr(proxy_server.proxy_logging_obj, "db_spend_update_writer", writer) - windows = DatabaseSharedWindows() - - async def charge(key: str): - with capture_model_windows(True) as captured: - await windows.add(key, "1d", WINDOW_START.timestamp(), 0.25) - await asyncio.sleep(0) - return captured().transactions - - first, second = await asyncio.gather(charge("first"), charge("second")) - assert tuple(row["entity_id"] for row in first) == ("first",) - assert tuple(row["entity_id"] for row in second) == ("second",) - await windows.add("legacy", "1d", WINDOW_START.timestamp(), 0.5) - queued = await writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() - assert tuple((row["entity_id"], row["spend"]) for row in queued) == (("legacy", 0.5),) diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index 9384a3f398b..f1fd02981d6 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -945,26 +945,3 @@ async def test_concurrent_first_charges_on_one_replica_both_count(): shared_windows=_WindowSpendDatabase().replica(), ) assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4) - - -@pytest.mark.asyncio -async def test_every_completed_poll_of_a_job_carries_its_window_for_the_ledger(): - from litellm.proxy.hooks.model_budget_shared_windows import captured_model_windows - - database: Final = _WindowSpendDatabase() - limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter( - dual_cache=DualCache(), - shared_windows=DatabaseSharedWindows( - prisma_client_source=lambda: SimpleNamespace(db=database, writer_db=database), # pyright: ignore[reportArgumentType] # duck-typed prisma double - ), - defer_window_spend=True, - ) - polls: Final = tuple({**_event("avideo_status", BATCH_COST), POLLED_VIDEO_STATUS_KWARG: "completed"} for _ in "ab") - for poll in polls: - await _log_success(limiter, poll, response_obj=VideoObject(id="video_first", object="video", status="completed")) - - carried: Final = tuple( - sorted((window["entity_id"], window["spend"]) for window in captured_model_windows(poll)) for poll in polls - ) - assert carried[0] and carried[0] == carried[1] - assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 9d75c3e2c6f..db5be07fb6f 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -721,37 +721,20 @@ def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch) @pytest.mark.asyncio @pytest.mark.parametrize("setup_count", [1, 2]) -async def test_cost_tracking_publishes_captured_window_charges_once_after_repeated_setup( +async def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set( monkeypatch: pytest.MonkeyPatch, setup_count: int ) -> None: - from datetime import datetime, timezone - from types import SimpleNamespace + from datetime import datetime import litellm - from litellm.caching.dual_cache import DualCache - from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows - from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter - from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import build_window_spend_transaction + from litellm.integrations.shadow_eval_logger import ShadowEvalLogger from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer from litellm.types.utils import ModelResponse, Usage from litellm.utils import Rules, function_setup - now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - key_hash: Final = "a" * 64 producer: Final = AsyncMock(spec=SpendEventProducer) - prisma: Final = MagicMock() - prisma.writer_db = prisma.db - prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock( - return_value=SimpleNamespace(window_start=now, spend=0.0) - ) - monkeypatch.setattr(ps, "prisma_client", prisma) - limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter( - dual_cache=DualCache(), - clock=lambda: now.timestamp(), - shared_windows=DatabaseSharedWindows(prisma_client_source=lambda: prisma), - ) - monkeypatch.setattr(ps, "model_max_budget_limiter", limiter) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) monkeypatch.setattr(ps, "spend_event_producer", None) monkeypatch.setattr(ps, "build_spend_event_producer", lambda settings, fallback: producer) for name in ( @@ -764,10 +747,13 @@ async def test_cost_tracking_publishes_captured_window_charges_once_after_repeat ): monkeypatch.setattr(litellm, name, []) - litellm.logging_callback_manager.add_litellm_callback(limiter) for _ in range(setup_count): cost_tracking() + assert len(litellm.callbacks) == 2 + assert sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks) == 1 + assert len(litellm._async_success_callback) == 1 + now: Final = datetime(2026, 1, 1) logging_obj, _ = function_setup( "acompletion", Rules(), @@ -776,17 +762,7 @@ async def test_cost_tracking_publishes_captured_window_charges_once_after_repeat messages=[{"role": "user", "content": "hello"}], litellm_call_id="single-accounting-event", ) - logging_obj.update_environment_variables( - litellm_params={ - "acompletion": True, - "metadata": { - "user_api_key": key_hash, - "user_api_key_model_max_budget": {"*": {"max_budget": 100, "budget_duration": "1d"}}, - }, - }, - optional_params={}, - ) - logging_obj.model_call_details["response_cost"] = 0.25 + logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={}) response: Final = ModelResponse( id="chatcmpl-single-accounting-event", model="gpt-4o-mini", @@ -797,16 +773,7 @@ async def test_cost_tracking_publishes_captured_window_charges_once_after_repeat event: Final = decode_spend_event(producer.publish.call_args.args[0]) assert not isinstance(event, SpendEventDecodeError) assert event.response_cost == logging_obj.model_call_details["response_cost"] - assert event.response_cost == 0.25 - assert event.accounting_model_windows == ( - build_window_spend_transaction( - entity_type="model_max_budget", - entity_id=f"virtual_key_spend:{key_hash}:*:1d", - window_duration="1d", - window_start=now, - spend=0.25, - ), - ) + assert event.response_cost > 0 def test_cost_tracking_no_op_when_prisma_missing(monkeypatch): diff --git a/tests/unit/proxy/spend_tracking/test_spend_event.py b/tests/unit/proxy/spend_tracking/test_spend_event.py index a27accfbdf0..e1912d36d91 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_event.py +++ b/tests/unit/proxy/spend_tracking/test_spend_event.py @@ -214,29 +214,3 @@ def test_batch_retrieves_and_video_status_polls_stay_in_process(): ) is False ) - - -def test_collector_round_trip_preserves_captured_model_budget_windows(): - from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import build_window_spend_transaction - from litellm.proxy.hooks.model_budget_shared_windows import CapturedModelWindows, captured_model_windows - - transaction = build_window_spend_transaction( - entity_type="model_max_budget", - entity_id="model-window", - window_duration="1d", - window_start=datetime(2026, 1, 1), - spend=0.125, - ) - kwargs = {**_success_kwargs(), "accounting_model_windows": CapturedModelWindows((transaction,))} - event = decode_spend_event(_build(kwargs, _response())) - assert not isinstance(event, SpendEventDecodeError) - assert captured_model_windows(spend_event_callback_args(event).kwargs) == (transaction,) - - -def test_collector_does_not_trust_caller_supplied_model_window_dicts(): - from litellm.proxy.hooks.model_budget_shared_windows import captured_model_windows - - kwargs = {**_success_kwargs(), "accounting_model_windows": {"transactions": [{"spend": -100}]}} - event = decode_spend_event(_build(kwargs, _response())) - assert not isinstance(event, SpendEventDecodeError) - assert captured_model_windows(spend_event_callback_args(event).kwargs) == ()