refactor: add black formatting

This commit is contained in:
Krrish Dholakia 2023-12-25 14:10:38 +05:30
parent b87d630b0a
commit 4905929de3
156 changed files with 19723 additions and 10869 deletions

View file

@ -2,6 +2,7 @@
import sys, os, time
import traceback, asyncio
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
@ -10,6 +11,7 @@ from litellm import Router
from concurrent.futures import ThreadPoolExecutor
from collections import defaultdict
from dotenv import load_dotenv
load_dotenv()
# every time we load the router we should have 4 clients:
@ -18,6 +20,7 @@ load_dotenv()
# Async + Stream
# Sync + Stream
def test_init_clients():
litellm.set_verbose = True
try:
@ -32,7 +35,7 @@ def test_init_clients():
"api_base": os.getenv("AZURE_API_BASE"),
"timeout": 0.01,
"stream_timeout": 0.000_001,
"max_retries": 7
"max_retries": 7,
},
},
]
@ -42,7 +45,7 @@ def test_init_clients():
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"]
@ -55,6 +58,7 @@ def test_init_clients():
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
# test_init_clients()
@ -80,20 +84,22 @@ def test_init_clients_basic():
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()
def test_timeouts_router():
"""
Test the timeouts of the router with multiple clients. This HASas to raise a timeout error
Test the timeouts of the router with multiple clients. This HASas to raise a timeout error
"""
import openai
litellm.set_verbose = True
try:
print("testing init 4 clients with diff timeouts")
@ -111,28 +117,32 @@ def test_timeouts_router():
},
]
router = Router(model_list=model_list)
print("PASSED !")
async def test():
try:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": "hello, write a 20 pg essay"
}
{"role": "user", "content": "hello, write a 20 pg essay"}
],
)
except Exception as e:
raise e
asyncio.run(test())
except openai.APITimeoutError as e:
print("Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e)
print(
"Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
)
print(type(e))
pass
except Exception as e:
pytest.fail(f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}")
pytest.fail(
f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
)
# test_timeouts_router()
@ -142,7 +152,7 @@ def test_stream_timeouts_router():
Test the stream timeouts router. See if it selected the correct client with stream timeout
"""
import openai
litellm.set_verbose = True
try:
print("testing init 4 clients with diff timeouts")
@ -154,37 +164,35 @@ def test_stream_timeouts_router():
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"timeout": 200, # regular calls will not timeout, stream calls will
"timeout": 200, # regular calls will not timeout, stream calls will
"stream_timeout": 0.000_001,
},
},
]
router = Router(model_list=model_list)
print("PASSED !")
selected_client = router._get_client(
deployment=router.model_list[0],
kwargs={
"model": "gpt-3.5-turbo",
"messages": [
{
"role": "user",
"content": "hello, write a 20 pg essay"
}
],
"stream": True
"messages": [{"role": "user", "content": "hello, write a 20 pg essay"}],
"stream": True,
},
client_type=None
client_type=None,
)
print("Select client timeout", selected_client.timeout)
assert selected_client.timeout == 0.000_001
assert selected_client.timeout == 0.000_001
except openai.APITimeoutError as e:
print("Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e)
print(
"Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
)
print(type(e))
pass
except Exception as e:
pytest.fail(f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}")
pytest.fail(
f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
)
test_stream_timeouts_router()