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

feat(router.py): support async image generation on router

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-20 17:24:20 +05:30
parent f355e03515
commit 4040f60feb
2 changed files with 87 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

35
litellm/tests/test_router.py
View file

@ -423,6 +423,41 @@ def test_function_calling_on_router():
# test_function_calling_on_router()
### IMAGE GENERATION
async def test_aimg_gen_on_router():
litellm.set_verbose = True
try:
model_list = [
{
"model_name": "dall-e-3",
"litellm_params": {
"model": "dall-e-3",
},
},
{
"model_name": "dall-e-3",
"litellm_params": {
"model": "azure/dall-e-3-test",
"api_version": "2023-12-01-preview",
"api_base": os.getenv("AZURE_SWEDEN_API_BASE"),
"api_key": os.getenv("AZURE_SWEDEN_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)
router.reset()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
asyncio.run(test_aimg_gen_on_router())
###
def test_aembedding_on_router():
litellm.set_verbose = True
try: