phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) testing model alias map

Browse source
This commit is contained in:
ishaan-jaff 2023-11-15 17:40:15 -08:00
parent 6f0b34c771
commit a7c31d54fa
1 changed files with 18 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

20
litellm/tests/test_model_alias_map.py
View file

@ -9,23 +9,29 @@ sys.path.insert(
) # Adds the parent directory to the system path ) # Adds the parent directory to the system path
import litellm import litellm
from litellm import embedding, completion from litellm import embedding, completion
import pytest
litellm.set_verbose = True litellm.set_verbose = True
model_alias_map = { model_alias_map = {
"llama2": "replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf" "good-model": "anyscale/meta-llama/Llama-2-7b-chat-hf"
} }
litellm.model_alias_map = model_alias_map litellm.model_alias_map = model_alias_map
try: def test_model_alias_map():
completion( try:
"llama2", response = completion(
"good-model",
messages=[{"role": "user", "content": "Hey, how's it going?"}], messages=[{"role": "user", "content": "Hey, how's it going?"}],
top_p=0.1, top_p=0.1,
temperature=0.01, temperature=0.01,
num_beams=4,
max_tokens=60, max_tokens=60,
) )
except Exception as e: print(response.model)
print(e) assert "Llama-2-7b-chat-hf" in response.model
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_model_alias_map()