fix undo changes cost tracking

This commit is contained in:
Ishaan Jaff 2024-11-12 17:45:12 -08:00
parent 84603763cc
commit 0a73fa6d01

View file

@ -171,7 +171,6 @@ def cost_per_token( # noqa: PLR0915
model_with_provider = model_with_provider_and_region model_with_provider = model_with_provider_and_region
else: else:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
model_without_prefix = model model_without_prefix = model
model_parts = model.split("/", 1) model_parts = model.split("/", 1)
if len(model_parts) > 1: if len(model_parts) > 1:
@ -516,7 +515,6 @@ 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="",
@ -530,7 +528,6 @@ 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,
@ -625,8 +622,6 @@ def completion_cost( # noqa: PLR0915
model = _select_model_name_for_cost_calc( model = _select_model_name_for_cost_calc(
model=model, model=model,
completion_response=completion_response, 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:
@ -858,14 +853,23 @@ 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(
completion_response=response_object,
model=model,
call_type=call_type,
custom_llm_provider=custom_llm_provider,
optional_params=optional_params,
)
else:
if custom_pricing is True: # override defaults if custom pricing is set
base_model = model
# base_model defaults to None if not set on model_info # base_model defaults to None if not set on model_info
response_cost = completion_cost( response_cost = completion_cost(
completion_response=response_object, completion_response=response_object,
call_type=call_type, call_type=call_type,
model=model, model=base_model,
base_model=base_model,
custom_llm_provider=custom_llm_provider, custom_llm_provider=custom_llm_provider,
custom_pricing=custom_pricing,
) )
return response_cost return response_cost
except Exception as e: except Exception as e: