diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 67d90d027ec..4ca577e9b69 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -1,32 +1,22 @@ """ -OpenRouter Image Generation Support +OpenRouter image generation through POST {api_base}/images -OpenRouter provides image generation through chat completion endpoints. -Models like google/gemini-2.5-flash-image return images in the message content. - -Response format: +Response shape: { - "choices": [{ - "message": { - "content": "Here is a beautiful sunset for you! ", - "role": "assistant", - "images": [{ - "image_url": {"url": "data:image/png;base64,..."}, - "index": 0, - "type": "image_url" - }] - } - }], + "created": 1790994420, + "data": [{"b64_json": "...", "media_type": "image/png"}], "usage": { - "completion_tokens": 1299, - "prompt_tokens": 6, - "total_tokens": 1305, - "completion_tokens_details": {"image_tokens": 1290}, - "cost": 0.0387243 + "prompt_tokens": 18, + "completion_tokens": 272, + "total_tokens": 290, + "cost": 0.002212, + "cost_details": {"upstream_inference_cost": 0.002212, ...}, + "completion_tokens_details": {"image_tokens": 272} } } """ +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx @@ -55,23 +45,23 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +OPENROUTER_API_BASE: Final = "https://openrouter.ai/api/v1" +IMAGES_PATH: Final = "/images" +LEGACY_CHAT_COMPLETIONS_SUFFIX: Final = "/chat/completions" +QUALITY_ALIASES: Final = MappingProxyType({"standard": "low", "hd": "high"}) +LEGACY_IMAGE_CONFIG_FIELDS: Final = MappingProxyType({"aspect_ratio": "aspect_ratio", "image_size": "resolution"}) +NON_BODY_PARAMS: Final = frozenset( + {"model", "prompt", "messages", "modalities", "stream", "image_config", "extra_headers"} +) + class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): """ - Configuration for OpenRouter image generation via chat completions. - - OpenRouter uses chat completion endpoints for image generation, - so we need to transform image generation requests to chat format - and extract images from chat responses. + OpenRouter image generation through the dedicated /images endpoint, which serves both + image-only models (openai/gpt-image-*) and image+text models (google/gemini-*-image) """ def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - """ - Get supported OpenAI parameters for OpenRouter image generation. - - Since OpenRouter uses chat completions for image generation, - we support standard image generation params. - """ return [ "size", "quality", @@ -86,104 +76,16 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): drop_params: bool, ) -> dict: """ - Map image generation params to OpenRouter chat completion format. - - Maps OpenAI parameters to OpenRouter's image_config format: - - size -> image_config.aspect_ratio - - quality -> image_config.image_size + size and n pass through as is: /images takes explicit pixel sizes and normalizes them per + provider. quality is native on /images, so only the dall-e-3 names are translated """ supported_params: Final = self.get_supported_openai_params(model) - - for key, value in non_default_params.items(): - if key in supported_params: - if key == "size": - # Map OpenAI size to OpenRouter aspect_ratio - aspect_ratio = self._map_size_to_aspect_ratio(value) - if "image_config" not in optional_params: - optional_params["image_config"] = {} - optional_params["image_config"]["aspect_ratio"] = aspect_ratio - elif key == "quality": - # Map OpenAI quality to OpenRouter image_size - image_size = self._map_quality_to_image_size(value) - if image_size: - if "image_config" not in optional_params: - optional_params["image_config"] = {} - optional_params["image_config"]["image_size"] = image_size - else: - # Pass through other supported params (like n) - optional_params[key] = value - elif not drop_params: - # If not supported and drop_params is False, pass through - optional_params[key] = value - - return optional_params - - def _map_size_to_aspect_ratio(self, size: str) -> str: - """ - Map OpenAI size format to OpenRouter aspect_ratio format. - - OpenAI sizes: - - 1024x1024 (square) - - 1536x1024 (landscape) - - 1024x1536 (portrait) - - 1792x1024 (wide landscape, dall-e-3) - - 1024x1792 (tall portrait, dall-e-3) - - 256x256, 512x512 (dall-e-2) - - auto (default) - - OpenRouter aspect_ratios: - - 1:1 → 1024×1024 (default) - - 2:3 → 832×1248 - - 3:2 → 1248×832 - - 3:4 → 864×1184 - - 4:3 → 1184×864 - - 4:5 → 896×1152 - - 5:4 → 1152×896 - - 9:16 → 768×1344 - - 16:9 → 1344×768 - - 21:9 → 1536×672 - """ - size_to_aspect_ratio: Final = { - # Square formats - "256x256": "1:1", - "512x512": "1:1", - "1024x1024": "1:1", - # Landscape formats - "1536x1024": "3:2", # 1.5:1 ratio, closest to 3:2 - "1792x1024": "16:9", # 1.75:1 ratio, closest to 16:9 - # Portrait formats - "1024x1536": "2:3", # 0.67:1 ratio, closest to 2:3 - "1024x1792": "9:16", # 0.57:1 ratio, closest to 9:16 - # Default - "auto": "1:1", + mapped_params: Final = { + key: QUALITY_ALIASES.get(value, value) if key == "quality" else value + for key, value in non_default_params.items() + if (key in supported_params or not drop_params) and (key, value) != ("size", "auto") } - return size_to_aspect_ratio.get(size, "1:1") - - def _map_quality_to_image_size(self, quality: str) -> str | None: - """ - Map OpenAI quality to OpenRouter image_size format. - - OpenAI quality values: - - auto (default) - automatically select best quality - - high, medium, low - for GPT image models - - hd, standard - for dall-e-3 - - OpenRouter image_size values (Gemini only): - - 1K → Standard resolution (default) - - 2K → Higher resolution - - 4K → Highest resolution - """ - quality_to_image_size: Final = { - # OpenAI quality mappings - "low": "1K", - "standard": "1K", - "medium": "2K", - "high": "4K", - "hd": "4K", - # Auto defaults to standard - "auto": "1K", - } - return quality_to_image_size.get(quality) + return {**optional_params, **mapped_params} def _set_usage_and_cost( self, @@ -244,19 +146,10 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - """ - Get the complete URL for OpenRouter image generation. - - OpenRouter uses chat completions endpoint for image generation. - Default: https://openrouter.ai/api/v1/chat/completions - """ - if api_base: - if not api_base.endswith("/chat/completions"): - api_base = api_base.rstrip("/") - return f"{api_base}/chat/completions" - return api_base - - return "https://openrouter.ai/api/v1/chat/completions" + base_url: Final = (api_base or OPENROUTER_API_BASE).rstrip("/") + if base_url.endswith(IMAGES_PATH): + return base_url + return base_url.removesuffix(LEGACY_CHAT_COMPLETIONS_SUFFIX) + IMAGES_PATH def validate_environment( self, @@ -285,30 +178,21 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): headers: dict, ) -> dict: """ - Transform image generation request to OpenRouter chat completion format. - - Args: - model: The model name - prompt: The image generation prompt - optional_params: Optional parameters (including image_config) - litellm_params: LiteLLM parameters - headers: Request headers - - Returns: - dict: Request body in chat completion format with image_config + image_config is the request shape of the older chat-based path. Its fields map onto the + /images names so existing configs keep working, and explicit top-level values win """ - request_body: Final = { + legacy_image_config: Final = optional_params.get("image_config") or {} + return { "model": model, - "messages": [{"role": "user", "content": prompt}], + "prompt": prompt, + **{ + LEGACY_IMAGE_CONFIG_FIELDS[key]: value + for key, value in legacy_image_config.items() + if key in LEGACY_IMAGE_CONFIG_FIELDS + }, + **{key: value for key, value in optional_params.items() if key not in NON_BODY_PARAMS}, } - # These will be passed through to OpenRouter - for key, value in optional_params.items(): - if key not in ["model", "messages", "modalities"]: - request_body[key] = value - - return request_body - def transform_image_generation_response( self, model: str, @@ -322,83 +206,21 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: - """ - Transform OpenRouter chat completion response to ImageResponse format. - - Extracts images from the message content and maps usage/cost information. - - Args: - model: The model name - raw_response: Raw HTTP response from OpenRouter - model_response: ImageResponse object to populate - logging_obj: Logging object - request_data: Original request data - optional_params: Optional parameters - litellm_params: LiteLLM parameters - encoding: Encoding - api_key: API key - json_mode: JSON mode flag - - Returns: - ImageResponse: Populated image response - """ try: response_json: Final = raw_response.json() - except Exception as e: + except ValueError as e: raise OpenRouterException( message=f"Error parsing OpenRouter response: {e}", status_code=raw_response.status_code, headers=raw_response.headers, - ) + ) from e - if not model_response.data: - model_response.data = [] - - try: - choices: Final = response_json.get("choices", []) - - for choice in choices: - message = choice.get("message", {}) - images = message.get("images", []) - - for image_data in images: - image_url_obj = image_data.get("image_url", {}) - image_url = image_url_obj.get("url") - - if image_url: - if image_url.startswith("data:"): - # Extract base64 data - # Format: data:image/png;base64, - parts = image_url.split(",", 1) - b64_data = parts[1] if len(parts) > 1 else None - - model_response.data.append( - ImageObject( - b64_json=b64_data, - url=None, - revised_prompt=None, - ) - ) - else: - model_response.data.append( - ImageObject( - b64_json=None, - url=image_url, - revised_prompt=None, - ) - ) - - # Extract and set usage and cost information - self._set_usage_and_cost(model_response, response_json, model) - - return model_response - - except Exception as e: - raise OpenRouterException( - message=f"Error transforming OpenRouter image generation response: {e}", - status_code=500, - headers={}, - ) + image_response: Final = ImageResponse( + created=response_json.get("created"), + data=[ImageObject(b64_json=item.get("b64_json")) for item in response_json.get("data") or []], + ) + self._set_usage_and_cost(image_response, response_json, model) + return image_response def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: """Get the appropriate error class for OpenRouter errors.""" diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index eb27d3fe810..30b63365796 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2009,7 +2009,7 @@ "messages": true, "responses": true, "embeddings": true, - "image_generations": false, + "image_generations": true, "audio_transcriptions": false, "audio_speech": false, "moderations": false, diff --git a/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py b/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py index e45270fb5e3..71c8e31b892 100644 --- a/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py +++ b/tests/unit/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py @@ -1,582 +1,359 @@ import json +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest - +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.openrouter.common_utils import OpenRouterException from litellm.llms.openrouter.image_generation.transformation import ( OpenRouterImageGenerationConfig, ) -from litellm.llms.openrouter.common_utils import OpenRouterException -from litellm.types.utils import ImageResponse +from litellm.types.utils import ImageResponse, ImageUsage, ImageUsageInputTokensDetails + +CONFIG: Final = OpenRouterImageGenerationConfig() +IMAGE_ONLY_MODEL: Final = "openai/gpt-image-1-mini" +HYBRID_MODEL: Final = "google/gemini-2.5-flash-image" +PROMPT: Final = "a small red apple on a white table, simple flat illustration" +IMAGES_URL: Final = "https://openrouter.ai/api/v1/images" + +# usage object returned by a real POST https://openrouter.ai/api/v1/images call for openai/gpt-image-1-mini +# (quality low, 1024x1024) on 2026-10-03 +OPENROUTER_IMAGES_USAGE: Final = { + "prompt_tokens": 18, + "completion_tokens": 272, + "total_tokens": 290, + "cost": 0.002212, + "is_byok": False, + "prompt_tokens_details": {"cached_tokens": 0}, + "cost_details": { + "upstream_inference_cost": 0.002212, + "upstream_inference_prompt_cost": 3.6e-05, + "upstream_inference_completions_cost": 0.002176, + }, + "completion_tokens_details": {"reasoning_tokens": 0, "image_tokens": 272}, +} -class TestOpenRouterImageGenerationTransformation: - def setup_method(self): - """Set up test fixtures before each test method.""" - self.config = OpenRouterImageGenerationConfig() - self.model = "google/gemini-2.5-flash-image" - self.logging_obj = MagicMock() +def _images_response(*b64_images: str, created: int = 1790994427) -> dict[str, object]: + return { + "created": created, + "data": [{"b64_json": image, "media_type": "image/png"} for image in b64_images], + "usage": OPENROUTER_IMAGES_USAGE, + } - def test_get_supported_openai_params(self): - """Test that get_supported_openai_params returns correct parameters.""" - supported_params = self.config.get_supported_openai_params(self.model) - assert "size" in supported_params - assert "quality" in supported_params - assert "n" in supported_params - assert len(supported_params) == 3 +def _transform_response(raw_response: httpx.Response) -> ImageResponse: + return CONFIG.transform_image_generation_response( + model=IMAGE_ONLY_MODEL, + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) - def test_map_size_to_aspect_ratio_square(self): - """Test mapping square sizes to aspect ratio.""" - assert self.config._map_size_to_aspect_ratio("256x256") == "1:1" - assert self.config._map_size_to_aspect_ratio("512x512") == "1:1" - assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" - def test_map_size_to_aspect_ratio_landscape(self): - """Test mapping landscape sizes to aspect ratio.""" - assert self.config._map_size_to_aspect_ratio("1536x1024") == "3:2" - assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" +class RequestRecorder: + """httpx.MockTransport handler that keeps every request it was called with""" - def test_map_size_to_aspect_ratio_portrait(self): - """Test mapping portrait sizes to aspect ratio.""" - assert self.config._map_size_to_aspect_ratio("1024x1536") == "2:3" - assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" + def __init__(self, response_payload: object, status_code: int = 200) -> None: + self.response_payload = response_payload + self.status_code = status_code + self.requests: list[httpx.Request] = [] - def test_map_size_to_aspect_ratio_auto(self): - """Test mapping auto size to default aspect ratio.""" - assert self.config._map_size_to_aspect_ratio("auto") == "1:1" + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests.append(request) + return httpx.Response(status_code=self.status_code, json=self.response_payload) - def test_map_size_to_aspect_ratio_unknown(self): - """Test mapping unknown size defaults to 1:1.""" - assert self.config._map_size_to_aspect_ratio("999x999") == "1:1" - def test_map_quality_to_image_size_low(self): - """Test mapping low quality values to 1K.""" - assert self.config._map_quality_to_image_size("low") == "1K" - assert self.config._map_quality_to_image_size("standard") == "1K" - assert self.config._map_quality_to_image_size("auto") == "1K" +def _client(recorder: RequestRecorder) -> HTTPHandler: + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(recorder))) - def test_map_quality_to_image_size_medium(self): - """Test mapping medium quality to 2K.""" - assert self.config._map_quality_to_image_size("medium") == "2K" - def test_map_quality_to_image_size_high(self): - """Test mapping high quality values to 4K.""" - assert self.config._map_quality_to_image_size("high") == "4K" - assert self.config._map_quality_to_image_size("hd") == "4K" +def test_get_supported_openai_params(): + assert CONFIG.get_supported_openai_params(IMAGE_ONLY_MODEL) == ["size", "quality", "n"] - def test_map_quality_to_image_size_unknown(self): - """Test mapping unknown quality returns None.""" - assert self.config._map_quality_to_image_size("unknown") is None - def test_map_openai_params_size_only(self): - """Test that map_openai_params correctly maps size parameter.""" - non_default_params = {"size": "1024x1024"} - optional_params = {} +@pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + (None, IMAGES_URL), + ("https://openrouter.ai/api/v1", IMAGES_URL), + ("https://openrouter.ai/api/v1/", IMAGES_URL), + ("https://openrouter.ai/api/v1/chat/completions", IMAGES_URL), + ("https://gateway.example.com/openrouter/v1", "https://gateway.example.com/openrouter/v1/images"), + ("https://gateway.example.com/api/v1/images", "https://gateway.example.com/api/v1/images"), + ], +) +def test_get_complete_url_points_at_the_images_endpoint(api_base: str | None, expected_url: str): + url = CONFIG.get_complete_url( + api_base=api_base, + api_key="sk-test", + model=IMAGE_ONLY_MODEL, + optional_params={}, + litellm_params={}, + ) - result = self.config.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model=self.model, - drop_params=False, + assert url == expected_url + + +@pytest.mark.parametrize( + ("non_default_params", "expected_params"), + [ + ({"size": "1536x1024"}, {"size": "1536x1024"}), + ({"size": "auto"}, {}), + ({"quality": "low"}, {"quality": "low"}), + ({"quality": "medium"}, {"quality": "medium"}), + ({"quality": "high"}, {"quality": "high"}), + ({"quality": "auto"}, {"quality": "auto"}), + ({"quality": "standard"}, {"quality": "low"}), + ({"quality": "hd"}, {"quality": "high"}), + ({"n": 2}, {"n": 2}), + ], +) +def test_map_openai_params_sends_size_quality_and_n_as_images_fields( + non_default_params: dict[str, object], expected_params: dict[str, object] +): + mapped = CONFIG.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=IMAGE_ONLY_MODEL, + drop_params=False, + ) + + assert mapped == expected_params + + +@pytest.mark.parametrize( + ("drop_params", "expected_params"), + [ + (False, {"size": "1024x1024", "unsupported_param": "value"}), + (True, {"size": "1024x1024"}), + ], +) +def test_map_openai_params_unsupported_param_follows_drop_params(drop_params: bool, expected_params: dict[str, object]): + mapped = CONFIG.map_openai_params( + non_default_params={"size": "1024x1024", "unsupported_param": "value"}, + optional_params={}, + model=IMAGE_ONLY_MODEL, + drop_params=drop_params, + ) + + assert mapped == expected_params + + +def test_map_openai_params_keeps_params_already_in_optional_params(): + mapped = CONFIG.map_openai_params( + non_default_params={"n": 1}, + optional_params={"resolution": "2K"}, + model=HYBRID_MODEL, + drop_params=False, + ) + + assert mapped == {"resolution": "2K", "n": 1} + + +@patch("litellm.llms.openrouter.image_generation.transformation.get_secret_str") +def test_validate_environment_with_api_key(mock_get_secret: MagicMock): + result = CONFIG.validate_environment( + headers={}, + model=HYBRID_MODEL, + messages=[], + optional_params={}, + litellm_params={}, + api_key="test_api_key", + ) + + assert result["Authorization"] == "Bearer test_api_key" + mock_get_secret.assert_not_called() + + +@patch("litellm.llms.openrouter.image_generation.transformation.get_secret_str") +def test_validate_environment_with_secret_key(mock_get_secret: MagicMock): + mock_get_secret.return_value = "secret_api_key" + + result = CONFIG.validate_environment( + headers={}, + model=HYBRID_MODEL, + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + assert result["Authorization"] == "Bearer secret_api_key" + mock_get_secret.assert_called_once_with("OPENROUTER_API_KEY") + + +def test_transform_request_body_holds_only_images_fields(): + body = CONFIG.transform_image_generation_request( + model=IMAGE_ONLY_MODEL, + prompt=PROMPT, + optional_params={ + "size": "1024x1024", + "quality": "low", + "n": 1, + "modalities": ["image", "text"], + "stream": True, + "extra_headers": {"Authorization": "Bearer sk-test"}, + }, + litellm_params={}, + headers={}, + ) + + assert body == {"model": IMAGE_ONLY_MODEL, "prompt": PROMPT, "size": "1024x1024", "quality": "low", "n": 1} + + +@pytest.mark.parametrize( + ("optional_params", "expected_fields"), + [ + ( + {"image_config": {"aspect_ratio": "16:9", "image_size": "4K"}}, + {"aspect_ratio": "16:9", "resolution": "4K"}, + ), + ( + {"image_config": {"aspect_ratio": "16:9", "image_size": "4K"}, "aspect_ratio": "1:1", "resolution": "2K"}, + {"aspect_ratio": "1:1", "resolution": "2K"}, + ), + ], +) +def test_transform_request_maps_legacy_image_config_and_explicit_fields_win( + optional_params: dict[str, object], expected_fields: dict[str, object] +): + body = CONFIG.transform_image_generation_request( + model=HYBRID_MODEL, + prompt=PROMPT, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert body == {"model": HYBRID_MODEL, "prompt": PROMPT, **expected_fields} + + +def test_transform_response_returns_every_image_in_order(): + response = _transform_response(httpx.Response(200, json=_images_response("aW1hZ2Ux", "aW1hZ2Uy"))) + + assert [(image.b64_json, image.url) for image in response.data] == [("aW1hZ2Ux", None), ("aW1hZ2Uy", None)] + + +def test_transform_response_copies_the_openrouter_created_timestamp(): + response = _transform_response(httpx.Response(200, json=_images_response("aW1hZ2Ux", created=1790994427))) + + assert response.created == 1790994427 + + +def test_transform_response_with_zero_created_keeps_a_real_timestamp(): + response = _transform_response(httpx.Response(200, json=_images_response("aW1hZ2Ux", created=0))) + + assert response.created > 0 + + +def test_transform_response_reports_openrouter_usage_and_cost(): + response = _transform_response(httpx.Response(200, json=_images_response("aW1hZ2Ux"))) + + assert response.usage == ImageUsage( + input_tokens=18, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=18), + output_tokens=272, + total_tokens=290, + ) + assert response._hidden_params["additional_headers"] == { + "llm_provider-x-litellm-response-cost": OPENROUTER_IMAGES_USAGE["cost"] + } + assert response._hidden_params["response_cost_details"] == OPENROUTER_IMAGES_USAGE["cost_details"] + assert response._hidden_params["model"] == IMAGE_ONLY_MODEL + + +def test_transform_response_with_non_json_body_raises_openrouter_exception(): + with pytest.raises(OpenRouterException, match="Error parsing OpenRouter response") as exc_info: + _transform_response(httpx.Response(502, content=b"bad gateway")) + + assert exc_info.value.status_code == 502 + assert isinstance(exc_info.value.__cause__, json.JSONDecodeError) + + +def test_get_error_class(): + error = CONFIG.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, OpenRouterException) + assert "Test error" in str(error) + assert error.status_code == 400 + + +def test_image_only_model_is_sent_to_the_images_endpoint_and_charged_the_openrouter_cost(): + recorder = RequestRecorder(_images_response("aW1hZ2Ux")) + + response = litellm.image_generation( + model=f"openrouter/{IMAGE_ONLY_MODEL}", + prompt=PROMPT, + size="1024x1024", + quality="low", + n=1, + api_key="sk-test", + client=_client(recorder), + ) + + (request,) = recorder.requests + assert str(request.url) == IMAGES_URL + assert request.headers["Authorization"] == "Bearer sk-test" + assert json.loads(request.content) == { + "model": IMAGE_ONLY_MODEL, + "prompt": PROMPT, + "size": "1024x1024", + "quality": "low", + "n": 1, + } + assert [image.b64_json for image in response.data] == ["aW1hZ2Ux"] + assert response._hidden_params["response_cost"] == OPENROUTER_IMAGES_USAGE["cost"] + + +def test_hybrid_image_text_model_uses_the_same_images_endpoint(): + recorder = RequestRecorder(_images_response("aW1hZ2Ux")) + + litellm.image_generation( + model=f"openrouter/{HYBRID_MODEL}", + prompt=PROMPT, + api_key="sk-test", + client=_client(recorder), + ) + + (request,) = recorder.requests + assert str(request.url) == IMAGES_URL + assert json.loads(request.content) == {"model": HYBRID_MODEL, "prompt": PROMPT} + + +def test_legacy_chat_completions_api_base_still_reaches_the_images_endpoint(): + recorder = RequestRecorder(_images_response("aW1hZ2Ux")) + + litellm.image_generation( + model=f"openrouter/{IMAGE_ONLY_MODEL}", + prompt=PROMPT, + api_key="sk-test", + api_base="https://openrouter.ai/api/v1/chat/completions", + client=_client(recorder), + ) + + (request,) = recorder.requests + assert str(request.url) == IMAGES_URL + + +def test_openrouter_error_response_surfaces_as_not_found_error(): + recorder = RequestRecorder({"error": {"code": 404, "message": "Resource not found"}}, status_code=404) + + with pytest.raises(litellm.NotFoundError, match="Resource not found"): + litellm.image_generation( + model=f"openrouter/{IMAGE_ONLY_MODEL}", + prompt=PROMPT, + api_key="sk-test", + client=_client(recorder), ) - - assert "image_config" in result - assert result["image_config"]["aspect_ratio"] == "1:1" - - def test_map_openai_params_quality_only(self): - """Test that map_openai_params correctly maps quality parameter.""" - non_default_params = {"quality": "high"} - optional_params = {} - - result = self.config.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model=self.model, - drop_params=False, - ) - - assert "image_config" in result - assert result["image_config"]["image_size"] == "4K" - - def test_map_openai_params_size_and_quality(self): - """Test that map_openai_params correctly maps both size and quality.""" - non_default_params = {"size": "1792x1024", "quality": "hd"} - optional_params = {} - - result = self.config.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model=self.model, - drop_params=False, - ) - - assert "image_config" in result - assert result["image_config"]["aspect_ratio"] == "16:9" - assert result["image_config"]["image_size"] == "4K" - - def test_map_openai_params_with_n_parameter(self): - """Test that map_openai_params correctly passes through n parameter.""" - non_default_params = {"size": "1024x1024", "n": 2} - optional_params = {} - - result = self.config.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model=self.model, - drop_params=False, - ) - - assert "image_config" in result - assert result["image_config"]["aspect_ratio"] == "1:1" - assert result["n"] == 2 - - def test_map_openai_params_unsupported_param_drop_false(self): - """Test that unsupported params are passed through when drop_params=False.""" - non_default_params = {"size": "1024x1024", "unsupported_param": "value"} - optional_params = {} - - result = self.config.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model=self.model, - drop_params=False, - ) - - assert "image_config" in result - assert result["unsupported_param"] == "value" - - def test_map_openai_params_unsupported_param_drop_true(self): - """Test that unsupported params are dropped when drop_params=True.""" - non_default_params = {"size": "1024x1024", "unsupported_param": "value"} - optional_params = {} - - result = self.config.map_openai_params( - non_default_params=non_default_params, - optional_params=optional_params, - model=self.model, - drop_params=True, - ) - - assert "image_config" in result - assert "unsupported_param" not in result - - def test_get_complete_url_default(self): - """Test that get_complete_url returns default OpenRouter URL.""" - result = self.config.get_complete_url( - api_base=None, - api_key="test_key", - model=self.model, - optional_params={}, - litellm_params={}, - ) - - assert result == "https://openrouter.ai/api/v1/chat/completions" - - def test_get_complete_url_with_custom_base(self): - """Test that get_complete_url uses custom api_base.""" - custom_base = "https://custom.openrouter.ai/api/v1" - - result = self.config.get_complete_url( - api_base=custom_base, - api_key="test_key", - model=self.model, - optional_params={}, - litellm_params={}, - ) - - assert result == f"{custom_base}/chat/completions" - - def test_get_complete_url_with_base_already_complete(self): - """Test that get_complete_url doesn't duplicate /chat/completions.""" - custom_base = "https://custom.openrouter.ai/api/v1/chat/completions" - - result = self.config.get_complete_url( - api_base=custom_base, - api_key="test_key", - model=self.model, - optional_params={}, - litellm_params={}, - ) - - assert result == custom_base - - @patch("litellm.llms.openrouter.image_generation.transformation.get_secret_str") - def test_validate_environment_with_api_key(self, mock_get_secret): - """Test that validate_environment correctly sets authorization header.""" - headers = {} - api_key = "test_api_key" - - result = self.config.validate_environment( - headers=headers, - model=self.model, - messages=[], - optional_params={}, - litellm_params={}, - api_key=api_key, - ) - - assert result["Authorization"] == f"Bearer {api_key}" - mock_get_secret.assert_not_called() - - @patch("litellm.llms.openrouter.image_generation.transformation.get_secret_str") - def test_validate_environment_with_secret_key(self, mock_get_secret): - """Test that validate_environment uses secret API key when api_key is None.""" - mock_get_secret.return_value = "secret_api_key" - headers = {} - - result = self.config.validate_environment( - headers=headers, - model=self.model, - messages=[], - optional_params={}, - litellm_params={}, - api_key=None, - ) - - assert result["Authorization"] == "Bearer secret_api_key" - mock_get_secret.assert_called_once_with("OPENROUTER_API_KEY") - - def test_transform_image_generation_request_basic(self): - """Test that transform_image_generation_request creates correct request body.""" - prompt = "A beautiful sunset over mountains" - optional_params = {} - - result = self.config.transform_image_generation_request( - model=self.model, - prompt=prompt, - optional_params=optional_params, - litellm_params={}, - headers={}, - ) - - assert result["model"] == self.model - assert result["messages"] == [{"role": "user", "content": prompt}] - assert "modalities" not in result # modalities should not be added by default - - def test_transform_image_generation_request_with_image_config(self): - """Test that transform_image_generation_request includes image_config.""" - prompt = "A beautiful sunset" - optional_params = { - "image_config": {"aspect_ratio": "16:9", "image_size": "4K"}, - "n": 2, - } - - result = self.config.transform_image_generation_request( - model=self.model, - prompt=prompt, - optional_params=optional_params, - litellm_params={}, - headers={}, - ) - - assert result["model"] == self.model - assert result["messages"] == [{"role": "user", "content": prompt}] - assert result["image_config"]["aspect_ratio"] == "16:9" - assert result["image_config"]["image_size"] == "4K" - assert result["n"] == 2 - - def test_transform_image_generation_response_with_base64_images(self): - """Test that transform_image_generation_response correctly extracts base64 images.""" - response_data = { - "choices": [ - { - "message": { - "content": "Here is your image!", - "role": "assistant", - "images": [ - { - "image_url": { - "url": "data:image/png;base64,iVBORw0KGgoAAAANS" - }, - "index": 0, - "type": "image_url", - } - ], - } - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 1300, - "total_tokens": 1310, - "completion_tokens_details": {"image_tokens": 1290}, - "cost": 0.0387243, - }, - "model": "google/gemini-2.5-flash-image", - } - - mock_response = MagicMock() - mock_response.json.return_value = response_data - mock_response.status_code = 200 - mock_response.headers = {} - - model_response = ImageResponse(data=[]) - - result = self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].b64_json == "iVBORw0KGgoAAAANS" - assert result.data[0].url is None - - def test_transform_image_generation_response_with_url_images(self): - """Test that transform_image_generation_response correctly extracts URL images.""" - response_data = { - "choices": [ - { - "message": { - "content": "Here is your image!", - "role": "assistant", - "images": [ - { - "image_url": {"url": "https://example.com/image.png"}, - "index": 0, - "type": "image_url", - } - ], - } - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 1300, - "total_tokens": 1310, - }, - "model": "google/gemini-2.5-flash-image", - } - - mock_response = MagicMock() - mock_response.json.return_value = response_data - mock_response.status_code = 200 - mock_response.headers = {} - - model_response = ImageResponse(data=[]) - - result = self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 1 - assert result.data[0].url == "https://example.com/image.png" - assert result.data[0].b64_json is None - - def test_transform_image_generation_response_with_usage_and_cost(self): - """Test that transform_image_generation_response correctly extracts usage and cost.""" - response_data = { - "choices": [ - { - "message": { - "content": "Here is your image!", - "role": "assistant", - "images": [ - { - "image_url": {"url": "data:image/png;base64,abc123"}, - "index": 0, - "type": "image_url", - } - ], - } - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 1300, - "total_tokens": 1310, - "completion_tokens_details": {"image_tokens": 1290}, - "cost": 0.0387243, - "cost_details": {"input_cost": 0.001, "output_cost": 0.037}, - }, - "model": "google/gemini-2.5-flash-image", - } - - mock_response = MagicMock() - mock_response.json.return_value = response_data - mock_response.status_code = 200 - mock_response.headers = {} - - model_response = ImageResponse(data=[]) - - result = self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - # Check usage - assert result.usage is not None - assert result.usage.input_tokens == 10 - assert result.usage.output_tokens == 1290 - assert result.usage.total_tokens == 1310 - assert result.usage.input_tokens_details.text_tokens == 10 - assert result.usage.input_tokens_details.image_tokens == 0 - - # Check cost - assert hasattr(result, "_hidden_params") - assert "additional_headers" in result._hidden_params - assert ( - result._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] - == 0.0387243 - ) - - # Check cost details - assert "response_cost_details" in result._hidden_params - assert result._hidden_params["response_cost_details"]["input_cost"] == 0.001 - assert result._hidden_params["response_cost_details"]["output_cost"] == 0.037 - - # Check model - assert result._hidden_params["model"] == "google/gemini-2.5-flash-image" - - def test_transform_image_generation_response_multiple_images(self): - """Test that transform_image_generation_response handles multiple images.""" - response_data = { - "choices": [ - { - "message": { - "content": "Here are your images!", - "role": "assistant", - "images": [ - { - "image_url": { - "url": "data:image/png;base64,image1data" - }, - "index": 0, - "type": "image_url", - }, - { - "image_url": { - "url": "data:image/png;base64,image2data" - }, - "index": 1, - "type": "image_url", - }, - ], - } - } - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 2600, - "total_tokens": 2610, - }, - "model": "google/gemini-2.5-flash-image", - } - - mock_response = MagicMock() - mock_response.json.return_value = response_data - mock_response.status_code = 200 - mock_response.headers = {} - - model_response = ImageResponse(data=[]) - - result = self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert len(result.data) == 2 - assert result.data[0].b64_json == "image1data" - assert result.data[1].b64_json == "image2data" - - def test_transform_image_generation_response_json_error(self): - """Test that transform_image_generation_response raises error on invalid JSON.""" - mock_response = MagicMock() - mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0) - mock_response.status_code = 500 - mock_response.headers = {} - - model_response = ImageResponse(data=[]) - - with pytest.raises(OpenRouterException) as exc_info: - self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert "Error parsing OpenRouter response" in str(exc_info.value) - assert exc_info.value.status_code == 500 - - def test_transform_image_generation_response_transformation_error(self): - """Test that transform_image_generation_response handles transformation errors.""" - response_data = { - "choices": [ - { - "message": { - "content": "Here is your image!", - "role": "assistant", - "images": "invalid_format", # Invalid format - } - } - ] - } - - mock_response = MagicMock() - mock_response.json.return_value = response_data - mock_response.status_code = 200 - mock_response.headers = {} - - model_response = ImageResponse(data=[]) - - with pytest.raises(OpenRouterException) as exc_info: - self.config.transform_image_generation_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=self.logging_obj, - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - ) - - assert "Error transforming OpenRouter image generation response" in str( - exc_info.value - ) - - def test_get_error_class(self): - """Test that get_error_class returns OpenRouterException.""" - error = self.config.get_error_class( - error_message="Test error", - status_code=400, - headers={"Content-Type": "application/json"}, - ) - - assert isinstance(error, OpenRouterException) - assert "Test error" in str(error) - assert error.status_code == 400