diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index dc2723c4267..ca4aa00e601 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -51,7 +51,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_request_model_access_groups, 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, @@ -106,6 +106,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, @@ -275,20 +342,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=user_api_key_dict.end_user_id, - team_id=user_api_key_dict.team_id, - kwargs=request_data, - 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, + None, + ), ) + 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=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, + ) + @log_db_metrics async def _PROXY_track_cost_callback( self, 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 28376be64b6..2add2636615 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -93,6 +93,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(): """ @@ -2877,3 +2936,94 @@ def test_autonomous_agent_cost_tracking_needs_no_human_or_virtual_key(agent_id: assert _should_track_cost_callback( user_api_key=None, user_id=None, team_id=None, end_user_id=None, call_type="acompletion", agent_id=agent_id ) is expected +||||||| parent of b217e5fa70 (fixup! feat(proxy): write one spend row per billed timed-out retry attempt (#47416)) + + +class _SharedCounter: + def __init__(self) -> None: + self.totals: dict[str, float] = {} + + async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: + self.totals[key] = self.totals.get(key, 0) + value + return self.totals[key] + + +def _grounded_gemini_kwargs(response_cost: float) -> dict: + return { + "model": "vertex_ai/gemini-test", + "call_type": "acompletion", + "response_cost": response_cost, + "litellm_params": {"metadata": {"user_api_key": "sk-test", "user_api_key_user_id": "user-1"}}, + "standard_logging_object": { + "response_cost": response_cost, + "request_tags": None, + "cache_hit": None, + "custom_llm_provider": "vertex_ai", + "cost_breakdown": { + "input_cost": 0.0, + "output_cost": 0.0, + "tool_usage_cost": 0.5, + "total_cost": response_cost, + }, + "model_map_information": { + "model_map_key": "vertex_ai/gemini-test", + "model_map_value": { + "litellm_provider": "vertex_ai", + "mode": "chat", + "web_search_billing_unit": "per_query", + "search_context_cost_per_query": {"search_context_size_medium": 0.1}, + "search_free_requests": {"count": 100, "period": "month", "shared_group": "test"}, + }, + }, + }, + "stream": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("opted_in", [True, False]) +async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in(monkeypatch, opted_in): + monkeypatch.setattr(litellm, "apply_search_free_tier", opted_in, raising=False) + counter = _SharedCounter() + kwargs = _grounded_gemini_kwargs(response_cost=1.0) + response = ModelResponse( + usage=Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=5), + ) + ) + + with ( + patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ), + patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True) + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.internal_usage_cache.dual_cache.redis_cache = counter + + await _ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=kwargs, completion_response=response, start_time=datetime.now(), end_time=datetime.now() + ) + await asyncio.sleep(0) + + mock_proxy_logging.failed_tracking_alert.assert_not_called() + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + + breakdown = kwargs["standard_logging_object"]["cost_breakdown"] + if opted_in: + assert kwargs["response_cost"] == pytest.approx(1.0 - 5 * 0.1) + assert breakdown["original_cost"] == 1.0 and breakdown["discount_amount"] == pytest.approx(5 * 0.1) + assert sum(counter.totals.values()) == 5 + else: + assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown + assert counter.totals == {}