mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
test(test_amazing_vertex_completion.py): pass pdf as base64 to vertex ai
This commit is contained in:
Krrish Dholakia
parent
f15cb2bdc8
commit
c89aa1e626
2 changed files with 37 additions and 156 deletions
|
|
@ -593,6 +593,43 @@ async def test_gemini_pro_vision(provider, sync_mode):
|
|||
# test_gemini_pro_vision()
|
||||
|
||||
|
||||
def test_completion_function_plus_pdf():
|
||||
litellm.set_verbose = True
|
||||
load_vertex_ai_credentials()
|
||||
try:
|
||||
import base64
|
||||
|
||||
import requests
|
||||
|
||||
# URL of the file
|
||||
url = "https://storage.googleapis.com/cloud-samples-data/generative-ai/pdf/2403.05530.pdf"
|
||||
|
||||
# Download the file
|
||||
response = requests.get(url)
|
||||
file_data = response.content
|
||||
|
||||
encoded_file = base64.b64encode(file_data).decode("utf-8")
|
||||
|
||||
image_content = [
|
||||
{"type": "text", "text": "What's this file about?"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:application/pdf;base64,{encoded_file}"},
|
||||
},
|
||||
]
|
||||
image_message = {"role": "user", "content": image_content}
|
||||
|
||||
response = completion(
|
||||
model="vertex_ai_beta/gemini-1.5-flash-preview-0514",
|
||||
messages=[image_message],
|
||||
stream=False,
|
||||
)
|
||||
|
||||
print(response)
|
||||
except litellm.InternalServerError as e:
|
||||
pytest.fail("Got={}".format(str(e)))
|
||||
|
||||
|
||||
def encode_image(image_path):
|
||||
import base64
|
||||
|
||||
|
|
@ -1430,123 +1467,6 @@ def test_tool_name_conversion():
|
|||
)
|
||||
|
||||
|
||||
# Extra gemini Vision tests for completion + stream, async, async + stream
|
||||
# if we run into issues with gemini, we will also add these to our ci/cd pipeline
|
||||
# def test_gemini_pro_vision_stream():
|
||||
# try:
|
||||
# litellm.set_verbose = False
|
||||
# litellm.num_retries=0
|
||||
# print("streaming response from gemini-pro-vision")
|
||||
# resp = litellm.completion(
|
||||
# model = "vertex_ai/gemini-pro-vision",
|
||||
# messages=[
|
||||
# {
|
||||
# "role": "user",
|
||||
# "content": [
|
||||
# {
|
||||
# "type": "text",
|
||||
# "text": "Whats in this image?"
|
||||
# },
|
||||
# {
|
||||
# "type": "image_url",
|
||||
# "image_url": {
|
||||
# "url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
# }
|
||||
# }
|
||||
# ]
|
||||
# }
|
||||
# ],
|
||||
# stream=True
|
||||
# )
|
||||
# print(resp)
|
||||
# for chunk in resp:
|
||||
# print(chunk)
|
||||
# except Exception as e:
|
||||
# import traceback
|
||||
# traceback.print_exc()
|
||||
# raise e
|
||||
# test_gemini_pro_vision_stream()
|
||||
|
||||
|
||||
def test_gemini_pro_vision_async():
|
||||
try:
|
||||
litellm.set_verbose = True
|
||||
litellm.num_retries = 0
|
||||
|
||||
async def test():
|
||||
load_vertex_ai_credentials()
|
||||
resp = await litellm.acompletion(
|
||||
model="vertex_ai/gemini-pro-vision",
|
||||
messages=[
|
||||
{"role": "system", "content": ""},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Whats in this image?"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
)
|
||||
print("async response gemini pro vision")
|
||||
print(resp)
|
||||
|
||||
asyncio.run(test())
|
||||
except litellm.RateLimitError:
|
||||
pass
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
raise e
|
||||
|
||||
|
||||
# test_gemini_pro_vision_async()
|
||||
|
||||
|
||||
# def test_gemini_pro_vision_async_stream():
|
||||
# try:
|
||||
# litellm.set_verbose = True
|
||||
# litellm.num_retries=0
|
||||
# async def test():
|
||||
# resp = await litellm.acompletion(
|
||||
# model = "vertex_ai/gemini-pro-vision",
|
||||
# messages=[
|
||||
# {
|
||||
# "role": "user",
|
||||
# "content": [
|
||||
# {
|
||||
# "type": "text",
|
||||
# "text": "Whats in this image?"
|
||||
# },
|
||||
# {
|
||||
# "type": "image_url",
|
||||
# "image_url": {
|
||||
# "url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
# }
|
||||
# }
|
||||
# ]
|
||||
# }
|
||||
# ],
|
||||
# stream=True
|
||||
# )
|
||||
# print("async response gemini pro vision")
|
||||
# print(resp)
|
||||
# for chunk in resp:
|
||||
# print(chunk)
|
||||
# asyncio.run(test())
|
||||
# except Exception as e:
|
||||
# import traceback
|
||||
# traceback.print_exc()
|
||||
# raise e
|
||||
# test_gemini_pro_vision_async()
|
||||
|
||||
|
||||
def test_prompt_factory():
|
||||
messages = [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -848,45 +848,6 @@ def test_completion_function_plus_image(model):
|
|||
print(response)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model", ["gemini/gemini-1.5-pro"] # "claude-3-sonnet-20240229",
|
||||
)
|
||||
def test_completion_function_plus_pdf(model):
|
||||
litellm.set_verbose = True
|
||||
try:
|
||||
import base64
|
||||
|
||||
import requests
|
||||
|
||||
# URL of the file
|
||||
url = "https://storage.googleapis.com/cloud-samples-data/generative-ai/pdf/2403.05530.pdf"
|
||||
|
||||
# Download the file
|
||||
response = requests.get(url)
|
||||
file_data = response.content
|
||||
|
||||
encoded_file = base64.b64encode(file_data).decode("utf-8")
|
||||
|
||||
image_content = [
|
||||
{"type": "text", "text": "What's this file about?"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:application/pdf;base64,{encoded_file}"},
|
||||
},
|
||||
]
|
||||
image_message = {"role": "user", "content": image_content}
|
||||
|
||||
response = completion(
|
||||
model=model,
|
||||
messages=[image_message],
|
||||
stream=False,
|
||||
)
|
||||
|
||||
print(response)
|
||||
except litellm.InternalServerError:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"provider",
|
||||
["azure", "azure_ai"],
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue