From f77fcbcad37bc25b39bbe1f6fd254a95554fed2b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 14:28:26 -0700 Subject: [PATCH] rename _is_model_gemini_spec_model --- litellm/llms/vertex_ai/common_utils.py | 2 +- .../gemini/vertex_and_google_ai_studio_gemini.py | 4 ++-- .../gemini/test_vertex_and_google_ai_studio.py | 11 ++++------- 3 files changed, 7 insertions(+), 10 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 2361939d81..9c9db2f047 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_gemini_spec_model(model): + if litellm.VertexGeminiConfig._is_model_gemini_spec_model(model): supports_system_message = True except Exception as e: verbose_logger.warning( 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 8aecfffd86..2931b1b3a8 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 @@ -434,12 +434,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Returns: str: The model name to use in the request to Vertex AI """ - if VertexGeminiConfig._is_model_gemini_gemini_spec_model(model): + if VertexGeminiConfig._is_model_gemini_spec_model(model): return VertexGeminiConfig._get_model_name_from_gemini_spec_model(model) return model @staticmethod - def _is_model_gemini_gemini_spec_model(model: Optional[str]) -> bool: + def _is_model_gemini_spec_model(model: Optional[str]) -> bool: """ Returns true if user is trying to call custom model in `/gemini` request/response format """ diff --git a/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py b/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py index 9fc9be7e0f..18d965200c 100644 --- a/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py +++ b/tests/litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio.py @@ -45,18 +45,15 @@ def test_get_model_for_vertex_ai_url(): assert result == "ft-uuid-123" -def test_is_model_gemini_gemini_spec_model(): +def test_is_model_gemini_spec_model(): # Test case 1: None input - assert VertexGeminiConfig._is_model_gemini_gemini_spec_model(None) == False + assert VertexGeminiConfig._is_model_gemini_spec_model(None) == False # Test case 2: Regular model name - assert VertexGeminiConfig._is_model_gemini_gemini_spec_model("gemini-pro") == False + assert VertexGeminiConfig._is_model_gemini_spec_model("gemini-pro") == False # Test case 3: Gemini spec model - assert ( - VertexGeminiConfig._is_model_gemini_gemini_spec_model("gemini/custom-model") - == True - ) + assert VertexGeminiConfig._is_model_gemini_spec_model("gemini/custom-model") == True def test_get_model_name_from_gemini_spec_model():