mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(accounting): preserve prepared usage observation on failure
This commit is contained in:
parent
33cf5c172f
commit
c8bbb1e2d9
2 changed files with 57 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue