(test) init router clients

This commit is contained in:
ishaan-jaff 2023-12-04 16:59:48 -08:00
parent cba98cf530
commit 886b52d448

View file

@ -1,4 +1,4 @@
# this tests if the router is intiaized correctly # this tests if the router is initialized correctly
import sys, os, time import sys, os, time
import traceback, asyncio import traceback, asyncio
import pytest import pytest
@ -12,40 +12,79 @@ from collections import defaultdict
from dotenv import load_dotenv from dotenv import load_dotenv
load_dotenv() load_dotenv()
# every time we load the router we should have 4 clients:
# everytime we load the router we should have 4 clients:
# Async # Async
# Sync # Sync
# Async + Stream # Async + Stream
# Sync + Stream # Sync + Stream
def test_init_clients(): def test_init_clients():
litellm.set_verbose = True litellm.set_verbose = True
try: try:
model_list = [ print("testing init 4 clients with diff timeouts")
{ model_list = [
"model_name": "gpt-3.5-turbo", {
"litellm_params": { "model_name": "gpt-3.5-turbo",
"model": "azure/chatgpt-v-2", "litellm_params": {
"api_key": os.getenv("AZURE_API_KEY"), "model": "azure/chatgpt-v-2",
"api_version": os.getenv("AZURE_API_VERSION"), "api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE") "api_version": os.getenv("AZURE_API_VERSION"),
}, "api_base": os.getenv("AZURE_API_BASE"),
}, "timeout": 0.01,
] "stream_timeout": 0.000_001,
"max_retries": 7
},
},
]
router = Router(model_list=model_list)
for elem in router.model_list:
assert elem["client"] is not None
assert elem["async_client"] is not None
assert elem["stream_client"] is not None
assert elem["stream_async_client"] is not None
# check if timeout for stream/non stream clients is set correctly
async_client = elem["async_client"]
stream_async_client = elem["stream_async_client"]
assert async_client.timeout == 0.01
assert stream_async_client.timeout == 0.000_001
print("PASSED !")
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
test_init_clients()
router = Router(model_list=model_list) def test_init_clients_basic():
print(router.model_list) litellm.set_verbose = True
for elem in router.model_list: try:
print(elem) print("Test basic client init")
assert elem["client"] is not None model_list = [
assert elem["async_client"] is not None {
assert elem["stream_client"] is not None "model_name": "gpt-3.5-turbo",
assert elem["stream_async_client"] is not None "litellm_params": {
"model": "azure/chatgpt-v-2",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
},
]
router = Router(model_list=model_list)
for elem in router.model_list:
assert elem["client"] is not None
assert elem["async_client"] is not None
assert elem["stream_client"] is not None
assert elem["stream_async_client"] is not None
print("PASSED !")
# see if we can init clients without timeout or max retries set
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
test_init_clients_basic()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
# test_init_clients()