From 447502b409ebd68c10f4c46f3dcafa0c8763683d Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Tue, 21 Apr 2026 15:03:40 +0530
Subject: [PATCH] fix(image_edit): read vertex_project/location from
 litellm_params in Imagen get_complete_url

VertexAIImagenImageEditConfig.get_complete_url was resolving vertex_project
and vertex_location only from env vars and global settings, ignoring
litellm_params. Users supplying project/location exclusively via YAML
config would get a ValueError or wrong URL even after auth headers were fixed.

Mirrors the pattern already used by VertexAIGeminiImageEditConfig and
image_generation counterpart (safe_get_vertex_ai_project/location).

Also fixes api_key type hint in MockImageEditConfig (str -> Optional[str])
and adds a test covering get_complete_url credential resolution.

Made-with: Cursor
---
 .../vertex_imagen_transformation.py           | 10 +++++--
 .../images/test_image_edit_utils.py           | 30 ++++++++++++++++++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py
index 11126b2a3e..9c0b07b827 100644
--- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py
+++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py
@@ -137,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(
diff --git a/tests/test_litellm/images/test_image_edit_utils.py b/tests/test_litellm/images/test_image_edit_utils.py
index 1a13dd0671..2146c1fab0 100644
--- a/tests/test_litellm/images/test_image_edit_utils.py
+++ b/tests/test_litellm/images/test_image_edit_utils.py
@@ -29,7 +29,7 @@ class MockImageEditConfig(BaseImageEditConfig):
         self,
         headers: dict,
         model: str,
-        api_key: str = None,
+        api_key: Optional[str] = None,
         litellm_params: Optional[dict] = None,
         api_base: Optional[str] = None,
     ) -> dict:
@@ -341,6 +341,34 @@ class TestImageEditHandlerCredentialsForwarding:
             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