From 93daf5cbac2af4cc632b47acd6f3c90cd8275b60 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 12:16:18 -0700 Subject: [PATCH] _get_model_name_from_gemini_spec_model --- litellm/llms/vertex_ai/common_utils.py | 4 +++- .../gemini/vertex_and_google_ai_studio_gemini.py | 15 ++++++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index bbaab31401..b5bad6b858 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -31,7 +31,7 @@ def get_supports_system_message( ) # Vertex Models called in the `/gemini` request/response format also support system messages - if litellm.VertexGeminiConfig._is_model_gemini_spec_model(model): + if litellm.VertexGeminiConfig._is_model_gemini_gemini_spec_model(model): supports_system_message = True except Exception as e: verbose_logger.warning( @@ -75,6 +75,8 @@ def _get_vertex_url( ) -> Tuple[str, str]: url: Optional[str] = None endpoint: Optional[str] = None + if litellm.VertexGeminiConfig._is_model_gemini_gemini_spec_model(model): + model = litellm.VertexGeminiConfig._get_model_name_from_gemini_spec_model(model) if mode == "chat": ### SET RUNTIME ENDPOINT ### endpoint = "generateContent" diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index fdb8c0e558..0e48c690ba 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -420,7 +420,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ] @staticmethod - def _is_model_gemini_spec_model(model: Optional[str]) -> bool: + def _is_model_gemini_gemini_spec_model(model: Optional[str]) -> bool: """ Returns true if user is trying to call custom model in `/gemini` request/response format """ @@ -430,6 +430,19 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return True return False + @staticmethod + def _get_model_name_from_gemini_spec_model(model: str) -> str: + """ + Returns the model name if model="vertex_ai/gemini/" + + Example: + - model = "gemini/1234567890" + - returns "1234567890" + """ + if "gemini/" in model: + return model.split("/")[-1] + return model + def get_flagged_finish_reasons(self) -> Dict[str, str]: """ Return Dictionary of finish reasons which indicate response was flagged