fix get litellm params

This commit is contained in:
Ishaan Jaffer 2025-11-26 13:20:05 -08:00
parent 0437be13d3
commit 9aef28fe70
2 changed files with 12 additions and 4 deletions

View file

@ -118,6 +118,10 @@ def get_litellm_params(
"bucket_name": kwargs.get("bucket_name"),
"vertex_credentials": kwargs.get("vertex_credentials"),
"vertex_project": kwargs.get("vertex_project"),
"vertex_location": kwargs.get("vertex_location"),
"vertex_ai_project": kwargs.get("vertex_ai_project"),
"vertex_ai_location": kwargs.get("vertex_ai_location"),
"vertex_ai_credentials": kwargs.get("vertex_ai_credentials"),
"use_litellm_proxy": use_litellm_proxy,
"litellm_request_debug": litellm_request_debug,
"aws_region_name": kwargs.get("aws_region_name"),

View file

@ -122,8 +122,10 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
"""
Get the complete URL for Vertex AI Gemini generateContent API
"""
vertex_project = self._resolve_vertex_project()
vertex_location = self._resolve_vertex_location()
# First check litellm_params (where vertex_ai_project/vertex_ai_location are passed)
# then fall back to environment variables and other sources
vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project()
vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location()
if not vertex_project or not vertex_location:
raise ValueError("vertex_project and vertex_location are required for Vertex AI")
@ -151,8 +153,10 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
api_base: Optional[str] = None,
) -> dict:
headers = headers or {}
vertex_project = self._resolve_vertex_project()
vertex_credentials = self._resolve_vertex_credentials()
# First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed)
# then fall back to environment variables and other sources
vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project()
vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials()
access_token, _ = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,