rename _is_model_gemini_spec_model

This commit is contained in:
Ishaan Jaff 2025-03-26 14:28:26 -07:00
parent 64c2a46647
commit f77fcbcad3
3 changed files with 7 additions and 10 deletions

View file

@ -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(

View file

@ -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
"""

View file

@ -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():