Merge pull request #26160 from BerriAI/litellm_vertex_image_edit_credentials_fix

fix(image_edit): forward litellm_params to validate_environment for Vertex AI credentials

litellm/llms/azure/image_edit/transformation.py

@@ -14,6 +14,8 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = (
             api_key

litellm/llms/azure_ai/image_edit/flux2_transformation.py

@@ -65,6 +65,8 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate Azure AI Foundry environment and set up authentication

litellm/llms/azure_ai/image_edit/transformation.py

@@ -25,6 +25,8 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate Azure AI Foundry environment and set up authentication

litellm/llms/base_llm/image_edit/transformation.py

@@ -67,6 +67,8 @@ class BaseImageEditConfig(ABC):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         return {}
 

litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py

@@ -483,6 +483,8 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         if headers is None:
             headers = {}

litellm/llms/bedrock/image_edit/stability_transformation.py

@@ -372,6 +372,8 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate environment for Bedrock Stability image edit.

litellm/llms/black_forest_labs/image_edit/transformation.py

@@ -123,6 +123,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate environment and set up headers for Black Forest Labs.

litellm/llms/custom_httpx/llm_http_handler.py

@@ -5515,6 +5515,8 @@ class BaseLLMHTTPHandler:
             api_key=litellm_params.api_key,
             headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
             model=model,
+            litellm_params=dict(litellm_params),
+            api_base=litellm_params.api_base,
         )
 
         if extra_headers:
@@ -5611,6 +5613,8 @@ class BaseLLMHTTPHandler:
             api_key=litellm_params.api_key,
             headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
             model=model,
+            litellm_params=dict(litellm_params),
+            api_base=litellm_params.api_base,
         )
 
         if extra_headers:

litellm/llms/gemini/image_edit/transformation.py

@@ -54,6 +54,8 @@ class GeminiImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         final_api_key: Optional[str] = api_key or get_secret_str("GEMINI_API_KEY")
         if not final_api_key:

litellm/llms/litellm_proxy/image_edit/transformation.py

@@ -8,7 +8,12 @@ class LiteLLMProxyImageEditConfig(OpenAIImageEditConfig):
     """Configuration for image edit requests routed through LiteLLM Proxy."""
 
     def validate_environment(
-        self, headers: dict, model: str, api_key: Optional[str] = None
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = api_key or get_secret_str("LITELLM_PROXY_API_KEY")
         headers.update({"Authorization": f"Bearer {api_key}"})

litellm/llms/openai/image_edit/transformation.py

@@ -165,6 +165,8 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = (
             api_key

litellm/llms/openrouter/image_edit/transformation.py

@@ -116,6 +116,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         api_key = api_key or litellm.api_key or get_secret_str("OPENROUTER_API_KEY")
         if not api_key:

litellm/llms/recraft/image_edit/transformation.py

@@ -81,6 +81,8 @@ class RecraftImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         final_api_key: Optional[str] = api_key or get_secret_str("RECRAFT_API_KEY")
         if not final_api_key:

litellm/llms/stability/image_edit/transformations.py

@@ -149,6 +149,8 @@ class StabilityImageEditConfig(BaseImageEditConfig):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         """
         Validate environment and set up headers for Stability AI.

litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py

@@ -103,10 +103,24 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
         headers: dict,
         model: str,
         api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         headers = headers or {}
-        vertex_project = self._resolve_vertex_project()
-        vertex_credentials = self._resolve_vertex_credentials()
+        litellm_params = litellm_params or {}
+
+        _api_base = litellm_params.get("api_base") or api_base
+        if _api_base is not None:
+            return headers
+
+        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,
@@ -123,8 +137,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
         """
         Get the complete URL for Vertex AI Imagen predict API
         """
-        vertex_project = self._resolve_vertex_project()
-        vertex_location = self._resolve_vertex_location()
+        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(

tests/test_litellm/images/test_image_edit_utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -26,7 +26,12 @@ class MockImageEditConfig(BaseImageEditConfig):
         return "https://example.com/api"
 
     def validate_environment(
-        self, headers: dict, model: str, api_key: str = None
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
     ) -> dict:
         return headers
 
@@ -262,3 +267,141 @@ class TestImageEditCustomPricing:
     def test_custom_pricing_not_detected_without_model_info(self):
         litellm_params = {"litellm_call_id": "test-call-id"}
         assert use_custom_pricing_for_model(litellm_params) is False
+
+
+class TestImageEditHandlerCredentialsForwarding:
+    """
+    Regression tests for Vertex AI image_edit credentials bug.
+
+    image_edit handler must forward litellm_params to validate_environment,
+    so that credentials passed via YAML config (vertex_ai_project,
+    vertex_ai_credentials, etc.) reach the auth layer instead of falling
+    through to Application Default Credentials.
+    """
+
+    def test_vertex_gemini_image_edit_reads_credentials_from_litellm_params(self):
+        """
+        VertexAIGeminiImageEditConfig.validate_environment should read
+        vertex_ai_project/vertex_ai_credentials from litellm_params first.
+        """
+        from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import (
+            VertexAIGeminiImageEditConfig,
+        )
+
+        config = VertexAIGeminiImageEditConfig()
+
+        litellm_params = {
+            "vertex_ai_project": "test-project-from-params",
+            "vertex_ai_credentials": "/path/to/creds.json",
+        }
+
+        with patch.object(
+            config, "_ensure_access_token", return_value=("token", "project")
+        ) as mock_ensure:
+            config.validate_environment(
+                headers={},
+                model="test-model",
+                litellm_params=litellm_params,
+            )
+
+            mock_ensure.assert_called_once()
+            call_kwargs = mock_ensure.call_args[1]
+
+            assert call_kwargs["credentials"] == "/path/to/creds.json"
+            assert call_kwargs["project_id"] == "test-project-from-params"
+
+    def test_vertex_imagen_image_edit_reads_credentials_from_litellm_params(self):
+        """
+        VertexAIImagenImageEditConfig.validate_environment should read
+        vertex_ai_project/vertex_ai_credentials from litellm_params first.
+        """
+        from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
+            VertexAIImagenImageEditConfig,
+        )
+
+        config = VertexAIImagenImageEditConfig()
+
+        litellm_params = {
+            "vertex_ai_project": "test-project-from-params",
+            "vertex_ai_credentials": "/path/to/creds.json",
+        }
+
+        with patch.object(
+            config, "_ensure_access_token", return_value=("token", "project")
+        ) as mock_ensure:
+            config.validate_environment(
+                headers={},
+                model="test-model",
+                litellm_params=litellm_params,
+            )
+
+            mock_ensure.assert_called_once()
+            call_kwargs = mock_ensure.call_args[1]
+
+            assert call_kwargs["credentials"] == "/path/to/creds.json"
+            assert call_kwargs["project_id"] == "test-project-from-params"
+
+    def test_vertex_imagen_get_complete_url_reads_project_and_location_from_litellm_params(
+        self,
+    ):
+        """
+        VertexAIImagenImageEditConfig.get_complete_url should read
+        vertex_ai_project and vertex_ai_location from litellm_params,
+        not only from env vars / global settings.
+        """
+        from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
+            VertexAIImagenImageEditConfig,
+        )
+
+        config = VertexAIImagenImageEditConfig()
+
+        litellm_params = {
+            "vertex_ai_project": "param-project",
+            "vertex_ai_location": "us-east1",
+        }
+
+        url = config.get_complete_url(
+            model="vertex_ai/imagegeneration@002",
+            api_base=None,
+            litellm_params=litellm_params,
+        )
+
+        assert "param-project" in url
+        assert "us-east1" in url
+
+    def test_validate_environment_signature_includes_litellm_params(self):
+        """
+        All image_edit config validate_environment methods should accept
+        litellm_params to allow credentials to be forwarded from the handler.
+        """
+        import inspect
+
+        from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import (
+            VertexAIGeminiImageEditConfig,
+        )
+        from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
+            VertexAIImagenImageEditConfig,
+        )
+        from litellm.llms.openai.image_edit.transformation import (
+            OpenAIImageEditConfig,
+        )
+
+        configs = [
+            VertexAIGeminiImageEditConfig(),
+            VertexAIImagenImageEditConfig(),
+            OpenAIImageEditConfig(),
+            MockImageEditConfig(),
+        ]
+
+        for config in configs:
+            sig = inspect.signature(config.validate_environment)
+            params = list(sig.parameters.keys())
+
+            assert "litellm_params" in params, (
+                f"{config.__class__.__name__}.validate_environment "
+                "missing litellm_params parameter"
+            )
+            assert "api_base" in params, (
+                f"{config.__class__.__name__}.validate_environment "
+                "missing api_base parameter"
+            )