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

vtx test fix gemini-2.5-flash-lite

Browse source
This commit is contained in:
Ishaan Jaff 2025-08-07 00:11:10 -07:00
parent 65cf982d5b
commit deb1dd93e2
1 changed files with 70 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
tests/local_testing/test_amazing_vertex_completion.py
View file

@ -504,6 +504,75 @@ async def test_async_vertexai_streaming_response():
pytest.fail(f"An exception occurred: {e}")
# asyncio.run(test_async_vertexai_streaming_response())
@pytest.mark.parametrize("provider", ["vertex_ai"]) # "vertex_ai_beta"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.flaky(retries=3, delay=1)
@pytest.mark.asyncio
async def test_gemini_pro_vision(provider, sync_mode):
try:
load_vertex_ai_credentials()
litellm.set_verbose = True
litellm.num_retries = 3
if sync_mode:
resp = litellm.completion(
model="{}/gemini-2.5-flash-lite".format(provider),
messages=[
{"role": "system", "content": "Be a good bot"},
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "gs://cloud-samples-data/generative-ai/image/boats.jpeg"
},
},
],
},
],
)
else:
resp = await litellm.acompletion(
model="{}/gemini-2.5-flash-lite".format(provider),
messages=[
{"role": "system", "content": "Be a good bot"},
{
"role": "user",
"content": [
{"type": "text", "text": "Whats in this image?"},
{
"type": "image_url",
"image_url": {
"url": "gs://cloud-samples-data/generative-ai/image/boats.jpeg"
},
},
],
},
],
)
print(resp)
prompt_tokens = resp.usage.prompt_tokens
# DO Not DELETE this ASSERT
# Google counts the prompt tokens for us, we should ensure we use the tokens from the orignal response
assert prompt_tokens == 267 # the gemini api returns 267 to us
except litellm.RateLimitError as e:
pass
except Exception as e:
if "500 Internal error encountered.'" in str(e):
pass
else:
pytest.fail(f"An exception occurred - {str(e)}")
# test_gemini_pro_vision()
@pytest.mark.parametrize("load_pdf", [False]) # True,
@pytest.mark.flaky(retries=3, delay=1)
@ -536,7 +605,7 @@ def test_completion_function_plus_pdf(load_pdf):
image_message = {"role": "user", "content": image_content}
response = completion(
model="vertex_ai_beta/gemini-1.5-flash-preview-0514",
model="vertex_ai_beta/gemini-2.5-flash-lite",
messages=[image_message],
stream=False,
)