File size: 950 Bytes
395201c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
#### What this tests ####
# This tests the model alias mapping - if user passes in an alias, and has set an alias, set it to the actual value
import sys, os
import traceback
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm import embedding, completion
import pytest
litellm.set_verbose = True
model_alias_map = {
"good-model": "anyscale/meta-llama/Llama-2-7b-chat-hf"
}
litellm.model_alias_map = model_alias_map
def test_model_alias_map():
try:
response = completion(
"good-model",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
top_p=0.1,
temperature=0.01,
max_tokens=10,
)
print(response.model)
assert "Llama-2-7b-chat-hf" in response.model
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_model_alias_map() |