phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #1690 from BerriAI/litellm_custom_pricing_fix

Browse source 
fix(main.py): register both model name and model name with provider
This commit is contained in:
Krish Dholakia 2024-01-30 13:56:38 -08:00 • committed by GitHub
commit 8ef7e9ad20
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 28 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

18
litellm/main.py
View file

@ -592,26 +592,37 @@ def completion(
### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
if input_cost_per_token is not None and output_cost_per_token is not None:
print_verbose(f"Registering model={model} in model cost map")
litellm.register_model(
{
f"{custom_llm_provider}/{model}": {
"input_cost_per_token": input_cost_per_token,
"output_cost_per_token": output_cost_per_token,
"litellm_provider": custom_llm_provider,
},
model: {
"input_cost_per_token": input_cost_per_token,
"output_cost_per_token": output_cost_per_token,
"litellm_provider": custom_llm_provider,
}
},
}
)
if (
elif (
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(
{
f"{custom_llm_provider}/{model}": {
"input_cost_per_second": input_cost_per_second,
"output_cost_per_second": output_cost_per_second,
"litellm_provider": custom_llm_provider,
},
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 ###
@ -3270,6 +3281,7 @@ async def ahealth_check(
## Set verbose to true -> ```litellm.set_verbose = True```
def print_verbose(print_statement):
try:
verbose_logger.debug(print_statement)
if litellm.set_verbose:
print(print_statement) # noqa
except: