litellm/tests/unit/images/test_image_edit_utils.py

from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock, patch

import pytest

import litellm
from litellm.images.utils import ImageEditRequestUtils
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.types.images.main import ImageEditOptionalRequestParams


class MockImageEditConfig(BaseImageEditConfig):
    def get_supported_openai_params(self, model: str) -> List[str]:
        return ["size", "quality"]

    def map_openai_params(
        self,
        image_edit_optional_params: ImageEditOptionalRequestParams,
        model: str,
        drop_params: bool,
    ) -> Dict[str, Any]:
        return dict(image_edit_optional_params)

    def get_complete_url(self, model: str, api_base: str, litellm_params: dict) -> str:
        return "https://example.com/api"

    def validate_environment(
        self,
        headers: dict,
        model: str,
        api_key: Optional[str] = None,
        litellm_params: Optional[dict] = None,
        api_base: Optional[str] = None,
    ) -> dict:
        return headers

    def transform_image_edit_request(self, *args, **kwargs):
        return {}, []

    def transform_image_edit_response(self, *args, **kwargs):
        return MagicMock()


class TestImageEditRequestUtilsDropParams:
    def setup_method(self):
        self.config = MockImageEditConfig()
        self.model = "test-model"
        self._original_drop_params = getattr(litellm, "drop_params", None)

    def teardown_method(self):
        if self._original_drop_params is None:
            if hasattr(litellm, "drop_params"):
                delattr(litellm, "drop_params")
        else:
            litellm.drop_params = self._original_drop_params

    def test_unsupported_params_raises_without_drop(self):
        litellm.drop_params = False
        optional_params: ImageEditOptionalRequestParams = {
            "size": "1024x1024",
            "unsupported_param": "value",
        }

        with pytest.raises(litellm.UnsupportedParamsError) as exc_info:
            ImageEditRequestUtils.get_optional_params_image_edit(
                model=self.model,
                image_edit_provider_config=self.config,
                image_edit_optional_params=optional_params,
            )

        assert "unsupported_param" in str(exc_info.value)

    def test_drop_params_global_setting(self):
        litellm.drop_params = True
        optional_params: ImageEditOptionalRequestParams = {
            "size": "1024x1024",
            "unsupported_param": "value",
        }

        result = ImageEditRequestUtils.get_optional_params_image_edit(
            model=self.model,
            image_edit_provider_config=self.config,
            image_edit_optional_params=optional_params,
        )

        assert "size" in result
        assert "unsupported_param" not in result

    def test_drop_params_explicit_parameter(self):
        litellm.drop_params = False
        optional_params: ImageEditOptionalRequestParams = {
            "size": "1024x1024",
            "unsupported_param": "value",
        }

        result = ImageEditRequestUtils.get_optional_params_image_edit(
            model=self.model,
            image_edit_provider_config=self.config,
            image_edit_optional_params=optional_params,
            drop_params=True,
        )

        assert "size" in result
        assert "unsupported_param" not in result

    def test_additional_drop_params(self):
        litellm.drop_params = False
        optional_params: ImageEditOptionalRequestParams = {
            "size": "1024x1024",
            "quality": "high",
        }

        result = ImageEditRequestUtils.get_optional_params_image_edit(
            model=self.model,
            image_edit_provider_config=self.config,
            image_edit_optional_params=optional_params,
            additional_drop_params=["quality"],
        )

        assert "size" in result
        assert "quality" not in result

    def test_drop_params_false_with_global_true(self):
        litellm.drop_params = True
        optional_params: ImageEditOptionalRequestParams = {
            "size": "1024x1024",
            "unsupported_param": "value",
        }

        result = ImageEditRequestUtils.get_optional_params_image_edit(
            model=self.model,
            image_edit_provider_config=self.config,
            image_edit_optional_params=optional_params,
            drop_params=False,
        )

        assert "size" in result
        assert "unsupported_param" not in result

    def test_supported_params_pass_through(self):
        litellm.drop_params = False
        optional_params: ImageEditOptionalRequestParams = {
            "size": "1024x1024",
            "quality": "high",
        }

        result = ImageEditRequestUtils.get_optional_params_image_edit(
            model=self.model,
            image_edit_provider_config=self.config,
            image_edit_optional_params=optional_params,
        )

        assert result["size"] == "1024x1024"
        assert result["quality"] == "high"

    def test_additional_drop_params_with_unsupported_and_drop_true(self):
        litellm.drop_params = True
        optional_params: ImageEditOptionalRequestParams = {
            "size": "1024x1024",
            "quality": "high",
            "unsupported_param": "value",
        }

        result = ImageEditRequestUtils.get_optional_params_image_edit(
            model=self.model,
            image_edit_provider_config=self.config,
            image_edit_optional_params=optional_params,
            additional_drop_params=["quality"],
        )

        assert "size" in result
        assert "quality" not in result
        assert "unsupported_param" not in result


class TestImageEditCustomPricing:
    """
    Regression tests for https://github.com/BerriAI/litellm/issues/22244

    image_edit must forward model_info and metadata into litellm_params
    when calling update_environment_variables, so that custom pricing
    detection works after PR #20679 stripped custom pricing fields from
    the shared backend model key.
    """

    def test_image_edit_passes_model_info_to_logging(self):
        """
        When the router provides model_info with custom pricing fields,
        image_edit should include model_info and metadata in litellm_params.
        """
        from litellm.images.main import image_edit

        custom_model_info = {
            "id": "test-deployment-id",
            "input_cost_per_image": 0.00676128,
            "mode": "image_generation",
        }
        custom_metadata = {
            "model_info": custom_model_info,
        }

        captured_litellm_params = {}

        mock_logging_obj = MagicMock()
        mock_logging_obj.model_call_details = {}

        original_update = mock_logging_obj.update_from_kwargs

        def capturing_update(**update_kwargs):
            captured_litellm_params.update(update_kwargs.get("litellm_params", {}))
            inner_kwargs = update_kwargs.get("kwargs", {})
            if "metadata" in inner_kwargs:
                captured_litellm_params["metadata"] = inner_kwargs["metadata"]
            return original_update(**update_kwargs)

        mock_logging_obj.update_from_kwargs = capturing_update

        with (
            patch(
                "litellm.images.main.get_llm_provider",
                return_value=("test-model", "openai", None, None),
            ),
            patch(
                "litellm.images.main.ProviderConfigManager.get_provider_image_edit_config",
                return_value=MagicMock(),
            ),
            patch(
                "litellm.images.main._get_ImageEditRequestUtils",
                return_value=MagicMock(
                    get_requested_image_edit_optional_param=MagicMock(return_value={}),
                    get_optional_params_image_edit=MagicMock(return_value={}),
                ),
            ),
            patch("litellm.images.main.base_llm_http_handler") as mock_handler,
        ):
            mock_handler.image_edit_handler.return_value = MagicMock()

            try:
                image_edit(
                    image=b"fake-image-data",
                    prompt="test prompt",
                    model="openai/test-model",
                    litellm_logging_obj=mock_logging_obj,
                    model_info=custom_model_info,
                    metadata=custom_metadata,
                )
            except Exception:
                pass

        assert "model_info" in captured_litellm_params
        assert captured_litellm_params["model_info"] == custom_model_info
        assert "metadata" in captured_litellm_params
        assert captured_litellm_params["metadata"] == custom_metadata

    def test_custom_pricing_detected_from_model_info_in_metadata(self):
        litellm_params = {
            "metadata": {
                "model_info": {
                    "id": "deployment-id",
                    "input_cost_per_image": 0.00676128,
                },
            },
        }
        assert use_custom_pricing_for_model(litellm_params) is True

    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"
            )