From 078e2d341bd99e075e532bc2b688f044907a2f57 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 17 Mar 2025 15:12:01 -0700
Subject: [PATCH] feat(cost_calculator.py): support reading litellm response
 cost header in client sdk

allows consistent cost tracking when sdk is calling proxy
---
 litellm/cost_calculator.py            | 27 +++++++++++++++++++++-
 tests/litellm/test_cost_calculator.py | 32 +++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 tests/litellm/test_cost_calculator.py

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 58600ea14f..e17a94c87e 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -808,6 +808,23 @@ def completion_cost(  # noqa: PLR0915
         raise e
 
 
+def get_response_cost_from_hidden_params(
+    hidden_params: Union[dict, BaseModel]
+) -> Optional[float]:
+    if isinstance(hidden_params, BaseModel):
+        _hidden_params_dict = hidden_params.model_dump()
+    else:
+        _hidden_params_dict = hidden_params
+
+    additional_headers = _hidden_params_dict.get("additional_headers", {})
+    if additional_headers and "x-litellm-response-cost" in additional_headers:
+        response_cost = additional_headers["x-litellm-response-cost"]
+        if response_cost is None:
+            return None
+        return float(additional_headers["x-litellm-response-cost"])
+    return None
+
+
 def response_cost_calculator(
     response_object: Union[
         ModelResponse,
@@ -844,7 +861,7 @@ def response_cost_calculator(
     base_model: Optional[str] = None,
     custom_pricing: Optional[bool] = None,
     prompt: str = "",
-) -> Optional[float]:
+) -> float:
     """
     Returns
     - float or None: cost of response
@@ -856,6 +873,14 @@ def response_cost_calculator(
         else:
             if isinstance(response_object, BaseModel):
                 response_object._hidden_params["optional_params"] = optional_params
+
+                if hasattr(response_object, "_hidden_params"):
+                    provider_response_cost = get_response_cost_from_hidden_params(
+                        response_object._hidden_params
+                    )
+                    if provider_response_cost is not None:
+                        return provider_response_cost
+
             response_cost = completion_cost(
                 completion_response=response_object,
                 model=model,
diff --git a/tests/litellm/test_cost_calculator.py b/tests/litellm/test_cost_calculator.py
new file mode 100644
index 0000000000..9c9f6d9043
--- /dev/null
+++ b/tests/litellm/test_cost_calculator.py
@@ -0,0 +1,32 @@
+import json
+import os
+import sys
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+
+from unittest.mock import MagicMock, patch
+
+from pydantic import BaseModel
+
+from litellm.cost_calculator import response_cost_calculator
+
+
+def test_cost_calculator():
+    class MockResponse(BaseModel):
+        _hidden_params = {"additional_headers": {"x-litellm-response-cost": 1000}}
+
+    result = response_cost_calculator(
+        response_object=MockResponse(),
+        model="",
+        custom_llm_provider=None,
+        call_type="",
+        optional_params={},
+        cache_hit=None,
+        base_model=None,
+    )
+
+    assert result == 1000