From aa0de83083e8ca4896d918138d2eacc9bc15c281 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Jan 2024 12:47:46 -0800 Subject: [PATCH] fix(main.py): register both model name and model name with provider --- litellm/main.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 05b352fa6b..511a40ea11 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -591,26 +591,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 ### @@ -3265,6 +3276,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: