feat(cost_calculator.py): support reading litellm response cost header in client sdk

allows consistent cost tracking when sdk is calling proxy
This commit is contained in:
Krrish Dholakia 2025-03-17 15:12:01 -07:00
parent ce9b6f49bb
commit 078e2d341b
2 changed files with 58 additions and 1 deletions

View File

@ -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,

View File

@ -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