mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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 <noreply@anthropic.com>
This commit is contained in:
Joshua Valluru
parent
8564c9bd63
commit
ca09ee0f22
13 changed files with 25 additions and 410 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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),)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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) == ()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue