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 # 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 supports_system_message = True
except Exception as e: except Exception as e:
verbose_logger.warning( verbose_logger.warning(

View file

@ -434,12 +434,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
Returns: Returns:
str: The model name to use in the request to Vertex AI 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 VertexGeminiConfig._get_model_name_from_gemini_spec_model(model)
return model return model
@staticmethod @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 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" 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 # 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 # 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 # Test case 3: Gemini spec model
assert ( assert VertexGeminiConfig._is_model_gemini_spec_model("gemini/custom-model") == True
VertexGeminiConfig._is_model_gemini_gemini_spec_model("gemini/custom-model")
== True
)
def test_get_model_name_from_gemini_spec_model(): def test_get_model_name_from_gemini_spec_model():