From c89aa1e626d5f461fb61e7c3040e50cc3a39a7f3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 8 Jul 2024 10:01:47 -0700 Subject: [PATCH] test(test_amazing_vertex_completion.py): pass pdf as base64 to vertex ai --- .../tests/test_amazing_vertex_completion.py | 154 +++++------------- litellm/tests/test_completion.py | 39 ----- 2 files changed, 37 insertions(+), 156 deletions(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 9c11a42484c..35ba35d8fd0 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -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 = [ { diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index ec6a32d8bcc..0598c52dfee 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -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"],