feat(utils.py): support custom cost tracking per second

https://github.com/BerriAI/litellm/issues/1374
This commit is contained in:
Krrish Dholakia 2024-01-22 11:15:08 -08:00
parent 44f756efb5
commit 276a685a59
4 changed files with 74 additions and 31 deletions

View file

@ -457,6 +457,8 @@ def completion(
### CUSTOM MODEL COST ###
input_cost_per_token = kwargs.get("input_cost_per_token", None)
output_cost_per_token = kwargs.get("output_cost_per_token", None)
input_cost_per_second = kwargs.get("input_cost_per_second", None)
output_cost_per_second = kwargs.get("output_cost_per_second", None)
### CUSTOM PROMPT TEMPLATE ###
initial_prompt_value = kwargs.get("initial_prompt_value", None)
roles = kwargs.get("roles", None)
@ -596,6 +598,19 @@ def completion(
}
}
)
if (
input_cost_per_second is not None
): # time based pricing just needs cost in place
output_cost_per_second = output_cost_per_second or 0.0
litellm.register_model(
{
model: {
"input_cost_per_second": input_cost_per_second,
"output_cost_per_second": output_cost_per_second,
"litellm_provider": custom_llm_provider,
}
}
)
### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ###
custom_prompt_dict = {} # type: ignore
if (