From 00d628544ce9ef405c52fd78ef2437d7417afaf6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 08:59:55 -0800 Subject: [PATCH 1/9] fix(utils.py): fix vertex ai finish reason handling --- litellm/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 7de5199b46f..0e718d31a31 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8538,7 +8538,11 @@ class CustomStreamWrapper: if hasattr(chunk, "candidates") == True: try: completion_obj["content"] = chunk.text - if hasattr(chunk.candidates[0], "finish_reason"): + if ( + hasattr(chunk.candidates[0], "finish_reason") + and chunk.candidates[0].finish_reason.name + != "FINISH_REASON_UNSPECIFIED" + ): # every non-final chunk in vertex ai has this model_response.choices[0].finish_reason = ( map_finish_reason( chunk.candidates[0].finish_reason.name From 02095886816e6aab7eec1db171838c995eeca061 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 09:00:20 -0800 Subject: [PATCH 2/9] =?UTF-8?q?bump:=20version=201.27.5=20=E2=86=92=201.27?= =?UTF-8?q?.6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2dd37a4848a..367c6583c88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.27.5" +version = "1.27.6" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -74,7 +74,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.27.5" +version = "1.27.6" version_files = [ "pyproject.toml:^version" ] From e48fff47dd35f05a801d7b24a2d631937973d0b5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 09:18:46 -0800 Subject: [PATCH 3/9] test(test_custom_callback_input.py): assert async success called only once during vertex ai streaming --- litellm/llms/vertex_ai.py | 3 + litellm/tests/test_custom_callback_input.py | 75 +++++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py index fdbc1625e80..f4447a9e91d 100644 --- a/litellm/llms/vertex_ai.py +++ b/litellm/llms/vertex_ai.py @@ -1000,12 +1000,15 @@ async def async_streaming( if stream: response = TextStreamer(completion_response) + logging_obj.post_call(input=prompt, api_key=None, original_response=response) + streamwrapper = CustomStreamWrapper( completion_stream=response, model=model, custom_llm_provider="vertex_ai", logging_obj=logging_obj, ) + return streamwrapper diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index 5da46ffeeac..ca1fe19a964 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -600,6 +600,81 @@ async def test_async_chat_sagemaker_stream(): pytest.fail(f"An exception occurred: {str(e)}") +## Test Vertex AI + Async +import json +import tempfile + + +def load_vertex_ai_credentials(): + # Define the path to the vertex_key.json file + print("loading vertex ai credentials") + filepath = os.path.dirname(os.path.abspath(__file__)) + vertex_key_path = filepath + "/vertex_key.json" + + # Read the existing content of the file or create an empty dictionary + try: + with open(vertex_key_path, "r") as file: + # Read the file content + print("Read vertexai file path") + content = file.read() + + # If the file is empty or not valid JSON, create an empty dictionary + if not content or not content.strip(): + service_account_key_data = {} + else: + # Attempt to load the existing JSON content + file.seek(0) + service_account_key_data = json.load(file) + except FileNotFoundError: + # If the file doesn't exist, create an empty dictionary + service_account_key_data = {} + + # Update the service_account_key_data with environment variables + private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") + private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") + private_key = private_key.replace("\\n", "\n") + service_account_key_data["private_key_id"] = private_key_id + service_account_key_data["private_key"] = private_key + + # Create a temporary file + with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: + # Write the updated content to the temporary file + json.dump(service_account_key_data, temp_file, indent=2) + + # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) + + +@pytest.mark.asyncio +async def test_async_chat_vertex_ai_stream(): + try: + load_vertex_ai_credentials() + customHandler = CompletionCustomHandler() + litellm.callbacks = [customHandler] + # test streaming + response = await litellm.acompletion( + model="gemini-pro", + messages=[ + { + "role": "user", + "content": f"Hi 👋 - i'm async vertex_ai {uuid.uuid4()}", + } + ], + stream=True, + ) + print(f"response: {response}") + async for chunk in response: + print(f"chunk: {chunk}") + continue + print(f"customHandler.states: {customHandler.states}") + assert ( + customHandler.states.count("async_success") == 1 + ) # pre, post, success, pre, post, failure + assert len(customHandler.states) >= 3 # pre, post, success + except Exception as e: + pytest.fail(f"An exception occurred: {str(e)}") + + # Text Completion From 1ac1f0a8365b1b44942b87b3a90102cb32761941 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 09:25:25 -0800 Subject: [PATCH 4/9] test(test_batch_completions.py): change flaky model --- litellm/tests/test_batch_completions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_batch_completions.py b/litellm/tests/test_batch_completions.py index 55e3084b4f4..485f51757b9 100644 --- a/litellm/tests/test_batch_completions.py +++ b/litellm/tests/test_batch_completions.py @@ -24,7 +24,7 @@ from litellm import ( def test_batch_completions(): messages = [[{"role": "user", "content": "write a short poem"}] for _ in range(3)] - model = "j2-mid" + model = "gpt-3.5-turbo" litellm.set_verbose = True try: result = batch_completion( @@ -44,7 +44,7 @@ def test_batch_completions(): pytest.fail(f"An error occurred: {e}") -test_batch_completions() +# test_batch_completions() def test_batch_completions_models(): From 144254bd9dc13db04a6dbd7d69597e76598a5f72 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 09:32:37 -0800 Subject: [PATCH 5/9] test(test_amazing_vertex_completion.py): skip gemini 1.5 as we don't have access to it --- litellm/tests/test_amazing_vertex_completion.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 76ebde7aefd..14dfa14b683 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -206,6 +206,8 @@ async def test_async_vertexai_response(): "code-gecko@latest", "code-bison@001", "text-bison@001", + "gemini-1.5-pro", + "gemini-1.5-pro-vision", ]: # our account does not have access to this model continue From f460da4770f7784e130ae6968147a5763f155ff6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 10:38:49 -0800 Subject: [PATCH 6/9] test: testing fixes --- litellm/tests/test_amazing_vertex_completion.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 14dfa14b683..bca241cd703 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -109,6 +109,8 @@ def test_vertex_ai(): "code-gecko@latest", "code-bison@001", "text-bison@001", + "gemini-1.5-pro", + "gemini-1.5-pro-vision", ]: # our account does not have access to this model continue From 28f4b5809c60f171ce3639e9e47bffec72c9603e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 10:42:05 -0800 Subject: [PATCH 7/9] test(test_amazing_vertex_completion.py): fix test --- litellm/main.py | 1 - litellm/tests/test_amazing_vertex_completion.py | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index ed8dddf05b2..33fad52cce8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -10,7 +10,6 @@ import os, openai, sys, json, inspect, uuid, datetime, threading from typing import Any, Literal, Union from functools import partial - import dotenv, traceback, random, asyncio, time, contextvars from copy import deepcopy import httpx diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index bca241cd703..35beb75fe6f 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -250,6 +250,8 @@ async def test_async_vertexai_streaming_response(): "code-gecko@latest", "code-bison@001", "text-bison@001", + "gemini-1.5-pro", + "gemini-1.5-pro-vision", ]: # our account does not have access to this model continue From a1c6e6d52be444ee718f9c5f0b456e9758986730 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 10:44:24 -0800 Subject: [PATCH 8/9] build(main.py): trigger new build --- litellm/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/main.py b/litellm/main.py index 33fad52cce8..ed8dddf05b2 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -10,6 +10,7 @@ import os, openai, sys, json, inspect, uuid, datetime, threading from typing import Any, Literal, Union from functools import partial + import dotenv, traceback, random, asyncio, time, contextvars from copy import deepcopy import httpx From 6cce9213d80f0f057df7e632fbf63019727cf213 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 26 Feb 2024 10:47:01 -0800 Subject: [PATCH 9/9] fix(main.py): refactor --- litellm/main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index ed8dddf05b2..33fad52cce8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -10,7 +10,6 @@ import os, openai, sys, json, inspect, uuid, datetime, threading from typing import Any, Literal, Union from functools import partial - import dotenv, traceback, random, asyncio, time, contextvars from copy import deepcopy import httpx