phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat(main.py): add async image generation support

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-20 16:58:15 +05:30
parent b3962e483f
commit f59b9436be
1 changed files with 12 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

21
litellm/tests/test_image_generation.py
View file

@ -8,7 +8,7 @@ import logging
logging.basicConfig(level=logging.DEBUG)
load_dotenv()
import os
import asyncio
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
@ -35,12 +35,15 @@ def test_image_generation_azure_dall_e_3():
print(f"response: {response}")
assert len(response.data) > 0
# test_image_generation_azure_dall_e_3()
# @pytest.mark.asyncio
# async def test_async_image_generation_openai():
# response = litellm.image_generation(prompt="A cute baby sea otter", model="dall-e-3")
# print(f"response: {response}")
@pytest.mark.asyncio
async def test_async_image_generation_openai():
response = litellm.image_generation(prompt="A cute baby sea otter", model="dall-e-3")
print(f"response: {response}")
assert len(response.data) > 0
# @pytest.mark.asyncio
# async def test_async_image_generation_azure():
# response = litellm.image_generation(prompt="A cute baby sea otter", model="azure/dall-e-3")
# print(f"response: {response}")
# asyncio.run(test_async_image_generation_openai())
@pytest.mark.asyncio
async def test_async_image_generation_azure():
response = await litellm.aimage_generation(prompt="A cute baby sea otter", model="azure/dall-e-3-test")
print(f"response: {response}")