phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-28 04:04:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(vertex_httpx.py): fix supports system message check for vertex_ai_beta

Browse source 
Fixes https://github.com/BerriAI/litellm/issues/4283
This commit is contained in:
Krrish Dholakia 2024-06-19 13:17:22 -07:00
parent 7966266d4f
commit 0a96865c48
2 changed files with 8 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/llms/vertex_httpx.py
View file

@ -709,8 +709,11 @@ class VertexLLM(BaseLLM):
## TRANSFORMATION ## ## TRANSFORMATION ##
try: 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( 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: except Exception as e:
verbose_logger.error( verbose_logger.error(

5
litellm/tests/test_amazing_vertex_completion.py
View file

@ -529,6 +529,7 @@ async def test_gemini_pro_vision(provider, sync_mode):
resp = litellm.completion( resp = litellm.completion(
model="{}/gemini-1.5-flash-preview-0514".format(provider), model="{}/gemini-1.5-flash-preview-0514".format(provider),
messages=[ messages=[
{"role": "system", "content": "Be a good bot"},
{ {
"role": "user", "role": "user",
"content": [ "content": [
@ -540,7 +541,7 @@ async def test_gemini_pro_vision(provider, sync_mode):
}, },
}, },
], ],
} },
], ],
) )
else: else:
@ -568,6 +569,8 @@ async def test_gemini_pro_vision(provider, sync_mode):
# DO Not DELETE this ASSERT # DO Not DELETE this ASSERT
# Google counts the prompt tokens for us, we should ensure we use the tokens from the orignal response # 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 prompt_tokens == 263 # the gemini api returns 263 to us
assert False
except litellm.RateLimitError as e: except litellm.RateLimitError as e:
pass pass
except Exception as e: except Exception as e: