fix(accounting): preserve prepared usage observation on failure

This commit is contained in:
Joshua Valluru 2026-10-29 23:26:36 -07:00
parent 33cf5c172f
commit c8bbb1e2d9
2 changed files with 57 additions and 0 deletions

View file

@ -1744,6 +1744,7 @@ def _failure_fields_to_lift(
if _prompt_estimate is not None
else ()
)
_recorded_observation: Final[object] = _model_call_details.get("_accounting_usage_observed")
_entries: Final = (
("start_time", _model_call_details.get("start_time")),
("first_api_call_start_time", _first_handoff),
@ -1764,6 +1765,8 @@ def _failure_fields_to_lift(
(
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

View file

@ -4158,3 +4158,57 @@ def test_is_client_error_exception_existing_behaviour_unchanged() -> None:
ProxyException(message="x", type="x", param=None, code="429")
) is True
assert _is_client_error_exception(ValueError("x")) is False
@pytest.mark.asyncio
@pytest.mark.parametrize("tokens", [None, 0, 3])
async def test_failure_capture_preserves_prepared_usage_observation(tokens: int | None) -> None:
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
timestamp: Final = real_datetime.datetime(2026, 1, 1)
committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
committed.put_nowait(write)
logger: Final = Logging(
model=FAILURE_USAGE_MODEL,
messages=[],
stream=False,
call_type="acompletion",
start_time=timestamp,
litellm_call_id="post-response-failure",
function_id="post-response-failure",
)
logger.update_environment_variables(
model=FAILURE_USAGE_MODEL, custom_llm_provider="openai", litellm_params={}, optional_params={}
)
logger.model_call_details["response_cost"] = 0.2
runtime: Final = UsageCaptureRuntime()
runtime.start(persist, UsageCaptureSettings())
try:
prepared: Final = logger.prepare_response_cost(
{"id": "provider", "usage": None if tokens is None else {"prompt_tokens": tokens, "total_tokens": tokens}},
timestamp,
)
assert prepared is not None
lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True)
assert lifted["combined_usage_object"] is logger.model_call_details["combined_usage_object"]
assert lifted["response_cost"] == 0.2
assert runtime.meter is not None
capture_callback(
runtime.meter,
{**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"},
None,
timestamp,
timestamp,
)
finally:
await runtime.stop(1)
committed.get_nowait()
envelope: Final = json.loads(committed.get_nowait().envelope)
assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == tokens
assert envelope["usage_source"] == ("unknown" if tokens is None else "measured")