scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

fix(get_api_base): fix get_api_base to handle model with alias

Browse source
This commit is contained in:
Krrish Dholakia 2024-05-09 08:01:17 -07:00
parent 4c8787f896
commit 4cfd988529
2 changed files with 10 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/tests/test_model_alias_map.py
View file

@ -16,7 +16,7 @@ litellm.set_verbose = True
model_alias_map = {"good-model": "anyscale/meta-llama/Llama-2-7b-chat-hf"}
def test_model_alias_map():
def test_model_alias_map(caplog):
try:
litellm.model_alias_map = model_alias_map
response = completion(
@ -27,9 +27,15 @@ def test_model_alias_map():
max_tokens=10,
)
print(response.model)
captured_logs = [rec.levelname for rec in caplog.records]
for log in captured_logs:
assert "ERROR" not in log
assert "Llama-2-7b-chat-hf" in response.model
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_model_alias_map()
# test_model_alias_map()

2
litellm/utils.py
View file

@ -5934,6 +5934,8 @@ def get_api_base(model: str, optional_params: dict) -> Optional[str]:
if _optional_params.api_base is not None:
return _optional_params.api_base
if litellm.model_alias_map and model in litellm.model_alias_map:
model = litellm.model_alias_map[model]
try:
model, custom_llm_provider, dynamic_api_key, dynamic_api_base = (
get_llm_provider(