diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index cf0f98bb46d..943d82148c1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -271,6 +271,7 @@ if TYPE_CHECKING: from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation from litellm.proxy.spend_tracking.usage_capture import CapturedUsage + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -642,9 +643,34 @@ class UsageObservation(Enum): ESTIMATED = "estimated" +@dataclass(frozen=True, slots=True) +class UsageCaptureObservation: + usage_json: str | None + provenance: UsageObservation + response_cost: float | None = None + + @classmethod + def from_usage( + cls, usage: Usage | None, provenance: UsageObservation, response_cost: float | None = None + ) -> "UsageCaptureObservation | UsageCaptureError": + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError + + try: + serialized: Final = ( + None if usage is None or provenance is UsageObservation.UNKNOWN else usage.model_dump_json() + ) + except Exception as error: # noqa: BLE001 # observation failures must not interrupt legacy retry handling + return UsageCaptureError(str(error)) + else: + return cls(serialized, provenance, response_cost) + + class Logging(LiteLLMLoggingBaseClass): usage_capture_callback: ClassVar[ - Callable[["Logging", object, datetime.datetime, datetime.datetime, bool], None] | None + Callable[ + ["Logging", object, datetime.datetime, datetime.datetime, bool], "CapturedUsage | UsageCaptureError | None" + ] + | None ] = None def capture_usage_before_logging( @@ -2592,7 +2618,9 @@ class Logging(LiteLLMLoggingBaseClass): }, usage, ) - self.record_partial_usage_for_failure(usage, raw_cost) + self.record_partial_usage_for_failure( + usage, raw_cost, observation=self.model_call_details["_accounting_usage_observed"] + ) return self.prepared_cost def record_finalized_cost(self, finalized: FinalizedCost) -> None: @@ -2605,10 +2633,14 @@ class Logging(LiteLLMLoggingBaseClass): def record_usage_pricing(self, pricing: object) -> None: self.model_call_details["_accounting_pricing_snapshot"] = pricing - def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None: + def record_partial_usage_for_failure( + self, usage: Usage, response_cost: float, *, observation: UsageObservation = UsageObservation.KNOWN + ) -> None: """Stash what an interrupted stream already consumed so the failure log bills it instead of zero.""" self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + if type(self).usage_capture_callback is not None: + self.model_call_details["_accounting_usage_observed"] = observation def record_failure_after_earlier_rounds( self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float @@ -2704,38 +2736,50 @@ class Logging(LiteLLMLoggingBaseClass): def capture_failed_attempt_usage(self, usage: Usage) -> None: from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY + from litellm.proxy.spend_tracking.usage_capture import CapturedUsage capture: Final = type(self).usage_capture_callback if capture is None: return attempt: Final = self.attempt_scoped_copy() + attempt.prepared_cost = None + attempt.finalized_cost = None identity: Final = ( f"{self.litellm_call_id}-attempt-{len(self.model_call_details.get('failed_attempts_usage', ()))}" ) params: Final = _LITELLM_PARAMS.validate_python(self.model_call_details.get("litellm_params") or {}) metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(params) + observed: Final[object] = self.model_call_details.get("_accounting_usage_observed") + cost: Final[object] = self.model_call_details.get("response_cost") attempt.model_call_details = { **attempt.model_call_details, "litellm_call_id": identity, "parent_request_id": self.litellm_call_id, "_accounting_capture": None, + "_accounting_prior_captures": (), + "_accounting_usage_snapshot": UsageCaptureObservation.from_usage( + usage, + observed if isinstance(observed, UsageObservation) else UsageObservation.KNOWN, + float(cost) if isinstance(cost, (int, float)) else None, + ), "litellm_params": { **params, "metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, "litellm_metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, }, } - capture( + captured: Final = capture( attempt, {"id": identity, "usage": usage.model_dump(), "status": "failed"}, self.start_time, datetime.datetime.now(datetime.timezone.utc), False, ) - self.model_call_details["_accounting_prior_attempt_ids"] = ( - *self.model_call_details.get("_accounting_prior_attempt_ids", ()), - identity, - ) + if isinstance(captured, CapturedUsage): + self.model_call_details["_accounting_prior_captures"] = ( + *self.model_call_details.get("_accounting_prior_captures", ()), + captured, + ) def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 03dbe148de0..1b9d2a3d20d 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -22,7 +22,11 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, ) -from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup, UsageObservation +from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + UsageCaptureObservation, + UsageObservation, +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping @@ -346,6 +350,10 @@ def _attempt_spend_row_kwargs( "combined_usage_object": attempt.usage, "response_cost": attempt.response_cost, "_accounting_usage_observed": UsageObservation.ESTIMATED, + "_accounting_usage_snapshot": UsageCaptureObservation.from_usage( + attempt.usage, UsageObservation.ESTIMATED, attempt.response_cost + ), + "_accounting_prior_captures": (), } row.pop("billed_failed_attempts", None) row.pop("_accounting_capture", None) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1b736f52330..93205a2acf1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -472,7 +472,6 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l log_llm_api_exception, open_sse_before_first_byte, proxy_exception_from_route_error, - request_litellm_call_id, resolve_litellm_call_id, should_return_raw_model_name, sse_error_payload, @@ -13034,8 +13033,10 @@ async def completion( """ global user_temperature, user_request_timeout, user_max_tokens, user_api_base data = {} + base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line + base_llm_response_processor.data = data if user_api_key_dict is not None: if data.get("metadata") is None: data["metadata"] = {} @@ -13049,7 +13050,6 @@ async def completion( data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None: data["metadata"]["agent_id"] = user_api_key_dict.agent_id - base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) return await base_llm_response_processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, @@ -13151,9 +13151,12 @@ async def completion( return _response except Exception as e: await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=base_llm_response_processor.data, + dispatched_call_type=base_llm_response_processor.route_type, ) - litellm_call_id: Final = request_litellm_call_id(data) + litellm_call_id: Final = base_llm_response_processor.litellm_call_id log_llm_api_exception(e, litellm_call_id) if isinstance(e, ProxyException): raise with_litellm_call_id(e, litellm_call_id) diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index fb6414c38eb..795bb131029 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -6,6 +6,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, replace from datetime import datetime, timezone from decimal import Decimal +from itertools import chain from typing import TYPE_CHECKING, Final from pydantic import BaseModel, JsonValue, TypeAdapter @@ -129,6 +130,7 @@ class UsageCaptureContext: provider_timestamps: str = "{}" usage_source: str = "measured" included_attempt_ids: tuple[str, ...] = () + related_attempt_ids: tuple[str, ...] = () def _database_usage_write(event: InboundUsageEvent | OutboundUsageEvent, envelope: Mapping[str, object]) -> UsageWrite: @@ -175,6 +177,7 @@ def usage_writes( "usage": usage, "usage_source": "unknown" if usage is None else context.usage_source, "included_attempt_ids": context.included_attempt_ids, + "related_attempt_ids": context.related_attempt_ids, "pricing": _JSON_OBJECT.validate_json(context.pricing), "charge_eligible": context.charge_eligible, "request_count_eligible": context.request_count_eligible, @@ -267,7 +270,12 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.internal_call_metadata import is_background_response, is_unbilled_non_inference_call - from litellm.litellm_core_utils.litellm_logging import Logging, UsageObservation, mask_api_base_credentials + from litellm.litellm_core_utils.litellm_logging import ( + Logging, + UsageCaptureObservation, + UsageObservation, + mask_api_base_credentials, + ) from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups @@ -324,6 +332,7 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res ) allocation: Final = allocation_value if isinstance(allocation_value, SearchAllowanceAllocation) else None pricing_snapshot: Final = kwargs.get("_accounting_pricing_snapshot") + observation: Final = kwargs.get("_accounting_usage_snapshot") captured_pricing: Final = { "model_map_information": ( _JSON_OBJECT.validate_json(pricing_snapshot.model_map_information) @@ -353,7 +362,9 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res "service_tier": raw_response.get("service_tier"), "region_name": hidden.get("region_name"), }, - "reported_cost": standard.get( + "reported_cost": observation.response_cost + if isinstance(observation, UsageCaptureObservation) + else standard.get( "response_cost", kwargs.get("response_cost", _fields(get_hidden_params(response)).get("response_cost")) ), "adjustments": ( @@ -413,10 +424,9 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res if kwargs.get("_accounting_usage_observed") is UsageObservation.ESTIMATED else "measured", included_attempt_ids=TypeAdapter(tuple[str, ...]).validate_python( - kwargs.get( - "_accounting_included_attempt_ids", () if retry else kwargs.get("_accounting_prior_attempt_ids", ()) - ) + kwargs.get("_accounting_included_attempt_ids", ()) ), + related_attempt_ids=tuple(write.event.event_id for write in _prior_writes(kwargs)), provider_timestamps=_json( { key: raw_response.get(key) @@ -463,7 +473,13 @@ class _CaptureBuilder: self.writes: tuple[UsageWrite, ...] = () def observe(self, payload: SpendLogsPayload, usage: str | None) -> None: - parsed: Final = _JSON_OBJECT.validate_json(usage) if usage is not None else {} + from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation, UsageObservation + + snapshot: Final = self.kwargs.get("_accounting_usage_snapshot") + if isinstance(snapshot, UsageCaptureError): + raise snapshot + effective: Final = snapshot.usage_json if isinstance(snapshot, UsageCaptureObservation) else usage + parsed: Final = _JSON_OBJECT.validate_json(effective) if effective is not None else {} batch: Final = _fields(parsed.get("batch_usage_by_model")) unknown: Final = (bool(batch) or bool(parsed.get("batch_unknown_usage_requests"))) and all( value is None for value in batch.values() @@ -476,10 +492,32 @@ class _CaptureBuilder: } ) if unknown - else usage + else effective ) context: Final = capture_context(payload, self.kwargs, self.response) - self.writes = usage_writes(payload, observed, replace(context, usage_source="unknown") if unknown else context) + provenance: Final = ( + "unknown" + if unknown + else "estimated" + if isinstance(snapshot, UsageCaptureObservation) and snapshot.provenance is UsageObservation.ESTIMATED + else context.usage_source + ) + self.writes = usage_writes(payload, observed, replace(context, usage_source=provenance)) + + +def _prior_writes(kwargs: Mapping[str, object]) -> tuple[UsageWrite, ...]: + prior: Final = kwargs.get("_accounting_prior_captures", ()) + return ( + tuple( + chain.from_iterable( + item.writes + for item in TypeAdapter(tuple[object, ...]).validate_python(prior) + if isinstance(item, CapturedUsage) + ) + ) + if isinstance(prior, tuple) + else () + ) def capture_callback( @@ -499,7 +537,8 @@ def capture_callback( payload: Final = get_logging_payload( dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe ) - return CapturedUsage(payload, builder.writes, meter.submit(builder.writes)) + dependencies: Final = _prior_writes(kwargs) + return CapturedUsage(payload, builder.writes, meter.submit((*dependencies, *builder.writes))) def captured_callback_kwargs( @@ -559,7 +598,7 @@ def capture_logging_usage( start_time: datetime, end_time: datetime, freeze_only: bool = False, -) -> None: +) -> CapturedUsage | UsageCaptureError | None: from litellm._internal_context import internal_parent_request_id, is_internal_call from litellm.batches.batch_utils import batch_cost_is_final from litellm.litellm_core_utils.litellm_logging import Logging @@ -572,10 +611,7 @@ def capture_logging_usage( return try: if not internal and isinstance(logging_obj.model_call_details.get("_accounting_capture"), CapturedUsage): - logging_obj.record_usage_capture( - capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time) - ) - return + return capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time) normalized_result: Final = logging_obj.usage_capture_response(response) prepared: Final = logging_obj.prepared_cost normalized: Final = ( @@ -616,8 +652,11 @@ def capture_logging_usage( ) if not internal: logging_obj.record_usage_capture(captured) - except Exception: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging + except Exception as error: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging verbose_proxy_logger.exception("Usage capture failed at the execution boundary") + return UsageCaptureError(str(error)) + else: + return captured @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c3150c6f467..a69d5c70196 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1744,7 +1744,39 @@ def _failure_fields_to_lift( if _prompt_estimate is not None else () ) + from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation + _recorded_observation: Final[object] = _model_call_details.get("_accounting_usage_observed") + _current_usage: Final = ( + _current_failure_usage_to_lift( + _model_call_details, + request_data, + _model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None, + _billed, + _deployment_id, + _base_model, + dispatched_call_type, + ) + if Logging.usage_capture_callback is not None + and (_billed or _model_call_details.get("combined_usage_object") is not None) + else None + ) + _current_observation: Final = ( + _recorded_observation + if isinstance(_recorded_observation, UsageObservation) + else UsageObservation.UNKNOWN + if _current_usage is None + else UsageObservation.KNOWN + if _model_call_details.get("combined_usage_object") is not None + else UsageObservation.ESTIMATED + ) + _capture_observation: Final = UsageCaptureObservation.from_usage( + _current_usage[0] if _current_usage is not None and isinstance(_current_usage[0], Usage) else None, + _current_observation, + float(_current_usage[1]) + if _current_usage is not None and isinstance(_current_usage[1], (int, float)) + else None, + ) _entries: Final = ( ("start_time", _model_call_details.get("start_time")), ("first_api_call_start_time", _first_handoff), @@ -1760,22 +1792,9 @@ def _failure_fields_to_lift( "_accounting_allocation", _logging_obj.cost_allowance_allocation if isinstance(_logging_obj, Logging) else None, ), - ( - "_accounting_usage_observed", - ( - None - if Logging.usage_capture_callback is None - else _recorded_observation - if isinstance(_recorded_observation, UsageObservation) - else UsageObservation.UNKNOWN - if _usage_to_lift is None - else UsageObservation.KNOWN - if _model_call_details.get("combined_usage_object") is not None - or _model_call_details.get("failed_attempts_usage") - else UsageObservation.ESTIMATED - ), - ), - ("_accounting_included_attempt_ids", _model_call_details.get("_accounting_prior_attempt_ids")), + ("_accounting_usage_observed", _current_observation if Logging.usage_capture_callback is not None else None), + ("_accounting_usage_snapshot", _capture_observation if Logging.usage_capture_callback is not None else None), + ("_accounting_prior_captures", _model_call_details.get("_accounting_prior_captures")), ) return MappingProxyType({key: value for key, value in _entries if value is not None}) diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py index 46025bdb1cd..85ed9046e8e 100644 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ b/tests/unit/proxy/spend_tracking/test_usage_capture.py @@ -133,7 +133,10 @@ class StoredUsage: existing: Final = {row.event.event_id: row for row in self.rows} if any(existing.get(row.event.event_id, row) != row for row in writes): raise AccountingConflictError("Same identity has different content") - self.rows += tuple(row for row in writes if row.event.event_id not in existing) + unique: Final = {row.event.event_id: row for row in writes} + if any(unique[row.event.event_id] != row for row in writes): + raise AccountingConflictError("Submission contains conflicting content") + self.rows += tuple(row for identity, row in unique.items() if identity not in existing) @pytest.mark.asyncio @@ -537,7 +540,8 @@ async def test_failed_attempt_keeps_its_usage_without_claiming_the_callers_succe assert failed["pricing"]["reported_cost"] == 0.4 assert all(json.loads(row.envelope)["included_attempt_ids"] == [] for row in database.rows[:2]) successful: Final = json.loads(database.rows[-1].envelope) - assert successful["included_attempt_ids"] == ["caller-attempt-0", "caller-attempt-1"] + assert successful["related_attempt_ids"] == ["caller-attempt-0", "caller-attempt-1"] + assert successful["included_attempt_ids"] == [] assert successful["usage"]["total_tokens"] == 3 assert successful["charge_eligible"] is True assert database.rows[2].event.outcome is Outcome.SUCCEEDED @@ -1146,3 +1150,155 @@ async def test_prepared_native_google_capture_preserves_reported_usage( ) == tokens assert envelope["usage_source"] == ("unknown" if tokens is None else "measured") assert response == native_response.json() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rejected", [False, True]) +async def test_retry_links_only_admitted_usage_and_preserves_later_measurement(rejected: bool) -> None: + from litellm.litellm_core_utils.litellm_logging import UsageObservation + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings(max_pending_records=4)) + logger: Final = _logger() + assert runtime.meter is not None + try: + if rejected: + runtime.meter.submit(tuple(_write(f"occupied-{index}") for index in range(4))) + logger.model_call_details["_accounting_usage_observed"] = UsageObservation.UNKNOWN + logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) + logger.carry_failed_attempt_usage() + await runtime.meter.flush() + logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) + finally: + await runtime.stop(1) + successful: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "success")) + assert successful.get("related_attempt_ids", successful["included_attempt_ids"]) == ( + [] if rejected else ["caller-attempt-0"] + ) + assert successful["usage"]["prompt_tokens"] == 3 + if not rejected: + failed: Final = json.loads( + next(row.envelope for row in database.rows if row.event.event_id == "caller-attempt-0") + ) + assert failed["usage"]["prompt_tokens"] == 5 + assert failed["usage_source"] == "measured" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("current_tokens", [None, 0, 3]) +async def test_final_failure_captures_current_attempt_without_legacy_retry_aggregate( + current_tokens: int | None, +) -> None: + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback + from litellm.proxy.utils import _failure_fields_to_lift + from litellm.types.utils import Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + try: + logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) + logger.carry_failed_attempt_usage() + if current_tokens is not None: + logger.record_partial_usage_for_failure( + Usage(prompt_tokens=current_tokens, total_tokens=current_tokens), 0.2 + ) + lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) + assert lifted["combined_usage_object"].total_tokens == 5 + (current_tokens or 0) + assert runtime.meter is not None + capture_callback( + runtime.meter, + {**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"}, + None, + NOW, + NOW, + ) + finally: + await runtime.stop(1) + final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) + assert (None if final["usage"] is None else final["usage"]["total_tokens"]) == current_tokens + assert final["usage_source"] == ("unknown" if current_tokens is None else "measured") + assert final["pricing"]["reported_cost"] == (None if current_tokens is None else 0.2) + assert final["included_attempt_ids"] == [] + assert final["related_attempt_ids"] == ["caller-attempt-0"] + + +@pytest.mark.asyncio +async def test_estimated_final_attempt_keeps_its_provenance_after_measured_retry() -> None: + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback + from litellm.proxy.utils import _failure_fields_to_lift + from litellm.types.utils import Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + logger.messages = [{"role": "user", "content": "hello"}] + logger.model_call_details["messages"] = logger.messages + try: + logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) + logger.carry_failed_attempt_usage() + logger.model_call_details["current_attempt_api_call_start_time"] = NOW + lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) + assert runtime.meter is not None + capture_callback(runtime.meter, {**logger.model_call_details, **lifted}, None, NOW, NOW) + finally: + await runtime.stop(1) + final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) + assert final["usage_source"] == "estimated" + assert final["usage"]["total_tokens"] > 0 + assert lifted["combined_usage_object"].total_tokens == final["usage"]["total_tokens"] + 5 + assert final["pricing"]["reported_cost"] == 0 + + +@pytest.mark.asyncio +async def test_parent_receipt_cannot_acknowledge_when_a_linked_attempt_conflicts() -> None: + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + try: + logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) + logger.carry_failed_attempt_usage() + assert runtime.meter is not None + await runtime.meter.flush() + other: Final = _logger() + other.record_partial_usage_for_failure(Usage(prompt_tokens=8), 0.8) + other.carry_failed_attempt_usage() + other.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) + captured: Final = other.model_call_details["_accounting_capture"] + await runtime.meter.flush() + with pytest.raises(UsageCaptureError, match="identity"): + await captured.receipt + finally: + await runtime.stop(1) + assert tuple(row.event.event_id for row in database.rows) == ("caller-attempt-0",) + + +@pytest.mark.asyncio +async def test_unserializable_observation_does_not_interrupt_retry_or_legacy_usage() -> None: + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + usage: Final = Usage(prompt_tokens=5, total_tokens=5, opaque_usage=object()) + try: + logger.record_partial_usage_for_failure(usage, 0.4) + assert logger.model_call_details["combined_usage_object"] is usage + logger.carry_failed_attempt_usage() + assert logger.model_call_details["failed_attempts_usage"] == ((usage, 0.4),) + logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) + finally: + await runtime.stop(1) + assert tuple(row.event.event_id for row in database.rows) == ("caller:inbound", "success") + assert json.loads(database.rows[-1].envelope)["related_attempt_ids"] == [] diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 47032e30e62..b0087f66837 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -4465,3 +4465,79 @@ async def test_sync_ack_preserves_async_request_support( else: assert "hi" in response.text assert not persisted.empty(), "Terminal usage must still be captured" + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_ack", [False, True]) +async def test_text_completion_provider_rejection_captures_unknown_usage( + sync_ack: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + from fastapi import Request + + from litellm.accounting import Outcome + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.db.accounting_ledger import UsageWrite + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings + + persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() + + async def persist(writes: tuple[UsageWrite, ...]) -> None: + for write in writes: + persisted.put_nowait(write) + + async def authorize(request: Request) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path) + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key") + respx_mock.post().respond( + 404, + json={ + "error": {"message": "The model is unavailable", "type": "invalid_request_error", "code": "model_not_found"} + }, + ) + runtime: Final = UsageCaptureRuntime() + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": sync_ack}) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "instruct", + "litellm_params": {"model": "openai/gpt-3.5-turbo-instruct", "api_key": "test-provider-key"}, + } + ], + num_retries=0, + ), + ) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache())) + monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()]) + monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize) + runtime.start(persist, UsageCaptureSettings()) + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test" + ) as client: + response: Final = await client.post( + "/v1/completions", json={"model": "instruct", "prompt": "hi", "max_tokens": 1} + ) + finally: + await runtime.stop(1) + assert response.status_code == 404, response.text + assert len(respx_mock.calls) == 1 + assert persisted.qsize() == 2 + inbound: Final = persisted.get_nowait() + outbound: Final = persisted.get_nowait() + assert outbound.event.outcome is Outcome.FAILED + assert outbound.event.call_type.value == "atext_completion" + assert outbound.event.metrics is None + assert inbound.event.request_id == outbound.event.request_id + envelope: Final = json.loads(outbound.envelope) + assert envelope["usage"] is None + assert envelope["usage_source"] == "unknown" + assert envelope["pricing"]["reported_cost"] is None