diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a0c0a90be5f..d272a235758 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -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"], diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index 5b94d4b4e5b..7c7b7ea5c4b 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -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")), + ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5494d0d82b7..1fcb06b80e8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -38,7 +38,7 @@ from litellm._logging import ( verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import batch_cost_is_final, handle_completed_batch +from litellm.batches.batch_utils import handle_completed_batch from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.caching.redis_batch import flush_post_call_redis_batches @@ -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): diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d9fb236632b..35e7ae1c579 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,6 +34,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.batch_utils import BATCH_COST_REQUEST_ID_SUFFIX from litellm.litellm_core_utils.core_helpers import finite_cost, single_log_line from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -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 diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 002aeb145a5..636345a7ef7 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -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.""" diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 1c10826124e..889fb9fb373 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,7 +19,12 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.batch_utils import 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 ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 10b6ec2f661..90887422567 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -4,6 +4,7 @@ from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar from datetime import datetime, timezone +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast @@ -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.""" diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index c0d81abeba7..18ce18da936 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -21,8 +21,8 @@ from typing import ( from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.batches.batch_utils import batch_cost_is_final from litellm.constants import MAX_FILE_LIST_LIMIT +from litellm.litellm_core_utils.batch_utils import batch_cost_is_final from litellm.proxy._types import ProxyException from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.table_repositories import ( diff --git a/litellm/proxy/spend_tracking/spend_event.py b/litellm/proxy/spend_tracking/spend_event.py index 031e84896b2..d054519cb9c 100644 --- a/litellm/proxy/spend_tracking/spend_event.py +++ b/litellm/proxy/spend_tracking/spend_event.py @@ -28,6 +28,7 @@ from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import LiteLLMBatch, Usage +from litellm.types.videos.utils import is_video_status_call_type SPEND_EVENT_VERSION: Final = 1 CACHE_OFF_KEY: Final = "Cache OFF" @@ -205,9 +206,10 @@ class SpendEventDecodeError: reason: str -def is_offloadable_success(response_obj: object) -> bool: - """Batch retrieves and interaction polls branch on the concrete response class, so they stay in-process.""" - return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) +def is_offloadable_success(response_obj: object, call_type: object) -> bool: + """Batch retrieves and interaction polls branch on the concrete response class, and a video status poll on + the status and video id the event leaves out, so they stay in-process.""" + return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES) and not is_video_status_call_type(call_type) def _json_fallback(value: object) -> str: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index eb8c28fcbd7..85a7f870eb4 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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)}" diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 1d502dd8d95..98dd39b40e3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,7 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.batch_utils import 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) diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index 0fc2b8ede1b..240e36ecdf6 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -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, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 57ae9f1905f..42ee6178bf0 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -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 diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 16b8cf3f0e1..0f85f1e3993 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -22,7 +22,6 @@ from typing import Final, Literal import httpx import pytest import respx -from openai.types.batch import BatchRequestCounts import litellm @@ -2028,61 +2027,6 @@ def test_extract_credentials_keeps_openai_workload_identity_params(): } -# --------------------------------------------------------------------------- # -# batch_cost_is_final -# --------------------------------------------------------------------------- # - - -def _retrieved_batch( - status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None -) -> LiteLLMBatch: - return LiteLLMBatch( - id="batch_abc", - completion_window="24h", - created_at=1, - endpoint="/v1/chat/completions", - input_file_id="file-in", - object="batch", - status="validating", - output_file_id=output_file_id, - request_counts=counts, - ).model_copy(update={"status": status}) - - -class TestBatchCostIsFinal: - """Every retrieve of one batch writes the same spend row, so the first retrieve - that prices it decides the row for good. A poll before the output exists must - therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" - - @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) - def test_in_flight_batch_is_not_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status)) is False - - @pytest.mark.parametrize("status", ["completed", "complete"]) - def test_completed_with_output_is_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True - - def test_completed_without_output_and_unknown_counts_is_not_final(self): - assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False - - def test_completed_without_output_and_zero_counts_is_not_final(self): - counts = BatchRequestCounts(total=0, completed=0, failed=0) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - - def test_completed_without_output_but_successful_lines_is_not_final(self): - counts = BatchRequestCounts(total=2, completed=2, failed=0) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - - @pytest.mark.parametrize("status", ["completed", "complete"]) - def test_completed_without_output_and_every_line_failed_is_final(self, status): - counts = BatchRequestCounts(total=2, completed=0, failed=2) - assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True - - @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) - def test_other_terminal_statuses_are_final(self, status): - assert bu.batch_cost_is_final(_retrieved_batch(status)) is True - - def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: @@ -2965,7 +2909,7 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( @@ -3060,7 +3004,6 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( wrong token sum, or mispriced model fails this test. """ from litellm.batches.batch_utils import handle_completed_batch - from litellm.types.utils import LiteLLMBatch batch = LiteLLMBatch( id="batch-real-cost-123", @@ -3105,7 +3048,7 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( @@ -3187,7 +3130,7 @@ async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown card renders blank for every batch. Regression for the split being dropped.""" from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes mock_batch = LiteLLMBatch( id="batch-breakdown-1", @@ -3243,7 +3186,7 @@ async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batc from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch, SpecialEnums + from litellm.types.utils import CallTypes, SpecialEnums # Create a proper unified file ID by encoding the correct prefix unified_id_str = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}:test_file_789;unified_id:batch-789" @@ -3322,7 +3265,7 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): from unittest.mock import AsyncMock, patch from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import CallTypes, LiteLLMBatch + from litellm.types.utils import CallTypes # Mock batch result with completed status mock_batch = LiteLLMBatch( diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index cd37d635c43..45461b50761 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -5,14 +5,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 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index cf9ad4f91a3..a014d1d352d 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -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) diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 298b153f621..3a79df0370b 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1,10 +1,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", diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 45cc7eb8922..966e94fad80 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -8,6 +8,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(): """ diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index f5a2ccf7b76..f1fd02981d6 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -27,6 +27,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})} diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c8434007a12..c6366a40504 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -4,7 +4,7 @@ import logging from collections.abc import Mapping from datetime import datetime, timedelta, timezone from pathlib import Path -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, create_autospec, patch @@ -24,6 +24,7 @@ from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( ProxyDBLogger, + _charge_request_cost, _get_budget_reservation_from_metadata, _should_track_cost_callback, _update_database_and_spend_counters, @@ -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,))) diff --git a/tests/unit/proxy/spend_tracking/test_spend_event.py b/tests/unit/proxy/spend_tracking/test_spend_event.py index ff449235582..e1912d36d91 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_event.py +++ b/tests/unit/proxy/spend_tracking/test_spend_event.py @@ -194,9 +194,11 @@ def test_undecodable_line_is_an_error_value(): assert isinstance(decode_spend_event(b"not json\n"), SpendEventDecodeError) -def test_batch_retrieves_stay_in_process(): - assert is_offloadable_success(_response()) is True - assert is_offloadable_success(None) is True +def test_batch_retrieves_and_video_status_polls_stay_in_process(): + assert is_offloadable_success(_response(), "acompletion") is True + assert is_offloadable_success(None, None) is True + assert is_offloadable_success(None, "avideo_generation") is True + assert is_offloadable_success(None, "avideo_status") is False assert ( is_offloadable_success( LiteLLMBatch( @@ -207,7 +209,8 @@ def test_batch_retrieves_stay_in_process(): input_file_id="f", object="batch", status="completed", - ) + ), + "aretrieve_batch", ) is False ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 32f6eea91ac..448cc498c7a 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -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 diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index c0583566531..d24d0598d87 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -7,6 +7,7 @@ the callback raise before any spend was recorded, so those budgets never moved. """ from datetime import datetime, timezone +from functools import partial from typing import Final import litellm @@ -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