From ad8ca533ca2cd5fa4360d519183c1a2aaccc084a Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 10:39:29 -0700 Subject: [PATCH] feat(proxy): write one spend row per billed timed-out retry attempt (#47416) _failure_spend_rows yields the final row plus one re-keyed row per lifted billed_failed_attempts entry (-attempt-, the attempt's usage, cost, retry count and deployment id), each written through DBSpendUpdateWriter.update_database so it also deducts from key, team and user spend Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/proxy_track_cost_callback.py | 115 +++++++++++++++--- .../hooks/test_proxy_track_cost_callback.py | 59 +++++++++ 2 files changed, 158 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 38a76028ddc..1bfb04ab883 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -55,7 +55,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # sanitize_error_information_for_spend_logs, should_store_prompts_and_responses_in_spend_logs, ) -from litellm.proxy.utils import ProxyUpdateSpend +from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend from litellm.types.utils import ( CallTypes, LiteLLMBatch, @@ -153,6 +153,73 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _billed_failed_attempts(request_data: Mapping[str, object]) -> tuple[BilledFailedAttempt, ...]: + """The earlier billed attempts ``_failure_fields_to_lift`` stashed on the request, if any.""" + attempts: Final = request_data.get("billed_failed_attempts") + if not isinstance(attempts, tuple): + return () + lifted: Final = cast("tuple[object, ...]", attempts) # cast-ok: lifted by _failure_fields_to_lift, checked below + return tuple(attempt for attempt in lifted if isinstance(attempt, BilledFailedAttempt)) + + +def _failure_spend_rows( + request_data: Mapping[str, object], final_cost: float, late_guardrails_snapshot: object +) -> tuple[tuple[float, Mapping[str, object]], ...]: + """One ``(response_cost, kwargs)`` per spend row a failure writes: the final attempt under the + request id, then every earlier attempt the provider also billed under its own row.""" + final_row: Final[Mapping[str, object]] = ( + request_data + if late_guardrails_snapshot is None + else {**request_data, "standard_logging_object": late_guardrails_snapshot} + ) + return ( + (final_cost, final_row), + *( + (attempt.response_cost, _attempt_spend_row_kwargs(request_data, attempt, position)) + for position, attempt in enumerate(_billed_failed_attempts(request_data)) + ), + ) + + +def _attempt_spend_row_kwargs( + request_data: Mapping[str, object], attempt: BilledFailedAttempt, position: int +) -> Mapping[str, object]: + """The final attempt's row kwargs re-keyed for one earlier billed attempt: its own + ``request_id`` (``-attempt-``, since ``request_id`` is the spend-log primary key + and the retry count restarts after a fallback hop), its estimated usage and cost, its retry + count under ``attempted_retries`` and the deployment that billed it under ``model_info.id``.""" + suffix: Final = f"-attempt-{position}" + raw_params: Final = cast("Mapping[str, object]", request_data.get("litellm_params") or {}) # cast-ok: request dict + litellm_params: Final = dict(raw_params) + metadata: Final = { + **cast("Mapping[str, object]", litellm_params.get("metadata") or {}), # cast-ok: request dict + "attempted_retries": attempt.attempted_retries, + } + metadata.pop("previous_models", None) + if attempt.deployment_id is not None: + metadata["model_info"] = { + **cast("Mapping[str, object]", metadata.get("model_info") or {}), # cast-ok: router dict + "id": attempt.deployment_id, + } + litellm_params["metadata"] = metadata + row: Final = { + **request_data, + "litellm_params": litellm_params, + "combined_usage_object": attempt.usage, + "response_cost": attempt.response_cost, + } + row.pop("billed_failed_attempts", None) + call_id: Final = request_data.get("litellm_call_id") or litellm_params.get("litellm_call_id") + if isinstance(call_id, str): + row["litellm_call_id"] = f"{call_id}{suffix}" + standard_logging_object: Final = request_data.get("standard_logging_object") + if isinstance(standard_logging_object, dict): + payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict + if isinstance(payload.get("id"), str): + row["standard_logging_object"] = {**payload, "id": f"{payload['id']}{suffix}"} + return row + + class _ProxyDBLogger(CustomLogger): def __init__( self, @@ -330,24 +397,40 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.get("standard_logging_guardrail_information"), ) - await self._spend_writer().update_database( - token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), - response_cost=recovered_response_cost, - user_id=user_api_key_dict.user_id, - end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id, - team_id=user_api_key_dict.team_id, - kwargs=( - request_data - if failure_snapshot_with_late_guardrails is None - else {**request_data, "standard_logging_object": failure_snapshot_with_late_guardrails} - ), - completion_response=original_exception, + await self._write_failure_spend_rows( + user_api_key_dict=user_api_key_dict, + original_exception=original_exception, start_time=actual_start_time, - end_time=datetime.now(), - org_id=user_api_key_dict.org_id, - project_id=user_api_key_dict.project_id, + spend_rows=_failure_spend_rows( + cast("Mapping[str, object]", request_data), # cast-ok: request dict, copied per row below + recovered_response_cost, + failure_snapshot_with_late_guardrails, + ), ) + async def _write_failure_spend_rows( + self, + user_api_key_dict: UserAPIKeyAuth, + original_exception: Exception, + start_time: datetime, + spend_rows: Sequence[tuple[float, Mapping[str, object]]], + ) -> None: + """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair.""" + for row_cost, row_kwargs in spend_rows: + await self._spend_writer().update_database( + token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), + response_cost=row_cost, + user_id=user_api_key_dict.user_id, + end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id, + team_id=user_api_key_dict.team_id, + kwargs=dict(row_kwargs), + completion_response=original_exception, + start_time=start_time, + end_time=datetime.now(), + org_id=user_api_key_dict.org_id, + project_id=user_api_key_dict.project_id, + ) + async def _PROXY_track_cost_callback( self, kwargs, # kwargs to completion diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 1f0dd67d86b..1132d339f15 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -95,6 +95,65 @@ async def test_async_post_call_failure_hook(): assert metadata["original_key"] == "original_value" +@pytest.mark.asyncio +async def test_each_earlier_billed_attempt_gets_its_own_spend_row(): + """A timed-out request that the router retried writes one row per billed attempt: the final + attempt under the request id, every earlier Timeout attempt under ``-attempt-`` with its + own cost, attempt number and deployment.""" + from litellm.proxy.utils import BilledFailedAttempt + + logger = ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key", user_id="u1", team_id="t1") + usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + request_data = { + "model": "gpt-4.1", + "litellm_call_id": "call-1", + "standard_logging_object": {"id": "call-1", "model_id": "dep-final"}, + "metadata": { + "attempted_retries": 2, + "model_info": {"id": "dep-final", "input_cost_per_token": 0.5}, + "previous_models": ({"exception_type": "Timeout"}, {"exception_type": "Timeout"}), + }, + "combined_usage_object": usage, + "response_cost": 0.5, + "billed_failed_attempts": ( + BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25), + BilledFailedAttempt(attempted_retries=0, deployment_id="dep-final", usage=usage, response_cost=0.5), + ), + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=litellm.Timeout("Request timed out.", model="gpt-4.1", llm_provider="openai"), + user_api_key_dict=user_api_key_dict, + ) + + rows = [call.kwargs for call in mock_update_database.call_args_list] + assert [row["response_cost"] for row in rows] == [0.5, 0.25, 0.5] + final_row, attempt_0, attempt_1 = (row["kwargs"] for row in rows) + assert final_row["litellm_call_id"] == "call-1" + assert final_row["litellm_params"]["metadata"]["attempted_retries"] == 2 + assert attempt_0["litellm_call_id"] == "call-1-attempt-0" + assert attempt_0["standard_logging_object"]["id"] == "call-1-attempt-0" + assert attempt_0["litellm_params"]["metadata"]["attempted_retries"] == 0 + assert attempt_0["litellm_params"]["metadata"]["model_info"]["id"] == "dep-a" + assert "previous_models" not in attempt_0["litellm_params"]["metadata"] + assert "billed_failed_attempts" not in attempt_0 + assert attempt_0["combined_usage_object"] is usage + assert attempt_1["litellm_call_id"] == "call-1-attempt-1" + assert attempt_1["litellm_params"]["metadata"]["attempted_retries"] == 0 + assert attempt_1["litellm_params"]["metadata"]["model_info"]["id"] == "dep-final" + assert final_row["litellm_params"]["metadata"]["model_info"]["id"] == "dep-final" + for row in rows: + assert row["token"] == "test_api_key" + assert row["user_id"] == "u1" + assert row["team_id"] == "t1" + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_metadata(): """