fix(acompletion): support client side timeouts + raise exceptions correctly for async calls

This commit is contained in:
Krrish Dholakia 2023-11-17 15:39:39 -08:00
parent 43ee581d46
commit 02ed97d0b2
8 changed files with 142 additions and 81 deletions

View file

@ -5,8 +5,6 @@ import sys, os
import pytest
import traceback
import asyncio, logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
sys.path.insert(
0, os.path.abspath("../..")
@ -20,7 +18,8 @@ def test_sync_response():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = completion(model="gpt-3.5-turbo", messages=messages, api_key=os.environ["OPENAI_API_KEY"])
response = completion(model="gpt-3.5-turbo", messages=messages, timeout=5)
print(f"response: {response}")
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
# test_sync_response()
@ -30,7 +29,7 @@ def test_sync_response_anyscale():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = completion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages)
response = completion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, timeout=5)
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
@ -43,10 +42,13 @@ def test_async_response_openai():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(model="gpt-3.5-turbo", messages=messages)
response = await acompletion(model="gpt-3.5-turbo", messages=messages, timeout=5)
print(f"response: {response}")
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
print(e)
asyncio.run(test_get_response())
@ -56,17 +58,18 @@ def test_async_response_azure():
import asyncio
litellm.set_verbose = True
async def test_get_response():
user_message = "Hello, how are you?"
user_message = "What do you know?"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(model="azure/chatgpt-v-2", messages=messages)
response = await acompletion(model="azure/chatgpt-v-2", messages=messages, timeout=5)
print(f"response: {response}")
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
asyncio.run(test_get_response())
def test_async_anyscale_response():
import asyncio
litellm.set_verbose = True
@ -74,9 +77,11 @@ def test_async_anyscale_response():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages)
response = await acompletion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, timeout=5)
# response = await response
print(f"response: {response}")
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
@ -91,7 +96,7 @@ def test_get_response_streaming():
messages = [{"content": user_message, "role": "user"}]
try:
litellm.set_verbose = True
response = await acompletion(model="gpt-3.5-turbo", messages=messages, stream=True)
response = await acompletion(model="gpt-3.5-turbo", messages=messages, stream=True, timeout=5)
print(type(response))
import inspect
@ -108,12 +113,12 @@ def test_get_response_streaming():
assert isinstance(output, str), "output needs to be of type str"
assert len(output) > 0, "Length of output needs to be greater than 0."
print(f'output: {output}')
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
return response
asyncio.run(test_async_call())
# test_get_response_streaming()
def test_get_response_non_openai_streaming():
@ -123,7 +128,7 @@ def test_get_response_non_openai_streaming():
user_message = "Hello, how are you?"
messages = [{"content": user_message, "role": "user"}]
try:
response = await acompletion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, stream=True)
response = await acompletion(model="anyscale/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, stream=True, timeout=5)
print(type(response))
import inspect
@ -140,10 +145,11 @@ def test_get_response_non_openai_streaming():
assert output is not None, "output cannot be None."
assert isinstance(output, str), "output needs to be of type str"
assert len(output) > 0, "Length of output needs to be greater than 0."
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred: {e}")
return response
asyncio.run(test_async_call())
# test_get_response_non_openai_streaming()
test_get_response_non_openai_streaming()