From ed4d3ae6ccc9f7582cf5e6404140791435575a8a Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:17:44 -0700 Subject: [PATCH] fix(accounting): enforce batch charge ownership and recover retained pricing --- .../migration.sql | 3 +- .../litellm_proxy_extras/schema.prisma | 1 + litellm/cost_calculator.py | 4 +- litellm/litellm_core_utils/litellm_logging.py | 8 +- litellm/proxy/db/accounting_ledger.py | 13 +- litellm/proxy/schema.prisma | 1 + litellm/proxy/spend_tracking/cost_compiler.py | 48 ++++- litellm/proxy/spend_tracking/usage_capture.py | 40 +++- schema.prisma | 1 + .../database/test_accounting_ledger.py | 146 +++++++++++++- .../spend_tracking/test_cost_compiler.py | 181 ++++++++++++++++++ .../spend_tracking/test_usage_capture.py | 11 +- 12 files changed, 427 insertions(+), 30 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql index 04fbdfdc237..ef18da183ed 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql @@ -1,5 +1,6 @@ ALTER TABLE "LiteLLM_AccountingEvent" ADD COLUMN "compilation_state" TEXT NOT NULL DEFAULT 'pending', -ADD COLUMN "compilation_error" TEXT; +ADD COLUMN "compilation_error" TEXT, +ADD COLUMN "compilation_retry_at" TIMESTAMP(3); CREATE INDEX "LiteLLM_AccountingEvent_kind_compilation_state_recorded_at_idx" ON "LiteLLM_AccountingEvent"("kind", "compilation_state", "recorded_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 048e0e1e17a..45ea5cd20d0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2113,6 +2113,7 @@ model LiteLLM_AccountingEvent { payload Json compilation_state String @default("pending") compilation_error String? + compilation_retry_at DateTime? @@index([recorded_at]) @@index([kind, source_event_id]) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4a72e88ad3f..d50627cceb0 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2009,7 +2009,7 @@ def completion_cost( video_resolution=video_resolution, ) elif call_type in _SPEECH_CALL_TYPES: - prompt_characters = _speech_billed_characters( + prompt_characters = speech_billed_characters( completion_response if isinstance(completion_response, HttpxBinaryResponseContent) else None, prompt, ) @@ -2527,7 +2527,7 @@ def response_cost_calculator( raise e -def _speech_billed_characters(speech_response: HttpxBinaryResponseContent | None, prompt: str) -> int: +def speech_billed_characters(speech_response: HttpxBinaryResponseContent | None, prompt: str) -> int: reported: Final = speech_response.hidden_params.get("billed_characters") if speech_response is not None else None return reported if isinstance(reported, int) else litellm.utils.count_characters(text=prompt) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 41668370609..c87aaabc4b7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -65,6 +65,7 @@ from litellm.cost_calculator import ( get_usage_object, pricing_entry_for_cost_calc, select_model_name_for_cost_calc, + speech_billed_characters, ) from litellm.exceptions import ( BudgetExceededError, @@ -3069,8 +3070,6 @@ class Logging(LiteLLMLoggingBaseClass): return self.cost_billing_time def usage_capture_pricing(self, result: object) -> Mapping[str, JsonValue]: - from litellm.utils import count_characters - return TypeAdapter(dict[str, JsonValue]).validate_python( { "response_ms": getattr(result, "_response_ms", None), @@ -3086,7 +3085,10 @@ class Logging(LiteLLMLoggingBaseClass): if self.model_call_details.get("call_type") in (CallTypes.call_mcp_tool.value, CallTypes.send_message.value, CallTypes.asend_message.value) else None, - "prompt_characters": count_characters(self._prompt_for_cost_calculation()) + "prompt_characters": speech_billed_characters( + result if isinstance(result, HttpxBinaryResponseContent) else None, + self._prompt_for_cost_calculation(), + ) if self.model_call_details.get("call_type") in (CallTypes.speech.value, CallTypes.aspeech.value) else None, } diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py index b479cea6a0d..7ebdfb70506 100644 --- a/litellm/proxy/db/accounting_ledger.py +++ b/litellm/proxy/db/accounting_ledger.py @@ -164,14 +164,22 @@ FROM "LiteLLM_AccountingEvent" WHERE source_event_id = $1 AND kind = 'cost' _PENDING_COMPILATION: Final = """ SELECT payload::text AS payload FROM "LiteLLM_AccountingEvent" WHERE kind = 'outbound_usage' AND compilation_state = 'pending' +AND (compilation_retry_at IS NULL OR compilation_retry_at <= CURRENT_TIMESTAMP) ORDER BY recorded_at, event_id LIMIT $1::int """ _COMPILATION_RESULT: Final = """ -UPDATE "LiteLLM_AccountingEvent" SET compilation_state = $2, compilation_error = $3 +UPDATE "LiteLLM_AccountingEvent" SET compilation_state = $2, compilation_error = $3, compilation_retry_at = NULL WHERE event_id = $1 AND kind = 'outbound_usage' """ +_RETRY_COMPILATION: Final = """ +UPDATE "LiteLLM_AccountingEvent" +SET compilation_error = 'calculation_retryable', compilation_retry_at = CURRENT_TIMESTAMP + INTERVAL '30 seconds' +WHERE event_id = $1 AND kind = 'outbound_usage' AND compilation_state = 'pending' +""" + + async def _unpriced(db: Prisma, event_id: str, reason: str) -> Unpriced: await db.execute_raw( _COMPILATION_RESULT, event_id, "excluded" if reason == "not_charge_owner" else "unpriced", reason @@ -401,6 +409,9 @@ class AccountingLedger: priced: Final = price(source["payload"]) except ValueError: return await _unpriced(db, validated.event_id, "invalid_pricing_inputs") + except Exception: # noqa: BLE001 # retain failed pricing for bounded retries without starving newer events + await db.execute_raw(_RETRY_COMPILATION, validated.event_id) + return Unpriced(usage_event_id=validated.event_id, reason="calculation_retryable") if isinstance(priced, Unpriced): if priced.usage_event_id != validated.event_id: raise AccountingConflictError("Unpriced result does not match its source usage") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 048e0e1e17a..45ea5cd20d0 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2113,6 +2113,7 @@ model LiteLLM_AccountingEvent { payload Json compilation_state String @default("pending") compilation_error String? + compilation_retry_at DateTime? @@index([recorded_at]) @@index([kind, source_event_id]) diff --git a/litellm/proxy/spend_tracking/cost_compiler.py b/litellm/proxy/spend_tracking/cost_compiler.py index 0ad72735401..52452a79106 100644 --- a/litellm/proxy/spend_tracking/cost_compiler.py +++ b/litellm/proxy/spend_tracking/cost_compiler.py @@ -5,6 +5,7 @@ from datetime import datetime from decimal import Decimal from typing import Final +from httpx import Response from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter from litellm._internal_context import pinned_billing_time @@ -18,19 +19,21 @@ from litellm.accounting import ( PricingSnapshot, Unpriced, ) -from litellm.cost_calculator import batch_cost_calculator +from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost from litellm.litellm_core_utils.hidden_params import set_hidden_param from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_cost_calc.finalization import finalize_cost from litellm.litellm_core_utils.llm_cost_calc.pricing_context import RetainedPricing, retained_pricing +from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.proxy.db.accounting_ledger import AccountingLedger from litellm.proxy.spend_tracking.usage_meter import usage_delivery_targets -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, CostBreakdown, CostPerToken, LiteLLMRealtimeStreamLoggingObject, + ModelInfo, ModelResponse, StandardBuiltInToolsParams, Usage, @@ -55,6 +58,7 @@ class _UsageEnvelope(BaseModel): usage: Mapping[str, JsonValue] | None = None pricing: Mapping[str, JsonValue] charge_eligible: bool + parent_request_id: str | None = None finalized_cost: Mapping[str, JsonValue] | None = None allocation: Mapping[str, JsonValue] | None = None @@ -62,8 +66,19 @@ class _UsageEnvelope(BaseModel): def price_usage(envelope: str) -> PricingSnapshot | Unpriced: captured: Final = _UsageEnvelope.model_validate_json(envelope) event: Final = captured.event - if not captured.charge_eligible: + batch_summary: Final = event.call_type in {CallTypes.retrieve_batch, CallTypes.aretrieve_batch} and not ( + captured.parent_request_id is not None and event.event_id.startswith(f"{captured.parent_request_id}:row:") + ) + if not captured.charge_eligible or batch_summary: return Unpriced(usage_event_id=event.event_id, reason="not_charge_owner") + logged: Final = _fields(captured.pricing.get("logging")) + speech_characters: Final = logged.get("prompt_characters") + speech_usage: Final = ( + event.call_type in {CallTypes.speech, CallTypes.aspeech} + and isinstance(speech_characters, int) + and not isinstance(speech_characters, bool) + and speech_characters >= 0 + ) fixed_call: Final = event.call_type in {CallTypes.call_mcp_tool, CallTypes.send_message, CallTypes.asend_message} realtime_events: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( captured.pricing.get("realtime_events") or () @@ -71,17 +86,22 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: realtime_usage: Final = event.call_type is CallTypes.arealtime and any( bool(item.get("usage")) for item in realtime_events ) - if captured.usage is None and not fixed_call and not realtime_usage: + if captured.usage is None and not fixed_call and not realtime_usage and not speech_usage: return Unpriced(usage_event_id=event.event_id, reason="usage_unknown") pricing: Final = captured.pricing mapped: Final = _fields(pricing.get("model_map_information")) rates: Final = _fields(mapped.get("model_map_value")) + if ( + event.call_type in {CallTypes.speech, CallTypes.aspeech} + and rates.get("input_cost_per_character") is not None + and not speech_usage + ): + return Unpriced(usage_event_id=event.event_id, reason="usage_unknown") model: Final = next( (d.value for d in event.dimensions if d.name is DimensionName.MODEL), event.call_type.value if fixed_call else None, ) provider: Final = next((d.value for d in event.dimensions if d.name is DimensionName.PROVIDER), None) - logged: Final = _fields(pricing.get("logging")) if ( fixed_call and event.outcome is not Outcome.SUCCEEDED @@ -180,8 +200,14 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: }, response_ms=logged.get("response_ms"), ) + speech_response: Final = HttpxBinaryResponseContent(Response(200, content=b"")) if speech_usage else None + if speech_response is not None: + set_hidden_param(speech_response, "billed_characters", speech_characters) + speech_response.usage = replay_usage response: Final = ( - LiteLLMRealtimeStreamLoggingObject.model_validate( + speech_response + if speech_response is not None + else LiteLLMRealtimeStreamLoggingObject.model_validate( { "results": pricing.get("realtime_events", []), "usage": captured.usage or {}, @@ -251,7 +277,15 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: logger.cost_billing_time = billed_at with retained_pricing(context), pinned_billing_time(billed_at): batch_costs: Final = ( - batch_cost_calculator( + ocr_batch_cost( + model=model, + custom_llm_provider=provider, + usage_info=OCRUsageInfo.model_validate(captured.usage), + model_info=TypeAdapter(ModelInfo).validate_python(info), + ) + if event.call_type in {CallTypes.retrieve_batch, CallTypes.aretrieve_batch} + and any(name in (captured.usage or {}) for name in OCRUsageInfo.model_fields) + else batch_cost_calculator( replay_usage, model, provider, diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index 5c735eaa6d9..4fede891018 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -173,7 +173,15 @@ def usage_writes( ended_at=ended, outcome=outcome, dimensions=dimensions, - metrics=_usage_metrics(None if context.usage_source == "unknown" else usage), + metrics=_usage_metrics( + None + if context.usage_source == "unknown" + or ( + context.ownership == "aggregate" + and payload["call_type"] in {CallTypes.retrieve_batch.value, CallTypes.aretrieve_batch.value} + ) + else usage + ), ) envelope: Final = { "event": outbound.model_dump(mode="json"), @@ -218,6 +226,7 @@ def usage_writes( _INCLUDED_IN_PARENT: Final = object() +_BATCH_ROWS: Final = object() _PRICING_FIELDS: Final = frozenset( @@ -306,6 +315,10 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res "expired", "incomplete", } + batch_summary: Final = ( + payload["call_type"] in {CallTypes.retrieve_batch.value, CallTypes.aretrieve_batch.value} + and kwargs.get("_accounting_batch_rows") is not _BATCH_ROWS + ) deferred: Final = billed_id.endswith(("_batch_cost", "_background_cost")) and terminal management: Final = ( payload["call_type"] in {CallTypes.create_batch.value, CallTypes.acreate_batch.value} @@ -320,6 +333,8 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res if management else f"{billed_id}:included" if included_in_parent + else f"{billed_id}:summary" + if batch_summary and terminal else billed_id ) if not attempt or attempt == "None" or (cached and not caller): @@ -390,12 +405,12 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res parent_request_id=parent, retry_of_attempt_id=_text(request_metadata.get("retry_of_attempt_id")), job_id=_text(raw_response.get("id")) if deferred or management else None, - charge_eligible=not cached and not included_in_parent and not management, + charge_eligible=not cached and not included_in_parent and not management and not batch_summary, request_count_eligible=bool(caller) and not origin and not retry and not deferred and not included_in_parent, ownership="included_in_parent" if included_in_parent else "aggregate" - if kwargs.get("additional_response_cost") is not None + if batch_summary or kwargs.get("additional_response_cost") is not None else "attempt", pricing=_json(captured_pricing), finalized_cost=_json( @@ -564,13 +579,26 @@ def _batch_rows(response: object) -> tuple[Mapping[str, JsonValue], ...]: def _batch_root(write: UsageWrite, rows: tuple[Mapping[str, JsonValue], ...]) -> UsageWrite: - event: Final = write.event.model_copy(update={"metrics": None}) + from litellm.proxy.openai_files_endpoints.common_utils import ( + convert_b64_uid_to_unified_uid, + get_batch_id_from_unified_batch_id, + get_original_file_id, + ) + envelope: Final = _JSON_OBJECT.validate_json(write.envelope) + observed_job: Final = TypeAdapter(str).validate_python(envelope["job_id"]) + decoded: Final = convert_b64_uid_to_unified_uid(observed_job) + job: Final = ( + get_batch_id_from_unified_batch_id(decoded) if decoded != observed_job else get_original_file_id(observed_job) + ) + identity: Final = f"{job}_batch_cost" + event: Final = write.event.model_copy(update={"metrics": None, "event_id": identity, "attempt_id": identity}) return _database_usage_write( event, { **envelope, "event": event.model_dump(mode="json"), + "job_id": job, "usage": { "batch_row_count": len(rows), "batch_rows_sha256": sha256(_json(rows).encode()).hexdigest(), @@ -644,7 +672,7 @@ def _prepare_batch_capture( from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload rows: Final = tuple(sorted(_batch_rows(response), key=_json)) - builder: Final = _CaptureBuilder(kwargs, response) + builder: Final = _CaptureBuilder({**kwargs, "_accounting_batch_rows": _BATCH_ROWS}, response) payload: Final = get_logging_payload( dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe ) @@ -846,7 +874,7 @@ def capture_logging_usage( if isinstance(existing_pricing, PricingCapture) and not internal and not _batch_rows(normalized) else _capture_pricing(logging_obj, normalized_result, details) ) - if freeze_only or _batch_rows(normalized): + if freeze_only or logging_obj.call_type in {CallTypes.retrieve_batch.value, CallTypes.aretrieve_batch.value}: if not internal: logging_obj.record_usage_pricing(pricing) return diff --git a/schema.prisma b/schema.prisma index 048e0e1e17a..45ea5cd20d0 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2113,6 +2113,7 @@ model LiteLLM_AccountingEvent { payload Json compilation_state String @default("pending") compilation_error String? + compilation_retry_at DateTime? @@index([recorded_at]) @@index([kind, source_event_id]) diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py index efff0d9b16a..8703559d6aa 100644 --- a/tests/integration/database/test_accounting_ledger.py +++ b/tests/integration/database/test_accounting_ledger.py @@ -358,7 +358,9 @@ async def test_compilation_rejects_changed_source_before_reusing_a_charge(writer @pytest.mark.asyncio -async def test_compilation_failure_rolls_back_and_can_be_retried(writer: _Writer, ledger_database: str) -> None: +async def test_compilation_failure_retains_retry_without_a_partial_charge( + writer: _Writer, ledger_database: str +) -> None: _, store = writer event: Final = _usage() assert isinstance(event, OutboundUsageEvent) @@ -367,8 +369,9 @@ async def test_compilation_failure_rolls_back_and_can_be_retried(writer: _Writer def failed_pricing(envelope: str) -> PricingSnapshot: raise RuntimeError("pricing temporarily unavailable") - with pytest.raises(RuntimeError, match="pricing temporarily unavailable"): - await LoggingCostCompiler(store, failed_pricing, _TARGETS).compile(event) + assert await LoggingCostCompiler(store, failed_pricing, _TARGETS).compile(event) == Unpriced( + usage_event_id=event.event_id, reason="calculation_retryable" + ) assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(0,)] result: Final = await LoggingCostCompiler(store, lambda _: _cost(amount="0").snapshot, _TARGETS).compile(event) assert isinstance(result, CostEvent) @@ -872,12 +875,19 @@ async def test_restarted_compiler_recovers_transient_failures_without_revisiting return _cost().snapshot await LoggingCostCompiler(store, transient, _TARGETS).compile_pending() - assert await store.pending_compilation() == (delayed,) + assert await store.pending_compilation() == () + assert _rows( + ledger_database, + "SELECT compilation_state, compilation_error FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'delayed'", + ) == [("pending", "calculation_retryable")] assert _rows(ledger_database, "SELECT source_event_id FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [ (healthy.event_id,) ] async with connected_proxy_prisma(ledger_database) as restarted: recovered: Final = AccountingLedger.for_client(restarted) + await restarted.db.execute_raw( + 'UPDATE "LiteLLM_AccountingEvent" SET compilation_retry_at = NULL WHERE event_id = $1', delayed.event_id + ) await LoggingCostCompiler(recovered, lambda _: _cost().snapshot, _TARGETS).compile_pending() assert await recovered.pending_compilation() == () await LoggingCostCompiler(recovered, transient, _TARGETS).compile_pending() @@ -982,3 +992,131 @@ async def test_proxy_capture_and_background_compiler_converge_on_final_header_ch finally: release.set() await runtime.stop(5) + + +@pytest.mark.asyncio +async def test_failing_compilation_page_does_not_starve_later_usage(writer: _Writer, ledger_database: str) -> None: + _, store = writer + broken: Final = _usage("broken-price") + healthy: Final = _usage("healthy-price") + assert isinstance(broken, OutboundUsageEvent) + assert isinstance(healthy, OutboundUsageEvent) + await _insert(store, broken) + await _insert(store, healthy) + + def price(envelope: str) -> PricingSnapshot: + if OutboundUsageEvent.model_validate(json.loads(envelope)["event"]).event_id == broken.event_id: + raise RuntimeError("temporary calculator failure") + return _cost().snapshot + + await LoggingCostCompiler(store, price, _TARGETS).compile_pending(limit=1) + await LoggingCostCompiler(store, price, _TARGETS).compile_pending(limit=1) + assert _rows(ledger_database, "SELECT source_event_id FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [ + (healthy.event_id,) + ] + assert isinstance(await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(broken), CostEvent) + assert await store.pending_compilation() == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("summary_first", (False, True)) +async def test_batch_callbacks_commit_only_row_charges_across_restart( + writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch, summary_first: bool +) -> None: + import litellm + from litellm.proxy.spend_tracking.usage_capture import ( + UsageBatchWriter, + capture_batch_callback, + capture_callback, + capture_logging_usage, + ) + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from tests.unit.proxy.spend_tracking.test_usage_capture import _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + _, store = writer + rates: Final = { + "input_cost_per_token_batches": 0.003, + "output_cost_per_token_batches": 0.004, + "litellm_provider": "openai", + "mode": "chat", + } + monkeypatch.setitem(litellm.model_cost, "capture-model", rates) + usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + response: Final = { + "id": "managed-job", + "status": "completed", + "usage": { + "batch_usage_by_model": {"capture-model": usage}, + "batch_usage_by_request": [{"model": "capture-model", "usage": usage} for _ in range(2)], + }, + } + logger: Final = _logger("aretrieve_batch") + batches: Final = UsageBatchWriter(store.insert_usage_batch, store.load_usage_batch) + meter: Final = DurableUsageMeter(batches.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) + compiler: Final = LoggingCostCompiler(store) + capture_logging_usage(meter, logger, response, _AT, _AT, freeze_only=True) + kwargs: Final = {**logger.model_call_details, "litellm_logging_obj": logger} + + async def summary() -> None: + capture_callback( + meter, + kwargs, + { + "id": "provider-job", + "status": "completed", + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + }, + _AT, + _AT, + ) + await meter.flush() + await compiler.compile_pending() + + try: + if summary_first: + await summary() + await capture_batch_callback(meter, store.load_usage_batch, kwargs, response, _AT, _AT) + await compiler.compile_pending() + if not summary_first: + await summary() + before: Final = _rows( + ledger_database, 'SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent" ORDER BY event_id' + ) + costs: Final = _rows( + ledger_database, "SELECT payload::text FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" + ) + deliveries: Final = _rows( + ledger_database, + 'SELECT event_id, target, state FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target', + ) + assert len(costs) == 2 + snapshots: Final = tuple( + CostEvent.model_validate_json(json.dumps(json.loads(str(row[0]))["event"])).snapshot for row in costs + ) + assert sum(float(snapshot.total_cost) for snapshot in snapshots) == pytest.approx(2 * (10 * 0.003 + 5 * 0.004)) + monkeypatch.setitem(litellm.model_cost, "capture-model", {**rates, "input_cost_per_token_batches": 99.0}) + async with connected_proxy_prisma(ledger_database) as restarted: + ledger: Final = AccountingLedger.for_client(restarted) + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + for job_id in ("managed-job", encode_file_id_with_model("managed-job", "capture-model", "batch")): + await capture_batch_callback( + meter, ledger.load_usage_batch, kwargs, {**response, "id": job_id}, _AT, _AT + ) + await asyncio.gather(*(LoggingCostCompiler(ledger).compile_pending() for _ in range(2))) + assert ( + _rows(ledger_database, 'SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent" ORDER BY event_id') + == before + ) + assert ( + _rows( + ledger_database, + 'SELECT event_id, target, state FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target', + ) + == deliveries + ) + for row, snapshot in zip(costs, snapshots, strict=True): + assert price_usage(json.dumps(json.loads(str(row[0]))["pricing_inputs"])) == snapshot + finally: + await meter.stop(1) diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py index 677dc2e4e08..1b89f472248 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_compiler.py +++ b/tests/unit/proxy/spend_tracking/test_cost_compiler.py @@ -1167,3 +1167,184 @@ def test_header_split_retains_token_rates_when_catalog_mode_is_unspecified(mode: assert float(result.total_cost) == pytest.approx(0.18) assert float(result.input_cost or 0) == pytest.approx(0.1) assert float(result.output_cost or 0) == pytest.approx(0.08) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("logging_first", (False, True)) +async def test_batch_logging_and_rows_have_one_charge_owner( + monkeypatch: pytest.MonkeyPatch, logging_first: bool +) -> None: + import litellm + from litellm.batches.batch_utils import batch_cost_and_usage + from litellm.proxy.spend_tracking.usage_capture import ( + UsageBatchWriter, + capture_batch_callback, + capture_callback, + capture_logging_usage, + ) + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.types.utils import LiteLLMBatch + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + rates: Final = { + "input_cost_per_token_batches": 0.003, + "output_cost_per_token_batches": 0.004, + "input_cost_per_token": 0.006, + "output_cost_per_token": 0.008, + "litellm_provider": "openai", + "mode": "chat", + } + monkeypatch.setitem(litellm.model_cost, "capture-model", rates) + billed: Final = batch_cost_and_usage( + [ + { + "response": { + "status_code": 200, + "body": { + "model": "capture-model", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + } + } + for _ in range(2) + ], + "openai", + ) + response: Final = LiteLLMBatch( + id="owner-job", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="file", + completion_window="24h", + status="completed", + created_at=1, + completed_at=2, + usage=billed.usage, + ) + database: Final = StoredUsage() + writer: Final = UsageBatchWriter(database.insert, database.load) + meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) + logger: Final = _logger("aretrieve_batch") + capture_logging_usage(meter, logger, response, _AT, _AT, freeze_only=True) + try: + for _ in range(2): + if logging_first: + capture_callback( + meter, + {**logger.model_call_details, "litellm_logging_obj": logger}, + { + **response.model_dump(), + "id": "provider-job", + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + }, + _AT, + _AT, + ) + capture_logging_usage(meter, logger, response, _AT, _AT) + await capture_batch_callback( + meter, database.load, {**logger.model_call_details, "litellm_logging_obj": logger}, response, _AT, _AT + ) + if not logging_first: + capture_callback( + meter, + {**logger.model_call_details, "litellm_logging_obj": logger}, + { + **response.model_dump(), + "id": "provider-job", + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + }, + _AT, + _AT, + ) + capture_logging_usage(meter, logger, response, _AT, _AT) + await meter.flush() + costs: Final = tuple( + price_usage(row.envelope) for row in database.rows if isinstance(row.event, OutboundUsageEvent) + ) + priced: Final = tuple(cost for cost in costs if isinstance(cost, PricingSnapshot)) + assert len(priced) == 2, costs + assert sum(float(cost.total_cost) for cost in priced) == pytest.approx(2 * (10 * 0.003 + 5 * 0.004)) + finally: + await meter.stop(1) + + +@pytest.mark.parametrize("characters", (0, 17)) +def test_speech_replay_prices_retained_characters(characters: int) -> None: + original: Final = _retained_envelope( + rates={"input_cost_per_character": 0.003, "mode": "audio_speech"}, logging={"prompt_characters": characters} + ) + event: Final = OutboundUsageEvent.model_validate(original["event"]).model_copy( + update={"call_type": CallTypes.aspeech} + ) + result: Final = price_usage(json.dumps({**original, "event": event.model_dump(mode="json")})) + assert isinstance(result, PricingSnapshot), result + assert float(result.total_cost) == pytest.approx(characters * 0.003) + + +def test_batch_ocr_replay_prices_pages_and_annotations() -> None: + original: Final = _retained_envelope( + rates={"ocr_cost_per_page_batches": 0.003, "annotation_cost_per_page_batches": 0.004, "mode": "ocr"}, + usage={"pages_processed": 3, "pages_processed_annotation": 2}, + ) + event: Final = OutboundUsageEvent.model_validate(original["event"]).model_copy( + update={"call_type": CallTypes.aretrieve_batch} + ) + result: Final = price_usage( + json.dumps( + { + **original, + "event": {**event.model_dump(mode="json"), "event_id": "batch:row:0"}, + "parent_request_id": "batch", + } + ) + ) + assert isinstance(result, PricingSnapshot), result + assert float(result.total_cost) == pytest.approx(3 * 0.003 + 2 * 0.004) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reported_characters", (None, 5)) +async def test_speech_capture_replays_billed_quantities_without_input_text(reported_characters: int | None) -> None: + from httpx import Response + + from litellm.litellm_core_utils.hidden_params import set_hidden_param + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.types.llms.openai import HttpxBinaryResponseContent + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + logger: Final = _logger("aspeech") + logger.set_request_pricing("capture-model", "openai", {"input_cost_per_character": 0.003, "mode": "audio_speech"}) + logger.model_call_details["input"] = "privatespeechinput" + response: Final = HttpxBinaryResponseContent(Response(200, content=b"audio")) + if reported_characters is not None: + set_hidden_param(response, "billed_characters", reported_characters) + expected: Final = (reported_characters if reported_characters is not None else len("privatespeechinput")) * 0.003 + assert logger.response_cost_calculator(response) == pytest.approx(expected) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, response, _AT, _AT) + await meter.flush() + envelope: Final = next(row.envelope for row in database.rows if isinstance(row.event, OutboundUsageEvent)) + assert "privatespeechinput" not in envelope + priced: Final = price_usage(envelope) + assert isinstance(priced, PricingSnapshot), priced + assert float(priced.total_cost) == pytest.approx(expected) + finally: + await meter.stop(1) + + +@pytest.mark.parametrize("characters", (None, -1, True)) +def test_character_pricing_does_not_invent_zero_for_missing_or_invalid_quantity(characters: object) -> None: + original: Final = _retained_envelope( + rates={"input_cost_per_character": 0.003}, logging={"prompt_characters": characters} + ) + event: Final = OutboundUsageEvent.model_validate(original["event"]).model_copy( + update={"call_type": CallTypes.aspeech} + ) + assert price_usage(json.dumps({**original, "event": event.model_dump(mode="json")})) == Unpriced( + usage_event_id=event.event_id, reason="usage_unknown" + ) diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py index 1ef6e6d2d4b..c83d11250c9 100644 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ b/tests/unit/proxy/spend_tracking/test_usage_capture.py @@ -334,7 +334,7 @@ async def test_competing_writers_reuse_only_terminal_pricing_snapshots(terminal: @pytest.mark.asyncio -async def test_batch_poll_has_distinct_identity_and_only_terminal_usage_is_chargeable() -> None: +async def test_batch_poll_and_terminal_summary_are_not_charge_owners() -> None: from litellm.proxy.spend_tracking.usage_capture import capture_callback database: Final = StoredUsage() @@ -352,8 +352,8 @@ async def test_batch_poll_has_distinct_identity_and_only_terminal_usage_is_charg await meter.close() assert pending.writes[-1].event.event_id == "poll:status" assert json.loads(pending.writes[-1].envelope)["charge_eligible"] is False - assert completed.writes[-1].event.event_id == "batch_batch_cost" - assert json.loads(completed.writes[-1].envelope)["charge_eligible"] is True + assert completed.writes[-1].event.event_id == "batch_batch_cost:summary" + assert json.loads(completed.writes[-1].envelope)["charge_eligible"] is False assert len(completed.writes) == 1 @@ -1021,9 +1021,8 @@ async def test_batch_capture_provenance_tracks_observations_not_envelope_presenc envelope: Final = json.loads(database.rows[-1].envelope) assert envelope["usage_source"] == ("unknown" if measured is None else "measured") assert envelope["usage"]["batch_usage_by_model"] == usage["batch_usage_by_model"] - assert database.rows[-1].event.metrics == ( - None if measured is None else (Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(measured["total_tokens"])),) - ) + assert database.rows[-1].event.metrics is None + assert envelope["charge_eligible"] is False @pytest.mark.asyncio