mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-19 06:18:55 +00:00
Sameer Kankute
447502b409
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
408 lines
14 KiB
Python
408 lines
14 KiB
Python
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"
|
|
)
|