refactor: simplify URL construction for Gemini image generation

- Remove unnecessary model name prefix stripping
- Directly use the model name in the API URL construction

This change streamlines the URL generation process for the Google AI API, ensuring compatibility with model names without the 'gemini/' prefix.

Signed-off-by: Ifta Khairul Alam Adil <ikaadil007@gmail.com>
This commit is contained in:
Ifta Khairul Alam Adil
2025-08-27 11:53:47 +02:00
parent 0f5b31fd78
commit 2d515e72ff
@@ -88,18 +88,15 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
Google AI API format: https://generativelanguage.googleapis.com/v1beta/models/{model}:predict
"""
from litellm.llms.gemini.common_utils import GeminiModelInfo
complete_url: str = (
api_base
or get_secret_str("GEMINI_API_BASE")
or self.DEFAULT_BASE_URL
)
# Strip the "gemini/" prefix from the model name for the API call
base_model = GeminiModelInfo.get_base_model(model) or model
complete_url = complete_url.rstrip("/")
complete_url = f"{complete_url}/models/{base_model}:predict"
complete_url = f"{complete_url}/models/{model}:predict"
return complete_url
def validate_environment(