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

fix(router.py): add support for async image generation endpoints

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-21 14:38:44 +05:30
parent ae361230fd
commit c084f04a35
6 changed files with 109 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

25
litellm/tests/test_router.py
View file

@ -424,6 +424,7 @@ def test_function_calling_on_router():
# test_function_calling_on_router()
### IMAGE GENERATION
@pytest.mark.asyncio
async def test_aimg_gen_on_router():
litellm.set_verbose = True
try:
@ -442,14 +443,32 @@ async def test_aimg_gen_on_router():
"api_base": os.getenv("AZURE_SWEDEN_API_BASE"),
"api_key": os.getenv("AZURE_SWEDEN_API_KEY")
}
},
{
"model_name": "dall-e-2",
"litellm_params": {
"model": "azure/",
"api_version": "2023-06-01-preview",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY")
}
}
]
router = Router(model_list=model_list)
# response = await router.aimage_generation(
# model="dall-e-3",
# prompt="A cute baby sea otter"
# )
# print(response)
# assert len(response.data) > 0
response = await router.aimage_generation(
model="dall-e-3",
model="dall-e-2",
prompt="A cute baby sea otter"
)
print(response)
assert len(response.data) > 0
router.reset()
except Exception as e:
traceback.print_exc()
@ -489,7 +508,7 @@ def test_img_gen_on_router():
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
test_img_gen_on_router()
# test_img_gen_on_router()
###
def test_aembedding_on_router():
@ -625,7 +644,7 @@ async def test_mistral_on_router():
]
)
print(response)
asyncio.run(test_mistral_on_router())
# asyncio.run(test_mistral_on_router())
def test_openai_completion_on_router():
# [PROD Use Case] - Makes an acompletion call + async acompletion call, and sync acompletion call, sync completion + stream