_get_model_name_from_gemini_spec_model

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"

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/<unique_id>"
+
+        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