mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
feat(proxy): lift every earlier billed timed-out retry attempt onto the failed request (#47490)
* feat(proxy): lift every earlier billed timed-out retry attempt onto the failed request _billed_earlier_attempts reads the router's previous_models records, keeps the Timeout ones other than the final attempt (matched on deployment_id plus attempted_retries, since the count restarts after a fallback hop and the fallback handler re-logs the exhausting attempt) and prices each at the deployment that served it. post_call_failure_hook lifts the tuple as billed_failed_attempts; nothing writes it yet. previous_models joins the router-reserved metadata fields so a client-planted record never reaches it * fix(proxy): price earlier timed-out attempts from the prompt estimate only The final row may lift a mid-stream recovered usage (prompt plus delivered chunks) or a failed_attempts_usage sum; earlier attempts never produced output, so they are priced from _estimated_prompt_usage, and when failed_attempts_usage already bills them into the final row none is lifted again --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
047cb4a8e2
commit
7a4a255d1d
4 changed files with 282 additions and 15 deletions
|
|
@ -467,6 +467,7 @@ _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
|
|||
{
|
||||
"attempted_fallbacks",
|
||||
"original_model_group",
|
||||
"previous_models",
|
||||
"request_retry_count",
|
||||
CLIENT_OUTPUT_CEILING_METADATA_KEY,
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
|
||||
|
|
|
|||
|
|
@ -1031,9 +1031,7 @@ def _failure_usage_to_lift(
|
|||
at the rate the router registered under ``deployment_id`` (``Logging.get_router_model_id``)
|
||||
when the deployment declares one, else the cost map's rate for ``base_model`` (an Azure
|
||||
deployment's ``model_info.base_model``) or the deployment model; 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
|
||||
failure stays at zero cost. Returns the
|
||||
(combined_usage_object, response_cost) pair to lift, or None."""
|
||||
recovered_usage: Final = model_call_details.get("combined_usage_object")
|
||||
failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ()))
|
||||
|
|
@ -1047,6 +1045,22 @@ def _failure_usage_to_lift(
|
|||
# The proxy-only error handler stamps this flag on a client disconnect too, even though the
|
||||
# request had already reached the provider, so a billed failure keeps its estimate.
|
||||
return None
|
||||
estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body)
|
||||
if estimated_usage is None:
|
||||
return None
|
||||
if not provider_billed:
|
||||
return estimated_usage, 0.0
|
||||
return estimated_usage, _estimated_input_cost(
|
||||
model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
|
||||
)
|
||||
|
||||
|
||||
def _estimated_prompt_usage(
|
||||
model_call_details: Mapping[str, object], request_body: Mapping[str, object]
|
||||
) -> Usage | None:
|
||||
"""The input-only Usage estimate for the prompt this request sent, or None when the call type has no
|
||||
estimable input. 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."""
|
||||
if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES:
|
||||
return None
|
||||
optional_params: Final = model_call_details.get("optional_params")
|
||||
|
|
@ -1056,17 +1070,118 @@ def _failure_usage_to_lift(
|
|||
else None
|
||||
)
|
||||
system_input: Final = dispatched_system or request_body.get("system") or request_body.get("instructions")
|
||||
estimated_usage: Final = _estimate_dispatched_failure_usage(
|
||||
return _estimate_dispatched_failure_usage(
|
||||
model=str(model_call_details.get("model") or ""),
|
||||
request_input=model_call_details.get("messages"),
|
||||
system_input=system_input,
|
||||
)
|
||||
if estimated_usage is None:
|
||||
return None
|
||||
if not provider_billed:
|
||||
return estimated_usage, 0.0
|
||||
return estimated_usage, _estimated_input_cost(
|
||||
model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BilledFailedAttempt:
|
||||
"""An earlier retry attempt the provider also billed, written as its own spend row."""
|
||||
|
||||
attempted_retries: int
|
||||
deployment_id: str | None
|
||||
usage: Usage
|
||||
response_cost: float
|
||||
|
||||
|
||||
_NO_RETRY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _retry_metadata_bucket(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]:
|
||||
"""The request metadata bucket the router writes retry state into, ``litellm_metadata`` when the
|
||||
route carries one (/messages, /responses) and ``metadata`` otherwise, same order as the deployment id."""
|
||||
if litellm_params is None:
|
||||
return _NO_RETRY_METADATA
|
||||
for metadata_key in ("litellm_metadata", "metadata"):
|
||||
metadata = litellm_params.get(metadata_key)
|
||||
if isinstance(metadata, dict):
|
||||
return cast("Mapping[str, object]", metadata) # cast-ok: untyped router dict, keys read with .get
|
||||
return _NO_RETRY_METADATA
|
||||
|
||||
|
||||
def _attempt_number(record: Mapping[str, object], index: int) -> int:
|
||||
"""The attempt a retry record describes: its own ``attempted_retries`` count, else its position."""
|
||||
recorded: Final = record.get("attempted_retries")
|
||||
return recorded if isinstance(recorded, int) else index
|
||||
|
||||
|
||||
def _retry_attempt_records(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]:
|
||||
"""The router's ``RetryAttemptRecord`` breadcrumbs, one per failed attempt, read from the same bucket."""
|
||||
if litellm_params is None:
|
||||
return ()
|
||||
for metadata_key in ("litellm_metadata", "metadata"):
|
||||
metadata = litellm_params.get(metadata_key)
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
records = cast("Mapping[str, object]", metadata).get("previous_models") # cast-ok: untyped router dict
|
||||
if isinstance(records, (list, tuple)):
|
||||
return tuple(
|
||||
cast("Mapping[str, object]", record) # cast-ok: RetryAttemptRecord TypedDict, read with .get
|
||||
for record in cast("Sequence[object]", records) # cast-ok: untyped router tuple
|
||||
if isinstance(record, dict)
|
||||
)
|
||||
return ()
|
||||
|
||||
|
||||
def _record_deployment_id(record: Mapping[str, object]) -> str | None:
|
||||
deployment_id: Final = record.get("deployment_id")
|
||||
return deployment_id if isinstance(deployment_id, str) else None
|
||||
|
||||
|
||||
def _attempt_identity(record: Mapping[str, object]) -> tuple[str | None, int | None]:
|
||||
"""What tells one provider call from another across the request: the deployment that served it
|
||||
and the retry count within its model group. Two records with the same identity describe the
|
||||
same call (the fallback handler re-logs the attempt that exhausted the previous group)."""
|
||||
attempted_retries: Final = record.get("attempted_retries")
|
||||
return _record_deployment_id(record), attempted_retries if isinstance(attempted_retries, int) else None
|
||||
|
||||
|
||||
def _timed_out_earlier_records(
|
||||
litellm_params: Mapping[str, object] | None, final_identity: tuple[str | None, int | None]
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""The ``Timeout`` breadcrumbs for calls other than the final attempt, one per provider call."""
|
||||
by_identity: Final = {
|
||||
_attempt_identity(record): record
|
||||
for record in _retry_attempt_records(litellm_params)
|
||||
if record.get("exception_type") == litellm.Timeout.__name__
|
||||
}
|
||||
return tuple(record for identity, record in by_identity.items() if identity != final_identity)
|
||||
|
||||
|
||||
def _billed_earlier_attempts(
|
||||
model_call_details: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
estimated_usage: Usage,
|
||||
base_model: str | None,
|
||||
deployment_id: str | None,
|
||||
) -> tuple[BilledFailedAttempt, ...]:
|
||||
"""Every earlier retry attempt the provider also billed. The router records one
|
||||
``RetryAttemptRecord`` per failed attempt in ``metadata.previous_models``, the final attempt
|
||||
included, which the failure row itself bills. The final attempt is the record with the
|
||||
request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback
|
||||
hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a
|
||||
``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final
|
||||
attempt, at that deployment's input rate (else the cost map's rate for ``base_model`` or the
|
||||
model). Only the exception class survives on the record, so an earlier connect timeout is
|
||||
billed too; any other failure class is not."""
|
||||
final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries")
|
||||
final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None)
|
||||
return tuple(
|
||||
BilledFailedAttempt(
|
||||
attempted_retries=_attempt_number(record, index),
|
||||
deployment_id=_record_deployment_id(record),
|
||||
usage=estimated_usage,
|
||||
response_cost=_estimated_input_cost(
|
||||
model_call_details,
|
||||
estimated_usage.prompt_tokens,
|
||||
deployment_id=_record_deployment_id(record),
|
||||
base_model=base_model,
|
||||
),
|
||||
)
|
||||
for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity))
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1265,13 +1380,26 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
|
|||
return _EMPTY_LIFT
|
||||
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
|
||||
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
|
||||
_litellm_params: Final = getattr(_logging_obj, "litellm_params", None)
|
||||
_base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details)
|
||||
_deployment_id: Final = _logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None
|
||||
_usage_to_lift: Final = _failure_usage_to_lift(
|
||||
model_call_details=_model_call_details,
|
||||
request_body=request_data,
|
||||
dispatched=_first_handoff is not None,
|
||||
provider_billed=provider_billed,
|
||||
deployment_id=_logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None,
|
||||
base_model=get_base_model_from_metadata(model_call_details=_model_call_details),
|
||||
deployment_id=_deployment_id,
|
||||
base_model=_base_model,
|
||||
)
|
||||
_prompt_estimate: Final = (
|
||||
_estimated_prompt_usage(_model_call_details, request_data)
|
||||
if provider_billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage")
|
||||
else None
|
||||
)
|
||||
_earlier_attempts: Final = (
|
||||
_billed_earlier_attempts(_model_call_details, _litellm_params, _prompt_estimate, _base_model, _deployment_id)
|
||||
if _prompt_estimate is not None
|
||||
else ()
|
||||
)
|
||||
_entries: Final = (
|
||||
("start_time", _model_call_details.get("start_time")),
|
||||
|
|
@ -1280,6 +1408,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
|
|||
("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]),
|
||||
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
|
||||
("standard_logging_object", _model_call_details.get("standard_logging_object")),
|
||||
("billed_failed_attempts", _earlier_attempts or None),
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in _entries if value is not None})
|
||||
|
||||
|
|
|
|||
|
|
@ -3330,9 +3330,7 @@ def test_x_litellm_trace_id_keeps_the_body_session_beside_it_whatever_the_log_se
|
|||
assert data["litellm_trace_id"] == "req-trace-1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"]
|
||||
)
|
||||
@pytest.mark.parametrize("session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"])
|
||||
@pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"])
|
||||
def test_session_headers_still_beat_the_body_session_with_correlation_logging(
|
||||
monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool
|
||||
|
|
@ -8282,6 +8280,7 @@ _PLANTED_STAMPS = {
|
|||
"attempted_fallbacks": 99,
|
||||
"original_model_group": "spoofed-group",
|
||||
"request_retry_count": -100,
|
||||
"previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}],
|
||||
"_client_output_ceiling": {"api_base": "https://attacker.example"},
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9,
|
||||
"client_key": "client_value",
|
||||
|
|
@ -8316,6 +8315,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
|
|||
assert "original_model_group" not in updated["metadata"]
|
||||
assert "_client_output_ceiling" not in updated["metadata"]
|
||||
assert "request_retry_count" not in updated["metadata"]
|
||||
assert "previous_models" not in updated["metadata"]
|
||||
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"]
|
||||
assert updated["metadata"]["client_key"] == "client_value"
|
||||
|
||||
|
|
|
|||
|
|
@ -2135,6 +2135,143 @@ def test_a_timed_out_request_is_priced_at_the_rate_the_router_registered_for_the
|
|||
assert response_cost == pytest.approx(usage.prompt_tokens * 0.25)
|
||||
|
||||
|
||||
def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(monkeypatch):
|
||||
"""The router leaves one RetryAttemptRecord per failed attempt in metadata.previous_models, the
|
||||
final attempt included; the earlier Timeout ones are billed (each at the deployment that served
|
||||
it), a 429 attempt is not, and the final attempt is left to the failure row itself."""
|
||||
from litellm.proxy.utils import _billed_earlier_attempts
|
||||
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"dep-lit145-b",
|
||||
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0},
|
||||
)
|
||||
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
|
||||
public_rate, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=24)
|
||||
|
||||
attempts = _billed_earlier_attempts(
|
||||
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
|
||||
litellm_params={
|
||||
"metadata": {
|
||||
"attempted_retries": 3,
|
||||
"previous_models": (
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "RateLimitError", "attempted_retries": 1},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 2},
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 3},
|
||||
),
|
||||
}
|
||||
},
|
||||
estimated_usage=usage,
|
||||
base_model=None,
|
||||
deployment_id="dep-lit145-a",
|
||||
)
|
||||
|
||||
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [(0, "dep-lit145-a"), (2, "dep-lit145-b")]
|
||||
assert attempts[0].response_cost == pytest.approx(public_rate)
|
||||
assert attempts[0].response_cost > 0
|
||||
assert attempts[1].response_cost == pytest.approx(24 * 0.25)
|
||||
assert all(a.usage is usage for a in attempts)
|
||||
|
||||
|
||||
def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_streams_recovered_usage():
|
||||
"""A final attempt whose stream broke mid-flight lifts its recovered usage (prompt plus the chunks
|
||||
delivered); the earlier timed-out attempts never produced output, so they are priced from the
|
||||
input-only estimate and not from that recovered usage."""
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.proxy.utils import _estimated_prompt_usage, _failure_fields_to_lift
|
||||
|
||||
logging_obj = _dispatched_failure_logging_obj(
|
||||
FAILURE_USAGE_MODEL,
|
||||
get_litellm_params(
|
||||
metadata={
|
||||
"attempted_retries": 1,
|
||||
"model_info": {"id": "dep-lit145-final"},
|
||||
"previous_models": (
|
||||
{"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
),
|
||||
}
|
||||
),
|
||||
)
|
||||
recovered = litellm.Usage(prompt_tokens=500, completion_tokens=40, total_tokens=540)
|
||||
logging_obj.model_call_details["combined_usage_object"] = recovered
|
||||
logging_obj.model_call_details["response_cost"] = 0.01
|
||||
|
||||
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
|
||||
|
||||
estimate = _estimated_prompt_usage(logging_obj.model_call_details, {})
|
||||
assert lifted["combined_usage_object"] is recovered
|
||||
assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [
|
||||
(estimate.prompt_tokens, 0)
|
||||
]
|
||||
assert estimate.prompt_tokens < recovered.prompt_tokens
|
||||
|
||||
|
||||
def test_attempts_already_summed_into_the_final_row_are_not_lifted_again():
|
||||
"""``Logging.carry_failed_attempt_usage`` moves a cut stream's usage into ``failed_attempts_usage``
|
||||
and the final row already bills every entry of it, so no earlier attempt gets a row of its own."""
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.proxy.utils import _failure_fields_to_lift
|
||||
|
||||
logging_obj = _dispatched_failure_logging_obj(
|
||||
FAILURE_USAGE_MODEL,
|
||||
get_litellm_params(
|
||||
metadata={
|
||||
"attempted_retries": 1,
|
||||
"model_info": {"id": "dep-lit145-final"},
|
||||
"previous_models": (
|
||||
{"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
),
|
||||
}
|
||||
),
|
||||
)
|
||||
logging_obj.model_call_details["failed_attempts_usage"] = (
|
||||
(litellm.Usage(prompt_tokens=24, completion_tokens=3, total_tokens=27), 0.001),
|
||||
)
|
||||
|
||||
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
|
||||
|
||||
assert lifted["combined_usage_object"].prompt_tokens == 24
|
||||
assert "billed_failed_attempts" not in lifted
|
||||
|
||||
|
||||
def test_a_timed_out_attempt_from_the_group_a_fallback_left_is_still_billed():
|
||||
"""After a fallback hop the router restarts ``attempted_retries`` at 0 for the new model group and
|
||||
re-logs the attempt that exhausted the previous group, so the final attempt (dep-b, retry 1) must be
|
||||
told from the previous group's attempt with the same count (dep-a, retry 1), and the re-logged
|
||||
record must bill once."""
|
||||
from litellm.proxy.utils import _billed_earlier_attempts
|
||||
|
||||
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
|
||||
|
||||
attempts = _billed_earlier_attempts(
|
||||
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
|
||||
litellm_params={
|
||||
"metadata": {
|
||||
"attempted_retries": 1,
|
||||
"model_info": {"id": "dep-lit145-b"},
|
||||
"previous_models": (
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 1},
|
||||
),
|
||||
}
|
||||
},
|
||||
estimated_usage=usage,
|
||||
base_model=None,
|
||||
deployment_id="dep-lit145-b",
|
||||
)
|
||||
|
||||
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [
|
||||
(0, "dep-lit145-a"),
|
||||
(1, "dep-lit145-a"),
|
||||
(0, "dep-lit145-b"),
|
||||
]
|
||||
assert all(a.response_cost > 0 for a in attempts)
|
||||
|
||||
|
||||
def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object:
|
||||
"""A Logging object as the failure hook sees it: litellm_params in the real get_litellm_params()
|
||||
shape (every key present, ``litellm_metadata`` None on chat completions) and a first-handoff
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue