fix response_cost_calculator

This commit is contained in:
Ishaan Jaff 2024-11-12 16:19:43 -08:00
parent 2415e7d0a6
commit 1c2f7f792b

View file

@ -515,6 +515,7 @@ def _infer_call_type(
def completion_cost( # noqa: PLR0915 def completion_cost( # noqa: PLR0915
completion_response=None, completion_response=None,
model: Optional[str] = None, model: Optional[str] = None,
base_model: Optional[str] = None,
prompt="", prompt="",
messages: List = [], messages: List = [],
completion="", completion="",
@ -528,6 +529,7 @@ def completion_cost( # noqa: PLR0915
quality=None, quality=None,
n=None, # number of images n=None, # number of images
### CUSTOM PRICING ### ### CUSTOM PRICING ###
custom_pricing: Optional[bool] = None,
custom_cost_per_token: Optional[CostPerToken] = None, custom_cost_per_token: Optional[CostPerToken] = None,
custom_cost_per_second: Optional[float] = None, custom_cost_per_second: Optional[float] = None,
optional_params: Optional[dict] = None, optional_params: Optional[dict] = None,
@ -620,7 +622,10 @@ def completion_cost( # noqa: PLR0915
f"completion_response response ms: {getattr(completion_response, '_response_ms', None)} " f"completion_response response ms: {getattr(completion_response, '_response_ms', None)} "
) )
model = _select_model_name_for_cost_calc( model = _select_model_name_for_cost_calc(
model=model, completion_response=completion_response model=model,
completion_response=completion_response,
base_model=base_model,
custom_pricing=custom_pricing,
) )
hidden_params = getattr(completion_response, "_hidden_params", None) hidden_params = getattr(completion_response, "_hidden_params", None)
if hidden_params is not None: if hidden_params is not None:
@ -852,24 +857,18 @@ def response_cost_calculator(
else: else:
if isinstance(response_object, BaseModel): if isinstance(response_object, BaseModel):
response_object._hidden_params["optional_params"] = optional_params response_object._hidden_params["optional_params"] = optional_params
if isinstance(response_object, ImageResponse):
response_cost = completion_cost( if custom_pricing is True: # override defaults if custom pricing is set
completion_response=response_object, base_model = model
model=model, # base_model defaults to None if not set on model_info
call_type=call_type, response_cost = completion_cost(
custom_llm_provider=custom_llm_provider, completion_response=response_object,
optional_params=optional_params, call_type=call_type,
) model=model,
else: base_model=base_model,
if custom_pricing is True: # override defaults if custom pricing is set custom_llm_provider=custom_llm_provider,
base_model = model custom_pricing=custom_pricing,
# base_model defaults to None if not set on model_info )
response_cost = completion_cost(
completion_response=response_object,
call_type=call_type,
model=base_model,
custom_llm_provider=custom_llm_provider,
)
return response_cost return response_cost
except Exception as e: except Exception as e:
raise e raise e