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

fix(vertex_httpx.py): add sync vertex image gen support

Browse source 
Fixes https://github.com/BerriAI/litellm/issues/4623
This commit is contained in:
Krrish Dholakia 2024-07-09 13:33:54 -07:00
parent 7c30347c57
commit 789d2dab15
5 changed files with 151 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

21
litellm/tests/test_image_generation.py
View file

@ -190,21 +190,26 @@ async def test_aimage_generation_bedrock_with_optional_params():
pytest.fail(f"An exception occurred - {str(e)}")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_aimage_generation_vertex_ai():
async def test_aimage_generation_vertex_ai(sync_mode):
from test_amazing_vertex_completion import load_vertex_ai_credentials
litellm.set_verbose = True
load_vertex_ai_credentials()
data = {
"prompt": "An olympic size swimming pool",
"model": "vertex_ai/imagegeneration@006",
"vertex_ai_project": "adroit-crow-413218",
"vertex_ai_location": "us-central1",
"n": 1,
}
try:
response = await litellm.aimage_generation(
prompt="An olympic size swimming pool",
model="vertex_ai/imagegeneration@006",
vertex_ai_project="adroit-crow-413218",
vertex_ai_location="us-central1",
n=1,
)
if sync_mode:
response = litellm.image_generation(**data)
else:
response = await litellm.aimage_generation(**data)
assert response.data is not None
assert len(response.data) > 0