fix(proxy): price the estimated input tokens on a read-timed-out request's spend row (#47435)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Backport of #47435 to stable/1.105.x.
Cherry-picked from 7b76b2b15f (main).
This commit is contained in:
devin-ai-integration[bot] 2026-10-21 00:00:05 +00:00
parent da2116999b
commit be578726ed
2 changed files with 154 additions and 13 deletions

View file

@ -46,6 +46,7 @@ from typing import (
overload,
)
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm import _custom_logger_compatible_callbacks_literal
@ -939,12 +940,17 @@ def _failure_usage_to_lift(
model_call_details: Mapping[str, object],
request_body: Mapping[str, object],
dispatched: bool,
provider_billed: bool = False,
) -> tuple[object, object] | None:
"""A stream that broke mid-flight still billed the provider for the chunks
already delivered; the streaming handler stashes that recovered usage and
cost in model_call_details, so prefer it. Otherwise a request that was
dispatched to a provider and failed without upstream usage gets an
estimated input-side Usage with zero cost. The raw request body backfills
estimated input-side Usage. When the failure is one the provider bills
anyway (``provider_billed``: a read timeout, where the provider had the
prompt and kept working after LiteLLM hung up) that
usage is priced at the model's input rate; any other post-dispatch failure
stays at zero cost. The raw request body backfills
the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a
chat-completions provider) without filling optional_params. Returns the
(combined_usage_object, response_cost) pair to lift, or None."""
@ -969,7 +975,27 @@ def _failure_usage_to_lift(
)
if estimated_usage is None:
return None
return estimated_usage, 0.0
if not provider_billed:
return estimated_usage, 0.0
return estimated_usage, _estimated_input_cost(model_call_details, estimated_usage.prompt_tokens)
def _estimated_input_cost(model_call_details: Mapping[str, object], prompt_tokens: int) -> float:
"""Input price for the estimated prompt at the cost map's rate; an unmapped model keeps the
row at $0 rather than dropping the failure row."""
try:
model: Final = str(model_call_details.get("model") or "")
estimated_input_cost, _ = litellm.cost_per_token(
model=model,
custom_llm_provider=str(model_call_details.get("custom_llm_provider") or "") or None,
prompt_tokens=prompt_tokens,
)
except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written
verbose_proxy_logger.debug(
"failure spend estimate: could not price %s: %s", model_call_details.get("model"), exc
)
return 0.0
return estimated_input_cost
_EMPTY_LIFT: Final = MappingProxyType({})
@ -1083,7 +1109,28 @@ def _call_type_for_route(route: str | None) -> str | None:
_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
_EXCEPTION_CHAIN_MAX_DEPTH: Final = 16
def _exception_chain(exception: BaseException) -> tuple[BaseException, ...]:
chain: Final[list[BaseException]] = [exception] # mutable-ok: built once, returned frozen
for _ in range(_EXCEPTION_CHAIN_MAX_DEPTH):
following = chain[-1].__cause__ or chain[-1].__context__
if following is None or following in chain:
break
chain.append(following)
return tuple(chain)
def _timed_out_reading_from_provider(exception: Exception) -> bool:
"""A read timeout means the request was sent and the provider was working on it; a connect
timeout (``httpx.ConnectTimeout``) never reached the provider, so it is not billed."""
return isinstance(exception, litellm.Timeout) and any(
isinstance(cause, httpx.ReadTimeout) for cause in _exception_chain(exception)
)
def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: bool = False) -> Mapping[str, object]:
"""Failure-path callbacks run after ``litellm_logging_obj`` is popped from
request_data (it is not serialisable), so the caller merges these fields
onto request_data first: the request start and first-handoff instants for
@ -1099,6 +1146,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str,
model_call_details=_model_call_details,
request_body=request_data,
dispatched=_first_handoff is not None,
provider_billed=provider_billed,
)
_entries: Final = (
("start_time", _model_call_details.get("start_time")),
@ -3269,7 +3317,11 @@ class ProxyLogging:
original_exception=original_exception,
)
request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data))
request_data.update(
await offload_token_count(_failure_fields_to_lift)(
request_data, provider_billed=_timed_out_reading_from_provider(original_exception)
)
)
# Remove before callbacks iterate — not serialisable
request_data.pop("litellm_logging_obj", None)

View file

@ -2,6 +2,7 @@ import datetime as real_datetime
import smtplib
from typing import Final
import httpx
import pytest
from fastapi import HTTPException
@ -687,11 +688,13 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
"""A non-stream request that failed after dispatch (timeout, provider
error) consumed provider-billed input tokens but recovered no usage.
post_call_failure_hook must estimate the input side onto request_data so
the spend log's failure row records what was sent instead of zero, while
never charging spend for the failure (LIT-5690).
the spend log's failure row records what was sent instead of zero
(LIT-5690). On a timeout that estimate is priced at the model's input
rate since the provider bills the prompt it accepted (LIT-145); other
post-dispatch failures keep zero spend.
"""
async def _run(self, request_data):
async def _run(self, request_data, original_exception: Exception | None = None):
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import UserAPIKeyAuth
@ -701,17 +704,81 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()):
await proxy_logging_obj.post_call_failure_hook(
request_data=request_data,
original_exception=Exception("boom"),
original_exception=original_exception or Exception("boom"),
user_api_key_dict=UserAPIKeyAuth(),
)
def _timeout(self, model: str, cause: Exception) -> Exception:
try:
try:
raise cause
except Exception as inner:
raise litellm.Timeout("Request timed out.", model=model, llm_provider="openai") from inner
except litellm.Timeout as timed_out:
return timed_out
def _logging_obj(self, model_call_details):
logging_obj = MagicMock()
logging_obj.model_call_details = model_call_details
return logging_obj
@pytest.mark.asyncio
async def test_dispatched_failure_estimates_input_tokens_with_zero_cost(self):
async def test_timed_out_request_estimates_input_tokens_priced_at_input_rate(self):
from datetime import datetime
from litellm.types.utils import Usage
request_data = {
"litellm_logging_obj": self._logging_obj(
{
"first_api_call_start_time": datetime.now(),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
}
),
"metadata": {},
"response_cost": 123.0,
}
await self._run(
request_data, original_exception=self._timeout("gpt-3.5-turbo", httpx.ReadTimeout("no bytes in time"))
)
estimated = request_data["combined_usage_object"]
assert isinstance(estimated, Usage)
assert estimated.prompt_tokens > 0
assert estimated.completion_tokens == 0
assert estimated.total_tokens == estimated.prompt_tokens
expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens)
assert expected_input_cost > 0
assert request_data["response_cost"] == pytest.approx(expected_input_cost)
@pytest.mark.asyncio
async def test_connect_timeout_never_reached_the_provider_stays_zero_cost(self):
from datetime import datetime
from litellm.types.utils import Usage
request_data = {
"litellm_logging_obj": self._logging_obj(
{
"first_api_call_start_time": datetime.now(),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
}
),
"metadata": {},
}
await self._run(
request_data, original_exception=self._timeout("gpt-3.5-turbo", httpx.ConnectTimeout("no route"))
)
assert isinstance(request_data["combined_usage_object"], Usage)
assert request_data["response_cost"] == 0.0
@pytest.mark.asyncio
async def test_dispatched_non_timeout_failure_estimates_input_tokens_with_zero_cost(self):
from datetime import datetime
from litellm.types.utils import Usage
@ -733,8 +800,6 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
estimated = request_data["combined_usage_object"]
assert isinstance(estimated, Usage)
assert estimated.prompt_tokens > 0
assert estimated.completion_tokens == 0
assert estimated.total_tokens == estimated.prompt_tokens
assert request_data["response_cost"] == 0.0
@pytest.mark.asyncio
@ -1901,7 +1966,28 @@ def test_usage_the_stream_already_recovered_beats_an_estimate():
assert lifted == (recovered, 0.25)
def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost():
def test_a_timed_out_request_on_an_unmapped_model_keeps_its_row_at_zero_cost():
from litellm.proxy.utils import _failure_usage_to_lift
lifted = _failure_usage_to_lift(
model_call_details={
"call_type": "acompletion",
"model": "my-unmapped-model",
"custom_llm_provider": "openai",
"messages": ONE_USER_MESSAGE,
},
request_body={},
dispatched=True,
provider_billed=True,
)
assert lifted is not None
usage, response_cost = lifted
assert usage.prompt_tokens > 0
assert response_cost == 0.0
def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rate():
from litellm.proxy.utils import _failure_usage_to_lift
lifted = _failure_usage_to_lift(
@ -1912,13 +1998,16 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost():
},
request_body={},
dispatched=True,
provider_billed=True,
)
assert lifted is not None
usage, response_cost = lifted
assert usage.prompt_tokens > 0
assert usage.completion_tokens == 0
assert response_cost == 0.0
expected_input_cost, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=usage.prompt_tokens)
assert expected_input_cost > 0
assert response_cost == pytest.approx(expected_input_cost)
@pytest.mark.parametrize(