From 81c77f33b8e5f425ec47dec6365a8c4888a23f80 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 18 Jul 2024 18:16:00 -0700 Subject: [PATCH] fix test amazing vertex medlm --- .../tests/test_amazing_vertex_completion.py | 66 +++++++------------ 1 file changed, 22 insertions(+), 44 deletions(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 6a381022e..b8ba54cb4 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -36,6 +36,20 @@ litellm.cache = None user_message = "Write a short poem about the sky" messages = [{"content": user_message, "role": "user"}] +VERTEX_MODELS_TO_NOT_TEST = [ + "medlm-medium", + "medlm-large", + "code-gecko", + "code-gecko@001", + "code-gecko@002", + "code-gecko@latest", + "codechat-bison@latest", + "code-bison@001", + "text-bison@001", + "gemini-1.5-pro", + "gemini-1.5-pro-preview-0215", +] + def get_vertex_ai_creds_json() -> dict: # Define the path to the vertex_key.json file @@ -327,17 +341,7 @@ def test_vertex_ai(): test_models += litellm.vertex_language_models # always test gemini-pro for model in test_models: try: - if model in [ - "code-gecko", - "code-gecko@001", - "code-gecko@002", - "code-gecko@latest", - "codechat-bison@latest", - "code-bison@001", - "text-bison@001", - "gemini-1.5-pro", - "gemini-1.5-pro-preview-0215", - ] or ( + if model in VERTEX_MODELS_TO_NOT_TEST or ( "gecko" in model or "32k" in model or "ultra" in model or "002" in model ): # our account does not have access to this model @@ -382,17 +386,7 @@ def test_vertex_ai_stream(): test_models += litellm.vertex_language_models # always test gemini-pro for model in test_models: try: - if model in [ - "code-gecko", - "code-gecko@001", - "code-gecko@002", - "code-gecko@latest", - "codechat-bison@latest", - "code-bison@001", - "text-bison@001", - "gemini-1.5-pro", - "gemini-1.5-pro-preview-0215", - ] or ( + if model in VERTEX_MODELS_TO_NOT_TEST or ( "gecko" in model or "32k" in model or "ultra" in model or "002" in model ): # our account does not have access to this model @@ -437,17 +431,9 @@ async def test_async_vertexai_response(): test_models += litellm.vertex_language_models # always test gemini-pro for model in test_models: print(f"model being tested in async call: {model}") - if model in [ - "code-gecko", - "code-gecko@001", - "code-gecko@002", - "code-gecko@latest", - "codechat-bison@latest", - "code-bison@001", - "text-bison@001", - "gemini-1.5-pro", - "gemini-1.5-pro-preview-0215", - ] or ("gecko" in model or "32k" in model or "ultra" in model or "002" in model): + if model in VERTEX_MODELS_TO_NOT_TEST or ( + "gecko" in model or "32k" in model or "ultra" in model or "002" in model + ): # our account does not have access to this model continue try: @@ -484,17 +470,9 @@ async def test_async_vertexai_streaming_response(): test_models = random.sample(test_models, 1) test_models += litellm.vertex_language_models # always test gemini-pro for model in test_models: - if model in [ - "code-gecko", - "code-gecko@001", - "code-gecko@002", - "code-gecko@latest", - "codechat-bison@latest", - "code-bison@001", - "text-bison@001", - "gemini-1.5-pro", - "gemini-1.5-pro-preview-0215", - ] or ("gecko" in model or "32k" in model or "ultra" in model or "002" in model): + if model in VERTEX_MODELS_TO_NOT_TEST or ( + "gecko" in model or "32k" in model or "ultra" in model or "002" in model + ): # our account does not have access to this model continue try: