From 46634af06fef50a73425b8028192eb23fe98e5ae Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Jul 2024 18:15:00 -0700 Subject: [PATCH] fix(utils.py): fix model registeration to model cost map Fixes https://github.com/BerriAI/litellm/issues/4972 --- litellm/cost_calculator.py | 10 +++---- litellm/proxy/_new_secret_config.yaml | 10 ++----- litellm/tests/test_completion_cost.py | 41 +++++++++++++++++++++++++++ litellm/types/utils.py | 2 ++ litellm/utils.py | 24 +++++++++++++++- 5 files changed, 73 insertions(+), 14 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b680fc8a543..7b8bfb0d98c 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -106,7 +106,6 @@ def cost_per_token( Returns: tuple: A tuple containing the cost in USD dollars for prompt tokens and completion tokens, respectively. """ - args = locals() if model is None: raise Exception("Invalid arg. Model cannot be none.") ## CUSTOM PRICING ## @@ -117,6 +116,7 @@ def cost_per_token( custom_cost_per_second=custom_cost_per_second, custom_cost_per_token=custom_cost_per_token, ) + if response_cost is not None: return response_cost[0], response_cost[1] @@ -495,9 +495,9 @@ def completion_cost( completion_tokens = completion_response.get("usage", {}).get( "completion_tokens", 0 ) - total_time = completion_response.get("_response_ms", 0) + total_time = getattr(completion_response, "_response_ms", 0) verbose_logger.debug( - f"completion_response response ms: {completion_response.get('_response_ms')} " + f"completion_response response ms: {getattr(completion_response, '_response_ms', None)} " ) model = model or completion_response.get( "model", None @@ -659,9 +659,7 @@ def completion_cost( call_type=call_type, ) _final_cost = prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar - print_verbose( - f"final cost: {_final_cost}; prompt_tokens_cost_usd_dollar: {prompt_tokens_cost_usd_dollar}; completion_tokens_cost_usd_dollar: {completion_tokens_cost_usd_dollar}" - ) + return _final_cost except Exception as e: raise e diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index d13fb3f37af..0bd00067a28 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,9 +1,5 @@ model_list: - - model_name: "*" + - model_name: claude-3-haiku-20240307 litellm_params: - model: "*" - -litellm_settings: - cache: true - cache_params: - type: redis \ No newline at end of file + model: anthropic/claude-3-haiku-20240307 + max_tokens: 4096 \ No newline at end of file diff --git a/litellm/tests/test_completion_cost.py b/litellm/tests/test_completion_cost.py index 53cbaa31dc8..fc0fdc32355 100644 --- a/litellm/tests/test_completion_cost.py +++ b/litellm/tests/test_completion_cost.py @@ -966,3 +966,44 @@ def test_completion_cost_tts(model): ) assert cost > 0 + + +def test_completion_cost_anthropic(): + """ + model_name: claude-3-haiku-20240307 + litellm_params: + model: anthropic/claude-3-haiku-20240307 + max_tokens: 4096 + """ + router = litellm.Router( + model_list=[ + { + "model_name": "claude-3-haiku-20240307", + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "max_tokens": 4096, + }, + } + ] + ) + data = { + "model": "claude-3-haiku-20240307", + "prompt_tokens": 21, + "completion_tokens": 20, + "response_time_ms": 871.7040000000001, + "custom_llm_provider": "anthropic", + "region_name": None, + "prompt_characters": 0, + "completion_characters": 0, + "custom_cost_per_token": None, + "custom_cost_per_second": None, + "call_type": "acompletion", + } + + input_cost, output_cost = cost_per_token(**data) + + assert input_cost > 0 + assert output_cost > 0 + + print(input_cost) + print(output_cost) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3f7b16a2ad9..199edef6b45 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -40,6 +40,8 @@ class ModelInfo(TypedDict, total=False): Model info for a given model, this is information found in litellm.model_prices_and_context_window.json """ + key: Required[str] # the key in litellm.model_cost which is returned + max_tokens: Required[Optional[int]] max_input_tokens: Required[Optional[int]] max_output_tokens: Required[Optional[int]] diff --git a/litellm/utils.py b/litellm/utils.py index 1e9a0e87c9a..a32c67d0339 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2148,6 +2148,13 @@ def supports_parallel_function_calling(model: str): ####### HELPER FUNCTIONS ################ +def _update_dictionary(existing_dict: dict, new_dict: dict) -> dict: + for k, v in new_dict.items(): + existing_dict[k] = v + + return existing_dict + + def register_model(model_cost: Union[str, dict]): """ Register new / Override existing models (and their pricing) to specific providers. @@ -2170,8 +2177,17 @@ def register_model(model_cost: Union[str, dict]): loaded_model_cost = litellm.get_model_cost_map(url=model_cost) for key, value in loaded_model_cost.items(): + ## get model info ## + try: + existing_model = get_model_info(model=key) + model_cost_key = existing_model["key"] + except Exception: + existing_model = {} + model_cost_key = key ## override / add new keys to the existing model cost dictionary - litellm.model_cost.setdefault(key, {}).update(value) + litellm.model_cost.setdefault(model_cost_key, {}).update( + _update_dictionary(existing_model, value) + ) verbose_logger.debug(f"{key} added to model cost map") # add new model names to provider lists if value.get("litellm_provider") == "openai": @@ -4858,6 +4874,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod Returns: dict: A dictionary containing the following information: + key: Required[str] # the key in litellm.model_cost which is returned max_tokens: Required[Optional[int]] max_input_tokens: Required[Optional[int]] max_output_tokens: Required[Optional[int]] @@ -4959,6 +4976,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) return ModelInfo( + key=model, max_tokens=max_tokens, # type: ignore max_input_tokens=None, max_output_tokens=None, @@ -4979,6 +4997,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod 3. 'split_model' in litellm.model_cost. Checks "llama3-8b-8192" in litellm.model_cost if model="groq/llama3-8b-8192" """ if combined_model_name in litellm.model_cost: + key = combined_model_name _model_info = litellm.model_cost[combined_model_name] _model_info["supported_openai_params"] = supported_openai_params if ( @@ -4992,6 +5011,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod else: raise Exception elif model in litellm.model_cost: + key = model _model_info = litellm.model_cost[model] _model_info["supported_openai_params"] = supported_openai_params if ( @@ -5005,6 +5025,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod else: raise Exception elif split_model in litellm.model_cost: + key = split_model _model_info = litellm.model_cost[split_model] _model_info["supported_openai_params"] = supported_openai_params if ( @@ -5027,6 +5048,7 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod _model_info["supports_response_schema"] = True return ModelInfo( + key=key, max_tokens=_model_info.get("max_tokens", None), max_input_tokens=_model_info.get("max_input_tokens", None), max_output_tokens=_model_info.get("max_output_tokens", None),