fix(router.py): only do sync image gen fallbacks for now

The customhttptransport we use for dall-e-2 only works for sync httpx calls, not async. Will need to spend some time writing the async version

n
This commit is contained in:
Krrish Dholakia 2023-12-20 19:10:59 +05:30
parent 350389f501
commit 04bbd0649f
4 changed files with 82 additions and 9 deletions

View file

@ -455,7 +455,41 @@ async def test_aimg_gen_on_router():
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
asyncio.run(test_aimg_gen_on_router())
# asyncio.run(test_aimg_gen_on_router())
def test_img_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 = router.image_generation(
model="dall-e-3",
prompt="A cute baby sea otter"
)
print(response)
assert len(response.data) > 0
router.reset()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
test_img_gen_on_router()
###
def test_aembedding_on_router():