From 8ec813d2f28c8caa335a5f4ba4f29db0b5fc22b7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 21:34:33 -0800 Subject: [PATCH 01/10] fix(proxy): charge a video's provider-reported cost once, on the first poll that sees it completed A custom video provider can report a video's cost on the status poll that sees it completed. Every poll of a video shared the video id as its spend log request_id, so the paid poll's row was dropped as a duplicate and every later poll of the finished video charged the cost again. Video status rows now carry their own `_video_cost` request_id, and a poll that reports a cost goes through the same settle-once claim batches use, so the first poll that sees the video completed charges it and writes the row, and every later poll charges nothing. --- litellm/proxy/db/db_spend_update_writer.py | 73 ++++++++------ .../spend_tracking/spend_tracking_utils.py | 4 + litellm/types/videos/utils.py | 5 +- .../proxy/db/test_db_spend_update_writer.py | 98 ++++++++++++++++++- .../test_spend_tracking_utils.py | 12 +++ 5 files changed, 157 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d9fb236632b..c725b65a0e9 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -92,6 +92,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.repositories.prisma_protocols import BatchTable from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import CallTypes +from litellm.types.videos.utils import VIDEO_STATUS_CALL_TYPES if TYPE_CHECKING: from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction @@ -158,23 +159,34 @@ def _org_member_transaction_key(org_id: str, user_id: str) -> str: return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}" -def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: - return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" +def _is_settle_once_cost_row(payload: SpendLogsPayload) -> bool: + """A poll of a job whose cost is charged once, by the first poll that sees it, however often it is polled. + + A batch's cost is known only when a retrieve sees it finished (LIT-7048). A video's cost is + charged at create by the built-in providers, but a custom provider may report it on the + status poll that sees the video completed, and a poll that reports no cost charges nothing. + """ + if payload.get("status") != "success": + return False + call_type: Final = payload.get("call_type") + if call_type == CallTypes.aretrieve_batch.value: + return True + return call_type in VIDEO_STATUS_CALL_TYPES and payload["spend"] > 0 -_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) +_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) -def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: - """Reduce a batch's cost row to what tells the retrieves apart when logging is off. +def _cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: + """Reduce a job's cost row to what tells its polls apart when logging is off. - A proxy run with spend logs disabled still needs one row per batch to charge it once, + A proxy run with spend logs disabled still needs one row per job to charge it once, so the row is written either way, but it carries no request of its own: no metadata, no requester IP, no key, model, or token counts (LIT-7048). """ if disable_spend_logs is False: return payload - return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) + return MappingProxyType({field: value for field, value in payload.items() if field in _COST_CLAIM_FIELDS}) class _SpendIncrement(TypedDict): @@ -613,7 +625,7 @@ class DBSpendUpdateWriter: ) -> bool: """Record the request's spend, answering whether its cost still needs charging. - False only for a batch retrieve whose cost row another retrieve already wrote, + 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.proxy_server import ( @@ -755,24 +767,24 @@ class DBSpendUpdateWriter: disable_spend_logs: bool, llm_router: "Router | None" = None, ) -> bool: - if prisma_client is not None and _is_batch_cost_row(payload): - return await self._claim_batch_cost_spend_log( + if prisma_client is not None and _is_settle_once_cost_row(payload): + return await self._claim_settle_once_cost_row( payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs ) if disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client, llm_router=llm_router) return True - async def _claim_batch_cost_spend_log( + async def _claim_settle_once_cost_row( self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool ) -> bool: - """Write the batch's cost row now, or learn that another retrieve already did. + """Write the job's cost row now, or learn that another poll already did. - Every retrieve of one batch shares this row, so the insert that lands first owns + Every poll of one batch or video shares this row, so the insert that lands first owns the charge and every later one finds the row and charges nothing (LIT-7048). Only - a row that recorded a charge counts: a failed retrieve, a request whose client - picked the batch id as its call id, and the $0 row an older proxy left behind - while the batch was still running all leave the charge to be made. + a row that recorded a charge counts: a failed poll, a request whose client + picked the job id as its call id, and the $0 row a poll left behind + while the job was still running all leave the charge to be made. """ from litellm.repositories.table_repositories import SpendLogsRepository @@ -783,7 +795,7 @@ class DBSpendUpdateWriter: ) row: Final = spend_log_row_with_retained_metadata( - _batch_cost_row_to_write(payload, disable_spend_logs), configured_spend_logs_metadata_fields() + _cost_row_to_write(payload, disable_spend_logs), configured_spend_logs_metadata_fields() ) spend_logs: Final = SpendLogsRepository(prisma_client).table try: @@ -796,30 +808,29 @@ class DBSpendUpdateWriter: existing: Final = await spend_logs.find_unique(where={"request_id": request_id}) except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log verbose_proxy_logger.warning( - "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e + "Could not claim spend row %s for a job's cost, queueing it: %s", request_id, e ) await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) return True - if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": + if existing is None or existing.call_type != payload["call_type"] or existing.status != "success": verbose_proxy_logger.warning( - "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + "Spend row %s belongs to a %s request, so this job's cost is charged without a row of its own", request_id, getattr(existing, "call_type", None), ) return True if existing.spend > 0: - verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this job", request_id) return False - return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) + return await self._take_over_uncharged_cost_row(payload=payload, prisma_client=prisma_client, row=row) - async def _take_over_uncharged_batch_cost_row( + async def _take_over_uncharged_cost_row( self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] ) -> bool: - """Take the batch's cost row over from the poll that left it charging nothing. + """Take the job's cost row over from the poll that left it charging nothing. - A pre-upgrade proxy wrote that row every time it polled the batch while it was still - running, so the charge is still to be made and the row still has to end up carrying - it. The row stops matching the moment it carries a charge, so it is one retrieve that + A poll that saw the job still running wrote that row at $0, so the charge is still to be + made and the row still has to end up carrying it. The row stops matching the moment it carries a charge, so it is one poll that takes it over and charges, and every later one reads the charge and charges nothing. """ from litellm.repositories.table_repositories import SpendLogsRepository @@ -827,7 +838,7 @@ class DBSpendUpdateWriter: request_id: Final = payload["request_id"] if payload["spend"] <= 0: verbose_proxy_logger.debug( - "Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id + "Cost tracking skipped: this job costs nothing and spend row %s says so", request_id ) return False try: @@ -837,18 +848,18 @@ class DBSpendUpdateWriter: ), where={ "request_id": request_id, - "call_type": CallTypes.aretrieve_batch.value, + "call_type": payload["call_type"], "status": "success", "spend": 0.0, }, ) except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over verbose_proxy_logger.warning( - "Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e + "Could not take over spend row %s, leaving this job's cost to the next poll: %s", request_id, e ) return False if taken_over == 0: - verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this job", request_id) return False return True diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8b0ff718fec..69ae59b5ad1 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -99,6 +99,7 @@ from litellm.types.utils import ( VectorStoreSearchResponse, batch_chat_usage, ) +from litellm.types.videos.utils import VIDEO_STATUS_CALL_TYPES from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: @@ -344,11 +345,14 @@ def _get_spend_logs_metadata( BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" +VIDEO_COST_REQUEST_ID_SUFFIX: Final = "_video_cost" def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: if call_type == CallTypes.aretrieve_batch.value: return BATCH_COST_REQUEST_ID_SUFFIX + if call_type in VIDEO_STATUS_CALL_TYPES: + return VIDEO_COST_REQUEST_ID_SUFFIX if settles_background_response(call_type, metadata, kwargs): return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX return "" diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index 0fc2b8ede1b..3cb5ba50ef6 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -17,10 +17,9 @@ VIDEO_ID_PREFIX: Final = "video_" VIDEO_ID_SERVED_MODEL_PREFIX: Final = "served_model:" VIDEO_ID_DEPLOYMENT_PREFIX: Final = "deployment_id:" VIDEO_ID_REQUESTED_MODEL_PREFIX: Final = "requested_model:" -VIDEO_ID_CALL_TYPES: Final = frozenset( +VIDEO_STATUS_CALL_TYPES: Final = frozenset({"avideo_status", "video_status"}) +VIDEO_ID_CALL_TYPES: Final = VIDEO_STATUS_CALL_TYPES | frozenset( { - "avideo_status", - "video_status", "avideo_content", "video_content", "avideo_remix", 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 298b153f621..bc10e375eeb 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -4456,7 +4456,7 @@ async def _update_database_with( end_user_id=None, team_id=None, org_id=None, - kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"}, + kwargs={"model": "gpt-5.6-luna", "call_type": payload["call_type"]}, completion_response=None, start_time=datetime.now(timezone.utc), end_time=datetime.now(timezone.utc), @@ -4707,6 +4707,102 @@ async def test_update_database_queues_every_other_spend_row_for_the_next_flush(p assert db_writer._batch_database_updates.await_count == 1 +def _video_status_payload() -> dict: + return { + **_minimal_spend_payload(), + "request_id": "video_abc_video_cost", + "call_type": "avideo_status", + "status": "success", + } + + +class _SpendLogsTable: + """LiteLLM_SpendLogs keyed by its request_id primary key, as the writer's three queries see it.""" + + def __init__(self) -> None: + self.rows: dict[str, dict] = {} + + async def create_many(self, data: list[dict], skip_duplicates: bool) -> int: + fresh = {row["request_id"]: dict(row) for row in data if row["request_id"] not in self.rows} + self.rows.update(fresh) + return len(fresh) + + async def find_unique(self, where: dict) -> SimpleNamespace | None: + row = self.rows.get(where["request_id"]) + return None if row is None else SimpleNamespace(**row) + + async def update_many(self, data: dict, where: dict) -> int: + row = self.rows.get(where["request_id"]) + if row is None or any(row.get(field) != value for field, value in where.items()): + return 0 + row.update(data) + return 1 + + +def _spend_logs_table_prisma() -> MagicMock: + prisma = _tool_usage_prisma() + prisma.jsonify_object = lambda data: dict(data) + prisma.db.litellm_spendlogs = _SpendLogsTable() + return prisma + + +async def _flush_queued_spend_logs(prisma: MagicMock) -> None: + await prisma.db.litellm_spendlogs.create_many(data=prisma.spend_log_transactions, skip_duplicates=True) + prisma.spend_log_transactions.clear() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "flushes_before_completion", + [0, 1], + ids=["rendering_poll_still_queued", "rendering_poll_already_written"], +) +async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_completed( + flushes_before_completion: int, +): + """ + A custom video provider reports the video's cost on the status poll that sees it completed, + and clients poll a finished video again before they download it, so only that first completed + poll may charge, and its spend row has to carry the charge (LIT-10404). Every poll of one video + shares that row, the $0 one written while the video was still rendering included. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + rendering = await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.0) + for _ in range(flushes_before_completion): + await _flush_queued_spend_logs(prisma) + completed = await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.42) + repolled = await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.42) + await _flush_queued_spend_logs(prisma) + + assert (rendering, completed, repolled) == (True, True, False) + assert db_writer._batch_database_updates.await_count == 2 + rows = prisma.db.litellm_spendlogs.rows + assert {request_id: (row["call_type"], row["spend"]) for request_id, row in rows.items()} == { + "video_abc_video_cost": ("avideo_status", 0.42) + } + + +@pytest.mark.asyncio +async def test_update_database_queues_a_video_poll_that_reports_no_cost(): + """ + Built-in video providers charge at create and report nothing on the status poll, so those + polls keep the ordinary queued write and still count as requests in the daily rollups. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + + assert await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.0) is True + assert await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.0) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert [row["request_id"] for row in prisma.spend_log_transactions] == ["video_abc_video_cost"] * 2 + assert db_writer._batch_database_updates.await_count == 2 + + @pytest.mark.asyncio @pytest.mark.parametrize( "injected_deployment, attributed", diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index f73d33cc93e..4bc99ae91d4 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5055,6 +5055,18 @@ def test_get_spend_logs_id_prefers_the_response_id_for_batch_and_file_calls(call assert get_spend_logs_id(call_type, {"id": "batch_abc123"}, {"litellm_call_id": "call-id-1"}) == expected +@pytest.mark.parametrize("status_call_type", ["avideo_status", "video_status"]) +def test_get_spend_logs_id_keeps_a_video_status_row_apart_from_the_create_row(status_call_type): + """A video's create and its status polls answer with the same video id, and a custom + provider charges the video on the poll that sees it completed, so that charge needs a + row of its own: sharing the create row's id drops it on the skip_duplicates insert.""" + create_id = get_spend_logs_id("avideo_generation", {"id": "video_abc"}, {"litellm_call_id": "call-1"}) + status_id = get_spend_logs_id(status_call_type, {"id": "video_abc"}, {"litellm_call_id": "call-2"}) + + assert create_id == "video_abc" + assert status_id == "video_abc_video_cost" + + def test_get_logging_payload_gives_redacted_batch_and_file_rows_distinct_request_ids(): """End to end at the payload level: a batch retrieve and a file create whose bodies were both flattened to the same redaction placeholder must still produce two From 5b972297e14d728f6444107cd8dabce814b3b391 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:38:49 -0800 Subject: [PATCH 02/10] fix(proxy): charge a video's provider-reported cost once on model, router and session budgets too A video status poll gets the _video_cost row suffix only when it reports a cost, so a $0 poll keeps the bare video id and adds no row of its own, and the spend log writer claims a row by that suffix. The batch charge-once marker now covers videos, keyed on the video id under {spend_key}:video:{id}, with the batch marker key unchanged. Model and router budgets use it, and per-session budgets now charge a repolled batch or video once and ignore a cost that is not a finite number --- litellm/litellm_core_utils/batch_utils.py | 27 +++++-- litellm/proxy/db/db_spend_update_writer.py | 7 +- .../hooks/max_budget_per_session_limiter.py | 19 ++++- .../proxy/hooks/model_max_budget_limiter.py | 16 ++--- .../spend_tracking/spend_tracking_utils.py | 6 +- litellm/router_strategy/budget_limiter.py | 22 +++--- litellm/types/videos/utils.py | 7 ++ .../litellm_core_utils/test_batch_utils.py | 68 +++++++++++++----- .../proxy/db/test_db_spend_update_writer.py | 70 +++++++++++++++---- .../test_max_budget_per_session_limiter.py | 59 ++++++++++++++++ .../hooks/test_model_max_budget_limiter.py | 25 +++++++ .../test_spend_tracking_utils.py | 42 ++++++++--- .../router_strategy/test_budget_limiter.py | 30 ++++++++ 13 files changed, 324 insertions(+), 74 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index 5b94d4b4e5b..e11087299e0 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -3,8 +3,9 @@ import json from collections.abc import Callable, Generator, Iterable, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar +from dataclasses import dataclass from os import PathLike -from typing import Final, TypeVar +from typing import Final, Literal, TypeVar from openai.types import Batch from pydantic import TypeAdapter, ValidationError @@ -12,6 +13,8 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.videos.main import VideoObject +from litellm.types.videos.utils import is_video_status_call_type _DECODED_JSONL_ROW: Final = TypeAdapter(object) _DEPLOYMENT_FIELD: Final = TypeAdapter(Mapping[str, object]) @@ -289,12 +292,22 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES -def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str: - return f"{spend_key}:batch:{batch_id}" +@dataclass(frozen=True, slots=True) +class ChargeOnceJob: + kind: Literal["batch", "video"] + job_id: str -def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None: - """A finished batch reports its whole cost on every poll, so its id is charged once per counter.""" - if response_cost <= 0 or not is_batch_retrieve_call_type(call_type): +def job_charged_once_marker_key(spend_key: str, job: ChargeOnceJob) -> str: + return f"{spend_key}:{job.kind}:{job.job_id}" + + +def job_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> ChargeOnceJob | None: + """A finished batch or video reports its whole cost on every poll, so its id is charged once per counter.""" + if response_cost <= 0: return None - return response_obj.id if isinstance(response_obj, Batch) else None + if is_batch_retrieve_call_type(call_type) and isinstance(response_obj, Batch): + return ChargeOnceJob(kind="batch", job_id=response_obj.id) + if is_video_status_call_type(call_type) and isinstance(response_obj, VideoObject): + return ChargeOnceJob(kind="video", job_id=response_obj.id) + return None diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c725b65a0e9..e845da88002 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -92,7 +92,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.repositories.prisma_protocols import BatchTable from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import CallTypes -from litellm.types.videos.utils import VIDEO_STATUS_CALL_TYPES +from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES if TYPE_CHECKING: from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction @@ -164,14 +164,15 @@ def _is_settle_once_cost_row(payload: SpendLogsPayload) -> bool: A batch's cost is known only when a retrieve sees it finished (LIT-7048). A video's cost is charged at create by the built-in providers, but a custom provider may report it on the - status poll that sees the video completed, and a poll that reports no cost charges nothing. + status poll that sees the video completed. Only such a poll's row carries the video cost + suffix, so a poll that reports no cost keeps its own id and charges nothing. """ if payload.get("status") != "success": return False call_type: Final = payload.get("call_type") if call_type == CallTypes.aretrieve_batch.value: return True - return call_type in VIDEO_STATUS_CALL_TYPES and payload["spend"] > 0 + return call_type in VIDEO_STATUS_CALL_TYPES and payload["request_id"].endswith(VIDEO_COST_REQUEST_ID_SUFFIX) _COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 002aeb145a5..28eaffbf593 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -26,6 +26,8 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import log_redis_failure from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key, job_to_charge_once +from litellm.litellm_core_utils.core_helpers import finite_cost from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -44,6 +46,7 @@ _CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float]] = TypeAdapt _OPTIONAL_CACHE_AMOUNT_ADAPTER: Final[TypeAdapter[str | bytes | int | float | None]] = TypeAdapter( str | bytes | int | float | None, config=ConfigDict(strict=True, hide_input_in_errors=True) ) +_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) # Redis Lua script for atomic float increment with TTL. @@ -140,7 +143,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return None @with_service_target("session_budgets") - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None: """ After a successful LLM call, increment the session spend by the response cost. """ @@ -168,12 +171,16 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): if max_budget is None: return - response_cost: Final = kwargs.get("response_cost") or 0.0 + callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs) + response_cost: Final = finite_cost(callback_fields.get("response_cost")) or 0.0 if response_cost <= 0: return cache_key: Final = self._make_cache_key(str(session_id)) - await self._increment_spend(cache_key, float(response_cost)) + job: Final = job_to_charge_once(callback_fields.get("call_type"), response_obj, response_cost) + if job is not None and not await self._claim_job_charge(cache_key, job): + return + await self._increment_spend(cache_key, response_cost) verbose_proxy_logger.debug( "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", @@ -221,6 +228,12 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): def _make_cache_key(self, session_id: str) -> str: return f"{{session_budget:{session_id}}}:spend" + async def _claim_job_charge(self, cache_key: str, job: ChargeOnceJob) -> bool: + polls: Final = await self.internal_usage_cache.dual_cache.async_increment_cache( + key=job_charged_once_marker_key(spend_key=cache_key, job=job), value=1, ttl=self.ttl + ) + return polls == 1 + @with_service_target("session_budgets") async def _get_current_spend(self, cache_key: str) -> float: """Read current accumulated spend for a session.""" diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 1c10826124e..30b826a3245 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,7 +19,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once +from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key, job_to_charge_once from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions @@ -1208,7 +1208,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) return - batch_id: Final = batch_id_to_charge_once( + job: Final = job_to_charge_once( call_type=kwargs.get("call_type"), response_obj=response_obj, response_cost=response_cost, @@ -1219,7 +1219,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): entity_id=entity_id, resolved=resolved, response_cost=response_cost, - batch_id=batch_id, + job=job, ) if self.dual_cache.redis_cache is not None: @@ -1236,7 +1236,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): entity_id: str | None, resolved: ResolvedModelBudget, response_cost: float, - batch_id: str | None, + job: ChargeOnceJob | None, ) -> None: budget_duration: Final = resolved.budget_config.budget_duration if budget_duration is None: @@ -1247,9 +1247,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): budget_model=resolved.budget_model, budget_duration=budget_duration, ) - if batch_id is not None and not await self._claim_batch_charge( + if job is not None and not await self._claim_job_charge( spend_key=spend_key, - batch_id=batch_id, + job=job, ttl_seconds=duration_in_seconds(budget_duration), ): return @@ -1362,8 +1362,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): spend=response_cost, ) - async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: - marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) + async def _claim_job_charge(self, spend_key: str, job: ChargeOnceJob, ttl_seconds: int) -> bool: + marker_key: Final = job_charged_once_marker_key(spend_key=spend_key, job=job) polls: Final = await self.dual_cache.async_increment_cache( key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True ) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 69ae59b5ad1..eda19ce7433 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -44,6 +44,7 @@ from litellm.litellm_core_utils.classifier_logging import ( without_classifier_audit, ) from litellm.litellm_core_utils.core_helpers import ( + finite_cost, get_litellm_metadata_from_kwargs, proxy_stamped_kubernetes_pod_routing, proxy_stamped_used_client_oauth_token, @@ -99,7 +100,7 @@ from litellm.types.utils import ( VectorStoreSearchResponse, batch_chat_usage, ) -from litellm.types.videos.utils import VIDEO_STATUS_CALL_TYPES +from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: @@ -345,13 +346,12 @@ def _get_spend_logs_metadata( BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" -VIDEO_COST_REQUEST_ID_SUFFIX: Final = "_video_cost" def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: if call_type == CallTypes.aretrieve_batch.value: return BATCH_COST_REQUEST_ID_SUFFIX - if call_type in VIDEO_STATUS_CALL_TYPES: + if call_type in VIDEO_STATUS_CALL_TYPES and (finite_cost(kwargs.get("response_cost")) or 0.0) > 0: return VIDEO_COST_REQUEST_ID_SUFFIX if settles_background_response(call_type, metadata, kwargs): return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 1d502dd8d95..072d38717a3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,7 +36,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once +from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key, job_to_charge_once from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) @@ -52,7 +52,7 @@ from litellm.types.utils import BudgetConfig as GenericBudgetInfo DEFAULT_REDIS_SYNC_INTERVAL: Final = 1 _CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) -_BATCH_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) +_JOB_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) class _LiteLLMParamsDictView: @@ -545,13 +545,13 @@ class RouterBudgetLimiting(CustomLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is required") - batch_response, response_cost = _BATCH_CHARGE_VALUES.validate_python( + polled_response, response_cost = _JOB_CHARGE_VALUES.validate_python( (response_obj, _CALLBACK_FIELDS.validate_python(standard_logging_payload).get("response_cost", 0)) ) model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") - batch_id: Final = batch_id_to_charge_once(callback_fields.get("call_type"), batch_response, response_cost) + job: Final = job_to_charge_once(callback_fields.get("call_type"), polled_response, response_cost) budget_config: Final = ( self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None @@ -565,7 +565,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=spend_key, start_time_key=start_time_key, response_cost=response_cost, - batch_id=batch_id, + job=job, ) deployment_budget_config: Final = self._get_budget_config_for_deployment(model_id) @@ -578,7 +578,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=deployment_spend_key, start_time_key=deployment_start_time_key, response_cost=response_cost, - batch_id=batch_id, + job=job, ) request_tags: Final = get_tags_from_request_kwargs( @@ -596,7 +596,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=_tag_spend_key, start_time_key=_tag_start_time_key, response_cost=response_cost, - batch_id=batch_id, + job=job, ) async def _increment_spend_for_key( @@ -605,14 +605,14 @@ class RouterBudgetLimiting(CustomLogger): spend_key: str, start_time_key: str, response_cost: float, - batch_id: str | None = None, + job: ChargeOnceJob | None = None, ): if budget_config.budget_duration is None: return current_time: Final = self._clock() new_window_seconds: Final = budget_window_seconds(current_time, budget_config.budget_duration) - if batch_id is not None and not await self._claim_batch_charge(spend_key, batch_id, new_window_seconds): + if job is not None and not await self._claim_job_charge(spend_key, job, new_window_seconds): return budget_start: Final = await self._get_or_set_budget_start_time( @@ -639,8 +639,8 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, response_cost) - async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: - marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) + async def _claim_job_charge(self, spend_key: str, job: ChargeOnceJob, ttl_seconds: int) -> bool: + marker_key: Final = job_charged_once_marker_key(spend_key=spend_key, job=job) redis_cache: Final = self.dual_cache.redis_cache if redis_cache is not None: polls: Final = await redis_cache.async_increment(key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True) diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index 3cb5ba50ef6..c74f11b7abb 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -18,6 +18,7 @@ VIDEO_ID_SERVED_MODEL_PREFIX: Final = "served_model:" VIDEO_ID_DEPLOYMENT_PREFIX: Final = "deployment_id:" VIDEO_ID_REQUESTED_MODEL_PREFIX: Final = "requested_model:" VIDEO_STATUS_CALL_TYPES: Final = frozenset({"avideo_status", "video_status"}) +VIDEO_COST_REQUEST_ID_SUFFIX: Final = "_video_cost" VIDEO_ID_CALL_TYPES: Final = VIDEO_STATUS_CALL_TYPES | frozenset( { "avideo_content", @@ -30,6 +31,12 @@ VIDEO_ID_CALL_TYPES: Final = VIDEO_STATUS_CALL_TYPES | frozenset( "video_extension", } ) + + +def is_video_status_call_type(call_type: object) -> bool: + return isinstance(call_type, str) and call_type in VIDEO_STATUS_CALL_TYPES + + CHARACTER_ID_PREFIX: Final = "character_" CHARACTER_ID_TEMPLATE: Final = "litellm:custom_llm_provider:{};model_id:{};character_id:{}" diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index cd37d635c43..4a4194b184b 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -382,30 +382,62 @@ def setup_and_teardown(): asyncio.set_event_loop(None) -@pytest.mark.parametrize( - ("call_type", "cost", "eligible"), - (("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False), - ("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)), -) -def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None: +def _completed_batch(): from openai.types import Batch - from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once - - batch: Final = Batch( - id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", + return Batch( + id="job-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", input_file_id="file-budget", object="batch", status="completed", ) - assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None) - assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None -def test_batch_charge_markers_keep_counters_and_batches_independent() -> None: +def _completed_video(): + from litellm.types.videos.main import VideoObject + + return VideoObject(id="job-budget", object="video", status="completed") + + +@pytest.mark.parametrize( + ("call_type", "polled", "cost", "kind"), + ( + ("aretrieve_batch", _completed_batch, 0.25, "batch"), + ("retrieve_batch", _completed_batch, 0.25, "batch"), + ("avideo_status", _completed_video, 0.25, "video"), + ("video_status", _completed_video, 0.25, "video"), + ("acompletion", _completed_batch, 0.25, None), + ("avideo_status", _completed_batch, 0.25, None), + ("aretrieve_batch", _completed_video, 0.25, None), + ("avideo_generation", _completed_video, 0.25, None), + ("aretrieve_batch", _completed_batch, 0, None), + ("avideo_status", _completed_video, 0, None), + ("avideo_status", _completed_video, -0.25, None), + ), +) +def test_job_charge_identity_requires_a_cost_bearing_poll(call_type: str, polled, cost: float, kind) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + expected: Final = None if kind is None else ChargeOnceJob(kind=kind, job_id="job-budget") + assert job_to_charge_once(call_type, polled(), cost) == expected + assert job_to_charge_once(call_type, {"id": "job-budget"}, cost) is None + + +def test_job_charge_markers_keep_counters_jobs_and_kinds_independent() -> None: from litellm.caching.in_memory_cache import InMemoryCache - from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key + first_batch: Final = ChargeOnceJob(kind="batch", job_id="first-job") cache: Final = InMemoryCache() - cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1) - assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1 - assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None - assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None + cache.set_cache(job_charged_once_marker_key("provider-spend", first_batch), 1) + assert cache.get_cache(job_charged_once_marker_key("provider-spend", first_batch)) == 1 + assert cache.get_cache(job_charged_once_marker_key("deployment-spend", first_batch)) is None + assert cache.get_cache(job_charged_once_marker_key("provider-spend", ChargeOnceJob("batch", "second-job"))) is None + assert cache.get_cache(job_charged_once_marker_key("provider-spend", ChargeOnceJob("video", "first-job"))) is None + + +def test_batch_charge_marker_keeps_the_key_markers_already_written_use() -> None: + """Markers written before the video kind existed must still block a re-charge after an upgrade.""" + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key + + assert job_charged_once_marker_key("provider-spend", ChargeOnceJob("batch", "first-job")) == ( + "provider-spend:batch:first-job" + ) 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 bc10e375eeb..d9b06b6f10c 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -4707,10 +4707,10 @@ async def test_update_database_queues_every_other_spend_row_for_the_next_flush(p assert db_writer._batch_database_updates.await_count == 1 -def _video_status_payload() -> dict: +def _video_status_payload(request_id: str = "video_abc_video_cost") -> dict: return { **_minimal_spend_payload(), - "request_id": "video_abc_video_cost", + "request_id": request_id, "call_type": "avideo_status", "status": "success", } @@ -4763,14 +4763,14 @@ async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_c """ A custom video provider reports the video's cost on the status poll that sees it completed, and clients poll a finished video again before they download it, so only that first completed - poll may charge, and its spend row has to carry the charge (LIT-10404). Every poll of one video - shares that row, the $0 one written while the video was still rendering included. + poll may charge, and its spend row has to carry the charge (LIT-10404). The $0 poll written + while the video was still rendering keeps the bare video id, so it lands on the create row. """ db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() prisma = _spend_logs_table_prisma() - rendering = await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.0) + rendering = await _update_database_with(db_writer, prisma, _video_status_payload("video_abc"), response_cost=0.0) for _ in range(flushes_before_completion): await _flush_queued_spend_logs(prisma) completed = await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.42) @@ -4781,26 +4781,70 @@ async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_c assert db_writer._batch_database_updates.await_count == 2 rows = prisma.db.litellm_spendlogs.rows assert {request_id: (row["call_type"], row["spend"]) for request_id, row in rows.items()} == { - "video_abc_video_cost": ("avideo_status", 0.42) + "video_abc": ("avideo_status", 0.0), + "video_abc_video_cost": ("avideo_status", 0.42), } +async def _update_database_for_video_poll( + db_writer: DBSpendUpdateWriter, prisma: MagicMock, litellm_call_id: str, response_cost: float +) -> bool: + from litellm.types.videos.main import VideoObject + + with ( + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.prisma_client", prisma + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), + ): + charged = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={ + "model": "sora-2", + "call_type": "avideo_status", + "litellm_call_id": litellm_call_id, + "response_cost": response_cost, + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + completion_response=VideoObject(id="video_abc", object="video", status="completed"), + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=response_cost, + ) + await asyncio.sleep(0) + return charged + + @pytest.mark.asyncio -async def test_update_database_queues_a_video_poll_that_reports_no_cost(): +async def test_update_database_writes_no_row_of_its_own_for_a_video_poll_that_reports_no_cost(): """ Built-in video providers charge at create and report nothing on the status poll, so those - polls keep the ordinary queued write and still count as requests in the daily rollups. + polls keep the ordinary queued write under the create row's id and still count as requests + in the daily rollups, leaving the video with its one create row. """ db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() - prisma = _spend_logs_prisma(1, None) + prisma = _spend_logs_table_prisma() + await prisma.db.litellm_spendlogs.create_many( + data=[{**_minimal_spend_payload(), "request_id": "video_abc", "call_type": "avideo_generation", "spend": 0.3}], + skip_duplicates=True, + ) - assert await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.0) is True - assert await _update_database_with(db_writer, prisma, _video_status_payload(), response_cost=0.0) is True + assert await _update_database_for_video_poll(db_writer, prisma, "poll-1", response_cost=0.0) is True + assert await _update_database_for_video_poll(db_writer, prisma, "poll-2", response_cost=0.0) is True + await _flush_queued_spend_logs(prisma) - prisma.db.litellm_spendlogs.create_many.assert_not_called() - assert [row["request_id"] for row in prisma.spend_log_transactions] == ["video_abc_video_cost"] * 2 assert db_writer._batch_database_updates.await_count == 2 + rows = prisma.db.litellm_spendlogs.rows + assert {request_id: (row["call_type"], row["spend"]) for request_id, row in rows.items()} == { + "video_abc": ("avideo_generation", 0.3) + } @pytest.mark.asyncio diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 45cc7eb8922..18b1bce8e84 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -145,6 +145,65 @@ async def test_budget_per_session_independent_sessions(): assert result is None +@pytest.mark.asyncio +async def test_repeated_polls_of_a_finished_job_charge_the_session_once(): + from openai.types import Batch + + from litellm.types.videos.main import VideoObject + + handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + video = VideoObject(id="job-1", object="video", status="completed") + batch = Batch( + id="job-1", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", + input_file_id="file-1", object="batch", status="completed", + ) + polls = (("avideo_status", video, 0.4), ("aretrieve_batch", batch, 0.1), ("acompletion", None, 0.01)) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + for call_type, response_obj, cost in polls * 3: + await handler.async_log_success_event( + { + "call_type": call_type, + "response_cost": cost, + "litellm_params": {"metadata": {"session_id": "session-jobs", "agent_id": "agent-budget-123"}}, + }, + response_obj, + None, + None, + ) + + spend = await handler._get_current_spend(handler._make_cache_key("session-jobs")) + assert spend == pytest.approx(0.4 + 0.1 + 3 * 0.01) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), True]) +async def test_a_cost_that_is_not_a_finite_number_leaves_the_session_spend_alone(response_cost): + handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + cache_key = handler._make_cache_key("session-unpriced") + await handler._increment_spend(cache_key, 0.25) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + await handler.async_log_success_event( + { + "call_type": "acompletion", + "response_cost": response_cost, + "litellm_params": {"metadata": {"session_id": "session-unpriced", "agent_id": "agent-budget-123"}}, + }, + None, + None, + None, + ) + + assert await handler._get_current_spend(cache_key) == 0.25 + + @pytest.mark.asyncio async def test_no_agent_id_passes(): """ 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 f5a2ccf7b76..48b44a75cef 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -27,6 +27,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.router import Router from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import LiteLLMBatch, Usage +from litellm.types.videos.main import VideoObject KEY_HASH: Final = "key-hash-batch" USER_ID: Final = "user-batch" @@ -326,6 +327,30 @@ async def test_a_batch_polled_within_every_budget_window_is_never_charged_again( assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + +async def _poll_video( + limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, video_id: str, status: str, response_cost: float +) -> None: + await _log_success( + limiter, + _event("avideo_status", response_cost), + response_obj=VideoObject(id=video_id, object="video", status=status), + ) + + +@pytest.mark.asyncio +async def test_polls_of_a_finished_video_charge_each_per_model_budget_once(): + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + + await _poll_video(limiter, "video_first", "in_progress", response_cost=0) + for _ in range(3): + await _poll_video(limiter, "video_first", "completed", response_cost=BATCH_COST) + await _poll_video(limiter, "video_second", "completed", response_cost=BATCH_COST) + await _poll(limiter, _batch("video_first", "completed"), response_cost=BATCH_COST) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(3 * BATCH_COST) + assert await _spend(limiter, USER_SPEND_KEY) == pytest.approx(3 * BATCH_COST) + OPUS_GROUP: Final = "claude-opus" OPUS_GROUP_MEMBERS: Final = MappingProxyType( {"claude-opus-4-8": frozenset({OPUS_GROUP}), "claude-opus-5-5": frozenset({OPUS_GROUP})} diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 4bc99ae91d4..b5c62e86ee5 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5056,15 +5056,41 @@ def test_get_spend_logs_id_prefers_the_response_id_for_batch_and_file_calls(call @pytest.mark.parametrize("status_call_type", ["avideo_status", "video_status"]) -def test_get_spend_logs_id_keeps_a_video_status_row_apart_from_the_create_row(status_call_type): - """A video's create and its status polls answer with the same video id, and a custom - provider charges the video on the poll that sees it completed, so that charge needs a - row of its own: sharing the create row's id drops it on the skip_duplicates insert.""" - create_id = get_spend_logs_id("avideo_generation", {"id": "video_abc"}, {"litellm_call_id": "call-1"}) - status_id = get_spend_logs_id(status_call_type, {"id": "video_abc"}, {"litellm_call_id": "call-2"}) +@pytest.mark.parametrize( + ("response_cost", "status_id"), + [ + (0.42, "video_abc_video_cost"), + (0.0, "video_abc"), + (None, "video_abc"), + (float("nan"), "video_abc"), + ], + ids=["charged", "free", "unpriced", "non_finite"], +) +def test_get_logging_payload_gives_only_a_charged_video_poll_a_row_of_its_own(status_call_type, response_cost, status_id): + """A custom provider charges the video on the poll that sees it completed, so that poll needs + a row apart from the create row, which answers with the same video id and would drop it on the + skip_duplicates insert. A poll that charges nothing keeps the bare id, so it still lands on the + create row and adds no row of its own.""" + from litellm.types.videos.main import VideoObject - assert create_id == "video_abc" - assert status_id == "video_abc_video_cost" + video = VideoObject(id="video_abc", object="video", status="completed") + create, status = ( + get_logging_payload( + kwargs={ + "call_type": call_type, + "litellm_call_id": call_id, + "response_cost": cost, + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=video, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + for call_type, call_id, cost in (("avideo_generation", "call-1", 0.42), (status_call_type, "call-2", response_cost)) + ) + + assert create["request_id"] == "video_abc" + assert status["request_id"] == status_id def test_get_logging_payload_gives_redacted_batch_and_file_rows_distinct_request_ids(): diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index c0583566531..7fb43f17209 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -498,6 +498,36 @@ async def test_batch_polls_charge_each_router_budget_once(monkeypatch, call_type assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0) + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ("avideo_status", "video_status")) +async def test_video_polls_charge_each_router_budget_once(call_type): + from litellm.types.videos.main import VideoObject + + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "video-budget", + "litellm_params": {"model": "openai/sora-2", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + kwargs: Final = _success_kwargs(provider_in_litellm_params=None, provider_in_payload="openai", call_type=call_type) + video: Final = VideoObject(id="video-first", object="video", status="completed") + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") + + await asyncio.gather(*(limiter.async_log_success_event(kwargs, video, None, None) for _ in range(10))) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.25, 0.25) + + await limiter.async_log_success_event( + kwargs, VideoObject(id="video-second", object="video", status="completed"), None, None + ) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.5, 0.5) + @pytest.mark.asyncio async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(): from fakeredis import FakeServer, FakeRedis as FakeSyncRedis From 9590cb966e54775d9d05bc71ba7852b92af7393d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:56:22 -0800 Subject: [PATCH 03/10] fix(proxy): charge a finished video's own cost once, keyed by the id the client polled Guardrail cost on a running poll no longer takes the charge-once claim, only a completed poll charges, the session marker TTL refreshes like the others, and two deployments whose provider ids match each keep their own charge. --- litellm/litellm_core_utils/batch_utils.py | 61 +++++++++-- .../hooks/max_budget_per_session_limiter.py | 4 +- .../proxy/hooks/model_max_budget_limiter.py | 6 +- .../spend_tracking/spend_tracking_utils.py | 11 +- litellm/router_strategy/budget_limiter.py | 2 +- .../router_code_coverage.py | 4 + .../litellm_core_utils/test_batch_utils.py | 39 ++++++- .../proxy/db/test_db_spend_update_writer.py | 47 +++++++- .../test_max_budget_per_session_limiter.py | 101 ++++++++++++++++++ .../hooks/test_model_max_budget_limiter.py | 24 +++++ .../test_spend_tracking_utils.py | 18 ++++ .../router_strategy/test_budget_limiter.py | 40 +++++++ 12 files changed, 332 insertions(+), 25 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index e11087299e0..a1fa6e39e4b 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -11,6 +11,8 @@ from openai.types import Batch from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import finite_cost +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose from litellm.types.utils import CallTypes, LlmProviders from litellm.types.videos.main import VideoObject @@ -18,6 +20,7 @@ from litellm.types.videos.utils import is_video_status_call_type _DECODED_JSONL_ROW: Final = TypeAdapter(object) _DEPLOYMENT_FIELD: Final = TypeAdapter(Mapping[str, object]) +_PAYLOAD_FIELDS: Final = TypeAdapter(Mapping[str, object]) _ENDPOINT_PATHS: Final = TypeAdapter(tuple[str, ...]) BATCH_ENDPOINT: Final = "/v1/batch" @@ -302,12 +305,54 @@ def job_charged_once_marker_key(spend_key: str, job: ChargeOnceJob) -> str: return f"{spend_key}:{job.kind}:{job.job_id}" -def job_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> ChargeOnceJob | None: - """A finished batch or video reports its whole cost on every poll, so its id is charged once per counter.""" - if response_cost <= 0: +def job_cost_of_poll(response_cost: object, standard_logging_payload: object) -> float: + """What the polled job itself cost, since a guardrail on the poll charges every poll, running or finished.""" + guardrail_information: Final = ( + _PAYLOAD_FIELDS.validate_python(standard_logging_payload).get("guardrail_information") + if isinstance(standard_logging_payload, Mapping) + else None + ) + return (finite_cost(response_cost) or 0.0) - guardrail_information_cost(guardrail_information) + + +def _requested_video_id(kwargs: Mapping[str, object]) -> str | None: + litellm_params: Final = kwargs.get("litellm_params") + requested: Final = ( + _PAYLOAD_FIELDS.validate_python(litellm_params).get("video_id") if isinstance(litellm_params, Mapping) else None + ) + return requested if isinstance(requested, str) and requested else None + + +def charged_video_id( + call_type: object, status: object, video_id: str, kwargs: Mapping[str, object], response_cost: object +) -> str | None: + """The id a finished video's own cost is charged once under. + + That is the id the client polled, which names the deployment that made the video, since a + custom provider's own id can repeat across deployments. + """ + if not is_video_status_call_type(call_type) or status != "completed": return None - if is_batch_retrieve_call_type(call_type) and isinstance(response_obj, Batch): - return ChargeOnceJob(kind="batch", job_id=response_obj.id) - if is_video_status_call_type(call_type) and isinstance(response_obj, VideoObject): - return ChargeOnceJob(kind="video", job_id=response_obj.id) - return None + if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0: + return None + return _requested_video_id(kwargs) or video_id + + +def _charged_batch(kwargs: Mapping[str, object], response_obj: object, response_cost: float) -> ChargeOnceJob | None: + if not is_batch_retrieve_call_type(kwargs.get("call_type")) or not isinstance(response_obj, Batch): + return None + if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0: + return None + return ChargeOnceJob(kind="batch", job_id=response_obj.id) + + +def job_to_charge_once( + kwargs: Mapping[str, object], response_obj: object, response_cost: float +) -> ChargeOnceJob | None: + """A finished batch or video reports its whole cost on every poll, so it is charged once per counter.""" + if not isinstance(response_obj, VideoObject): + return _charged_batch(kwargs, response_obj, response_cost) + video_id: Final = charged_video_id( + kwargs.get("call_type"), response_obj.status, response_obj.id, kwargs, response_cost + ) + return None if video_id is None else ChargeOnceJob(kind="video", job_id=video_id) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 28eaffbf593..6b60a8b8de0 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -177,7 +177,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return cache_key: Final = self._make_cache_key(str(session_id)) - job: Final = job_to_charge_once(callback_fields.get("call_type"), response_obj, response_cost) + job: Final = job_to_charge_once(callback_fields, response_obj, response_cost) if job is not None and not await self._claim_job_charge(cache_key, job): return await self._increment_spend(cache_key, response_cost) @@ -230,7 +230,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): async def _claim_job_charge(self, cache_key: str, job: ChargeOnceJob) -> bool: polls: Final = await self.internal_usage_cache.dual_cache.async_increment_cache( - key=job_charged_once_marker_key(spend_key=cache_key, job=job), value=1, ttl=self.ttl + key=job_charged_once_marker_key(spend_key=cache_key, job=job), value=1, ttl=self.ttl, refresh_ttl=True ) return polls == 1 diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 30b826a3245..6414b7613b9 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1208,11 +1208,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) return - job: Final = job_to_charge_once( - call_type=kwargs.get("call_type"), - response_obj=response_obj, - response_cost=response_cost, - ) + job: Final = job_to_charge_once(kwargs=kwargs, response_obj=response_obj, response_cost=response_cost) for entity_type, entity_id, resolved in resolved_budgets: await self._charge_entity( entity_type=entity_type, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index eda19ce7433..03ef44bf34b 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -38,13 +38,13 @@ from litellm.constants import ( ) from litellm.integrations.otel.model.metadata import caller_session_id from litellm.integrations.otel.model.trace_controls import TraceControls +from litellm.litellm_core_utils.batch_utils import charged_video_id from litellm.litellm_core_utils.classifier_logging import ( CLASSIFIER_AUDIT_FIELDS, classifier_audit_fields, without_classifier_audit, ) from litellm.litellm_core_utils.core_helpers import ( - finite_cost, get_litellm_metadata_from_kwargs, proxy_stamped_kubernetes_pod_routing, proxy_stamped_used_client_oauth_token, @@ -100,7 +100,7 @@ from litellm.types.utils import ( VectorStoreSearchResponse, batch_chat_usage, ) -from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES +from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: @@ -351,8 +351,6 @@ BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: if call_type == CallTypes.aretrieve_batch.value: return BATCH_COST_REQUEST_ID_SUFFIX - if call_type in VIDEO_STATUS_CALL_TYPES and (finite_cost(kwargs.get("response_cost")) or 0.0) > 0: - return VIDEO_COST_REQUEST_ID_SUFFIX if settles_background_response(call_type, metadata, kwargs): return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX return "" @@ -375,6 +373,11 @@ def get_spend_logs_id( ) if resolved_id is None: return None + charged_video: Final = charged_video_id( + call_type, response_obj.get("status"), resolved_id, kwargs, kwargs.get("response_cost") + ) + if charged_video is not None: + return f"{charged_video}{VIDEO_COST_REQUEST_ID_SUFFIX}" return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 072d38717a3..bdcba8c8ae1 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -551,7 +551,7 @@ class RouterBudgetLimiting(CustomLogger): model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") - job: Final = job_to_charge_once(callback_fields.get("call_type"), polled_response, response_cost) + job: Final = job_to_charge_once(callback_fields, polled_response, response_cost) budget_config: Final = ( self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 1b1ee484afc..58d0e373eb0 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -175,6 +175,10 @@ ignored_function_names = [ "_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py "_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py "_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py + "job_cost_of_poll", # Read by job_to_charge_once and charged_video_id, tested through the guardrail-priced poll tests in test_batch_utils.py + "charged_video_id", # Read by job_to_charge_once and get_spend_logs_id, tested through both in test_batch_utils.py and test_spend_tracking_utils.py + "_requested_video_id", # Read by charged_video_id, tested through the polled-id video test in test_batch_utils.py + "_charged_batch", # Read by job_to_charge_once, tested through the batch charge identity tests in test_batch_utils.py "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index 4a4194b184b..d8b5bb856c6 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -417,8 +417,43 @@ def test_job_charge_identity_requires_a_cost_bearing_poll(call_type: str, polled from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once expected: Final = None if kind is None else ChargeOnceJob(kind=kind, job_id="job-budget") - assert job_to_charge_once(call_type, polled(), cost) == expected - assert job_to_charge_once(call_type, {"id": "job-budget"}, cost) is None + kwargs: Final = {"call_type": call_type, "standard_logging_object": {"response_cost": cost}} + assert job_to_charge_once(kwargs, polled(), cost) == expected + assert job_to_charge_once(kwargs, {"id": "job-budget"}, cost) is None + + +@pytest.mark.parametrize( + ("call_type", "polled", "kind"), + (("aretrieve_batch", _completed_batch, "batch"), ("avideo_status", _completed_video, "video")), +) +def test_only_the_jobs_own_cost_claims_its_charge_since_a_guardrail_charges_every_poll( + call_type: str, polled, kind +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + kwargs: Final = { + "call_type": call_type, + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + assert job_to_charge_once(kwargs, polled(), 0.01) is None + assert job_to_charge_once(kwargs, polled(), 0.43) == ChargeOnceJob(kind=kind, job_id="job-budget") + + +def test_a_video_is_charged_once_it_completes_under_the_id_the_client_polled() -> None: + """A provider may price a running video too, and a custom provider's own id can repeat across deployments.""" + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + from litellm.types.videos.main import VideoObject + + def poll(status: str, polled_id: str) -> ChargeOnceJob | None: + return job_to_charge_once( + {"call_type": "avideo_status", "litellm_params": {"video_id": polled_id}}, + VideoObject(id="job-1", object="video", status=status), + 0.42, + ) + + assert poll("in_progress", "video_deployment_a") is None + assert poll("completed", "video_deployment_a") == ChargeOnceJob(kind="video", job_id="video_deployment_a") + assert poll("completed", "video_deployment_b") == ChargeOnceJob(kind="video", job_id="video_deployment_b") def test_job_charge_markers_keep_counters_jobs_and_kinds_independent() -> None: 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 d9b06b6f10c..6ff4deb9e22 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -4787,7 +4787,12 @@ async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_c async def _update_database_for_video_poll( - db_writer: DBSpendUpdateWriter, prisma: MagicMock, litellm_call_id: str, response_cost: float + db_writer: DBSpendUpdateWriter, + prisma: MagicMock, + litellm_call_id: str, + response_cost: float, + status: str = "completed", + polled_video_id: str = "video_abc", ) -> bool: from litellm.types.videos.main import VideoObject @@ -4810,9 +4815,9 @@ async def _update_database_for_video_poll( "call_type": "avideo_status", "litellm_call_id": litellm_call_id, "response_cost": response_cost, - "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "litellm_params": {"metadata": {"user_api_key": "test-key"}, "video_id": polled_video_id}, }, - completion_response=VideoObject(id="video_abc", object="video", status="completed"), + completion_response=VideoObject(id="video_abc", object="video", status=status), start_time=datetime.now(timezone.utc), end_time=datetime.now(timezone.utc), response_cost=response_cost, @@ -4821,6 +4826,42 @@ async def _update_database_for_video_poll( return charged +@pytest.mark.asyncio +async def test_update_database_charges_the_completed_polls_cost_and_not_a_running_polls(): + """A provider that also prices a running video must still have the completed poll's cost billed, once.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + running = await _update_database_for_video_poll(db_writer, prisma, "poll-1", 0.1, status="in_progress") + completed = await _update_database_for_video_poll(db_writer, prisma, "poll-2", 0.42) + repolled = await _update_database_for_video_poll(db_writer, prisma, "poll-3", 0.42) + await _flush_queued_spend_logs(prisma) + + assert (running, completed, repolled) == (True, True, False) + assert prisma.db.litellm_spendlogs.rows["video_abc_video_cost"]["spend"] == 0.42 + + +@pytest.mark.asyncio +async def test_update_database_charges_each_deployments_video_though_their_provider_ids_match(): + """A custom provider's own ids can repeat across deployments, so the cost row keys off the id the client polled.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll(db_writer, prisma, call_id, 0.42, polled_video_id=polled_id) + for call_id, polled_id in (("poll-1", "video_dep_a"), ("poll-2", "video_dep_b"), ("poll-3", "video_dep_a")) + ] + + assert charged == [True, True, False] + assert { + request_id: row["spend"] + for request_id, row in prisma.db.litellm_spendlogs.rows.items() + if request_id.endswith("_video_cost") + } == {"video_dep_a_video_cost": 0.42, "video_dep_b_video_cost": 0.42} + + @pytest.mark.asyncio async def test_update_database_writes_no_row_of_its_own_for_a_video_poll_that_reports_no_cost(): """ diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 18b1bce8e84..ccfcb6ed4db 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -8,6 +8,8 @@ Tests that session-scoped budget tracking works correctly: - Requests without agent_id pass through """ +from collections.abc import Mapping +from types import MappingProxyType from unittest.mock import patch import logging @@ -179,6 +181,105 @@ async def test_repeated_polls_of_a_finished_job_charge_the_session_once(): assert spend == pytest.approx(0.4 + 0.1 + 3 * 0.01) +def _video_poll(cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs") -> dict[str, object]: + return { + "call_type": "avideo_status", + "response_cost": cost, + "standard_logging_object": { + "response_cost": cost, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], + }, + "litellm_params": {"metadata": {"session_id": session_id, "agent_id": "agent-budget-123"}}, + } + + +@pytest.mark.asyncio +async def test_a_guardrail_priced_poll_of_a_running_video_leaves_the_session_charge_to_the_finished_poll(): + from litellm.types.videos.main import VideoObject + + handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + await handler.async_log_success_event( + _video_poll(0.01, guardrail_cost=0.01), VideoObject(id="job-1", object="video", status="in_progress"), None, None + ) + for _ in range(3): + await handler.async_log_success_event( + _video_poll(0.41, guardrail_cost=0.01), VideoObject(id="job-1", object="video", status="completed"), None, None + ) + + spend = await handler._get_current_spend(handler._make_cache_key("session-jobs")) + assert spend == pytest.approx(0.01 + 0.41) + + +class _ClockedRedis: + """The Redis calls the session limiter makes, on a clock the test moves.""" + + def __init__(self) -> None: + self.now = 0.0 + self.entries: Mapping[str, tuple[float, float | None]] = MappingProxyType({}) + + def _live(self, key: str) -> tuple[float, float | None] | None: + entry = self.entries.get(key) + if entry is None or (entry[1] is not None and entry[1] <= self.now): + return None + return entry + + def _store(self, key: str, value: float, expires_at: float | None) -> float: + self.entries = MappingProxyType({**self.entries, key: (value, expires_at)}) + return value + + async def async_get_cache(self, key: str, **kwargs: object) -> float | None: + entry = self._live(key) + return None if entry is None else entry[0] + + async def async_increment( + self, key: str, value: float, ttl: int | None = None, parent_otel_span: object = None, refresh_ttl: bool = False + ) -> float: + live = self._live(key) + kept_expiry = None if live is None else live[1] + expires_at = kept_expiry if ttl is None or (live is not None and not refresh_ttl) else self.now + ttl + return self._store(key, value + (0.0 if live is None else live[0]), expires_at) + + def async_register_script(self, script: str): + async def increment_expiring_on_first_write(keys: list[str], args: list[object]) -> float: + live = self._live(keys[0]) + expires_at = self.now + float(str(args[1])) if live is None else live[1] + return self._store(keys[0], float(str(args[0])) + (0.0 if live is None else live[0]), expires_at) + + return increment_expiring_on_first_write + + +@pytest.mark.asyncio +async def test_a_finished_video_polled_without_pause_is_not_charged_again_once_the_session_window_rolls_over(): + from litellm.types.videos.main import VideoObject + + redis = _ClockedRedis() + handler = PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis)) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + video = VideoObject(id="job-1", object="video", status="completed") + chat = { + "call_type": "acompletion", + "response_cost": 0.01, + "litellm_params": {"metadata": {"session_id": "session-jobs", "agent_id": "agent-budget-123"}}, + } + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + for _ in range(3): + await handler.async_log_success_event(_video_poll(0.4), video, None, None) + redis.now += 0.6 * handler.ttl + await handler.async_log_success_event(chat, None, None, None) + + assert await handler._get_current_spend(handler._make_cache_key("session-jobs")) == pytest.approx(0.01) + + @pytest.mark.asyncio @pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), True]) async def test_a_cost_that_is_not_a_finite_number_leaves_the_session_spend_alone(response_cost): 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 48b44a75cef..4747b397f20 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -58,12 +58,14 @@ def _event( team_member_model_max_budget: Mapping[str, object] | None = None, user_id: str = USER_ID, team_id: str = TEAM_ID, + guardrail_information: tuple[Mapping[str, object], ...] = (), ) -> dict[str, object]: return { "call_type": call_type, "standard_logging_object": { "call_type": call_type, "response_cost": response_cost, + "guardrail_information": guardrail_information, "model": "openai/gpt-5.4-mini", "model_group": MODEL_GROUP, "metadata": { @@ -351,6 +353,28 @@ async def test_polls_of_a_finished_video_charge_each_per_model_budget_once(): assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(3 * BATCH_COST) assert await _spend(limiter, USER_SPEND_KEY) == pytest.approx(3 * BATCH_COST) + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "polled"), + ( + ("avideo_status", lambda status: VideoObject(id="job_first", object="video", status=status)), + ("aretrieve_batch", lambda status: _batch("job_first", status)), + ), +) +async def test_a_guardrail_priced_poll_of_a_running_job_leaves_its_charge_to_the_finished_poll(call_type, polled): + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + guardrail_cost: Final = 0.01 + priced: Final = ({"guardrail_name": "priced", "guardrail_cost": guardrail_cost},) + + await _log_success(limiter, _event(call_type, guardrail_cost, guardrail_information=priced), polled("in_progress")) + for _ in range(3): + await _log_success( + limiter, _event(call_type, BATCH_COST + guardrail_cost, guardrail_information=priced), polled("completed") + ) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(guardrail_cost + BATCH_COST + guardrail_cost) + OPUS_GROUP: Final = "claude-opus" OPUS_GROUP_MEMBERS: Final = MappingProxyType( {"claude-opus-4-8": frozenset({OPUS_GROUP}), "claude-opus-5-5": frozenset({OPUS_GROUP})} diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index b5c62e86ee5..db05159c5c1 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5093,6 +5093,24 @@ def test_get_logging_payload_gives_only_a_charged_video_poll_a_row_of_its_own(st assert status["request_id"] == status_id +@pytest.mark.parametrize( + ("response_cost", "status_id"), + [(0.01, "video_abc"), (0.43, "video_abc_video_cost")], + ids=["guardrail_only", "charged_with_guardrail"], +) +def test_a_video_poll_that_pays_only_its_guardrail_keeps_the_bare_id(response_cost, status_id): + """A guardrail charges every poll, so a running video's poll must not take the row the finished poll charges.""" + kwargs = { + "response_cost": response_cost, + "standard_logging_object": { + "response_cost": response_cost, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}], + }, + } + + assert get_spend_logs_id("avideo_status", {"id": "video_abc", "status": "completed"}, kwargs) == status_id + + def test_get_logging_payload_gives_redacted_batch_and_file_rows_distinct_request_ids(): """End to end at the payload level: a batch retrieve and a file create whose bodies were both flattened to the same redaction placeholder must still produce two diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 7fb43f17209..b0550814bce 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -7,6 +7,7 @@ the callback raise before any spend was recorded, so those budgets never moved. """ from datetime import datetime, timezone +from functools import partial from typing import Final import litellm @@ -38,6 +39,7 @@ def _success_kwargs( call_type: str = "aresponses", response_cost: float = 0.25, model_id: str = "deployment-1", + guardrail_information: tuple[dict[str, object], ...] = (), ) -> dict[str, object]: provider_params: Final[dict[str, str]] = ( {} if provider_in_litellm_params is None else {"custom_llm_provider": provider_in_litellm_params} @@ -51,6 +53,7 @@ def _success_kwargs( "response_cost": response_cost, "model_id": model_id, "custom_llm_provider": provider_in_payload, + "guardrail_information": guardrail_information, }, } @@ -528,6 +531,43 @@ async def test_video_polls_charge_each_router_budget_once(call_type): ) assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.5, 0.5) + +@pytest.mark.asyncio +async def test_a_guardrail_priced_poll_of_a_running_video_leaves_the_router_charge_to_the_finished_poll(): + from litellm.types.videos.main import VideoObject + + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "video-budget", + "litellm_params": {"model": "openai/sora-2", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + priced: Final = ({"guardrail_name": "priced", "guardrail_cost": 0.01},) + poll: Final = partial( + _success_kwargs, + provider_in_litellm_params=None, + provider_in_payload="openai", + call_type="avideo_status", + guardrail_information=priced, + ) + + await limiter.async_log_success_event( + poll(response_cost=0.01), VideoObject(id="video-first", object="video", status="in_progress"), None, None + ) + for _ in range(3): + await limiter.async_log_success_event( + poll(response_cost=0.26), VideoObject(id="video-first", object="video", status="completed"), None, None + ) + + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") + assert tuple([await cache.async_get_cache(key) for key in keys]) == pytest.approx((0.27, 0.27)) + @pytest.mark.asyncio async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(): from fakeredis import FakeServer, FakeRedis as FakeSyncRedis From 3b2027417b4238f591c4b3a208853e8fa3976632 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 04:10:03 -0800 Subject: [PATCH 04/10] fix(proxy): charge repeat poll guardrails, key video charges by canonical id and survive redaction --- litellm/litellm_core_utils/batch_utils.py | 59 ++++++++++++---- litellm/litellm_core_utils/litellm_logging.py | 3 + .../hooks/max_budget_per_session_limiter.py | 7 +- .../proxy/hooks/model_max_budget_limiter.py | 8 +++ .../proxy/hooks/proxy_track_cost_callback.py | 67 ++++++++++++++----- .../spend_tracking/spend_tracking_utils.py | 10 +-- litellm/router_strategy/budget_limiter.py | 7 ++ litellm/types/videos/utils.py | 44 ++++++------ .../router_code_coverage.py | 3 + .../litellm_core_utils/test_batch_utils.py | 47 +++++++++++-- .../test_litellm_logging.py | 42 ++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 1 + .../test_max_budget_per_session_limiter.py | 16 +++-- .../hooks/test_model_max_budget_limiter.py | 18 ++--- .../hooks/test_proxy_track_cost_callback.py | 42 +++++++++++- .../test_spend_tracking_utils.py | 33 +++++++-- .../router_strategy/test_budget_limiter.py | 20 ++++-- 17 files changed, 343 insertions(+), 84 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index a1fa6e39e4b..7f75619a464 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -16,7 +16,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_in from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose from litellm.types.utils import CallTypes, LlmProviders from litellm.types.videos.main import VideoObject -from litellm.types.videos.utils import is_video_status_call_type +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG, canonical_video_id, is_video_status_call_type _DECODED_JSONL_ROW: Final = TypeAdapter(object) _DEPLOYMENT_FIELD: Final = TypeAdapter(Mapping[str, object]) @@ -299,20 +299,41 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: class ChargeOnceJob: kind: Literal["batch", "video"] job_id: str + repeat_poll_cost: float = 0.0 def job_charged_once_marker_key(spend_key: str, job: ChargeOnceJob) -> str: return f"{spend_key}:{job.kind}:{job.job_id}" -def job_cost_of_poll(response_cost: object, standard_logging_payload: object) -> float: - """What the polled job itself cost, since a guardrail on the poll charges every poll, running or finished.""" +def _guardrail_cost_of_poll(standard_logging_payload: object) -> float: guardrail_information: Final = ( _PAYLOAD_FIELDS.validate_python(standard_logging_payload).get("guardrail_information") if isinstance(standard_logging_payload, Mapping) else None ) - return (finite_cost(response_cost) or 0.0) - guardrail_information_cost(guardrail_information) + return guardrail_information_cost(guardrail_information) + + +def job_cost_of_poll(response_cost: object, standard_logging_payload: object) -> float: + """What the polled job itself cost, since a guardrail on the poll charges every poll, running or finished.""" + return (finite_cost(response_cost) or 0.0) - _guardrail_cost_of_poll(standard_logging_payload) + + +def repeat_poll_cost(response_cost: object, standard_logging_payload: object) -> float: + """What a poll still costs once an earlier poll charged its job: the guardrails that ran on it.""" + return max(0.0, min(finite_cost(response_cost) or 0.0, _guardrail_cost_of_poll(standard_logging_payload))) + + +def pays_only_its_guardrail(call_type: object, kwargs: Mapping[str, object], response_cost: object) -> bool: + """A batch or video poll whose whole cost is its guardrails', charged like any other request.""" + if not is_batch_retrieve_call_type(call_type) and not is_video_status_call_type(call_type): + return False + return ( + job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) + <= 0 + < (finite_cost(response_cost) or 0.0) + ) def _requested_video_id(kwargs: Mapping[str, object]) -> str | None: @@ -324,35 +345,47 @@ def _requested_video_id(kwargs: Mapping[str, object]) -> str | None: def charged_video_id( - call_type: object, status: object, video_id: str, kwargs: Mapping[str, object], response_cost: object + call_type: object, kwargs: Mapping[str, object], fallback_video_id: str | None, response_cost: object ) -> str | None: """The id a finished video's own cost is charged once under. That is the id the client polled, which names the deployment that made the video, since a - custom provider's own id can repeat across deployments. + custom provider's own id can repeat across deployments. The status comes from the logging + kwargs because message redaction replaces the logged video object. """ - if not is_video_status_call_type(call_type) or status != "completed": + if not is_video_status_call_type(call_type) or kwargs.get(POLLED_VIDEO_STATUS_KWARG) != "completed": return None if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0: return None - return _requested_video_id(kwargs) or video_id + video_id: Final = _requested_video_id(kwargs) or fallback_video_id + return None if video_id is None else canonical_video_id(video_id) def _charged_batch(kwargs: Mapping[str, object], response_obj: object, response_cost: float) -> ChargeOnceJob | None: if not is_batch_retrieve_call_type(kwargs.get("call_type")) or not isinstance(response_obj, Batch): return None - if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0: + standard_logging_payload: Final = kwargs.get("standard_logging_object") + if job_cost_of_poll(response_cost, standard_logging_payload) <= 0: return None - return ChargeOnceJob(kind="batch", job_id=response_obj.id) + return ChargeOnceJob( + kind="batch", job_id=response_obj.id, repeat_poll_cost=repeat_poll_cost(response_cost, standard_logging_payload) + ) def job_to_charge_once( kwargs: Mapping[str, object], response_obj: object, response_cost: float ) -> ChargeOnceJob | None: """A finished batch or video reports its whole cost on every poll, so it is charged once per counter.""" - if not isinstance(response_obj, VideoObject): + call_type: Final = kwargs.get("call_type") + if not is_video_status_call_type(call_type): return _charged_batch(kwargs, response_obj, response_cost) video_id: Final = charged_video_id( - kwargs.get("call_type"), response_obj.status, response_obj.id, kwargs, response_cost + call_type, kwargs, response_obj.id if isinstance(response_obj, VideoObject) else None, response_cost + ) + if video_id is None: + return None + return ChargeOnceJob( + kind="video", + job_id=video_id, + repeat_poll_cost=repeat_poll_cost(response_cost, kwargs.get("standard_logging_object")), ) - return None if video_id is None else ChargeOnceJob(kind="video", job_id=video_id) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 37657a94f16..85c0838482c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -211,6 +211,7 @@ from litellm.types.utils import ( batch_chat_usage, ) from litellm.types.videos.main import VideoObject +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG from litellm.utils import get_base_model_from_metadata, print_verbose from ..integrations.argilla import ArgillaLogger @@ -3311,6 +3312,8 @@ class Logging(LiteLLMLoggingBaseClass): ) result = self.normalize_endpoint_logging_result(result) + if isinstance(result, VideoObject): + self.model_call_details[POLLED_VIDEO_STATUS_KWARG] = result.status logging_result: Final = self.normalize_logging_result(result=result) if cache_hit is None and _served_from_cache_at_no_charge(logging_result): diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 6b60a8b8de0..c0e76d758be 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -178,9 +178,12 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): cache_key: Final = self._make_cache_key(str(session_id)) job: Final = job_to_charge_once(callback_fields, response_obj, response_cost) - if job is not None and not await self._claim_job_charge(cache_key, job): + charge: Final = ( + response_cost if job is None or await self._claim_job_charge(cache_key, job) else job.repeat_poll_cost + ) + if charge <= 0: return - await self._increment_spend(cache_key, response_cost) + await self._increment_spend(cache_key, charge) verbose_proxy_logger.debug( "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 6414b7613b9..bd9fe7208ed 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1248,6 +1248,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): job=job, ttl_seconds=duration_in_seconds(budget_duration), ): + if job.repeat_poll_cost > 0: + await self._charge_entity( + entity_type=entity_type, + entity_id=entity_id, + resolved=resolved, + response_cost=job.repeat_poll_cost, + job=None, + ) return start_time_key: Final = model_budget_start_time_cache_key( entity_type=entity_type, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 980f0a5d410..c7b516a2a38 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -4,6 +4,7 @@ from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar from datetime import datetime, timezone +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast @@ -15,6 +16,7 @@ from litellm.batches.batch_utils import batch_cost_is_final from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import repeat_poll_cost from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, finite_cost, @@ -983,25 +985,28 @@ class _ProxyDBLogger(CustomLogger): agent_id=metadata.get("billing_agent_id") or metadata.get("agent_id"), ): ## UPDATE DATABASE - charged: Final = await _update_database_and_spend_counters( - proxy_logging_obj=proxy_logging_obj, - increment_spend_counters=increment_spend_counters, - user_api_key=user_api_key, - user_id=user_id, - end_user_id=end_user_id, - team_id=team_id, - org_id=org_id, + charged_cost: Final = await _charge_request_cost( + partial( + _update_database_and_spend_counters, + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key=user_api_key, + user_id=user_id, + end_user_id=end_user_id, + team_id=team_id, + org_id=org_id, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, + request_tags=tags, + model_access_groups=model_access_groups, + project_id=project_id, + ), kwargs=kwargs, - completion_response=completion_response, - start_time=start_time, - end_time=end_time, response_cost=response_cost, budget_reservation=budget_reservation, - request_tags=tags, - model_access_groups=model_access_groups, - project_id=project_id, ) - if not charged: + if charged_cost is None: return # update cache (fire-and-forget for backward compat: @@ -1011,7 +1016,7 @@ class _ProxyDBLogger(CustomLogger): token=user_api_key, user_id=user_id, end_user_id=end_user_id, - response_cost=response_cost, + response_cost=charged_cost, team_id=team_id, parent_otel_span=parent_otel_span, tags=tags, @@ -1022,7 +1027,7 @@ class _ProxyDBLogger(CustomLogger): token=user_api_key, key_alias=key_alias, end_user_id=end_user_id, - response_cost=response_cost, + response_cost=charged_cost, max_budget=end_user_max_budget, ) elif budget_reservation is not None: @@ -1316,6 +1321,34 @@ def _get_request_tags_for_cost_tracking( return None +class _ChargeSpend(Protocol): + async def __call__( + self, + *, + kwargs: dict[str, object], # mutable-ok: update_database takes the logging kwargs dict itself + response_cost: float, + budget_reservation: dict[str, object] | None, # mutable-ok: reconcile stamps the shared reservation dict + ) -> bool: ... + + +async def _charge_request_cost( + charge_spend: _ChargeSpend, + kwargs: dict[str, object], # mutable-ok: the first charge passes the caller's logging kwargs dict on unchanged + response_cost: float, + budget_reservation: dict[str, object] | None, # mutable-ok: reconcile stamps the shared reservation dict +) -> float | None: + """What the request charged, since a poll of a job an earlier poll charged still pays for its own guardrails.""" + if await charge_spend(kwargs=kwargs, response_cost=response_cost, budget_reservation=budget_reservation): + return response_cost + cost_after_job: Final = repeat_poll_cost(response_cost, kwargs.get("standard_logging_object")) + if cost_after_job <= 0: + return None + charged: Final = await charge_spend( + kwargs={**kwargs, "response_cost": cost_after_job}, response_cost=cost_after_job, budget_reservation=None + ) + return cost_after_job if charged else None + + class _IncrementSpendCounters(Protocol): """The ``increment_spend_counters`` coroutine :func:`_update_database_and_spend_counters` awaits.""" diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 03ef44bf34b..01f2b4d6a90 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -38,7 +38,7 @@ from litellm.constants import ( ) from litellm.integrations.otel.model.metadata import caller_session_id from litellm.integrations.otel.model.trace_controls import TraceControls -from litellm.litellm_core_utils.batch_utils import charged_video_id +from litellm.litellm_core_utils.batch_utils import charged_video_id, pays_only_its_guardrail from litellm.litellm_core_utils.classifier_logging import ( CLASSIFIER_AUDIT_FIELDS, classifier_audit_fields, @@ -373,11 +373,13 @@ def get_spend_logs_id( ) if resolved_id is None: return None - charged_video: Final = charged_video_id( - call_type, response_obj.get("status"), resolved_id, kwargs, kwargs.get("response_cost") - ) + response_cost: Final = kwargs.get("response_cost") + charged_video: Final = charged_video_id(call_type, kwargs, resolved_id, response_cost) if charged_video is not None: return f"{charged_video}{VIDEO_COST_REQUEST_ID_SUFFIX}" + call_id: Final = kwargs.get("litellm_call_id") + if pays_only_its_guardrail(call_type, kwargs, response_cost) and isinstance(call_id, str) and call_id: + return call_id return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index bdcba8c8ae1..602c5b7279b 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -613,6 +613,13 @@ class RouterBudgetLimiting(CustomLogger): current_time: Final = self._clock() new_window_seconds: Final = budget_window_seconds(current_time, budget_config.budget_duration) if job is not None and not await self._claim_job_charge(spend_key, job, new_window_seconds): + if job.repeat_poll_cost > 0: + await self._increment_spend_for_key( + budget_config=budget_config, + spend_key=spend_key, + start_time_key=start_time_key, + response_cost=job.repeat_poll_cost, + ) return budget_start: Final = await self._get_or_set_budget_start_time( diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index c74f11b7abb..240e36ecdf6 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -19,6 +19,7 @@ VIDEO_ID_DEPLOYMENT_PREFIX: Final = "deployment_id:" VIDEO_ID_REQUESTED_MODEL_PREFIX: Final = "requested_model:" VIDEO_STATUS_CALL_TYPES: Final = frozenset({"avideo_status", "video_status"}) VIDEO_COST_REQUEST_ID_SUFFIX: Final = "_video_cost" +POLLED_VIDEO_STATUS_KWARG: Final = "polled_video_status" VIDEO_ID_CALL_TYPES: Final = VIDEO_STATUS_CALL_TYPES | frozenset( { "avideo_content", @@ -100,6 +101,27 @@ def _trailing_part(parts: Sequence[str], prefix: str) -> str | None: return next((part.removeprefix(prefix) or None for part in parts[3:] if part.startswith(prefix)), None) +def _decoded_video_id_payload(encoded_video_id: str) -> str | None: + if not encoded_video_id.startswith(VIDEO_ID_PREFIX): + return None + try: + decoded_id: Final = base64.b64decode( + _add_base64_padding(encoded_video_id.removeprefix(VIDEO_ID_PREFIX)).encode("utf-8") + ).decode("utf-8") + except ValueError as e: + verbose_logger.debug("Error decoding video_id '%s': %s", encoded_video_id, e) + return None + return decoded_id if ";" in decoded_id else None + + +def canonical_video_id(video_id: str) -> str: + """One spelling per video, since a client may poll a LiteLLM video id with or without its base64 padding.""" + payload: Final = _decoded_video_id_payload(video_id) + if payload is None: + return video_id + return f"{VIDEO_ID_PREFIX}{base64.b64encode(payload.encode('utf-8')).decode('utf-8')}" + + def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: """Decode provider and model_id from encoded video_id.""" if not encoded_video_id: @@ -109,26 +131,8 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: video_id=encoded_video_id, ) - if not encoded_video_id.startswith(VIDEO_ID_PREFIX): - return DecodedVideoId( - custom_llm_provider=None, - model_id=None, - video_id=encoded_video_id, - ) - - try: - cleaned_id = encoded_video_id.replace(VIDEO_ID_PREFIX, "") - cleaned_id = _add_base64_padding(cleaned_id) - decoded_id: Final = base64.b64decode(cleaned_id.encode("utf-8")).decode("utf-8") - except Exception as e: - verbose_logger.debug("Error decoding video_id '%s': %s", encoded_video_id, e) - return DecodedVideoId( - custom_llm_provider=None, - model_id=None, - video_id=encoded_video_id, - ) - - if ";" not in decoded_id: + decoded_id: Final = _decoded_video_id_payload(encoded_video_id) + if decoded_id is None: return DecodedVideoId( custom_llm_provider=None, model_id=None, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 58d0e373eb0..6724cf176d7 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -178,6 +178,9 @@ ignored_function_names = [ "job_cost_of_poll", # Read by job_to_charge_once and charged_video_id, tested through the guardrail-priced poll tests in test_batch_utils.py "charged_video_id", # Read by job_to_charge_once and get_spend_logs_id, tested through both in test_batch_utils.py and test_spend_tracking_utils.py "_requested_video_id", # Read by charged_video_id, tested through the polled-id video test in test_batch_utils.py + "_guardrail_cost_of_poll", # Read by job_cost_of_poll and repeat_poll_cost, tested through the guardrail-priced poll tests + "repeat_poll_cost", # Read by job_to_charge_once and _charge_request_cost, tested through every limiter's guardrail poll test + "pays_only_its_guardrail", # Read by get_spend_logs_id, tested through its guardrail-only poll test in test_spend_tracking_utils.py "_charged_batch", # Read by job_to_charge_once, tested through the batch charge identity tests in test_batch_utils.py "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index d8b5bb856c6..cf643482a78 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.batch_utils import ( parse_jsonl_with_embedded_newlines, replace_model_in_jsonl, ) +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG, encode_video_id_with_provider from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -417,7 +418,11 @@ def test_job_charge_identity_requires_a_cost_bearing_poll(call_type: str, polled from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once expected: Final = None if kind is None else ChargeOnceJob(kind=kind, job_id="job-budget") - kwargs: Final = {"call_type": call_type, "standard_logging_object": {"response_cost": cost}} + kwargs: Final = { + "call_type": call_type, + POLLED_VIDEO_STATUS_KWARG: "completed", + "standard_logging_object": {"response_cost": cost}, + } assert job_to_charge_once(kwargs, polled(), cost) == expected assert job_to_charge_once(kwargs, {"id": "job-budget"}, cost) is None @@ -426,17 +431,20 @@ def test_job_charge_identity_requires_a_cost_bearing_poll(call_type: str, polled ("call_type", "polled", "kind"), (("aretrieve_batch", _completed_batch, "batch"), ("avideo_status", _completed_video, "video")), ) -def test_only_the_jobs_own_cost_claims_its_charge_since_a_guardrail_charges_every_poll( +def test_only_the_jobs_own_cost_is_charged_once_since_a_guardrail_charges_every_poll( call_type: str, polled, kind ) -> None: from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once kwargs: Final = { "call_type": call_type, + POLLED_VIDEO_STATUS_KWARG: "completed", "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, } assert job_to_charge_once(kwargs, polled(), 0.01) is None - assert job_to_charge_once(kwargs, polled(), 0.43) == ChargeOnceJob(kind=kind, job_id="job-budget") + assert job_to_charge_once(kwargs, polled(), 0.43) == ChargeOnceJob( + kind=kind, job_id="job-budget", repeat_poll_cost=0.01 + ) def test_a_video_is_charged_once_it_completes_under_the_id_the_client_polled() -> None: @@ -446,7 +454,7 @@ def test_a_video_is_charged_once_it_completes_under_the_id_the_client_polled() - def poll(status: str, polled_id: str) -> ChargeOnceJob | None: return job_to_charge_once( - {"call_type": "avideo_status", "litellm_params": {"video_id": polled_id}}, + {"call_type": "avideo_status", POLLED_VIDEO_STATUS_KWARG: status, "litellm_params": {"video_id": polled_id}}, VideoObject(id="job-1", object="video", status=status), 0.42, ) @@ -456,6 +464,37 @@ def test_a_video_is_charged_once_it_completes_under_the_id_the_client_polled() - assert poll("completed", "video_deployment_b") == ChargeOnceJob(kind="video", job_id="video_deployment_b") +@pytest.mark.parametrize("status", ["in_progress", "completed"]) +def test_a_video_poll_is_judged_by_its_logged_status_when_redaction_replaced_the_video(status: str) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + kwargs: Final = { + "call_type": "avideo_status", + POLLED_VIDEO_STATUS_KWARG: status, + "litellm_params": {"video_id": "video_deployment_a"}, + } + expected: Final = ChargeOnceJob(kind="video", job_id="video_deployment_a") if status == "completed" else None + + assert job_to_charge_once(kwargs, {"text": "redacted-by-litellm"}, 0.42) == expected + + +def test_a_video_polled_with_or_without_its_base64_padding_is_charged_as_one_job() -> None: + from litellm.litellm_core_utils.batch_utils import job_to_charge_once + + padded: Final = encode_video_id_with_provider("job-1", "custom_video", "deployment-a") + assert padded.endswith("=") + + def charged_job(polled_id: str): + return job_to_charge_once( + {"call_type": "avideo_status", POLLED_VIDEO_STATUS_KWARG: "completed", "litellm_params": {"video_id": polled_id}}, + None, + 0.42, + ) + + assert charged_job(padded.rstrip("=")) == charged_job(padded) + assert charged_job(padded) is not None + + def test_job_charge_markers_keep_counters_jobs_and_kinds_independent() -> None: from litellm.caching.in_memory_cache import InMemoryCache from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 13e181b2b92..0dee6d2b45f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14451,3 +14451,45 @@ def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_call (first_cost, second_cost) = costs assert second_cost == pytest.approx(2 * first_cost) assert shared_metadata == {"tag": "x"} + + +@pytest.mark.asyncio +async def test_a_redacted_video_poll_still_tells_the_spend_callbacks_the_video_finished( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + from litellm.types.videos.main import VideoObject + + class Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: tuple[tuple[dict[str, object], object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events = (*self.events, (dict(kwargs), response_obj)) + + recorder: Final = Recorder() + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + logging_obj: Final = LitellmLogging( + model="custom-video", + messages=[], + stream=False, + call_type="avideo_status", + start_time=datetime.datetime.now(), + litellm_call_id="poll-1", + function_id="poll-1", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}, "video_id": "video_abc"}, optional_params={} + ) + + await logging_obj.async_success_handler(result=VideoObject(id="video_abc", object="video", status="completed")) + + ((kwargs, response_obj),) = recorder.events + assert not isinstance(response_obj, VideoObject) + assert job_to_charge_once(kwargs, response_obj, 0.42) == ChargeOnceJob(kind="video", job_id="video_abc") 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 6ff4deb9e22..8e9d9ea843d 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -4815,6 +4815,7 @@ async def _update_database_for_video_poll( "call_type": "avideo_status", "litellm_call_id": litellm_call_id, "response_cost": response_cost, + "polled_video_status": status, "litellm_params": {"metadata": {"user_api_key": "test-key"}, "video_id": polled_video_id}, }, completion_response=VideoObject(id="video_abc", object="video", status=status), diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index ccfcb6ed4db..32f25892d3d 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -24,6 +24,7 @@ from litellm.proxy.hooks.max_budget_per_session_limiter import ( ) from litellm.proxy.utils import InternalUsageCache from litellm.types.agents import AgentResponse +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG def _make_mock_agent(max_budget_per_session: float) -> AgentResponse: @@ -170,6 +171,7 @@ async def test_repeated_polls_of_a_finished_job_charge_the_session_once(): { "call_type": call_type, "response_cost": cost, + POLLED_VIDEO_STATUS_KWARG: "completed", "litellm_params": {"metadata": {"session_id": "session-jobs", "agent_id": "agent-budget-123"}}, }, response_obj, @@ -181,10 +183,13 @@ async def test_repeated_polls_of_a_finished_job_charge_the_session_once(): assert spend == pytest.approx(0.4 + 0.1 + 3 * 0.01) -def _video_poll(cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs") -> dict[str, object]: +def _video_poll( + cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs", status: str = "completed" +) -> dict[str, object]: return { "call_type": "avideo_status", "response_cost": cost, + POLLED_VIDEO_STATUS_KWARG: status, "standard_logging_object": { "response_cost": cost, "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], @@ -194,7 +199,7 @@ def _video_poll(cost: float, guardrail_cost: float = 0.0, session_id: str = "ses @pytest.mark.asyncio -async def test_a_guardrail_priced_poll_of_a_running_video_leaves_the_session_charge_to_the_finished_poll(): +async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_session_for_the_video(): from litellm.types.videos.main import VideoObject handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) @@ -204,7 +209,10 @@ async def test_a_guardrail_priced_poll_of_a_running_video_leaves_the_session_cha ) as mock_registry: mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) await handler.async_log_success_event( - _video_poll(0.01, guardrail_cost=0.01), VideoObject(id="job-1", object="video", status="in_progress"), None, None + _video_poll(0.01, guardrail_cost=0.01, status="in_progress"), + VideoObject(id="job-1", object="video", status="in_progress"), + None, + None, ) for _ in range(3): await handler.async_log_success_event( @@ -212,7 +220,7 @@ async def test_a_guardrail_priced_poll_of_a_running_video_leaves_the_session_cha ) spend = await handler._get_current_spend(handler._make_cache_key("session-jobs")) - assert spend == pytest.approx(0.01 + 0.41) + assert spend == pytest.approx(0.01 + 0.41 + 2 * 0.01) class _ClockedRedis: 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 4747b397f20..6d387e39efe 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -28,6 +28,7 @@ from litellm.router import Router from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import LiteLLMBatch, Usage from litellm.types.videos.main import VideoObject +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG KEY_HASH: Final = "key-hash-batch" USER_ID: Final = "user-batch" @@ -335,7 +336,7 @@ async def _poll_video( ) -> None: await _log_success( limiter, - _event("avideo_status", response_cost), + {**_event("avideo_status", response_cost), POLLED_VIDEO_STATUS_KWARG: status}, response_obj=VideoObject(id=video_id, object="video", status=status), ) @@ -362,18 +363,19 @@ async def test_polls_of_a_finished_video_charge_each_per_model_budget_once(): ("aretrieve_batch", lambda status: _batch("job_first", status)), ), ) -async def test_a_guardrail_priced_poll_of_a_running_job_leaves_its_charge_to_the_finished_poll(call_type, polled): +async def test_every_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_job(call_type, polled): limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) guardrail_cost: Final = 0.01 priced: Final = ({"guardrail_name": "priced", "guardrail_cost": guardrail_cost},) - await _log_success(limiter, _event(call_type, guardrail_cost, guardrail_information=priced), polled("in_progress")) - for _ in range(3): - await _log_success( - limiter, _event(call_type, BATCH_COST + guardrail_cost, guardrail_information=priced), polled("completed") - ) + def poll(status: str, response_cost: float) -> dict[str, object]: + return {**_event(call_type, response_cost, guardrail_information=priced), POLLED_VIDEO_STATUS_KWARG: status} - assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(guardrail_cost + BATCH_COST + guardrail_cost) + await _log_success(limiter, poll("in_progress", guardrail_cost), polled("in_progress")) + for _ in range(3): + await _log_success(limiter, poll("completed", BATCH_COST + guardrail_cost), polled("completed")) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST + 4 * guardrail_cost) OPUS_GROUP: Final = "claude-opus" OPUS_GROUP_MEMBERS: Final = MappingProxyType( diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index a4e78a48a1a..efc514697f2 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -4,7 +4,7 @@ import logging from collections.abc import Mapping from datetime import datetime, timedelta, timezone from pathlib import Path -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, create_autospec, patch @@ -24,6 +24,7 @@ from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( ProxyDBLogger, + _charge_request_cost, _get_budget_reservation_from_metadata, _should_track_cost_callback, _update_database_and_spend_counters, @@ -4086,3 +4087,42 @@ async def test_capture_and_legacy_share_resolved_key_attribution( legacy: Final = TypeAdapter(dict[str, object]).validate_python(writer.rows[0]["kwargs"]) payload: Final = get_logging_payload(legacy, {"id": "tool"}, _CALL_START, _CALL_START) assert (payload["user"], payload["team_id"], payload["organization_id"]) == (owner, "team", "org") + + +class _RecordedCharges: + def __init__(self, outcomes: tuple[bool, ...]) -> None: + self.outcomes = outcomes + self.calls: tuple[tuple[object, float, dict[str, object] | None], ...] = () + + async def __call__( + self, *, kwargs: dict, response_cost: float, budget_reservation: dict[str, object] | None + ) -> bool: + self.calls = (*self.calls, (kwargs.get("response_cost"), response_cost, budget_reservation)) + return self.outcomes[len(self.calls) - 1] + + +_GUARDED_POLL: Final = MappingProxyType( + { + "response_cost": 0.43, + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("outcomes", "poll", "charged", "second_call"), + [ + ((True,), _GUARDED_POLL, 0.43, None), + ((False, True), _GUARDED_POLL, 0.01, (0.01, 0.01, None)), + ((False, False), _GUARDED_POLL, None, (0.01, 0.01, None)), + ((False,), {"response_cost": 0.43}, None, None), + ], + ids=["first_poll_pays_the_job", "repeat_poll_pays_its_guardrail", "guardrail_charge_not_written", "repeat_poll_free"], +) +async def test_a_poll_of_an_already_charged_job_still_charges_its_own_guardrails(outcomes, poll, charged, second_call): + charge_spend: Final = _RecordedCharges(outcomes) + reservation: Final = {"reserved": 0.43} + + assert await _charge_request_cost(charge_spend, dict(poll), 0.43, reservation) == charged + assert charge_spend.calls == ((0.43, 0.43, reservation), *(() if second_call is None else (second_call,))) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index db05159c5c1..a25d5f6c4b0 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5080,6 +5080,7 @@ def test_get_logging_payload_gives_only_a_charged_video_poll_a_row_of_its_own(st "call_type": call_type, "litellm_call_id": call_id, "response_cost": cost, + "polled_video_status": "completed", "litellm_params": {"metadata": {"user_api_key": "test-key"}}, }, response_obj=video, @@ -5094,21 +5095,41 @@ def test_get_logging_payload_gives_only_a_charged_video_poll_a_row_of_its_own(st @pytest.mark.parametrize( - ("response_cost", "status_id"), - [(0.01, "video_abc"), (0.43, "video_abc_video_cost")], - ids=["guardrail_only", "charged_with_guardrail"], + ("call_type", "job_id", "charged_id"), + [("avideo_status", "video_abc", "video_abc_video_cost"), ("aretrieve_batch", "batch_abc", "batch_abc_batch_cost")], ) -def test_a_video_poll_that_pays_only_its_guardrail_keeps_the_bare_id(response_cost, status_id): - """A guardrail charges every poll, so a running video's poll must not take the row the finished poll charges.""" +@pytest.mark.parametrize(("response_cost", "pays_the_job"), [(0.01, False), (0.43, True)]) +def test_a_job_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own( + call_type, job_id, charged_id, response_cost, pays_the_job +): + """A guardrail charges every poll, so a poll paying only that must neither land on the create row + nor take the row the job's own charge goes on, either of which the skip_duplicates insert drops.""" kwargs = { + "litellm_call_id": "poll-call-id", "response_cost": response_cost, + "polled_video_status": "completed", "standard_logging_object": { "response_cost": response_cost, "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}], }, } - assert get_spend_logs_id("avideo_status", {"id": "video_abc", "status": "completed"}, kwargs) == status_id + assert get_spend_logs_id(call_type, {"id": job_id, "status": "completed"}, kwargs) == ( + charged_id if pays_the_job else "poll-call-id" + ) + + +def test_a_video_polled_with_or_without_its_base64_padding_charges_one_row(): + from litellm.types.videos.utils import encode_video_id_with_provider + + padded = encode_video_id_with_provider("job-1", "custom_video", "deployment-a") + assert padded.endswith("=") + + def row_id(polled_id: str) -> str | None: + kwargs = {"response_cost": 0.42, "polled_video_status": "completed", "litellm_params": {"video_id": polled_id}} + return get_spend_logs_id("avideo_status", {"id": "job-1", "status": "completed"}, kwargs) + + assert row_id(padded.rstrip("=")) == row_id(padded) == f"{padded}_video_cost" def test_get_logging_payload_gives_redacted_batch_and_file_rows_distinct_request_ids(): diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index b0550814bce..c2df918e1d0 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -18,6 +18,7 @@ from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.utils import BudgetConfig +from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG import sys, os, asyncio, time, random @@ -519,7 +520,10 @@ async def test_video_polls_charge_each_router_budget_once(call_type): } ], ) - kwargs: Final = _success_kwargs(provider_in_litellm_params=None, provider_in_payload="openai", call_type=call_type) + kwargs: Final = { + **_success_kwargs(provider_in_litellm_params=None, provider_in_payload="openai", call_type=call_type), + POLLED_VIDEO_STATUS_KWARG: "completed", + } video: Final = VideoObject(id="video-first", object="video", status="completed") keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") @@ -533,7 +537,7 @@ async def test_video_polls_charge_each_router_budget_once(call_type): @pytest.mark.asyncio -async def test_a_guardrail_priced_poll_of_a_running_video_leaves_the_router_charge_to_the_finished_poll(): +async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_router_for_the_video(): from litellm.types.videos.main import VideoObject cache: Final = DualCache() @@ -558,15 +562,21 @@ async def test_a_guardrail_priced_poll_of_a_running_video_leaves_the_router_char ) await limiter.async_log_success_event( - poll(response_cost=0.01), VideoObject(id="video-first", object="video", status="in_progress"), None, None + {**poll(response_cost=0.01), POLLED_VIDEO_STATUS_KWARG: "in_progress"}, + VideoObject(id="video-first", object="video", status="in_progress"), + None, + None, ) for _ in range(3): await limiter.async_log_success_event( - poll(response_cost=0.26), VideoObject(id="video-first", object="video", status="completed"), None, None + {**poll(response_cost=0.26), POLLED_VIDEO_STATUS_KWARG: "completed"}, + VideoObject(id="video-first", object="video", status="completed"), + None, + None, ) keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") - assert tuple([await cache.async_get_cache(key) for key in keys]) == pytest.approx((0.27, 0.27)) + assert tuple([await cache.async_get_cache(key) for key in keys]) == pytest.approx((0.29, 0.29)) @pytest.mark.asyncio async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(): From d355c1e58b9bfc78da06dc5168fa6bc985ead2d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 05:46:13 -0800 Subject: [PATCH 05/10] fix(proxy): keep video status polls in-process for the collector and give guardrail-only job polls an unclaimable row id --- litellm/litellm_core_utils/batch_utils.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 11 ++-- .../proxy/hooks/proxy_track_cost_callback.py | 2 +- litellm/proxy/spend_tracking/spend_event.py | 8 ++- .../spend_tracking/spend_tracking_utils.py | 10 ++- .../proxy/db/test_db_spend_update_writer.py | 64 ++++++++++++++++++- .../hooks/test_proxy_track_cost_callback.py | 40 ++++++++---- .../proxy/spend_tracking/test_spend_event.py | 11 ++-- .../test_spend_tracking_utils.py | 2 +- 9 files changed, 118 insertions(+), 31 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index 7f75619a464..2a81c04eea5 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -24,6 +24,7 @@ _PAYLOAD_FIELDS: Final = TypeAdapter(Mapping[str, object]) _ENDPOINT_PATHS: Final = TypeAdapter(tuple[str, ...]) BATCH_ENDPOINT: Final = "/v1/batch" +BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" _AZURE_PROVIDERS: Final = frozenset({LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value}) _batch_create_routing: Final[ContextVar[bool]] = ContextVar("batch_create_routing", default=False) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e845da88002..83508a08de6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,6 +34,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.batch_utils import BATCH_COST_REQUEST_ID_SUFFIX from litellm.litellm_core_utils.core_helpers import finite_cost, single_log_line from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -164,15 +165,17 @@ def _is_settle_once_cost_row(payload: SpendLogsPayload) -> bool: A batch's cost is known only when a retrieve sees it finished (LIT-7048). A video's cost is charged at create by the built-in providers, but a custom provider may report it on the - status poll that sees the video completed. Only such a poll's row carries the video cost - suffix, so a poll that reports no cost keeps its own id and charges nothing. + status poll that sees the video completed. Only a row carrying the job's cost suffix settles + the job, so a poll that pays only its guardrails, under an id that client call ids cannot + turn into the job's, is charged like any other request. """ if payload.get("status") != "success": return False call_type: Final = payload.get("call_type") + request_id: Final = payload["request_id"] if call_type == CallTypes.aretrieve_batch.value: - return True - return call_type in VIDEO_STATUS_CALL_TYPES and payload["request_id"].endswith(VIDEO_COST_REQUEST_ID_SUFFIX) + return request_id.endswith(BATCH_COST_REQUEST_ID_SUFFIX) + return call_type in VIDEO_STATUS_CALL_TYPES and request_id.endswith(VIDEO_COST_REQUEST_ID_SUFFIX) _COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 62492d186f0..b4e842050ee 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -406,7 +406,7 @@ class _ProxyDBLogger(CustomLogger): async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: - if self.spend_event_producer is None or not is_offloadable_success(response_obj): + if self.spend_event_producer is None or not is_offloadable_success(response_obj, kwargs.get("call_type")): await self._PROXY_track_cost_callback(dict(kwargs), response_obj, start_time, end_time) return event: Final = build_spend_event( diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py index 031e84896b2..d054519cb9c 100644 --- a/litellm/proxy/spend_tracking/spend_event.py +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -28,6 +28,7 @@ from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import LiteLLMBatch, Usage +from litellm.types.videos.utils import is_video_status_call_type SPEND_EVENT_VERSION: Final = 1 CACHE_OFF_KEY: Final = "Cache OFF" @@ -205,9 +206,10 @@ class SpendEventDecodeError: reason: str -def is_offloadable_success(response_obj: object) -> bool: - """Batch retrieves and interaction polls branch on the concrete response class, so they stay in-process.""" - return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) +def is_offloadable_success(response_obj: object, call_type: object) -> bool: + """Batch retrieves and interaction polls branch on the concrete response class, and a video status poll on + the status and video id the event leaves out, so they stay in-process.""" + return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) and not is_video_status_call_type(call_type) def _json_fallback(value: object) -> str: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 94f856d5e98..0e675f8deb0 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -38,7 +38,11 @@ from litellm.constants import ( ) from litellm.integrations.otel.model.metadata import caller_session_id from litellm.integrations.otel.model.trace_controls import TraceControls -from litellm.litellm_core_utils.batch_utils import charged_video_id, pays_only_its_guardrail +from litellm.litellm_core_utils.batch_utils import ( + BATCH_COST_REQUEST_ID_SUFFIX, + charged_video_id, + pays_only_its_guardrail, +) from litellm.litellm_core_utils.classifier_logging import ( CLASSIFIER_AUDIT_FIELDS, classifier_audit_fields, @@ -345,7 +349,7 @@ def _get_spend_logs_metadata( return clean_metadata -BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" +JOB_POLL_GUARDRAIL_COST_REQUEST_ID_SUFFIX: Final = "_guardrail_cost" def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: @@ -379,7 +383,7 @@ def get_spend_logs_id( return f"{charged_video}{VIDEO_COST_REQUEST_ID_SUFFIX}" call_id: Final = kwargs.get("litellm_call_id") if pays_only_its_guardrail(call_type, kwargs, response_cost) and isinstance(call_id, str) and call_id: - return call_id + return f"{call_id}{JOB_POLL_GUARDRAIL_COST_REQUEST_ID_SUFFIX}" return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" 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 8e9d9ea843d..4dd2ee59fb0 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -4793,9 +4793,25 @@ async def _update_database_for_video_poll( response_cost: float, status: str = "completed", polled_video_id: str = "video_abc", + guardrail_cost: float = 0.0, + call_type: str = "avideo_status", ) -> bool: + from litellm.types.utils import LiteLLMBatch from litellm.types.videos.main import VideoObject + polled_job = ( + VideoObject(id="video_abc", object="video", status=status) + if call_type == "avideo_status" + else LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status=status, + ) + ) with ( patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam "litellm.proxy.proxy_server.prisma_client", prisma @@ -4812,13 +4828,27 @@ async def _update_database_for_video_poll( org_id=None, kwargs={ "model": "sora-2", - "call_type": "avideo_status", + "call_type": call_type, "litellm_call_id": litellm_call_id, "response_cost": response_cost, "polled_video_status": status, "litellm_params": {"metadata": {"user_api_key": "test-key"}, "video_id": polled_video_id}, + **( + { + "standard_logging_object": { + "response_cost": response_cost, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], + "metadata": {}, + "hidden_params": {}, + "request_tags": [], + "model_map_information": None, + } + } + if guardrail_cost + else {} + ), }, - completion_response=VideoObject(id="video_abc", object="video", status=status), + completion_response=polled_job, start_time=datetime.now(timezone.utc), end_time=datetime.now(timezone.utc), response_cost=response_cost, @@ -4843,6 +4873,36 @@ async def test_update_database_charges_the_completed_polls_cost_and_not_a_runnin assert prisma.db.litellm_spendlogs.rows["video_abc_video_cost"]["spend"] == 0.42 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "charge_row_id"), + [("avideo_status", "video_abc_video_cost"), ("aretrieve_batch", "batch_abc_batch_cost")], +) +async def test_a_poll_paying_only_its_guardrail_cannot_take_the_jobs_charge_whatever_call_id_the_client_sent( + call_type, charge_row_id +): + """A client picks its own x-litellm-call-id, so a guardrail-only poll naming the job's charge row, once or + twice, must still be charged each time and leave the finished job's charge to the poll that sees it done.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll( + db_writer, prisma, call_id, cost, status=status, guardrail_cost=0.01, call_type=call_type + ) + for call_id, cost, status in ( + (charge_row_id, 0.01, "in_progress"), + (charge_row_id, 0.01, "in_progress"), + ("poll-3", 0.43, "completed"), + ) + ] + await _flush_queued_spend_logs(prisma) + + assert charged == [True, True, True] + assert prisma.db.litellm_spendlogs.rows[charge_row_id]["spend"] == 0.43 + + @pytest.mark.asyncio async def test_update_database_charges_each_deployments_video_though_their_provider_ids_match(): """A custom provider's own ids can repeat across deployments, so the cost row keys off the id the client polled.""" diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 06107705754..6299b82ea7a 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -43,6 +43,7 @@ from litellm.proxy.utils import ( ) from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.videos.main import VideoObject @pytest.mark.asyncio @@ -3377,18 +3378,8 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski assert event.response_cost == 0.0125 -@pytest.mark.asyncio -async def test_async_log_success_event_keeps_batch_retrieves_in_process(): - producer = SpendEventProducer( - address=UnixAddress(path="/nonexistent/spend.sock"), - on_unavailable="drop", - buffer_size=10, - connect_timeout=1.0, - fallback=_no_fallback, - ) - logger = ProxyDBLogger(producer) - kwargs = {**_offload_kwargs(), "call_type": CallTypes.aretrieve_batch.value} - completed_batch = LiteLLMBatch( +def _completed_batch() -> LiteLLMBatch: + return LiteLLMBatch( id="batch_abc", completion_window="24h", created_at=1, @@ -3399,6 +3390,29 @@ async def test_async_log_success_event_keeps_batch_retrieves_in_process(): status="completed", ) + +def _completed_video() -> VideoObject: + return VideoObject(id="video_abc", object="video", status="completed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "polled_job"), + [(CallTypes.aretrieve_batch.value, _completed_batch), ("avideo_status", _completed_video)], +) +async def test_async_log_success_event_keeps_job_polls_in_process(call_type, polled_job): + """The spend event leaves out the polled video's status and id, which decide whether the poll charges the job.""" + producer = SpendEventProducer( + address=UnixAddress(path="/nonexistent/spend.sock"), + on_unavailable="drop", + buffer_size=10, + connect_timeout=1.0, + fallback=_no_fallback, + ) + logger = ProxyDBLogger(producer) + kwargs = {**_offload_kwargs(), "call_type": call_type} + completed_job = polled_job() + with ( patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam "litellm.proxy.proxy_server.proxy_logging_obj" @@ -3412,7 +3426,7 @@ async def test_async_log_success_event_keeps_batch_retrieves_in_process(): ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() - await logger.async_log_success_event(kwargs, completed_batch, datetime.now(), datetime.now()) + await logger.async_log_success_event(kwargs, completed_job, datetime.now(), datetime.now()) mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() assert producer.stats().queued == 0 diff --git a/tests/unit/proxy/spend_tracking/test_spend_event.py b/tests/unit/proxy/spend_tracking/test_spend_event.py index ff449235582..e1912d36d91 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_event.py +++ b/tests/unit/proxy/spend_tracking/test_spend_event.py @@ -194,9 +194,11 @@ def test_undecodable_line_is_an_error_value(): assert isinstance(decode_spend_event(b"not json\n"), SpendEventDecodeError) -def test_batch_retrieves_stay_in_process(): - assert is_offloadable_success(_response()) is True - assert is_offloadable_success(None) is True +def test_batch_retrieves_and_video_status_polls_stay_in_process(): + assert is_offloadable_success(_response(), "acompletion") is True + assert is_offloadable_success(None, None) is True + assert is_offloadable_success(None, "avideo_generation") is True + assert is_offloadable_success(None, "avideo_status") is False assert ( is_offloadable_success( LiteLLMBatch( @@ -207,7 +209,8 @@ def test_batch_retrieves_stay_in_process(): input_file_id="f", object="batch", status="completed", - ) + ), + "aretrieve_batch", ) is False ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index b797eaf5b29..34c2fdfbd51 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5115,7 +5115,7 @@ def test_a_job_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own( } assert get_spend_logs_id(call_type, {"id": job_id, "status": "completed"}, kwargs) == ( - charged_id if pays_the_job else "poll-call-id" + charged_id if pays_the_job else "poll-call-id_guardrail_cost" ) From ca8d40455795426bd6aff191f92af93d180f6262 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:30:43 -0800 Subject: [PATCH 06/10] fix(proxy): keep repeat batch polls free and group guardrail-only batch rows under their batch's trace --- litellm/litellm_core_utils/batch_utils.py | 16 +++++++------- .../proxy/hooks/proxy_track_cost_callback.py | 4 +++- .../spend_tracking/spend_tracking_utils.py | 16 ++++++++------ .../litellm_core_utils/test_batch_utils.py | 8 +++---- .../hooks/test_model_max_budget_limiter.py | 12 ++++++----- .../hooks/test_proxy_track_cost_callback.py | 10 ++++++++- .../test_spend_tracking_utils.py | 21 ++++++++++++------- 7 files changed, 55 insertions(+), 32 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index 2a81c04eea5..a54dfc60069 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -321,8 +321,11 @@ def job_cost_of_poll(response_cost: object, standard_logging_payload: object) -> return (finite_cost(response_cost) or 0.0) - _guardrail_cost_of_poll(standard_logging_payload) -def repeat_poll_cost(response_cost: object, standard_logging_payload: object) -> float: - """What a poll still costs once an earlier poll charged its job: the guardrails that ran on it.""" +def repeat_poll_cost(call_type: object, response_cost: object, standard_logging_payload: object) -> float: + """What a poll still costs once an earlier poll charged its job: a video poll's guardrails, which every video + poll has always paid, and nothing for a batch poll, whose repeats have always been free.""" + if not is_video_status_call_type(call_type): + return 0.0 return max(0.0, min(finite_cost(response_cost) or 0.0, _guardrail_cost_of_poll(standard_logging_payload))) @@ -365,12 +368,9 @@ def charged_video_id( def _charged_batch(kwargs: Mapping[str, object], response_obj: object, response_cost: float) -> ChargeOnceJob | None: if not is_batch_retrieve_call_type(kwargs.get("call_type")) or not isinstance(response_obj, Batch): return None - standard_logging_payload: Final = kwargs.get("standard_logging_object") - if job_cost_of_poll(response_cost, standard_logging_payload) <= 0: + if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0: return None - return ChargeOnceJob( - kind="batch", job_id=response_obj.id, repeat_poll_cost=repeat_poll_cost(response_cost, standard_logging_payload) - ) + return ChargeOnceJob(kind="batch", job_id=response_obj.id) def job_to_charge_once( @@ -388,5 +388,5 @@ def job_to_charge_once( return ChargeOnceJob( kind="video", job_id=video_id, - repeat_poll_cost=repeat_poll_cost(response_cost, kwargs.get("standard_logging_object")), + repeat_poll_cost=repeat_poll_cost(call_type, response_cost, kwargs.get("standard_logging_object")), ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b4e842050ee..e2e031e0169 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1332,7 +1332,9 @@ async def _charge_request_cost( """What the request charged, since a poll of a job an earlier poll charged still pays for its own guardrails.""" if await charge_spend(kwargs=kwargs, response_cost=response_cost, budget_reservation=budget_reservation): return response_cost - cost_after_job: Final = repeat_poll_cost(response_cost, kwargs.get("standard_logging_object")) + cost_after_job: Final = repeat_poll_cost( + kwargs.get("call_type"), response_cost, kwargs.get("standard_logging_object") + ) if cost_after_job <= 0: return None charged: Final = await charge_spend( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 0e675f8deb0..81c33af2c7e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -977,7 +977,9 @@ def get_logging_payload( metadata=metadata, standard_logging_payload=standard_logging_payload, omit_when_missing=_omits_session_id_when_missing(metadata), - batch_trace_session_id=_get_batch_trace_session_id(call_type=call_type, request_id=id), + batch_trace_session_id=_get_batch_trace_session_id( + call_type=call_type, request_id=id, batch_id=response_obj_dict.get("id") + ), ), request_duration_ms=_get_request_duration_ms(start_time, end_time), status=_get_status_for_spend_log( @@ -1035,13 +1037,15 @@ _BATCH_TRACE_CALL_TYPES: Final = frozenset( ) -def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -> str | None: - """A batch's create row and its poller-written cost row both derive their request id - from the same batch id (the cost row appends BATCH_COST_REQUEST_ID_SUFFIX), so using - that id as the session groups the batch lifecycle into one trace on the logs UI. The - poller builds its own logging context, so per-request trace ids can never link them.""" +def _get_batch_trace_session_id(call_type: str | None, request_id: str | None, batch_id: object) -> str | None: + """A batch's create row, its poller-written cost row and a poll's own guardrail row all name the + same batch, so using its id as the session groups the batch lifecycle into one trace on the logs + UI. The poller builds its own logging context, so per-request trace ids can never link them. A + redacted response has no batch id, so the cost row's request id stands in for it.""" if call_type not in _BATCH_TRACE_CALL_TYPES or not request_id: return None + if isinstance(batch_id, str) and batch_id: + return batch_id return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index cf643482a78..bf2a1674420 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -428,11 +428,11 @@ def test_job_charge_identity_requires_a_cost_bearing_poll(call_type: str, polled @pytest.mark.parametrize( - ("call_type", "polled", "kind"), - (("aretrieve_batch", _completed_batch, "batch"), ("avideo_status", _completed_video, "video")), + ("call_type", "polled", "kind", "repeat_cost"), + (("aretrieve_batch", _completed_batch, "batch", 0.0), ("avideo_status", _completed_video, "video", 0.01)), ) def test_only_the_jobs_own_cost_is_charged_once_since_a_guardrail_charges_every_poll( - call_type: str, polled, kind + call_type: str, polled, kind, repeat_cost ) -> None: from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once @@ -443,7 +443,7 @@ def test_only_the_jobs_own_cost_is_charged_once_since_a_guardrail_charges_every_ } assert job_to_charge_once(kwargs, polled(), 0.01) is None assert job_to_charge_once(kwargs, polled(), 0.43) == ChargeOnceJob( - kind=kind, job_id="job-budget", repeat_poll_cost=0.01 + kind=kind, job_id="job-budget", repeat_poll_cost=repeat_cost ) 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 6d387e39efe..100ce7cb7cc 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -357,13 +357,15 @@ async def test_polls_of_a_finished_video_charge_each_per_model_budget_once(): @pytest.mark.asyncio @pytest.mark.parametrize( - ("call_type", "polled"), + ("call_type", "polled", "guardrail_charges"), ( - ("avideo_status", lambda status: VideoObject(id="job_first", object="video", status=status)), - ("aretrieve_batch", lambda status: _batch("job_first", status)), + ("avideo_status", lambda status: VideoObject(id="job_first", object="video", status=status), 4), + ("aretrieve_batch", lambda status: _batch("job_first", status), 2), ), ) -async def test_every_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_job(call_type, polled): +async def test_only_the_finished_poll_pays_the_job_and_only_a_video_repoll_pays_its_guardrail( + call_type, polled, guardrail_charges +): limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) guardrail_cost: Final = 0.01 priced: Final = ({"guardrail_name": "priced", "guardrail_cost": guardrail_cost},) @@ -375,7 +377,7 @@ async def test_every_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the for _ in range(3): await _log_success(limiter, poll("completed", BATCH_COST + guardrail_cost), polled("completed")) - assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST + 4 * guardrail_cost) + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST + guardrail_charges * guardrail_cost) OPUS_GROUP: Final = "claude-opus" OPUS_GROUP_MEMBERS: Final = MappingProxyType( diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 6299b82ea7a..4062ebfb0cb 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -4149,6 +4149,7 @@ class _RecordedCharges: _GUARDED_POLL: Final = MappingProxyType( { + "call_type": "avideo_status", "response_cost": 0.43, "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, } @@ -4163,8 +4164,15 @@ _GUARDED_POLL: Final = MappingProxyType( ((False, True), _GUARDED_POLL, 0.01, (0.01, 0.01, None)), ((False, False), _GUARDED_POLL, None, (0.01, 0.01, None)), ((False,), {"response_cost": 0.43}, None, None), + ((False,), {**_GUARDED_POLL, "call_type": "aretrieve_batch"}, None, None), + ], + ids=[ + "first_poll_pays_the_job", + "repeat_poll_pays_its_guardrail", + "guardrail_charge_not_written", + "repeat_poll_free", + "repeat_batch_poll_free_whatever_its_guardrails_cost", ], - ids=["first_poll_pays_the_job", "repeat_poll_pays_its_guardrail", "guardrail_charge_not_written", "repeat_poll_free"], ) async def test_a_poll_of_an_already_charged_job_still_charges_its_own_guardrails(outcomes, poll, charged, second_call): charge_spend: Final = _RecordedCharges(outcomes) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 34c2fdfbd51..e9111d0bde7 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -375,20 +375,27 @@ def test_legacy_policy_keeps_trace_id_fallback(): def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): - """The create call's request id IS the batch id and the poller's cost row appends - _batch_cost to it, so deriving the session from the request id lands both rows in one - trace on the logs UI even though the poller builds a fresh logging context per cycle.""" + """The create row, the poller's cost row, a poll's own guardrail row and a redacted cost row + all land in one trace on the logs UI even though the poller builds a fresh logging context per cycle.""" from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id - create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") - cost_session: Final = _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost") - assert create_session == cost_session == "batch-uid-1" + sessions: Final = ( + _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1", batch_id="batch-uid-1"), + _get_batch_trace_session_id( + call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost", batch_id="batch-uid-1" + ), + _get_batch_trace_session_id( + call_type="aretrieve_batch", request_id="poll-1_guardrail_cost", batch_id="batch-uid-1" + ), + _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost", batch_id=None), + ) + assert set(sessions) == {"batch-uid-1"} def test_non_batch_call_types_derive_no_batch_session(): from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id - assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1") is None + assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1", batch_id="chatcmpl-1") is None def test_batch_session_outranks_the_per_request_trace_id(): From 8cb539c2b5d128b636ff64f3bf6bf7a136ed5a1d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:59:54 -0800 Subject: [PATCH 07/10] fix(proxy): settle a finished job whose own cost is below its poll's guardrail cost under the job's charge row --- litellm/litellm_core_utils/batch_utils.py | 12 ++ .../proxy/hooks/proxy_track_cost_callback.py | 22 ++- .../spend_tracking/spend_tracking_utils.py | 3 +- .../proxy/db/test_db_spend_update_writer.py | 156 +++++++++++++----- 4 files changed, 149 insertions(+), 44 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index a54dfc60069..bf8d9fe5b73 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -316,6 +316,18 @@ def _guardrail_cost_of_poll(standard_logging_payload: object) -> float: return guardrail_information_cost(guardrail_information) +def charged_poll_cost(kwargs: Mapping[str, object]) -> object: + """The poll's whole cost as the cost callback charges it, guardrails included, which the logging kwargs' own + response_cost leaves out.""" + standard_logging_payload: Final = kwargs.get("standard_logging_object") + logged_cost: Final = ( + _PAYLOAD_FIELDS.validate_python(standard_logging_payload).get("response_cost") + if isinstance(standard_logging_payload, Mapping) + else None + ) + return kwargs.get("response_cost") if logged_cost is None else logged_cost + + def job_cost_of_poll(response_cost: object, standard_logging_payload: object) -> float: """What the polled job itself cost, since a guardrail on the poll charges every poll, running or finished.""" return (finite_cost(response_cost) or 0.0) - _guardrail_cost_of_poll(standard_logging_payload) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index e2e031e0169..29e375c75de 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1332,13 +1332,27 @@ async def _charge_request_cost( """What the request charged, since a poll of a job an earlier poll charged still pays for its own guardrails.""" if await charge_spend(kwargs=kwargs, response_cost=response_cost, budget_reservation=budget_reservation): return response_cost - cost_after_job: Final = repeat_poll_cost( - kwargs.get("call_type"), response_cost, kwargs.get("standard_logging_object") - ) + standard_logging_payload: Final = kwargs.get("standard_logging_object") + cost_after_job: Final = repeat_poll_cost(kwargs.get("call_type"), response_cost, standard_logging_payload) if cost_after_job <= 0: return None charged: Final = await charge_spend( - kwargs={**kwargs, "response_cost": cost_after_job}, response_cost=cost_after_job, budget_reservation=None + kwargs={ + **kwargs, + "response_cost": cost_after_job, + **( + { + "standard_logging_object": { + **_OBJECT_MAPPING.validate_python(standard_logging_payload), + "response_cost": cost_after_job, + } + } + if isinstance(standard_logging_payload, Mapping) + else {} + ), + }, + response_cost=cost_after_job, + budget_reservation=None, ) return cost_after_job if charged else None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 81c33af2c7e..8866ac76696 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -40,6 +40,7 @@ from litellm.integrations.otel.model.metadata import caller_session_id from litellm.integrations.otel.model.trace_controls import TraceControls from litellm.litellm_core_utils.batch_utils import ( BATCH_COST_REQUEST_ID_SUFFIX, + charged_poll_cost, charged_video_id, pays_only_its_guardrail, ) @@ -377,7 +378,7 @@ def get_spend_logs_id( ) if resolved_id is None: return None - response_cost: Final = kwargs.get("response_cost") + response_cost: Final = charged_poll_cost(kwargs) charged_video: Final = charged_video_id(call_type, kwargs, resolved_id, response_cost) if charged_video is not None: return f"{charged_video}{VIDEO_COST_REQUEST_ID_SUFFIX}" 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 4dd2ee59fb0..daf33aa3796 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1,10 +1,11 @@ import asyncio import copy +import functools import json import logging import re -from collections.abc import AsyncIterator, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace @@ -4786,16 +4787,16 @@ async def test_update_database_charges_a_video_once_from_the_poll_that_sees_it_c } -async def _update_database_for_video_poll( - db_writer: DBSpendUpdateWriter, - prisma: MagicMock, +def _job_poll( litellm_call_id: str, response_cost: float, status: str = "completed", polled_video_id: str = "video_abc", guardrail_cost: float = 0.0, call_type: str = "avideo_status", -) -> bool: +) -> tuple[dict[str, object], object]: + """The logging kwargs a job poll reaches the spend writer with: litellm logs the provider's own cost as + response_cost and the guardrails' on top of it only in the standard logging payload.""" from litellm.types.utils import LiteLLMBatch from litellm.types.videos.main import VideoObject @@ -4812,6 +4813,47 @@ async def _update_database_for_video_poll( status=status, ) ) + kwargs: Final[dict[str, object]] = { + "model": "sora-2", + "call_type": call_type, + "litellm_call_id": litellm_call_id, + "response_cost": response_cost - guardrail_cost, + "polled_video_status": status, + "litellm_params": {"metadata": {"user_api_key": "test-key"}, "video_id": polled_video_id}, + **( + { + "standard_logging_object": { + "response_cost": response_cost, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], + "metadata": {}, + "hidden_params": {}, + "request_tags": [], + "model_map_information": None, + } + } + if guardrail_cost + else {} + ), + } + return kwargs, polled_job + + +def _spend_writer_call(db_writer: DBSpendUpdateWriter, polled_job: object) -> Callable[..., Awaitable[bool]]: + return functools.partial( + db_writer.update_database, + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + completion_response=polled_job, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + ) + + +@contextmanager +def _proxy_globals(prisma: MagicMock) -> Iterator[None]: with ( patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam "litellm.proxy.proxy_server.prisma_client", prisma @@ -4820,39 +4862,22 @@ async def _update_database_for_video_poll( "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" ), ): - charged = await db_writer.update_database( - token="test-token", - user_id="test-user", - end_user_id=None, - team_id=None, - org_id=None, - kwargs={ - "model": "sora-2", - "call_type": call_type, - "litellm_call_id": litellm_call_id, - "response_cost": response_cost, - "polled_video_status": status, - "litellm_params": {"metadata": {"user_api_key": "test-key"}, "video_id": polled_video_id}, - **( - { - "standard_logging_object": { - "response_cost": response_cost, - "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": guardrail_cost}], - "metadata": {}, - "hidden_params": {}, - "request_tags": [], - "model_map_information": None, - } - } - if guardrail_cost - else {} - ), - }, - completion_response=polled_job, - start_time=datetime.now(timezone.utc), - end_time=datetime.now(timezone.utc), - response_cost=response_cost, - ) + yield + + +async def _update_database_for_video_poll( + db_writer: DBSpendUpdateWriter, + prisma: MagicMock, + litellm_call_id: str, + response_cost: float, + status: str = "completed", + polled_video_id: str = "video_abc", + guardrail_cost: float = 0.0, + call_type: str = "avideo_status", +) -> bool: + kwargs, polled_job = _job_poll(litellm_call_id, response_cost, status, polled_video_id, guardrail_cost, call_type) + with _proxy_globals(prisma): + charged = await _spend_writer_call(db_writer, polled_job)(kwargs=kwargs, response_cost=response_cost) await asyncio.sleep(0) return charged @@ -4903,6 +4928,59 @@ async def test_a_poll_paying_only_its_guardrail_cannot_take_the_jobs_charge_what assert prisma.db.litellm_spendlogs.rows[charge_row_id]["spend"] == 0.43 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "charge_row_id"), + [("avideo_status", "video_abc_video_cost"), ("aretrieve_batch", "batch_abc_batch_cost")], +) +async def test_a_finished_job_cheaper_than_its_polls_guardrail_is_still_charged_once(call_type, charge_row_id): + """A batch of a few cheap requests costs less than a priced guardrail on its retrieve, and the finished job's + cost must still settle under its charge row once instead of being billed again on every poll.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll( + db_writer, prisma, call_id, 0.01002775, guardrail_cost=0.01, call_type=call_type + ) + for call_id in ("poll-1", "poll-2") + ] + await _flush_queued_spend_logs(prisma) + + assert charged == [True, False] + assert {request_id: row["spend"] for request_id, row in prisma.db.litellm_spendlogs.rows.items()} == { + charge_row_id: 0.01002775 + } + + +@pytest.mark.asyncio +async def test_a_repeat_poll_of_a_finished_video_charges_its_guardrail_on_a_row_of_its_own(): + from litellm.proxy.hooks.proxy_track_cost_callback import _charge_request_cost + + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + await _update_database_for_video_poll(db_writer, prisma, "poll-1", 0.43, guardrail_cost=0.01) + kwargs, polled_job = _job_poll("poll-2", 0.43, guardrail_cost=0.01) + + async def charge_spend( + *, kwargs: dict[str, object], response_cost: float, budget_reservation: dict[str, object] | None + ) -> bool: + return await _spend_writer_call(db_writer, polled_job)(kwargs=kwargs, response_cost=response_cost) + + with _proxy_globals(prisma): + repeat_charge = await _charge_request_cost(charge_spend, kwargs, 0.43, None) + await asyncio.sleep(0) + await _flush_queued_spend_logs(prisma) + + assert repeat_charge == 0.01 + assert {request_id: row["spend"] for request_id, row in prisma.db.litellm_spendlogs.rows.items()} == { + "video_abc_video_cost": 0.43, + "poll-2_guardrail_cost": 0.01, + } + + @pytest.mark.asyncio async def test_update_database_charges_each_deployments_video_though_their_provider_ids_match(): """A custom provider's own ids can repeat across deployments, so the cost row keys off the id the client polled.""" From 0f6052c74f9c0a6a431676c8432f2cdfeb463f62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 07:31:17 -0800 Subject: [PATCH 08/10] test(proxy): count charged_poll_cost as covered through the finished job guardrail poll tests --- tests/code_coverage_tests/router_code_coverage.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 46df08419e2..a5c5d15b419 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -182,6 +182,7 @@ ignored_function_names = [ "_guardrail_cost_of_poll", # Read by job_cost_of_poll and repeat_poll_cost, tested through the guardrail-priced poll tests "repeat_poll_cost", # Read by job_to_charge_once and _charge_request_cost, tested through every limiter's guardrail poll test "pays_only_its_guardrail", # Read by get_spend_logs_id, tested through its guardrail-only poll test in test_spend_tracking_utils.py + "charged_poll_cost", # Read by get_spend_logs_id, tested through the finished job guardrail poll tests in test_db_spend_update_writer.py "_charged_batch", # Read by job_to_charge_once, tested through the batch charge identity tests in test_batch_utils.py "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py From 49da7ec6db440120ec38401d05ba7184ec4105b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 07:58:51 -0800 Subject: [PATCH 09/10] fix(proxy): settle a finished batch that cost nothing itself on its first priced retrieve --- litellm/litellm_core_utils/batch_utils.py | 24 ++++++++++++-- .../spend_tracking/spend_tracking_utils.py | 6 +++- .../router_code_coverage.py | 1 + .../litellm_core_utils/test_batch_utils.py | 32 +++++++++++++++---- .../proxy/db/test_db_spend_update_writer.py | 22 +++++++++++++ .../test_spend_tracking_utils.py | 21 ++++++++++-- 6 files changed, 94 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index bf8d9fe5b73..d1e750f6890 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -341,10 +341,23 @@ def repeat_poll_cost(call_type: object, response_cost: object, standard_logging_ return max(0.0, min(finite_cost(response_cost) or 0.0, _guardrail_cost_of_poll(standard_logging_payload))) -def pays_only_its_guardrail(call_type: object, kwargs: Mapping[str, object], response_cost: object) -> bool: - """A batch or video poll whose whole cost is its guardrails', charged like any other request.""" +FINISHED_BATCH_STATUSES: Final = frozenset({"completed", "complete", "failed", "cancelled", "expired"}) + + +def is_finished_batch_poll(call_type: object, polled_status: object) -> bool: + """A finished batch reports the same result on every retrieve, so its first retrieve that costs anything settles + it, guardrails included, and its repeats are free.""" + return is_batch_retrieve_call_type(call_type) and polled_status in FINISHED_BATCH_STATUSES + + +def pays_only_its_guardrail( + call_type: object, kwargs: Mapping[str, object], response_cost: object, polled_status: object +) -> bool: + """A running batch's or any video's poll whose whole cost is its guardrails', charged like any other request.""" if not is_batch_retrieve_call_type(call_type) and not is_video_status_call_type(call_type): return False + if is_finished_batch_poll(call_type, polled_status): + return False return ( job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0 @@ -380,7 +393,12 @@ def charged_video_id( def _charged_batch(kwargs: Mapping[str, object], response_obj: object, response_cost: float) -> ChargeOnceJob | None: if not is_batch_retrieve_call_type(kwargs.get("call_type")) or not isinstance(response_obj, Batch): return None - if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0: + settled_cost: Final = ( + finite_cost(response_cost) or 0.0 + if is_finished_batch_poll(kwargs.get("call_type"), response_obj.status) + else job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) + ) + if settled_cost <= 0: return None return ChargeOnceJob(kind="batch", job_id=response_obj.id) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8866ac76696..3a0af4a01d1 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -383,7 +383,11 @@ def get_spend_logs_id( if charged_video is not None: return f"{charged_video}{VIDEO_COST_REQUEST_ID_SUFFIX}" call_id: Final = kwargs.get("litellm_call_id") - if pays_only_its_guardrail(call_type, kwargs, response_cost) and isinstance(call_id, str) and call_id: + if ( + pays_only_its_guardrail(call_type, kwargs, response_cost, response_obj.get("status")) + and isinstance(call_id, str) + and call_id + ): return f"{call_id}{JOB_POLL_GUARDRAIL_COST_REQUEST_ID_SUFFIX}" return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a5c5d15b419..b1debf2a2e2 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -183,6 +183,7 @@ ignored_function_names = [ "repeat_poll_cost", # Read by job_to_charge_once and _charge_request_cost, tested through every limiter's guardrail poll test "pays_only_its_guardrail", # Read by get_spend_logs_id, tested through its guardrail-only poll test in test_spend_tracking_utils.py "charged_poll_cost", # Read by get_spend_logs_id, tested through the finished job guardrail poll tests in test_db_spend_update_writer.py + "is_finished_batch_poll", # Read by pays_only_its_guardrail and _charged_batch, tested through the finished batch guardrail tests "_charged_batch", # Read by job_to_charge_once, tested through the batch charge identity tests in test_batch_utils.py "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index bf2a1674420..e724f9fecdf 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -383,15 +383,19 @@ def setup_and_teardown(): asyncio.set_event_loop(None) -def _completed_batch(): +def _batch(status: str): from openai.types import Batch return Batch( id="job-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", - input_file_id="file-budget", object="batch", status="completed", + input_file_id="file-budget", object="batch", status=status, ) +def _completed_batch(): + return _batch("completed") + + def _completed_video(): from litellm.types.videos.main import VideoObject @@ -428,11 +432,14 @@ def test_job_charge_identity_requires_a_cost_bearing_poll(call_type: str, polled @pytest.mark.parametrize( - ("call_type", "polled", "kind", "repeat_cost"), - (("aretrieve_batch", _completed_batch, "batch", 0.0), ("avideo_status", _completed_video, "video", 0.01)), + ("call_type", "polled", "guardrail_only_poll", "kind", "repeat_cost"), + ( + ("aretrieve_batch", _completed_batch, lambda: _batch("in_progress"), "batch", 0.0), + ("avideo_status", _completed_video, _completed_video, "video", 0.01), + ), ) def test_only_the_jobs_own_cost_is_charged_once_since_a_guardrail_charges_every_poll( - call_type: str, polled, kind, repeat_cost + call_type: str, polled, guardrail_only_poll, kind, repeat_cost ) -> None: from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once @@ -441,12 +448,25 @@ def test_only_the_jobs_own_cost_is_charged_once_since_a_guardrail_charges_every_ POLLED_VIDEO_STATUS_KWARG: "completed", "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, } - assert job_to_charge_once(kwargs, polled(), 0.01) is None + assert job_to_charge_once(kwargs, guardrail_only_poll(), 0.01) is None assert job_to_charge_once(kwargs, polled(), 0.43) == ChargeOnceJob( kind=kind, job_id="job-budget", repeat_poll_cost=repeat_cost ) +@pytest.mark.parametrize("status", ("completed", "failed", "cancelled", "expired")) +def test_a_finished_batch_that_cost_nothing_itself_is_still_charged_once_for_its_retrieves_guardrail( + status: str, +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once + + kwargs: Final = { + "call_type": "aretrieve_batch", + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + assert job_to_charge_once(kwargs, _batch(status), 0.01) == ChargeOnceJob(kind="batch", job_id="job-budget") + + def test_a_video_is_charged_once_it_completes_under_the_id_the_client_polled() -> None: """A provider may price a running video too, and a custom provider's own id can repeat across deployments.""" from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_to_charge_once 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 daf33aa3796..dff83bbeb33 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -4954,6 +4954,28 @@ async def test_a_finished_job_cheaper_than_its_polls_guardrail_is_still_charged_ } +@pytest.mark.asyncio +async def test_a_finished_batch_that_cost_nothing_itself_settles_its_retrieves_guardrail_once(): + """A batch whose every request failed costs nothing itself, and its finished retrieves must still settle under + the batch's charge row once, as they always have, instead of each paying the retrieve's guardrail again.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_table_prisma() + + charged = [ + await _update_database_for_video_poll( + db_writer, prisma, call_id, 0.01, guardrail_cost=0.01, call_type="aretrieve_batch" + ) + for call_id in ("poll-1", "poll-2", "poll-3") + ] + await _flush_queued_spend_logs(prisma) + + assert charged == [True, False, False] + assert {request_id: row["spend"] for request_id, row in prisma.db.litellm_spendlogs.rows.items()} == { + "batch_abc_batch_cost": 0.01 + } + + @pytest.mark.asyncio async def test_a_repeat_poll_of_a_finished_video_charges_its_guardrail_on_a_row_of_its_own(): from litellm.proxy.hooks.proxy_track_cost_callback import _charge_request_cost diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index e9111d0bde7..44d8cc3b4cf 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5110,7 +5110,9 @@ def test_a_job_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own( call_type, job_id, charged_id, response_cost, pays_the_job ): """A guardrail charges every poll, so a poll paying only that must neither land on the create row - nor take the row the job's own charge goes on, either of which the skip_duplicates insert drops.""" + nor take the row the job's own charge goes on, either of which the skip_duplicates insert drops. A batch + settles on its first finished retrieve, so its guardrail-only poll is a running one.""" + polled_status = "completed" if pays_the_job or call_type == "avideo_status" else "in_progress" kwargs = { "litellm_call_id": "poll-call-id", "response_cost": response_cost, @@ -5121,11 +5123,26 @@ def test_a_job_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own( }, } - assert get_spend_logs_id(call_type, {"id": job_id, "status": "completed"}, kwargs) == ( + assert get_spend_logs_id(call_type, {"id": job_id, "status": polled_status}, kwargs) == ( charged_id if pays_the_job else "poll-call-id_guardrail_cost" ) +def test_a_finished_batch_retrieve_paying_only_its_guardrail_settles_under_the_batch_row(): + kwargs = { + "litellm_call_id": "poll-call-id", + "response_cost": 0.01, + "standard_logging_object": { + "response_cost": 0.01, + "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}], + }, + } + + assert get_spend_logs_id("aretrieve_batch", {"id": "batch_abc", "status": "failed"}, kwargs) == ( + "batch_abc_batch_cost" + ) + + def test_a_video_polled_with_or_without_its_base64_padding_charges_one_row(): from litellm.types.videos.utils import encode_video_id_with_provider From d1d595ffc9053d57ffeaaf31d92a7f2a64976c6c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 15:40:57 -0800 Subject: [PATCH 10/10] fix(proxy): charge no budget for a batch retrieve before its cost is final, and read the retrieved batch from before redaction --- litellm/batches/batch_utils.py | 21 ---- litellm/litellm_core_utils/batch_utils.py | 74 +++++++++---- litellm/litellm_core_utils/litellm_logging.py | 10 +- .../hooks/max_budget_per_session_limiter.py | 10 +- .../proxy/hooks/model_max_budget_limiter.py | 10 +- .../proxy/hooks/proxy_track_cost_callback.py | 16 +-- .../openai_files_endpoints/common_utils.py | 2 +- .../spend_tracking/spend_tracking_utils.py | 22 ++-- litellm/proxy/spend_tracking/usage_capture.py | 2 +- litellm/router_strategy/budget_limiter.py | 10 +- .../router_code_coverage.py | 2 +- tests/unit/batches/test_batch_utils.py | 67 +----------- .../litellm_core_utils/test_batch_utils.py | 101 +++++++++++++++++- .../test_litellm_logging.py | 59 ++++++++++ .../proxy/db/test_db_spend_update_writer.py | 18 ++-- .../test_max_budget_per_session_limiter.py | 40 ++++++- .../hooks/test_model_max_budget_limiter.py | 5 +- .../hooks/test_proxy_track_cost_callback.py | 35 ++++-- .../test_spend_tracking_utils.py | 58 +++++----- .../router_strategy/test_budget_limiter.py | 49 +++++++++ 20 files changed, 413 insertions(+), 198 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a0c0a90be5f..522e05ebbd5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -38,7 +38,6 @@ class BatchCostUsageResult: completion_cost: float = 0.0 -_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) _DECODED_BATCH_LINE: Final = TypeAdapter(object) @@ -54,26 +53,6 @@ def _uses_native_vertex_output( return first_row is not None and is_native_vertex_batch_output_row(first_row) -_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) - - -def batch_cost_is_final(batch: Batch) -> bool: - """Whether this retrieve of the batch is the one to account its cost from. - - A batch still in flight has nothing to price, and a "completed" batch can report - no output_file_id for a moment before the output populates; pricing either records - $0 under the batch's single spend row and pins it there. Final means a completed - batch whose output file has arrived or whose counts prove no line succeeded, or - any other terminal status (failed, cancelled, expired). - """ - if batch.status not in _TERMINAL_BATCH_STATUSES: - return False - if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: - return True - request_counts: Final = batch.request_counts - return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 - - async def calculate_batch_cost_and_usage( file_content_dictionary: Sequence[dict[str, object]], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index d1e750f6890..7c7b7ea5c4b 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -23,7 +23,11 @@ _DEPLOYMENT_FIELD: Final = TypeAdapter(Mapping[str, object]) _PAYLOAD_FIELDS: Final = TypeAdapter(Mapping[str, object]) _ENDPOINT_PATHS: Final = TypeAdapter(tuple[str, ...]) +_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) +_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) BATCH_ENDPOINT: Final = "/v1/batch" +POLLED_BATCH_ID_KWARG: Final = "polled_batch_id" +POLLED_BATCH_COST_IS_FINAL_KWARG: Final = "polled_batch_cost_is_final" BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" _AZURE_PROVIDERS: Final = frozenset({LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value}) @@ -341,22 +345,52 @@ def repeat_poll_cost(call_type: object, response_cost: object, standard_logging_ return max(0.0, min(finite_cost(response_cost) or 0.0, _guardrail_cost_of_poll(standard_logging_payload))) -FINISHED_BATCH_STATUSES: Final = frozenset({"completed", "complete", "failed", "cancelled", "expired"}) +def batch_cost_is_final(batch: Batch) -> bool: + """Whether this retrieve of the batch is the one to account its cost from. - -def is_finished_batch_poll(call_type: object, polled_status: object) -> bool: - """A finished batch reports the same result on every retrieve, so its first retrieve that costs anything settles - it, guardrails included, and its repeats are free.""" - return is_batch_retrieve_call_type(call_type) and polled_status in FINISHED_BATCH_STATUSES - - -def pays_only_its_guardrail( - call_type: object, kwargs: Mapping[str, object], response_cost: object, polled_status: object -) -> bool: - """A running batch's or any video's poll whose whole cost is its guardrails', charged like any other request.""" - if not is_batch_retrieve_call_type(call_type) and not is_video_status_call_type(call_type): + A batch still in flight has nothing to price, and a "completed" batch can report + no output_file_id for a moment before the output populates; pricing either records + $0 under the batch's single spend row and pins it there. Final means a completed + batch whose output file has arrived or whose counts prove no line succeeded, or + any other terminal status (failed, cancelled, expired). + """ + if batch.status not in _TERMINAL_BATCH_STATUSES: return False - if is_finished_batch_poll(call_type, polled_status): + if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: + return True + request_counts: Final = batch.request_counts + return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 + + +@dataclass(frozen=True, slots=True) +class PolledBatch: + id: str + cost_is_final: bool + + +def _polled_batch(kwargs: Mapping[str, object], response_obj: object) -> PolledBatch | None: + """The retrieved batch, read from the logging kwargs when message redaction replaced the logged response.""" + if not is_batch_retrieve_call_type(kwargs.get("call_type")): + return None + if isinstance(response_obj, Batch): + return PolledBatch(id=response_obj.id, cost_is_final=batch_cost_is_final(response_obj)) + stamped_id: Final = kwargs.get(POLLED_BATCH_ID_KWARG) + stamped_cost_is_final: Final = kwargs.get(POLLED_BATCH_COST_IS_FINAL_KWARG) + if not isinstance(stamped_id, str) or not stamped_id or not isinstance(stamped_cost_is_final, bool): + return None + return PolledBatch(id=stamped_id, cost_is_final=stamped_cost_is_final) + + +def defers_batch_cost(kwargs: Mapping[str, object], response_obj: object) -> bool: + """A batch retrieve before the batch's cost is final charges nothing, guardrails included, since the retrieve + that settles the batch charges its whole cost once.""" + polled: Final = _polled_batch(kwargs, response_obj) + return polled is not None and not polled.cost_is_final + + +def pays_only_its_guardrail(call_type: object, kwargs: Mapping[str, object], response_cost: object) -> bool: + """A video poll whose whole cost is its guardrails', charged like any other request.""" + if not is_video_status_call_type(call_type): return False return ( job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) @@ -391,16 +425,10 @@ def charged_video_id( def _charged_batch(kwargs: Mapping[str, object], response_obj: object, response_cost: float) -> ChargeOnceJob | None: - if not is_batch_retrieve_call_type(kwargs.get("call_type")) or not isinstance(response_obj, Batch): + polled: Final = _polled_batch(kwargs, response_obj) + if polled is None or not polled.cost_is_final or response_cost <= 0: return None - settled_cost: Final = ( - finite_cost(response_cost) or 0.0 - if is_finished_batch_poll(kwargs.get("call_type"), response_obj.status) - else job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) - ) - if settled_cost <= 0: - return None - return ChargeOnceJob(kind="batch", job_id=response_obj.id) + return ChargeOnceJob(kind="batch", job_id=polled.id) def job_to_charge_once( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 03947450b1b..040f78776e6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -38,7 +38,7 @@ from litellm._logging import ( verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import batch_cost_is_final, handle_completed_batch +from litellm.batches.batch_utils import handle_completed_batch from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.caching.redis_batch import flush_post_call_redis_batches @@ -74,6 +74,11 @@ from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.langtrace import langtrace_trace_endpoint from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger +from litellm.litellm_core_utils.batch_utils import ( + POLLED_BATCH_COST_IS_FINAL_KWARG, + POLLED_BATCH_ID_KWARG, + batch_cost_is_final, +) from litellm.litellm_core_utils.classifier_logging import ( classifier_audit_fields, classifier_input_snapshot, @@ -3364,6 +3369,9 @@ class Logging(LiteLLMLoggingBaseClass): result = self._as_served_response(result) # rebind-ok: bridged chat calls log the served response if isinstance(result, VideoObject): self.model_call_details[POLLED_VIDEO_STATUS_KWARG] = result.status + if isinstance(result, LiteLLMBatch): + self.model_call_details[POLLED_BATCH_ID_KWARG] = result.id + self.model_call_details[POLLED_BATCH_COST_IS_FINAL_KWARG] = batch_cost_is_final(result) logging_result: Final = self.normalize_logging_result(result=result) if cache_hit is None and _served_from_cache_at_no_charge(logging_result): diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index c0e76d758be..00939cbffe7 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -26,7 +26,12 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import log_redis_failure from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key, job_to_charge_once +from litellm.litellm_core_utils.batch_utils import ( + ChargeOnceJob, + defers_batch_cost, + job_charged_once_marker_key, + job_to_charge_once, +) from litellm.litellm_core_utils.core_helpers import finite_cost from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -176,6 +181,9 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): if response_cost <= 0: return + if defers_batch_cost(callback_fields, response_obj): + return + cache_key: Final = self._make_cache_key(str(session_id)) job: Final = job_to_charge_once(callback_fields, response_obj, response_cost) charge: Final = ( diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index bd9fe7208ed..889fb9fb373 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,7 +19,12 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key, job_to_charge_once +from litellm.litellm_core_utils.batch_utils import ( + ChargeOnceJob, + defers_batch_cost, + job_charged_once_marker_key, + job_to_charge_once, +) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions @@ -1208,6 +1213,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) return + if defers_batch_cost(kwargs, response_obj): + return + job: Final = job_to_charge_once(kwargs=kwargs, response_obj=response_obj, response_cost=response_cost) for entity_type, entity_id, resolved in resolved_budgets: await self._charge_entity( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 29e375c75de..0333016f20e 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,11 +12,10 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger -from litellm.batches.batch_utils import batch_cost_is_final from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.batch_utils import repeat_poll_cost +from litellm.litellm_core_utils.batch_utils import defers_batch_cost, repeat_poll_cost from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, finite_cost, @@ -79,7 +78,6 @@ from litellm.proxy.utils import ( from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, - LiteLLMBatch, StandardLoggingPayload, ) from litellm.utils import get_end_user_id_for_cost_tracking @@ -906,16 +904,8 @@ class _ProxyDBLogger(CustomLogger): budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) if accounting is not None and (budget_reservation is not None or litellm.apply_search_free_tier): accounting.write_started = True - if ( - isinstance(completion_response, LiteLLMBatch) - and kwargs.get("call_type") == CallTypes.aretrieve_batch.value - and not batch_cost_is_final(completion_response) - ): - verbose_proxy_logger.debug( - "Cost tracking deferred for batch %s still in status %s", - completion_response.id, - completion_response.status, - ) + if defers_batch_cost(kwargs, completion_response): + verbose_proxy_logger.debug("Cost tracking deferred for a batch whose cost is not final yet") await _release_budget_reservation(budget_reservation=budget_reservation) return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index c0d81abeba7..18ce18da936 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -21,8 +21,8 @@ from typing import ( from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.batches.batch_utils import batch_cost_is_final from litellm.constants import MAX_FILE_LIST_LIMIT +from litellm.litellm_core_utils.batch_utils import batch_cost_is_final from litellm.proxy._types import ProxyException from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.table_repositories import ( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3a0af4a01d1..adeed15afea 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -383,11 +383,7 @@ def get_spend_logs_id( if charged_video is not None: return f"{charged_video}{VIDEO_COST_REQUEST_ID_SUFFIX}" call_id: Final = kwargs.get("litellm_call_id") - if ( - pays_only_its_guardrail(call_type, kwargs, response_cost, response_obj.get("status")) - and isinstance(call_id, str) - and call_id - ): + if pays_only_its_guardrail(call_type, kwargs, response_cost) and isinstance(call_id, str) and call_id: return f"{call_id}{JOB_POLL_GUARDRAIL_COST_REQUEST_ID_SUFFIX}" return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" @@ -982,9 +978,7 @@ def get_logging_payload( metadata=metadata, standard_logging_payload=standard_logging_payload, omit_when_missing=_omits_session_id_when_missing(metadata), - batch_trace_session_id=_get_batch_trace_session_id( - call_type=call_type, request_id=id, batch_id=response_obj_dict.get("id") - ), + batch_trace_session_id=_get_batch_trace_session_id(call_type=call_type, request_id=id), ), request_duration_ms=_get_request_duration_ms(start_time, end_time), status=_get_status_for_spend_log( @@ -1042,15 +1036,13 @@ _BATCH_TRACE_CALL_TYPES: Final = frozenset( ) -def _get_batch_trace_session_id(call_type: str | None, request_id: str | None, batch_id: object) -> str | None: - """A batch's create row, its poller-written cost row and a poll's own guardrail row all name the - same batch, so using its id as the session groups the batch lifecycle into one trace on the logs - UI. The poller builds its own logging context, so per-request trace ids can never link them. A - redacted response has no batch id, so the cost row's request id stands in for it.""" +def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -> str | None: + """A batch's create row and its poller-written cost row both derive their request id + from the same batch id (the cost row appends BATCH_COST_REQUEST_ID_SUFFIX), so using + that id as the session groups the batch lifecycle into one trace on the logs UI. The + poller builds its own logging context, so per-request trace ids can never link them.""" if call_type not in _BATCH_TRACE_CALL_TYPES or not request_id: return None - if isinstance(batch_id, str) and batch_id: - return batch_id return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index 19480a66678..69fc7380a26 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -799,7 +799,7 @@ def capture_logging_usage( freeze_only: bool = False, ) -> CapturedUsage | UsageCaptureError | None: from litellm._internal_context import internal_parent_request_id, is_internal_call - from litellm.batches.batch_utils import batch_cost_is_final + from litellm.litellm_core_utils.batch_utils import batch_cost_is_final from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import LiteLLMBatch diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 602c5b7279b..98dd39b40e3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,7 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, job_charged_once_marker_key, job_to_charge_once +from litellm.litellm_core_utils.batch_utils import ( + ChargeOnceJob, + defers_batch_cost, + job_charged_once_marker_key, + job_to_charge_once, +) from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) @@ -551,6 +556,9 @@ class RouterBudgetLimiting(CustomLogger): model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") + if defers_batch_cost(callback_fields, polled_response): + return + job: Final = job_to_charge_once(callback_fields, polled_response, response_cost) budget_config: Final = ( diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index b1debf2a2e2..42ee6178bf0 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -183,7 +183,7 @@ ignored_function_names = [ "repeat_poll_cost", # Read by job_to_charge_once and _charge_request_cost, tested through every limiter's guardrail poll test "pays_only_its_guardrail", # Read by get_spend_logs_id, tested through its guardrail-only poll test in test_spend_tracking_utils.py "charged_poll_cost", # Read by get_spend_logs_id, tested through the finished job guardrail poll tests in test_db_spend_update_writer.py - "is_finished_batch_poll", # Read by pays_only_its_guardrail and _charged_batch, tested through the finished batch guardrail tests + "_polled_batch", # Read by defers_batch_cost and _charged_batch, tested through the redacted batch retrieve tests "_charged_batch", # Read by job_to_charge_once, tested through the batch charge identity tests in test_batch_utils.py "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index b32afd51fd3..e46fd21e525 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -22,7 +22,6 @@ from typing import Final, Literal import httpx import pytest import respx -from openai.types.batch import BatchRequestCounts import litellm @@ -2028,61 +2027,6 @@ def test_extract_credentials_keeps_openai_workload_identity_params(): } -# --------------------------------------------------------------------------- # -# batch_cost_is_final -# --------------------------------------------------------------------------- # - - -def _retrieved_batch( - status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None -) -> LiteLLMBatch: - return LiteLLMBatch( - id="batch_abc", - completion_window="24h", - created_at=1, - endpoint="/v1/chat/completions", - input_file_id="file-in", - object="batch", - status="validating", - output_file_id=output_file_id, - request_counts=counts, - ).model_copy(update={"status": status}) - - -class TestBatchCostIsFinal: - """Every retrieve of one batch writes the same spend row, so the first retrieve - that prices it decides the row for good. A poll before the output exists must - therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" - - @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) - def test_in_flight_batch_is_not_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status)) is False - - @pytest.mark.parametrize("status", ["completed", "complete"]) - def test_completed_with_output_is_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True - - def test_completed_without_output_and_unknown_counts_is_not_final(self): - assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False - - def test_completed_without_output_and_zero_counts_is_not_final(self): - counts = BatchRequestCounts(total=0, completed=0, failed=0) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - - def test_completed_without_output_but_successful_lines_is_not_final(self): - counts = BatchRequestCounts(total=2, completed=2, failed=0) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - - @pytest.mark.parametrize("status", ["completed", "complete"]) - def test_completed_without_output_and_every_line_failed_is_final(self, status): - counts = BatchRequestCounts(total=2, completed=0, failed=2) - assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True - - @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) - def test_other_terminal_statuses_are_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status)) is True - - def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: @@ -2965,7 +2909,7 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( @@ -3060,7 +3004,6 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( wrong token sum, or mispriced model fails this test. """ from litellm.batches.batch_utils import handle_completed_batch - from litellm.types.utils import LiteLLMBatch batch = LiteLLMBatch( id="batch-real-cost-123", @@ -3105,7 +3048,7 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( @@ -3187,7 +3130,7 @@ async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown card renders blank for every batch. Regression for the split being dropped.""" from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes mock_batch = LiteLLMBatch( id="batch-breakdown-1", @@ -3243,7 +3186,7 @@ async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batc from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch, SpecialEnums + from litellm.types.utils import CallTypes, SpecialEnums # Create a proper unified file ID by encoding the correct prefix unified_id_str = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}:test_file_789;unified_id:batch-789" @@ -3322,7 +3265,7 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index e724f9fecdf..45461b50761 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -5,14 +5,21 @@ from io import BytesIO from typing import Dict, Final, List import pytest +from openai.types.batch import BatchRequestCounts import litellm from litellm.litellm_core_utils.batch_utils import ( + POLLED_BATCH_COST_IS_FINAL_KWARG, + POLLED_BATCH_ID_KWARG, + ChargeOnceJob, InMemoryFile, + batch_cost_is_final, + defers_batch_cost, get_router_metadata_variable_name, parse_jsonl_with_embedded_newlines, replace_model_in_jsonl, ) +from litellm.types.utils import LiteLLMBatch from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG, encode_video_id_with_provider from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -383,17 +390,18 @@ def setup_and_teardown(): asyncio.set_event_loop(None) -def _batch(status: str): +def _batch(status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None): from openai.types import Batch return Batch( id="job-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", - input_file_id="file-budget", object="batch", status=status, + input_file_id="file-budget", object="batch", status=status, output_file_id=output_file_id, + request_counts=counts, ) def _completed_batch(): - return _batch("completed") + return _batch("completed", output_file_id="file-out") def _completed_video(): @@ -464,7 +472,21 @@ def test_a_finished_batch_that_cost_nothing_itself_is_still_charged_once_for_its "call_type": "aretrieve_batch", "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, } - assert job_to_charge_once(kwargs, _batch(status), 0.01) == ChargeOnceJob(kind="batch", job_id="job-budget") + finished_batch: Final = _batch(status, output_file_id="file-out") + assert job_to_charge_once(kwargs, finished_batch, 0.01) == ChargeOnceJob(kind="batch", job_id="job-budget") + + +def test_a_completed_batch_still_waiting_on_its_output_does_not_take_the_charge_its_real_cost_needs() -> None: + """Its first retrieve can report completed before the output file exists, and settling on that guardrail-only + retrieve would leave the batch's real cost nothing to charge once the output arrives.""" + from litellm.litellm_core_utils.batch_utils import job_to_charge_once + + kwargs: Final = { + "call_type": "aretrieve_batch", + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + assert job_to_charge_once(kwargs, _batch("completed"), 0.01) is None + assert defers_batch_cost(kwargs, _batch("completed")) def test_a_video_is_charged_once_it_completes_under_the_id_the_client_polled() -> None: @@ -498,6 +520,27 @@ def test_a_video_poll_is_judged_by_its_logged_status_when_redaction_replaced_the assert job_to_charge_once(kwargs, {"text": "redacted-by-litellm"}, 0.42) == expected +@pytest.mark.parametrize( + ("cost_is_final", "deferred", "expected"), + [(False, True, None), (True, False, ChargeOnceJob(kind="batch", job_id="batch_abc"))], +) +def test_a_batch_retrieve_is_judged_by_its_logged_batch_when_redaction_replaced_it( + cost_is_final: bool, deferred: bool, expected: ChargeOnceJob | None +) -> None: + from litellm.litellm_core_utils.batch_utils import job_to_charge_once + + kwargs: Final = { + "call_type": "aretrieve_batch", + POLLED_BATCH_ID_KWARG: "batch_abc", + POLLED_BATCH_COST_IS_FINAL_KWARG: cost_is_final, + "standard_logging_object": {"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}]}, + } + redacted: Final = {"text": "redacted-by-litellm"} + + assert defers_batch_cost(kwargs, redacted) is deferred + assert job_to_charge_once(kwargs, redacted, 0.01) == expected + + def test_a_video_polled_with_or_without_its_base64_padding_is_charged_as_one_job() -> None: from litellm.litellm_core_utils.batch_utils import job_to_charge_once @@ -535,3 +578,53 @@ def test_batch_charge_marker_keeps_the_key_markers_already_written_use() -> None assert job_charged_once_marker_key("provider-spend", ChargeOnceJob("batch", "first-job")) == ( "provider-spend:batch:first-job" ) + + +def _retrieved_batch( + status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None +) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="validating", + output_file_id=output_file_id, + request_counts=counts, + ).model_copy(update={"status": status}) + + +class TestBatchCostIsFinal: + """Every retrieve of one batch writes the same spend row, so the first retrieve + that prices it decides the row for good. A poll before the output exists must + therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" + + @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) + def test_in_flight_batch_is_not_final(self, status): + assert batch_cost_is_final(_retrieved_batch(status)) is False + + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_with_output_is_final(self, status): + assert batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True + + def test_completed_without_output_and_unknown_counts_is_not_final(self): + assert batch_cost_is_final(_retrieved_batch("completed")) is False + + def test_completed_without_output_and_zero_counts_is_not_final(self): + counts = BatchRequestCounts(total=0, completed=0, failed=0) + assert batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + def test_completed_without_output_but_successful_lines_is_not_final(self): + counts = BatchRequestCounts(total=2, completed=2, failed=0) + assert batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_without_output_and_every_line_failed_is_final(self, status): + counts = BatchRequestCounts(total=2, completed=0, failed=2) + assert batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True + + @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) + def test_other_terminal_statuses_are_final(self, status): + assert batch_cost_is_final(_retrieved_batch(status)) is True diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 0dee6d2b45f..8d8c693143a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14493,3 +14493,62 @@ async def test_a_redacted_video_poll_still_tells_the_spend_callbacks_the_video_f ((kwargs, response_obj),) = recorder.events assert not isinstance(response_obj, VideoObject) assert job_to_charge_once(kwargs, response_obj, 0.42) == ChargeOnceJob(kind="video", job_id="video_abc") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status", "output_file_id", "deferred", "charged"), + [("in_progress", None, True, False), ("completed", "file-out", False, True)], +) +async def test_a_redacted_batch_retrieve_still_tells_the_spend_callbacks_whether_the_batch_cost_is_final( + monkeypatch: pytest.MonkeyPatch, status: str, output_file_id: str | None, deferred: bool, charged: bool +) -> None: + from litellm.litellm_core_utils.batch_utils import ChargeOnceJob, defers_batch_cost, job_to_charge_once + from litellm.types.utils import LiteLLMBatch, Usage + + class Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: tuple[tuple[dict[str, object], object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events = (*self.events, (dict(kwargs), response_obj)) + + recorder: Final = Recorder() + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=[], + stream=False, + call_type="aretrieve_batch", + start_time=datetime.datetime.now(), + litellm_call_id="retrieve-1", + function_id="retrieve-1", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + await logging_obj.async_success_handler( + result=LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + ), + batch_cost=0.42, + batch_usage=Usage(prompt_tokens=20, completion_tokens=18, total_tokens=38), + batch_models=["gpt-4o-mini"], + ) + + ((kwargs, response_obj),) = recorder.events + assert not isinstance(response_obj, LiteLLMBatch) + assert defers_batch_cost(kwargs, response_obj) is deferred + assert (job_to_charge_once(kwargs, response_obj, 0.42) == ChargeOnceJob(kind="batch", job_id="batch_abc")) is charged 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 dff83bbeb33..999f538cf85 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -4811,6 +4811,7 @@ def _job_poll( input_file_id="f", object="batch", status=status, + output_file_id="file-out", ) ) kwargs: Final[dict[str, object]] = { @@ -4899,15 +4900,12 @@ async def test_update_database_charges_the_completed_polls_cost_and_not_a_runnin @pytest.mark.asyncio -@pytest.mark.parametrize( - ("call_type", "charge_row_id"), - [("avideo_status", "video_abc_video_cost"), ("aretrieve_batch", "batch_abc_batch_cost")], -) -async def test_a_poll_paying_only_its_guardrail_cannot_take_the_jobs_charge_whatever_call_id_the_client_sent( - call_type, charge_row_id -): - """A client picks its own x-litellm-call-id, so a guardrail-only poll naming the job's charge row, once or - twice, must still be charged each time and leave the finished job's charge to the poll that sees it done.""" +async def test_a_poll_paying_only_its_guardrail_cannot_take_the_jobs_charge_whatever_call_id_the_client_sent(): + """A client picks its own x-litellm-call-id, so a guardrail-only video poll naming the job's charge row, once or + twice, must still be charged each time and leave the finished job's charge to the poll that sees it done. An + unfinished batch retrieve never reaches the writer, and its row is named by the batch id the provider sent.""" + call_type: Final = "avideo_status" + charge_row_id: Final = "video_abc_video_cost" db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() prisma = _spend_logs_table_prisma() @@ -4964,7 +4962,7 @@ async def test_a_finished_batch_that_cost_nothing_itself_settles_its_retrieves_g charged = [ await _update_database_for_video_poll( - db_writer, prisma, call_id, 0.01, guardrail_cost=0.01, call_type="aretrieve_batch" + db_writer, prisma, call_id, 0.01, status="failed", guardrail_cost=0.01, call_type="aretrieve_batch" ) for call_id in ("poll-1", "poll-2", "poll-3") ] diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 32f25892d3d..1f7af3122eb 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -10,6 +10,7 @@ Tests that session-scoped budget tracking works correctly: from collections.abc import Mapping from types import MappingProxyType +from typing import Final from unittest.mock import patch import logging @@ -158,7 +159,7 @@ async def test_repeated_polls_of_a_finished_job_charge_the_session_once(): video = VideoObject(id="job-1", object="video", status="completed") batch = Batch( id="job-1", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", - input_file_id="file-1", object="batch", status="completed", + input_file_id="file-1", object="batch", status="completed", output_file_id="file-out", ) polls = (("avideo_status", video, 0.4), ("aretrieve_batch", batch, 0.1), ("acompletion", None, 0.01)) @@ -183,6 +184,43 @@ async def test_repeated_polls_of_a_finished_job_charge_the_session_once(): assert spend == pytest.approx(0.4 + 0.1 + 3 * 0.01) +@pytest.mark.asyncio +async def test_a_batch_retrieve_before_its_cost_is_final_charges_the_session_nothing(): + from openai.types import Batch + + handler: Final = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache())) + + def polled(status: str, output_file_id: str | None = None) -> Batch: + return Batch( + id="batch-1", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", + input_file_id="file-1", object="batch", status=status, output_file_id=output_file_id, + ) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = _make_mock_agent(max_budget_per_session=5.0) + for response_obj, cost in ( + (polled("in_progress"), 0.01), + (polled("completed"), 0.01), + (polled("completed", "file-out"), 0.25), + (polled("completed", "file-out"), 0.25), + ): + await handler.async_log_success_event( + { + "call_type": "aretrieve_batch", + "response_cost": cost, + "litellm_params": {"metadata": {"session_id": "session-batch", "agent_id": "agent-budget-123"}}, + }, + response_obj, + None, + None, + ) + + spend: Final = await handler._get_current_spend(handler._make_cache_key("session-batch")) + assert spend == pytest.approx(0.25) + + def _video_poll( cost: float, guardrail_cost: float = 0.0, session_id: str = "session-jobs", status: str = "completed" ) -> dict[str, object]: 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 100ce7cb7cc..f1fd02981d6 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -49,6 +49,7 @@ def _batch(batch_id: str, status: str) -> LiteLLMBatch: input_file_id="file-batch", object="batch", status=status, + output_file_id="file-output", usage=Usage(prompt_tokens=20, completion_tokens=18, total_tokens=38), ) @@ -360,12 +361,14 @@ async def test_polls_of_a_finished_video_charge_each_per_model_budget_once(): ("call_type", "polled", "guardrail_charges"), ( ("avideo_status", lambda status: VideoObject(id="job_first", object="video", status=status), 4), - ("aretrieve_batch", lambda status: _batch("job_first", status), 2), + ("aretrieve_batch", lambda status: _batch("job_first", status), 1), ), ) async def test_only_the_finished_poll_pays_the_job_and_only_a_video_repoll_pays_its_guardrail( call_type, polled, guardrail_charges ): + """A running video poll pays its guardrail like any request, while a batch retrieve before its cost is final + charges nothing, as key spend defers it, so the retrieve that settles the batch pays the only guardrail.""" limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) guardrail_cost: Final = 0.01 priced: Final = ({"guardrail_name": "priced", "guardrail_cost": guardrail_cost},) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 4062ebfb0cb..ba908a62dc1 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1597,14 +1597,16 @@ def _retrieved_batch(status: str, output_file_id: str | None): @pytest.mark.asyncio @pytest.mark.parametrize( - ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"), + ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged", "redacted"), [ - ("aretrieve_batch", "in_progress", None, True, False, False), - ("aretrieve_batch", "completed", None, True, False, False), - ("aretrieve_batch", "completed", "file-out", False, True, False), - ("aretrieve_batch", "completed", "file-out", True, True, True), - ("aretrieve_batch", "failed", None, True, True, True), - ("acreate_batch", "validating", None, True, True, True), + ("aretrieve_batch", "in_progress", None, True, False, False, False), + ("aretrieve_batch", "completed", None, True, False, False, False), + ("aretrieve_batch", "completed", "file-out", False, True, False, False), + ("aretrieve_batch", "completed", "file-out", True, True, True, False), + ("aretrieve_batch", "failed", None, True, True, True, False), + ("acreate_batch", "validating", None, True, True, True, False), + ("aretrieve_batch", "in_progress", None, True, False, False, True), + ("aretrieve_batch", "completed", "file-out", True, True, True, True), ], ids=[ "retrieve_before_final", @@ -1613,20 +1615,31 @@ def _retrieved_batch(status: str, output_file_id: str | None): "retrieve_first_final", "retrieve_failed_batch", "create_before_final", + "redacted_retrieve_before_final", + "redacted_retrieve_first_final", ], ) async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate - call_type, status, output_file_id, row_claimed, spend_written, charged + call_type, status, output_file_id, row_claimed, spend_written, charged, redacted ): """ A poll before the batch is final used to pin its shared spend row at $0, and every completed retrieve after the first charged the key again (LIT-7048). Only retrieves are gated, since creating a batch is its own billable request, and a retrieve that - charges nothing hands its budget reservation back instead. + charges nothing hands its budget reservation back instead. Message redaction replaces + the logged batch, so the gate reads the batch logging stamped before redacting it. """ + from litellm.litellm_core_utils.batch_utils import POLLED_BATCH_COST_IS_FINAL_KWARG, POLLED_BATCH_ID_KWARG + logger = ProxyDBLogger() budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} - kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) + retrieved: Final = _retrieved_batch(status, output_file_id) + stamped: Final = ( + {POLLED_BATCH_ID_KWARG: retrieved.id, POLLED_BATCH_COST_IS_FINAL_KWARG: status != "in_progress"} + if redacted + else {} + ) + kwargs = {**_batch_retrieve_kwargs(call_type, reservation=budget_reservation), **stamped} with ( patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam @@ -1648,7 +1661,7 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # await logger._PROXY_track_cost_callback( kwargs=kwargs, - completion_response=_retrieved_batch(status, output_file_id), + completion_response={"text": "redacted-by-litellm"} if redacted else retrieved, start_time=datetime.now(), end_time=datetime.now(), ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 44d8cc3b4cf..b3aec659a9c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -375,27 +375,20 @@ def test_legacy_policy_keeps_trace_id_fallback(): def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): - """The create row, the poller's cost row, a poll's own guardrail row and a redacted cost row - all land in one trace on the logs UI even though the poller builds a fresh logging context per cycle.""" + """The create call's request id IS the batch id and the poller's cost row appends + _batch_cost to it, so deriving the session from the request id lands both rows in one + trace on the logs UI even though the poller builds a fresh logging context per cycle.""" from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id - sessions: Final = ( - _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1", batch_id="batch-uid-1"), - _get_batch_trace_session_id( - call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost", batch_id="batch-uid-1" - ), - _get_batch_trace_session_id( - call_type="aretrieve_batch", request_id="poll-1_guardrail_cost", batch_id="batch-uid-1" - ), - _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost", batch_id=None), - ) - assert set(sessions) == {"batch-uid-1"} + create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") + cost_session: Final = _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost") + assert create_session == cost_session == "batch-uid-1" def test_non_batch_call_types_derive_no_batch_session(): from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id - assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1", batch_id="chatcmpl-1") is None + assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1") is None def test_batch_session_outranks_the_per_request_trace_id(): @@ -5101,18 +5094,10 @@ def test_get_logging_payload_gives_only_a_charged_video_poll_a_row_of_its_own(st assert status["request_id"] == status_id -@pytest.mark.parametrize( - ("call_type", "job_id", "charged_id"), - [("avideo_status", "video_abc", "video_abc_video_cost"), ("aretrieve_batch", "batch_abc", "batch_abc_batch_cost")], -) @pytest.mark.parametrize(("response_cost", "pays_the_job"), [(0.01, False), (0.43, True)]) -def test_a_job_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own( - call_type, job_id, charged_id, response_cost, pays_the_job -): +def test_a_video_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own(response_cost, pays_the_job): """A guardrail charges every poll, so a poll paying only that must neither land on the create row - nor take the row the job's own charge goes on, either of which the skip_duplicates insert drops. A batch - settles on its first finished retrieve, so its guardrail-only poll is a running one.""" - polled_status = "completed" if pays_the_job or call_type == "avideo_status" else "in_progress" + nor take the row the video's own charge goes on, either of which the skip_duplicates insert drops.""" kwargs = { "litellm_call_id": "poll-call-id", "response_cost": response_cost, @@ -5123,12 +5108,17 @@ def test_a_job_poll_that_pays_only_its_guardrail_gets_a_row_of_its_own( }, } - assert get_spend_logs_id(call_type, {"id": job_id, "status": polled_status}, kwargs) == ( - charged_id if pays_the_job else "poll-call-id_guardrail_cost" + assert get_spend_logs_id("avideo_status", {"id": "video_abc", "status": "completed"}, kwargs) == ( + "video_abc_video_cost" if pays_the_job else "poll-call-id_guardrail_cost" ) -def test_a_finished_batch_retrieve_paying_only_its_guardrail_settles_under_the_batch_row(): +@pytest.mark.parametrize("status", ["failed", "completed"]) +def test_a_finished_batch_retrieve_paying_only_its_guardrail_charges_under_the_batch_row(status): + """The retrieve that settles a batch charges its whole cost once, guardrails included, so a batch whose own + cost is $0 must still land on the row that makes its repeat retrieves free.""" + from litellm.types.utils import LiteLLMBatch + kwargs = { "litellm_call_id": "poll-call-id", "response_cost": 0.01, @@ -5137,11 +5127,19 @@ def test_a_finished_batch_retrieve_paying_only_its_guardrail_settles_under_the_b "guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}], }, } - - assert get_spend_logs_id("aretrieve_batch", {"id": "batch_abc", "status": "failed"}, kwargs) == ( - "batch_abc_batch_cost" + polled_batch = LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id="file-out", ) + assert get_spend_logs_id("aretrieve_batch", polled_batch.model_dump(), kwargs) == "batch_abc_batch_cost" + def test_a_video_polled_with_or_without_its_base64_padding_charges_one_row(): from litellm.types.videos.utils import encode_video_id_with_provider diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index c2df918e1d0..d24d0598d87 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -466,6 +466,7 @@ def _completed_budget_batch(batch_id: str): input_file_id="file-budget", object="batch", status="completed", + output_file_id="file-budget-output", ) @@ -503,6 +504,54 @@ async def test_batch_polls_charge_each_router_budget_once(monkeypatch, call_type +@pytest.mark.asyncio +async def test_a_batch_retrieve_before_its_cost_is_final_charges_no_router_budget_even_for_its_guardrail(): + from litellm.types.utils import LiteLLMBatch + + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "batch-budget", + "litellm_params": {"model": "openai/batch-budget", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + retrieve: Final = partial( + _success_kwargs, + provider_in_litellm_params=None, + provider_in_payload="openai", + call_type="aretrieve_batch", + guardrail_information=({"guardrail_name": "priced", "guardrail_cost": 0.01},), + ) + + def polled(status: str, output_file_id: str | None = None) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-running", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-budget", + object="batch", + status=status, + output_file_id=output_file_id, + ) + + for response_cost, batch in ( + (0.01, polled("in_progress")), + (0.01, polled("completed")), + (0.26, polled("completed", "file-out")), + (0.26, polled("completed", "file-out")), + ): + await limiter.async_log_success_event(retrieve(response_cost=response_cost), batch, None, None) + + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") + assert tuple([await cache.async_get_cache(key) for key in keys]) == pytest.approx((0.26, 0.26)) + + @pytest.mark.asyncio @pytest.mark.parametrize("call_type", ("avideo_status", "video_status")) async def test_video_polls_charge_each_router_budget_once(call_type):