mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(utils.py): fix model registeration to model cost map
Fixes https://github.com/BerriAI/litellm/issues/4972
This commit is contained in:
Krrish Dholakia
parent
142f4fefd0
commit
46634af06f
5 changed files with 73 additions and 14 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
model: anthropic/claude-3-haiku-20240307
|
||||
max_tokens: 4096
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]]
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue