scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(proxy): charge a video's provider-reported cost once, on the first poll that sees it completed (#49824)

Browse source 
* 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.

* 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

* 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.

* fix(proxy): charge repeat poll guardrails, key video charges by canonical id and survive redaction

* fix(proxy): keep video status polls in-process for the collector and give guardrail-only job polls an unclaimable row id

* fix(proxy): keep repeat batch polls free and group guardrail-only batch rows under their batch's trace

* fix(proxy): settle a finished job whose own cost is below its poll's guardrail cost under the job's charge row

* test(proxy): count charged_poll_cost as covered through the finished job guardrail poll tests

* fix(proxy): settle a finished batch that cost nothing itself on its first priced retrieve

* fix(proxy): charge no budget for a batch retrieve before its cost is final, and read the retrieved batch from before redaction

* fix(proxy): charge sessions a poll's guardrail-inclusive cost, count a job cost row's guardrail run, and keep batch_cost_is_final importable from litellm.batches

* fix(proxy): claim a job and charge the session budget in one Redis call

The session limiter claimed the job marker and then raised spend in a second Redis call. A failure between the two left the job claimed while its charge sat only in local memory, so later polls skipped it. One script now does both, and a refused script claims and charges locally

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tessellark 2026-11-04 12:22:06 -08:00 • committed by GitHub
parent e69da69891
commit 8ab0321816
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 1881 additions and 267 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/batches/batch_utils.py
View file

@ -10,6 +10,7 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils import batch_utils as core_batch_utils
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
@ -38,8 +39,8 @@ class BatchCostUsageResult:
completion_cost: float = 0.0
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
_DECODED_BATCH_LINE: Final = TypeAdapter(object)
batch_cost_is_final: Final = core_batch_utils.batch_cost_is_final
def _uses_native_vertex_output(
@ -54,26 +55,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"],

164
litellm/litellm_core_utils/batch_utils.py
View file

@ -3,21 +3,32 @@ 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
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
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])
_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})
_batch_create_routing: Final[ContextVar[bool]] = ContextVar("batch_create_routing", default=False)
@ -289,12 +300,151 @@ 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
repeat_poll_cost: float = 0.0
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 _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 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)
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)))
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
@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
return response_obj.id if isinstance(response_obj, Batch) else 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"))
<= 0
< (finite_cost(response_cost) or 0.0)
)
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, 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. 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 kwargs.get(POLLED_VIDEO_STATUS_KWARG) != "completed":
return None
if job_cost_of_poll(response_cost, kwargs.get("standard_logging_object")) <= 0:
return None
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:
polled: Final = _polled_batch(kwargs, response_obj)
if polled is None or not polled.cost_is_final or response_cost <= 0:
return None
return ChargeOnceJob(kind="batch", job_id=polled.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."""
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(
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(call_type, response_cost, kwargs.get("standard_logging_object")),
)

13
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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
@ -79,6 +79,11 @@ from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.litellm_core_utils.agentic_followup_failure import ServedRound
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,
@ -220,6 +225,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,
@ -3308,6 +3314,11 @@ class Logging(LiteLLMLoggingBaseClass):
result = self.normalize_endpoint_logging_result(result)
if self.stream is not True:
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):

105
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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
@ -92,6 +93,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_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES
if TYPE_CHECKING:
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
@ -158,23 +160,53 @@ 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"
async def _track_guardrail_usage_of_cost_row(
payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool
) -> None:
"""A job's cost row skips the spend log queue, so its guardrail runs are counted here instead of at the flush."""
if disable_spend_logs:
return
from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage
try:
await process_spend_logs_guardrail_usage(prisma_client=prisma_client, logs_to_process=(payload,))
except Exception as e: # noqa: BLE001 # guardrail usage is best effort, as at the spend log flush
verbose_proxy_logger.warning(
"Guardrail usage tracking failed for spend row %s (non-fatal): %s", payload["request_id"], e
)
_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"})
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. 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 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)
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.
_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"})
A proxy run with spend logs disabled still needs one row per batch to charge it once,
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 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 +645,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 +787,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 +815,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:
@ -791,35 +823,40 @@ class DBSpendUpdateWriter:
data=[prisma_client.jsonify_object(row)],
skip_duplicates=True,
)
if claimed == 1:
return True
existing: Final = await spend_logs.find_unique(where={"request_id": request_id})
existing: Final = None if claimed == 1 else 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 claimed == 1:
await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs)
return True
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)
taken_over: Final = await self._take_over_uncharged_cost_row(
payload=payload, prisma_client=prisma_client, row=row
)
if taken_over:
await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs)
return taken_over
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 +864,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 +874,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

89
litellm/proxy/hooks/max_budget_per_session_limiter.py
View file

@ -26,6 +26,14 @@ 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,
charged_poll_cost,
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
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
@ -44,6 +52,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.
@ -63,6 +72,30 @@ end
return new_val
"""
MAX_BUDGET_SESSION_CHARGE_JOB_ONCE_SCRIPT: Final = """
local spend_key = KEYS[1]
local marker_key = KEYS[2]
local ttl = tonumber(ARGV[3])
local polls = redis.call('INCR', marker_key)
redis.call('EXPIRE', marker_key, ttl)
local amount = ARGV[2]
if polls == 1 then
amount = ARGV[1]
end
if tonumber(amount) <= 0 then
return '0'
end
local existed = redis.call('EXISTS', spend_key)
redis.call('INCRBYFLOAT', spend_key, amount)
if existed == 0 then
redis.call('EXPIRE', spend_key, ttl)
end
return amount
"""
# Default TTL for session budget counters (1 hour)
DEFAULT_MAX_BUDGET_PER_SESSION_TTL: Final = 3600
@ -87,12 +120,15 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
)
)
if self.internal_usage_cache.dual_cache.redis_cache is not None:
self.increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
MAX_BUDGET_SESSION_INCREMENT_SCRIPT
)
else:
self.increment_script = None
redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache
self.increment_script = (
redis_cache.async_register_script(MAX_BUDGET_SESSION_INCREMENT_SCRIPT) if redis_cache is not None else None
)
self.charge_job_once_script = (
redis_cache.async_register_script(MAX_BUDGET_SESSION_CHARGE_JOB_ONCE_SCRIPT)
if redis_cache is not None
else None
)
@with_service_target("session_budgets")
async def async_pre_call_hook(
@ -140,7 +176,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,15 +204,23 @@ 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(charged_poll_cost(callback_fields)) or 0.0
if response_cost <= 0:
return
if defers_batch_cost(callback_fields, response_obj):
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, response_obj, response_cost)
if job is None:
await self._increment_spend(cache_key, response_cost)
else:
await self._charge_job_once(cache_key, job, response_cost)
verbose_proxy_logger.debug(
"MaxBudgetPerSessionHandler: incremented session %s spend by %.6f",
"MaxBudgetPerSessionHandler: charged session %s for a call that cost %.6f",
session_id,
response_cost,
)
@ -221,6 +265,31 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
def _make_cache_key(self, session_id: str) -> str:
return f"{{session_budget:{session_id}}}:spend"
async def _charge_job_once(self, cache_key: str, job: ChargeOnceJob, response_cost: float) -> None:
marker_key: Final = job_charged_once_marker_key(spend_key=cache_key, job=job)
if self.charge_job_once_script is not None:
try:
await self.charge_job_once_script(
keys=[cache_key, marker_key],
args=[str(response_cost), str(job.repeat_poll_cost), self.ttl],
)
except Exception as e:
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"MaxBudgetPerSessionHandler: Redis job charge failed, falling back to in-memory",
e,
)
else:
return
polls: Final = await self.internal_usage_cache.dual_cache.in_memory_cache.async_increment(
key=marker_key, value=1, ttl=self.ttl
)
charge: Final = response_cost if polls == 1 else job.repeat_poll_cost
if charge > 0:
await self._in_memory_increment_spend(cache_key, charge)
@with_service_target("session_budgets")
async def _get_current_spend(self, cache_key: str) -> float:
"""Read current accumulated spend for a session."""

36
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -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 batch_charged_once_marker_key, batch_id_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,18 +1213,17 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
)
return
batch_id: Final = batch_id_to_charge_once(
call_type=kwargs.get("call_type"),
response_obj=response_obj,
response_cost=response_cost,
)
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(
entity_type=entity_type,
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 +1240,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,11 +1251,19 @@ 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),
):
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,
@ -1362,8 +1374,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
)

99
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -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
@ -11,10 +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 defers_batch_cost, repeat_poll_cost
from litellm.litellm_core_utils.core_helpers import (
budget_reservation_from_metadata,
finite_cost,
@ -75,7 +76,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
@ -415,7 +415,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(
@ -797,16 +797,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))
@ -868,25 +860,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:
@ -896,7 +891,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,
@ -907,7 +902,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:
@ -1201,6 +1196,50 @@ 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
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,
**(
{
"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
class _IncrementSpendCounters(Protocol):
"""The ``increment_spend_counters`` coroutine :func:`_update_database_and_spend_counters` awaits."""

2
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -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 (

8
litellm/proxy/spend_tracking/spend_event.py
View file

@ -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:

16
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -38,6 +38,12 @@ 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 (
BATCH_COST_REQUEST_ID_SUFFIX,
charged_poll_cost,
charged_video_id,
pays_only_its_guardrail,
)
from litellm.litellm_core_utils.classifier_logging import (
CLASSIFIER_AUDIT_FIELDS,
classifier_audit_fields,
@ -98,6 +104,7 @@ from litellm.types.utils import (
VectorStoreSearchResponse,
batch_chat_usage,
)
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:
@ -342,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:
@ -370,6 +377,13 @@ def get_spend_logs_id(
)
if resolved_id is None:
return None
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}"
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 f"{call_id}{JOB_POLL_GUARDRAIL_COST_REQUEST_ID_SUFFIX}"
return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}"

37
litellm/router_strategy/budget_limiter.py
View file

@ -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 batch_charged_once_marker_key, batch_id_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,
)
@ -52,7 +57,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 +550,16 @@ 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)
if defers_batch_cost(callback_fields, polled_response):
return
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
@ -565,7 +573,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 +586,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 +604,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 +613,21 @@ 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):
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(
@ -639,8 +654,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)

56
litellm/types/videos/utils.py
View file

@ -17,10 +17,11 @@ 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_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_status",
"video_status",
"avideo_content",
"video_content",
"avideo_remix",
@ -31,6 +32,12 @@ VIDEO_ID_CALL_TYPES: Final = 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:{}"
@ -94,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:
@ -103,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,

9
tests/code_coverage_tests/router_code_coverage.py
View file

@ -176,6 +176,15 @@ 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
"_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
"_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
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py

67
tests/unit/batches/test_batch_utils.py
View file

@ -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(

257
tests/unit/litellm_core_utils/test_batch_utils.py
View file

@ -5,14 +5,22 @@ 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
@ -382,30 +390,241 @@ 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 _batch(status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None):
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",
input_file_id="file-budget", object="batch", status="completed",
return Batch(
id="job-budget", 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,
request_counts=counts,
)
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_batch():
return _batch("completed", output_file_id="file-out")
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")
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
@pytest.mark.parametrize(
("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, guardrail_only_poll, kind, repeat_cost
) -> 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, 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}]},
}
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:
"""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", POLLED_VIDEO_STATUS_KWARG: 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")
@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
@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
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 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"
)
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

101
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -14453,6 +14453,107 @@ def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_call
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")
@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
_FAILED_ATTEMPT: Final = (Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), 0.5)

393
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -1,10 +1,12 @@
import asyncio
import copy
import functools
import itertools
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
@ -4456,7 +4458,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 +4709,391 @@ 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(request_id: str = "video_abc_video_cost") -> dict:
return {
**_minimal_spend_payload(),
"request_id": request_id,
"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). 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("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)
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": ("avideo_status", 0.0),
"video_abc_video_cost": ("avideo_status", 0.42),
}
def _indexed_guardrail_request_ids(prisma: MagicMock, index_table: str) -> tuple[str, ...]:
calls: Final = getattr(prisma.db, index_table).create_many.await_args_list
batches: Final = tuple(call.kwargs["data"] for call in calls)
return tuple(row["request_id"] for row in itertools.chain.from_iterable(batches))
@pytest.mark.asyncio
@pytest.mark.parametrize("disable_spend_logs", [False, True], ids=["spend_logs_on", "spend_logs_off"])
@pytest.mark.parametrize("left_uncharged", [False, True], ids=["claimed_row", "taken_over_row"])
async def test_a_job_cost_row_counts_its_polls_guardrail_run_in_guardrail_usage_once(
disable_spend_logs: bool, left_uncharged: bool
):
"""The job's cost row is written by its claim, not the spend log queue whose flush counts guardrail runs, so a
finished video's guardrail run must still reach guardrail usage, once, and not when spend logs are off."""
from litellm.repositories.table_repositories import (
DailyGuardrailMetricsRepository,
SpendLogGuardrailIndexRepository,
)
db_writer = DBSpendUpdateWriter()
db_writer._batch_database_updates = AsyncMock()
prisma = _spend_logs_table_prisma()
getattr(prisma.db, SpendLogGuardrailIndexRepository.table_name).create_many = AsyncMock(return_value=1)
getattr(prisma.db, DailyGuardrailMetricsRepository.table_name).upsert = AsyncMock()
if left_uncharged:
prisma.db.litellm_spendlogs.rows["video_abc_video_cost"] = {**_video_status_payload(), "spend": 0.0}
guarded_poll: Final = {
**_video_status_payload(),
"metadata": json.dumps(
{"guardrail_information": [{"guardrail_name": "priced", "guardrail_status": "success"}]}
),
}
charged = [
await _update_database_with(
db_writer, prisma, guarded_poll, disable_spend_logs=disable_spend_logs, response_cost=0.43
)
for _ in range(2)
]
assert charged == [True, False]
assert _indexed_guardrail_request_ids(prisma, SpendLogGuardrailIndexRepository.table_name) == (
() if disable_spend_logs else ("video_abc_video_cost",)
)
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",
) -> 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
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,
output_file_id="file-out",
)
)
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
),
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"
),
):
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
@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_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()
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
@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_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, status="failed", 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
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."""
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():
"""
Built-in video providers charge at create and report nothing on the status poll, so those
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_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_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)
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
@pytest.mark.parametrize(
"injected_deployment, attributed",

253
tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py
View file

@ -8,6 +8,9 @@ Tests that session-scoped budget tracking works correctly:
- Requests without agent_id pass through
"""
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from unittest.mock import patch
import logging
@ -22,6 +25,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:
@ -145,6 +149,255 @@ 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", output_file_id="file-out",
)
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,
POLLED_VIDEO_STATUS_KWARG: "completed",
"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
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(
video_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": video_cost,
POLLED_VIDEO_STATUS_KWARG: status,
"standard_logging_object": {
"response_cost": video_cost + guardrail_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
@pytest.mark.parametrize(
("video_cost", "guardrail_cost"),
[(0.40, 0.01), (0.05, 0.10)],
ids=["video_costs_more_than_its_guardrail", "video_costs_no_more_than_its_guardrail"],
)
async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pays_the_session_for_the_video(
video_cost: float, guardrail_cost: float
):
from litellm.types.videos.main import VideoObject
handler = PROXY_MaxBudgetPerSessionHandler(internal_usage_cache=InternalUsageCache(DualCache()))
completed_polls: Final = 3
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.0, guardrail_cost=guardrail_cost, status="in_progress"),
VideoObject(id="job-1", object="video", status="in_progress"),
None,
None,
)
for _ in range(completed_polls):
await handler.async_log_success_event(
_video_poll(video_cost, guardrail_cost=guardrail_cost),
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(video_cost + (1 + completed_polls) * guardrail_cost)
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({})
self.scripts_fail = False
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 _increment_expiring_on_first_write(self, key: str, amount: float, ttl: float) -> float:
live = self._live(key)
expires_at = self.now + ttl if live is None else live[1]
return self._store(key, amount + (0.0 if live is None else live[0]), expires_at)
def async_register_script(self, script: str):
async def run(keys: list[str], args: list[object]) -> float:
if self.scripts_fail:
raise ConnectionError("Redis dropped the script call")
if len(keys) == 1:
return self._increment_expiring_on_first_write(keys[0], float(str(args[0])), float(str(args[1])))
marker = self._live(keys[1])
polls = self._store(keys[1], 1.0 + (0.0 if marker is None else marker[0]), self.now + float(str(args[2])))
amount = float(str(args[0] if polls == 1 else args[1]))
if amount > 0:
self._increment_expiring_on_first_write(keys[0], amount, float(str(args[2])))
return amount
return run
@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
async def test_a_finished_video_whose_charge_missed_redis_is_charged_on_the_next_poll_once_redis_answers():
from litellm.types.videos.main import VideoObject
redis: Final = _ClockedRedis()
handler: Final = PROXY_MaxBudgetPerSessionHandler(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis)) # pyright: ignore[reportArgumentType] # duck-typed Redis double
)
video: Final = VideoObject(id="job-1", object="video", status="completed")
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)
redis.scripts_fail = True
await handler.async_log_success_event(_video_poll(0.4), video, None, None)
redis.scripts_fail = False
for _ in range(2):
await handler.async_log_success_event(_video_poll(0.4), video, None, None)
assert await handler._get_current_spend(handler._make_cache_key("session-jobs")) == pytest.approx(0.4)
@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():
"""

56
tests/unit/proxy/hooks/test_model_max_budget_limiter.py
View file

@ -27,6 +27,8 @@ 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
from litellm.types.videos.utils import POLLED_VIDEO_STATUS_KWARG
KEY_HASH: Final = "key-hash-batch"
USER_ID: Final = "user-batch"
@ -47,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),
)
@ -57,12 +60,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": {
@ -326,6 +331,57 @@ 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), POLLED_VIDEO_STATUS_KWARG: status},
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)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("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), 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},)
def poll(status: str, response_cost: float) -> dict[str, object]:
return {**_event(call_type, response_cost, guardrail_information=priced), POLLED_VIDEO_STATUS_KWARG: status}
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 + guardrail_charges * 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})}

125
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -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,
@ -42,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
@ -1595,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",
@ -1611,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
@ -1646,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(),
)
@ -3376,18 +3391,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,
@ -3398,6 +3403,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"
@ -3411,7 +3439,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
@ -3981,3 +4009,50 @@ async def test_async_post_call_failure_hook_disable_pre_routing_error_logs(
user_api_key_dict=user_api_key_dict,
)
assert len(writer.rows) == rows_written
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(
{
"call_type": "avideo_status",
"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),
((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",
],
)
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,)))

11
tests/unit/proxy/spend_tracking/test_spend_event.py
View file

@ -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
)

99
tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -5130,6 +5130,105 @@ 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"])
@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
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,
"polled_video_status": "completed",
"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
@pytest.mark.parametrize(("response_cost", "pays_the_job"), [(0.01, False), (0.43, True)])
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 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,
"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) == (
"video_abc_video_cost" if pays_the_job else "poll-call-id_guardrail_cost"
)
@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,
"standard_logging_object": {
"response_cost": 0.01,
"guardrail_information": [{"guardrail_name": "priced", "guardrail_cost": 0.01}],
},
}
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
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():
"""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

129
tests/unit/router_strategy/test_budget_limiter.py
View file

@ -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
@ -17,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
@ -38,6 +40,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 +54,7 @@ def _success_kwargs(
"response_cost": response_cost,
"model_id": model_id,
"custom_llm_provider": provider_in_payload,
"guardrail_information": guardrail_information,
},
}
@ -462,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",
)
@ -498,6 +503,130 @@ 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
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):
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),
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")
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_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()
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), 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), 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.29, 0.29))
@pytest.mark.asyncio
async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation():
from fakeredis import FakeServer, FakeRedis as FakeSyncRedis