scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

feat(proxy): write one spend row per billed timed-out retry attempt (#47416)

Browse source 
_failure_spend_rows yields the final row plus one re-keyed row per lifted billed_failed_attempts entry (<request_id>-attempt-<position>, 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>
This commit is contained in:
Kerry Lu 2026-10-20 10:39:29 -07:00 • committed by GitHub
parent 2484f5dd6d
commit ad8ca533ca
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 158 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

115
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -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`` (``<id>-attempt-<position>``, 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

59
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -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 ``<id>-attempt-<n>`` 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():
"""