scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix(utils.py): accept {custom_llm_provider}/{model_name} in get_model_info

Browse source 
fixes https://github.com/BerriAI/litellm/issues/3100
This commit is contained in:
Krrish Dholakia 2024-04-17 16:38:53 -07:00
parent 409bd5b4ab
commit 53df916f69
2 changed files with 35 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/tests/test_get_model_info.py Normal file
View file

@ -0,0 +1,25 @@
# What is this?
## Unit testing for the 'get_model_info()' function
import os, sys, traceback
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm import get_model_info
def test_get_model_info_simple_model_name():
"""
tests if model name given, and model exists in model info - the object is returned
"""
model = "claude-3-opus-20240229"
litellm.get_model_info(model)
def test_get_model_info_custom_llm_with_model_name():
"""
Tests if {custom_llm_provider}/{model_name} name given, and model exists in model info, the object is returned
"""
model = "anthropic/claude-3-opus-20240229"
litellm.get_model_info(model)

12
litellm/utils.py
View file

@ -6134,7 +6134,13 @@ def get_model_info(model: str):
"mode": "chat",
}
else:
raise Exception()
"""
Check if model in model cost map
"""
if model in litellm.model_cost:
return litellm.model_cost[model]
else:
raise Exception()
except:
raise Exception(
"This model isn't mapped yet. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json"
@ -10595,7 +10601,9 @@ def trim_messages(
if max_tokens is None:
# Check if model is valid
if model in litellm.model_cost:
max_tokens_for_model = litellm.model_cost[model].get("max_input_tokens", litellm.model_cost[model]["max_tokens"])
max_tokens_for_model = litellm.model_cost[model].get(
"max_input_tokens", litellm.model_cost[model]["max_tokens"]
)
max_tokens = int(max_tokens_for_model * trim_ratio)
else:
# if user did not specify max (input) tokens