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