fix(litellm_logging.py): use 1 cost calc function across response headers + logging integrations

Ensures consistent cost calculation when azure base models are used
This commit is contained in:
Krrish Dholakia 2024-08-01 10:26:59 -07:00
parent d02d3d9712
commit 10b571ca42
3 changed files with 40 additions and 24 deletions

View file

@ -12,6 +12,8 @@ import traceback
import uuid
from typing import Any, Callable, Dict, List, Literal, Optional
from pydantic import BaseModel
import litellm
from litellm import (
json_logs,
@ -503,6 +505,34 @@ class Logging:
)
)
def _response_cost_calculator(self, result: BaseModel):
"""
Calculate response cost using result + logging object variables.
used for consistent cost calculation across response headers + logging integrations.
"""
## RESPONSE COST ##
custom_pricing = use_custom_pricing_for_model(
litellm_params=self.litellm_params
)
response_cost = litellm.response_cost_calculator(
response_object=result,
model=self.model,
cache_hit=self.model_call_details.get("cache_hit", False),
custom_llm_provider=self.model_call_details.get(
"custom_llm_provider", None
),
base_model=_get_base_model_from_metadata(
model_call_details=self.model_call_details
),
call_type=self.call_type,
optional_params=self.optional_params,
custom_pricing=custom_pricing,
)
return response_cost
def _success_handler_helper_fn(
self, result=None, start_time=None, end_time=None, cache_hit=None
):
@ -537,20 +567,7 @@ class Logging:
litellm_params=self.litellm_params
)
self.model_call_details["response_cost"] = (
litellm.response_cost_calculator(
response_object=result,
model=self.model,
cache_hit=self.model_call_details.get("cache_hit", False),
custom_llm_provider=self.model_call_details.get(
"custom_llm_provider", None
),
base_model=_get_base_model_from_metadata(
model_call_details=self.model_call_details
),
call_type=self.call_type,
optional_params=self.optional_params,
custom_pricing=custom_pricing,
)
self._response_cost_calculator(result=result)
)
## HIDDEN PARAMS ##