scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test - test_aimage_generation_vertex_ai

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-20 08:14:43 -07:00
parent 1cecdc4690
commit 5ba5f15b56
1 changed files with 19 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/tests/test_image_generation.py
View file

@ -169,3 +169,22 @@ async def test_aimage_generation_bedrock_with_optional_params():
pass
else:
pytest.fail(f"An exception occurred - {str(e)}")
@pytest.mark.asyncio
async def test_aimage_generation_vertex_ai():
try:
response = await litellm.aimage_generation(
prompt="A cute baby sea otter",
model="vertex_ai/imagegeneration@006",
)
print(f"response: {response}")
except litellm.RateLimitError as e:
pass
except litellm.ContentPolicyViolationError:
pass # Azure randomly raises these errors - skip when they occur
except Exception as e:
if "Your task failed as a result of our safety system." in str(e):
pass
else:
pytest.fail(f"An exception occurred - {str(e)}")