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:
parent
ce9b6f49bb
commit
078e2d341b
@ -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,
|
||||
|
||||
32
tests/litellm/test_cost_calculator.py
Normal file
32
tests/litellm/test_cost_calculator.py
Normal 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
|
||||
Loading…
Reference in New Issue
Block a user