mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(accounting): enforce batch charge ownership and recover retained pricing
This commit is contained in:
Joshua Valluru
parent
1025640bdd
commit
ed4d3ae6cc
12 changed files with 427 additions and 30 deletions
|
|
@ -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");
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue