From 933c3c845e8b923092306237d57a7ae686c5e8d5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 19 Jun 2024 13:17:22 -0700 Subject: [PATCH] fix(vertex_httpx.py): fix supports system message check for vertex_ai_beta Fixes https://github.com/BerriAI/litellm/issues/4283 --- litellm/llms/vertex_httpx.py | 5 ++++- litellm/tests/test_amazing_vertex_completion.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index f412f670896..01fef7e98bc 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -709,8 +709,11 @@ class VertexLLM(BaseLLM): ## TRANSFORMATION ## try: + _custom_llm_provider = custom_llm_provider + if custom_llm_provider == "vertex_ai_beta": + _custom_llm_provider = "vertex_ai" supports_system_message = litellm.supports_system_messages( - model=model, custom_llm_provider=custom_llm_provider + model=model, custom_llm_provider=_custom_llm_provider ) except Exception as e: verbose_logger.error( diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 2e50b48cee6..4edc580a9b6 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -529,6 +529,7 @@ async def test_gemini_pro_vision(provider, sync_mode): resp = litellm.completion( model="{}/gemini-1.5-flash-preview-0514".format(provider), messages=[ + {"role": "system", "content": "Be a good bot"}, { "role": "user", "content": [ @@ -540,7 +541,7 @@ async def test_gemini_pro_vision(provider, sync_mode): }, }, ], - } + }, ], ) else: @@ -568,6 +569,8 @@ async def test_gemini_pro_vision(provider, sync_mode): # 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 == 263 # the gemini api returns 263 to us + + assert False except litellm.RateLimitError as e: pass except Exception as e: