fix(proxy): price a billed failure row at the deployment's own rate like the success path (#47437)

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

Backport of #47437 to stable/1.106.x
Cherry-picked from fb555b00dc (main)

litellm/proxy/utils.py

@@ -84,7 +84,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
 from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.proxy.model_listing import ModelInfoResponse
-from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage
+from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, CostPerToken, ModelInfo, Usage
 
 try:
     from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
@@ -150,7 +150,7 @@ from litellm.litellm_core_utils.core_helpers import (
     is_expected_client_error,
 )
 from litellm.litellm_core_utils.duration_parser import get_budget_window_start
-from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
 from litellm.litellm_core_utils.served_output_texts import (
@@ -268,6 +268,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionRe
 from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
 from litellm.utils import (
     add_custom_logger_callback_to_specific_event,
+    get_base_model_from_metadata,
 )
 
 if TYPE_CHECKING:
@@ -964,6 +965,8 @@ def _failure_usage_to_lift(
     request_body: Mapping[str, object],
     dispatched: bool,
     provider_billed: bool = False,
+    deployment_id: str | None = None,
+    base_model: str | None = None,
 ) -> 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
@@ -972,8 +975,11 @@ def _failure_usage_to_lift(
     estimated input-side Usage. When the failure is one the provider bills
     anyway (``provider_billed``: a read timeout or a client disconnect, 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
+    usage is priced at the model's input rate, the way the success path prices it:
+    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
     (combined_usage_object, response_cost) pair to lift, or None."""
@@ -1004,18 +1010,43 @@ def _failure_usage_to_lift(
         return None
     if not provider_billed:
         return estimated_usage, 0.0
-    return estimated_usage, _estimated_input_cost(model_call_details, estimated_usage.prompt_tokens)
+    return estimated_usage, _estimated_input_cost(
+        model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
+    )
 
 
-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."""
+def _deployment_cost_per_token(deployment_id: str | None, model: str) -> CostPerToken | None:
+    """The pricing the router registered for this deployment under its ``model_info.id``, read
+    through ``deployment_pricing_model_info`` exactly as the success path does; None means the
+    public cost map. The id is ``Logging.get_router_model_id()`` (the logging object's own
+    ``litellm_params``), never the flattened ``model_call_details`` copy a request body key can
+    overwrite."""
+    model_info: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=model)
+    if model_info is None:
+        return None
+    return CostPerToken(
+        input_cost_per_token=model_info["input_cost_per_token"] or 0.0,
+        output_cost_per_token=model_info["output_cost_per_token"] or 0.0,
+    )
+
+
+def _estimated_input_cost(
+    model_call_details: Mapping[str, object],
+    prompt_tokens: int,
+    deployment_id: str | None,
+    base_model: str | None,
+) -> float:
+    """Input price for the estimated prompt: the deployment's registered rate when it has one, else
+    the cost map's rate for ``base_model`` or the deployment model (an Azure deployment name is not
+    in the cost map, its ``base_model`` is). An unpriceable 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,
+            model=base_model or model,
             custom_llm_provider=str(model_call_details.get("custom_llm_provider") or "") or None,
             prompt_tokens=prompt_tokens,
+            custom_cost_per_token=_deployment_cost_per_token(deployment_id, model),
         )
     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(
@@ -1184,6 +1215,8 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
         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),
     )
     _entries: Final = (
         ("start_time", _model_call_details.get("start_time")),

tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py

@@ -2105,6 +2105,127 @@ def test_a_timed_out_request_on_an_unmapped_model_keeps_its_row_at_zero_cost():
     assert response_cost == 0.0
 
 
+def test_a_timed_out_request_is_priced_at_the_rate_the_router_registered_for_the_deployment(monkeypatch):
+    """Router._create_deployment registers a deployment's own ``input_cost_per_token`` under its
+    ``model_info.id``; the failure row prices with that entry through the same
+    ``deployment_pricing_model_info`` the success path uses."""
+    from litellm.proxy.utils import _failure_usage_to_lift
+
+    monkeypatch.setitem(
+        litellm.model_cost,
+        "dep-lit145",
+        {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0},
+    )
+
+    lifted = _failure_usage_to_lift(
+        model_call_details={
+            "call_type": "acompletion",
+            "model": FAILURE_USAGE_MODEL,
+            "messages": ONE_USER_MESSAGE,
+        },
+        request_body={},
+        dispatched=True,
+        provider_billed=True,
+        deployment_id="dep-lit145",
+    )
+
+    assert lifted is not None
+    usage, response_cost = lifted
+    assert usage.prompt_tokens > 0
+    assert response_cost == pytest.approx(usage.prompt_tokens * 0.25)
+
+
+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
+    instant, so the lift treats the request as dispatched."""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    logging_obj = Logging(
+        model=model,
+        messages=ONE_USER_MESSAGE,
+        stream=False,
+        call_type="acompletion",
+        start_time=real_datetime.datetime.now(),
+        litellm_call_id="call-lit145",
+        function_id="fn-lit145",
+    )
+    logging_obj.update_environment_variables(model=model, user="", optional_params={}, litellm_params=litellm_params)
+    logging_obj.model_call_details["first_api_call_start_time"] = real_datetime.datetime.now()
+    logging_obj.model_call_details["messages"] = ONE_USER_MESSAGE
+    logging_obj.model_call_details["call_type"] = "acompletion"
+    return logging_obj
+
+
+def test_the_deployment_rate_is_found_on_the_logging_objects_real_litellm_params(monkeypatch):
+    """``get_litellm_params()`` always carries a ``litellm_metadata`` key (None on chat
+    completions) next to the router-stamped ``metadata.model_info.id``; the failure row must
+    still price at the deployment's own rate, like ``Logging.get_router_model_id`` on the
+    success path."""
+    from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
+    from litellm.proxy.utils import _failure_fields_to_lift
+
+    monkeypatch.setitem(
+        litellm.model_cost,
+        "dep-lit145-real-shape",
+        {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 0.25},
+    )
+    litellm_params = get_litellm_params(metadata={"model_info": {"id": "dep-lit145-real-shape"}})
+    assert "litellm_metadata" in litellm_params and litellm_params["litellm_metadata"] is None
+    logging_obj = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, litellm_params)
+
+    lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
+
+    assert lifted["response_cost"] == pytest.approx(lifted["combined_usage_object"].prompt_tokens * 0.25)
+
+
+def test_an_azure_deployment_is_priced_at_its_base_model_like_the_success_path():
+    """An Azure deployment name is not in the cost map; the success path prices it through
+    ``model_info.base_model`` (``get_base_model_from_metadata``) and so must the failure row."""
+    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(
+        "azure/lit145-not-in-cost-map",
+        get_litellm_params(metadata={"model_info": {"base_model": "azure/gpt-4.1"}}),
+    )
+    logging_obj.model_call_details["custom_llm_provider"] = "azure"
+
+    lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
+
+    expected_cost, _ = litellm.cost_per_token(
+        model="azure/gpt-4.1", prompt_tokens=lifted["combined_usage_object"].prompt_tokens
+    )
+    assert expected_cost > 0
+    assert lifted["response_cost"] == pytest.approx(expected_cost)
+
+
+def test_a_litellm_params_key_in_the_request_body_cannot_zero_the_timeout_price():
+    """``model_call_details`` is flattened with the request's optional_params, so a body key
+    named ``litellm_params`` can land there; the deployment rate must come from the logging
+    object's own ``litellm_params`` instead."""
+    from litellm.proxy.utils import _failure_usage_to_lift
+
+    lifted = _failure_usage_to_lift(
+        model_call_details={
+            "call_type": "acompletion",
+            "model": FAILURE_USAGE_MODEL,
+            "messages": ONE_USER_MESSAGE,
+            "litellm_params": {"input_cost_per_token": 0},
+        },
+        request_body={},
+        dispatched=True,
+        provider_billed=True,
+        deployment_id=None,
+    )
+
+    assert lifted is not None
+    usage, response_cost = lifted
+    expected_cost, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=usage.prompt_tokens)
+    assert response_cost == pytest.approx(expected_cost)
+    assert response_cost > 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