phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test: testing fixes

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-09 10:22:58 +05:30 • committed by ishaan-jaff
parent 520cd7fa89
commit 10f76ec36c
4 changed files with 58 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/tests/test_completion.py
View file

@ -750,7 +750,7 @@ def test_completion_ollama_hosted():
messages=messages, messages=messages,
max_tokens=10, max_tokens=10,
num_retries=3, num_retries=3,
timeout=90, timeout=20,
api_base="https://test-ollama-endpoint.onrender.com", api_base="https://test-ollama-endpoint.onrender.com",
) )
# Add any assertions here to check the response # Add any assertions here to check the response

2
litellm/tests/test_custom_logger.py
View file

@ -183,7 +183,7 @@ def test_azure_completion_stream():
# checks if the model response available in the async + stream callbacks is equal to the received response # checks if the model response available in the async + stream callbacks is equal to the received response
customHandler2 = MyCustomHandler() customHandler2 = MyCustomHandler()
litellm.callbacks = [customHandler2] litellm.callbacks = [customHandler2]
litellm.set_verbose = True litellm.set_verbose = False
messages = [ messages = [
{"role": "system", "content": "You are a helpful assistant."}, {"role": "system", "content": "You are a helpful assistant."},
{ {

91
litellm/tests/test_image_generation.py
View file

@ -19,61 +19,80 @@ import litellm
def test_image_generation_openai(): def test_image_generation_openai():
litellm.set_verbose = True try:
response = litellm.image_generation( litellm.set_verbose = True
prompt="A cute baby sea otter", model="dall-e-3", num_retries=3 response = litellm.image_generation(
) prompt="A cute baby sea otter", model="dall-e-3"
print(f"response: {response}") )
assert len(response.data) > 0 print(f"response: {response}")
assert len(response.data) > 0
except litellm.RateLimitError as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_image_generation_openai() # test_image_generation_openai()
def test_image_generation_azure(): def test_image_generation_azure():
response = litellm.image_generation( try:
prompt="A cute baby sea otter", response = litellm.image_generation(
model="azure/", prompt="A cute baby sea otter", model="azure/", api_version="2023-06-01-preview"
api_version="2023-06-01-preview", )
num_retries=3, print(f"response: {response}")
) assert len(response.data) > 0
print(f"response: {response}") except litellm.RateLimitError as e:
assert len(response.data) > 0 pass
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_image_generation_azure() # test_image_generation_azure()
def test_image_generation_azure_dall_e_3(): def test_image_generation_azure_dall_e_3():
litellm.set_verbose = True try:
response = litellm.image_generation( litellm.set_verbose = True
prompt="A cute baby sea otter", response = litellm.image_generation(
model="azure/dall-e-3-test", prompt="A cute baby sea otter",
api_version="2023-12-01-preview", model="azure/dall-e-3-test",
api_base=os.getenv("AZURE_SWEDEN_API_BASE"), api_version="2023-12-01-preview",
api_key=os.getenv("AZURE_SWEDEN_API_KEY"), api_base=os.getenv("AZURE_SWEDEN_API_BASE"),
num_retries=3, api_key=os.getenv("AZURE_SWEDEN_API_KEY"),
) )
print(f"response: {response}") print(f"response: {response}")
assert len(response.data) > 0 assert len(response.data) > 0
except litellm.RateLimitError as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_image_generation_azure_dall_e_3() # test_image_generation_azure_dall_e_3()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_image_generation_openai(): async def test_async_image_generation_openai():
response = litellm.image_generation( try:
prompt="A cute baby sea otter", model="dall-e-3", num_retries=3 response = litellm.image_generation(
) prompt="A cute baby sea otter", model="dall-e-3"
print(f"response: {response}") )
assert len(response.data) > 0 print(f"response: {response}")
assert len(response.data) > 0
except litellm.RateLimitError as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# asyncio.run(test_async_image_generation_openai()) # asyncio.run(test_async_image_generation_openai())
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_image_generation_azure(): async def test_async_image_generation_azure():
response = await litellm.aimage_generation( try:
prompt="A cute baby sea otter", model="azure/dall-e-3-test", num_retries=3 response = await litellm.aimage_generation(
) prompt="A cute baby sea otter", model="azure/dall-e-3-test"
print(f"response: {response}") )
print(f"response: {response}")
except litellm.RateLimitError as e:
pass
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")

2
litellm/tests/test_streaming.py
View file

@ -307,7 +307,7 @@ def test_completion_ollama_hosted_stream():
messages=messages, messages=messages,
max_tokens=10, max_tokens=10,
num_retries=3, num_retries=3,
timeout=90, timeout=20,
api_base="https://test-ollama-endpoint.onrender.com", api_base="https://test-ollama-endpoint.onrender.com",
stream=True, stream=True,
) )