diff --git a/docs/my-website/docs/image_edits.md b/docs/my-website/docs/image_edits.md index a843833454..f1cfc0ed8e 100644 --- a/docs/my-website/docs/image_edits.md +++ b/docs/my-website/docs/image_edits.md @@ -16,7 +16,7 @@ LiteLLM provides image editing functionality that maps to OpenAI's `/images/edit | Supported operations | Create image edits | Single and multiple images supported | | Supported LiteLLM SDK Versions | 1.63.8+ | Gemini support requires 1.79.3+ | | Supported LiteLLM Proxy Versions | 1.71.1+ | Gemini support requires 1.79.3+ | -| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. Stability AI and Bedrock Stability support various image editing operations. | +| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **OpenRouter**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. OpenRouter routes image edits through chat completions. Stability AI and Bedrock Stability support various image editing operations. | #### ⚡️See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/) @@ -244,6 +244,47 @@ response = litellm.image_edit( print(response) ``` + + + + +#### Basic Image Edit +```python showLineNumbers title="OpenRouter Image Edit" +import os +from litellm import image_edit + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=open("original_image.png", "rb"), + prompt="Add aurora borealis to the night sky", +) + +print(response) +``` + +#### Multiple Images Edit +```python showLineNumbers title="OpenRouter Multiple Images Edit" +import os +from litellm import image_edit + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=[ + open("scene.png", "rb"), + open("style_reference.png", "rb"), + ], + prompt="Blend the reference style into the scene", + size="1536x1024", # mapped to aspect_ratio 3:2 + quality="high", # mapped to image_size 4K +) + +print(response) +``` + @@ -398,6 +439,34 @@ curl -X POST "http://0.0.0.0:4000/v1/images/edits" \ -F "size=1024x1024" ``` + + + + +1. Add the OpenRouter image edit model to your `config.yaml`: +```yaml showLineNumbers title="OpenRouter Proxy Configuration" +model_list: + - model_name: openrouter-image-edit + litellm_params: + model: openrouter/google/gemini-2.5-flash-image + api_key: os.environ/OPENROUTER_API_KEY +``` + +2. Start the LiteLLM proxy server: +```bash showLineNumbers title="Start LiteLLM Proxy Server" +litellm --config /path/to/config.yaml +``` + +3. Make an image edit request: +```bash showLineNumbers title="OpenRouter Proxy Image Edit" +curl -X POST "http://0.0.0.0:4000/v1/images/edits" \ + -H "Authorization: Bearer " \ + -F "model=openrouter-image-edit" \ + -F "image=@original_image.png" \ + -F "prompt=Make the sky a vibrant purple sunset" \ + -F "size=1024x1024" +``` + diff --git a/docs/my-website/docs/providers/openrouter.md b/docs/my-website/docs/providers/openrouter.md index 38eb998c98..4c79c41cfd 100644 --- a/docs/my-website/docs/providers/openrouter.md +++ b/docs/my-website/docs/providers/openrouter.md @@ -210,3 +210,90 @@ response = image_generation( # Cost is available in the response metadata print(f"Request cost: ${response._hidden_params['additional_headers']['llm_provider-x-litellm-response-cost']}") ``` + +## Image Edit + +OpenRouter supports image editing through select models like Google Gemini image models. LiteLLM routes image edit requests to OpenRouter's chat completions endpoint with the source image sent as a base64 data URL and `modalities: ["image", "text"]`. + +### Supported Models + +| Model | Description | +|-------|-------------| +| `openrouter/google/gemini-2.5-flash-image` | Gemini 2.5 Flash with image editing | + +See all available image models on [OpenRouter's model list](https://openrouter.ai/models?modality=image). + +### Supported Parameters + +| Parameter | OpenRouter Mapping | Notes | +|-----------|--------------------|-------| +| `size` | `image_config.aspect_ratio` | `1024x1024` → `1:1`, `1536x1024` → `3:2`, `1024x1536` → `2:3`, `1792x1024` → `16:9`, `1024x1792` → `9:16` | +| `quality` | `image_config.image_size` | `low`/`standard` → `1K`, `medium` → `2K`, `high`/`hd` → `4K` | +| `n` | `n` | Number of images | + +:::note +`quality=high` (4K) is only supported by `google/gemini-3-pro-image-preview` and `google/gemini-3.1-flash-image-preview`. The `google/gemini-2.5-flash-image` model supports up to `medium` (2K). +::: + +### Usage + +```python +from litellm import image_edit +import os + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +# Basic image edit +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=open("original_image.png", "rb"), + prompt="Make the sky a vibrant purple sunset", +) + +print(response) +``` + +### Advanced Usage with Parameters + +```python +from litellm import image_edit +import os + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +# Edit with size and quality parameters +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=open("photo.png", "rb"), + prompt="Add northern lights to the sky", + size="1536x1024", # Maps to aspect_ratio 3:2 + quality="high", # Maps to image_size 4K +) + +# Access the edited image +image_data = response.data[0] +if image_data.b64_json: + import base64 + with open("edited.png", "wb") as f: + f.write(base64.b64decode(image_data.b64_json)) +``` + +### Multiple Images Edit + +```python +from litellm import image_edit +import os + +os.environ["OPENROUTER_API_KEY"] = "your-api-key" + +response = image_edit( + model="openrouter/google/gemini-2.5-flash-image", + image=[ + open("scene.png", "rb"), + open("style_reference.png", "rb"), + ], + prompt="Blend the reference style into the scene", +) + +print(response) +``` diff --git a/litellm/llms/openrouter/image_edit/__init__.py b/litellm/llms/openrouter/image_edit/__init__.py new file mode 100644 index 0000000000..6edd133f27 --- /dev/null +++ b/litellm/llms/openrouter/image_edit/__init__.py @@ -0,0 +1,11 @@ +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + +from .transformation import OpenRouterImageEditConfig + +__all__ = [ + "OpenRouterImageEditConfig", +] + + +def get_openrouter_image_edit_config(model: str) -> BaseImageEditConfig: + return OpenRouterImageEditConfig() diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py new file mode 100644 index 0000000000..19931f0024 --- /dev/null +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -0,0 +1,367 @@ +""" +OpenRouter Image Edit Support + +OpenRouter provides image editing through chat completion endpoints. +The source image is sent as a base64 data URL in the message content, +and the response contains edited images in the message's images array. + +Request format: +{ + "model": "google/gemini-2.5-flash-image", + "messages": [{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}}, + {"type": "text", "text": "Edit this image by..."} + ] + }], + "modalities": ["image", "text"] +} + +Response format: +{ + "choices": [{ + "message": { + "content": "Here is the edited image.", + "role": "assistant", + "images": [{ + "image_url": {"url": "data:image/png;base64,..."}, + "type": "image_url" + }] + } + }], + "usage": { + "completion_tokens": 1299, + "prompt_tokens": 300, + "total_tokens": 1599, + "completion_tokens_details": {"image_tokens": 1290}, + "cost": 0.0387243 + } +} +""" + +import base64 +from io import BufferedReader, BytesIO +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast + +import httpx +from httpx._types import RequestFiles + +import litellm +from litellm.images.utils import ImageEditRequestUtils +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.openrouter.common_utils import OpenRouterException +from litellm.secret_managers.main import get_secret_str +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import FileTypes, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class OpenRouterImageEditConfig(BaseImageEditConfig): + """ + Configuration for OpenRouter image editing via chat completions. + + OpenRouter uses the chat completions endpoint for image editing. + The source image is sent as a base64 data URL in the message content, + and the response contains edited images in the message's images array. + """ + + def get_supported_openai_params(self, model: str) -> list: + return ["size", "quality", "n"] + + def map_openai_params( + self, + image_edit_optional_params: ImageEditOptionalRequestParams, + model: str, + drop_params: bool, + ) -> Dict: + supported_params = self.get_supported_openai_params(model) + mapped_params: Dict[str, Any] = {} + + for key, value in image_edit_optional_params.items(): + if key in supported_params: + if key == "size": + if "image_config" not in mapped_params: + mapped_params["image_config"] = {} + mapped_params["image_config"]["aspect_ratio"] = self._map_size_to_aspect_ratio(value) + elif key == "quality": + image_size = self._map_quality_to_image_size(value) + if image_size: + if "image_config" not in mapped_params: + mapped_params["image_config"] = {} + mapped_params["image_config"]["image_size"] = image_size + else: + mapped_params[key] = value + + return mapped_params + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + ) -> dict: + api_key = ( + api_key + or litellm.api_key + or get_secret_str("OPENROUTER_API_KEY") + ) + if not api_key: + raise ValueError("OPENROUTER_API_KEY is not set") + headers.update( + { + "Authorization": f"Bearer {api_key}", + } + ) + return headers + + def use_multipart_form_data(self) -> bool: + """OpenRouter uses JSON requests, not multipart/form-data.""" + return False + + def get_complete_url( + self, + model: str, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + base_url = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" + base_url = base_url.rstrip("/") + if not base_url.endswith("/chat/completions"): + return f"{base_url}/chat/completions" + return base_url + + def transform_image_edit_request( + self, + model: str, + prompt: Optional[str], + image: Optional[FileTypes], + image_edit_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Tuple[Dict, RequestFiles]: + content_parts: List[Dict[str, Any]] = [] + + # Add source image(s) as base64 data URLs + if image is not None: + images = image if isinstance(image, list) else [image] + for img in images: + if img is None: + continue + mime_type = ImageEditRequestUtils.get_image_content_type(img) + image_bytes = self._read_image_bytes(img) + b64_data = base64.b64encode(image_bytes).decode("utf-8") + content_parts.append( + { + "type": "image_url", + "image_url": { + "url": f"data:{mime_type};base64,{b64_data}" + }, + } + ) + + # Add the text prompt + if prompt: + content_parts.append({"type": "text", "text": prompt}) + + request_body: Dict[str, Any] = { + "model": model, + "messages": [ + { + "role": "user", + "content": content_parts, + } + ], + "modalities": ["image", "text"], + } + + # Add mapped optional params (image_config, n, etc.) + for key, value in image_edit_optional_request_params.items(): + if key not in ("model", "messages", "modalities"): + request_body[key] = value + + empty_files = cast(RequestFiles, []) + return request_body, empty_files + + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ImageResponse: + try: + response_json = raw_response.json() + except Exception as e: + raise OpenRouterException( + message=f"Error parsing OpenRouter response: {str(e)}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + model_response = ImageResponse() + model_response.data = [] + + try: + choices = 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 from data URL + 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, + ) + ) + + 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 edit response: {str(e)}", + status_code=500, + headers={}, + ) + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + return OpenRouterException( + message=error_message, + status_code=status_code, + headers=headers, + ) + + # Private helper methods + + def _map_size_to_aspect_ratio(self, size: str) -> str: + """ + Map OpenAI size format to OpenRouter aspect_ratio format. + + Uses the same mapping as image generation since OpenRouter + handles both through the same chat completions endpoint. + """ + size_to_aspect_ratio = { + "256x256": "1:1", + "512x512": "1:1", + "1024x1024": "1:1", + "1536x1024": "3:2", + "1792x1024": "16:9", + "1024x1536": "2:3", + "1024x1792": "9:16", + "auto": "1:1", + } + return size_to_aspect_ratio.get(size, "1:1") + + def _map_quality_to_image_size(self, quality: str) -> Optional[str]: + """ + Map OpenAI quality to OpenRouter image_size format. + + Uses the same mapping as image generation since OpenRouter + handles both through the same chat completions endpoint. + """ + quality_to_image_size = { + "low": "1K", + "standard": "1K", + "medium": "2K", + "high": "4K", + "hd": "4K", + "auto": "1K", + } + return quality_to_image_size.get(quality) + + def _set_usage_and_cost( + self, + model_response: ImageResponse, + response_json: dict, + model: str, + ) -> None: + """Extract and set usage and cost information from OpenRouter response.""" + usage_data = response_json.get("usage", {}) + if usage_data: + prompt_tokens = usage_data.get("prompt_tokens", 0) + total_tokens = usage_data.get("total_tokens", 0) + + completion_tokens_details = usage_data.get("completion_tokens_details", {}) + image_tokens = completion_tokens_details.get("image_tokens", 0) + + # For image edit, input may include image tokens + input_image_tokens = 0 + prompt_tokens_details = usage_data.get("prompt_tokens_details", {}) + if prompt_tokens_details: + input_image_tokens = prompt_tokens_details.get("image_tokens", 0) + + model_response.usage = ImageUsage( + input_tokens=prompt_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + image_tokens=input_image_tokens, + text_tokens=prompt_tokens - input_image_tokens, + ), + output_tokens=image_tokens, + total_tokens=total_tokens, + ) + + cost = usage_data.get("cost") + if cost is not None: + if not hasattr(model_response, "_hidden_params"): + model_response._hidden_params = {} + if "additional_headers" not in model_response._hidden_params: + model_response._hidden_params["additional_headers"] = {} + model_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(cost) + + cost_details = usage_data.get("cost_details", {}) + if cost_details: + if "response_cost_details" not in model_response._hidden_params: + model_response._hidden_params["response_cost_details"] = {} + model_response._hidden_params["response_cost_details"].update(cost_details) + + model_response._hidden_params["model"] = response_json.get("model", model) + + def _read_image_bytes(self, image: FileTypes) -> bytes: + """Read raw bytes from various image input types.""" + if isinstance(image, bytes): + return image + if isinstance(image, BytesIO): + current_pos = image.tell() + image.seek(0) + data = image.read() + image.seek(current_pos) + return data + if isinstance(image, BufferedReader): + current_pos = image.tell() + image.seek(0) + data = image.read() + image.seek(current_pos) + return data + raise ValueError("Unsupported image type for OpenRouter image edit.") diff --git a/litellm/utils.py b/litellm/utils.py index 7988910b08..07e855d613 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8791,6 +8791,12 @@ class ProviderConfigManager: ) return BedrockStabilityImageEditConfig() + elif LlmProviders.OPENROUTER == provider: + from litellm.llms.openrouter.image_edit import ( + get_openrouter_image_edit_config, + ) + + return get_openrouter_image_edit_config(model) return None @staticmethod diff --git a/tests/test_litellm/llms/openrouter/image_edit/__init__.py b/tests/test_litellm/llms/openrouter/image_edit/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py b/tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py new file mode 100644 index 0000000000..924e45dbf3 --- /dev/null +++ b/tests/test_litellm/llms/openrouter/image_edit/test_openrouter_image_edit_transformation.py @@ -0,0 +1,540 @@ +import base64 +import json +import os +import sys +from io import BytesIO +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.openrouter.common_utils import OpenRouterException +from litellm.llms.openrouter.image_edit.transformation import ( + OpenRouterImageEditConfig, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse + + +class TestOpenRouterImageEditTransformation: + def setup_method(self): + """Set up test fixtures before each test method.""" + self.config = OpenRouterImageEditConfig() + self.model = "google/gemini-2.5-flash-image" + self.logging_obj = MagicMock() + self.sample_image_bytes = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100 + + 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 test_use_multipart_form_data_returns_false(self): + """Test that OpenRouter uses JSON, not multipart/form-data.""" + assert self.config.use_multipart_form_data() is False + + # Parameter mapping tests + + def test_map_openai_params_size(self): + """Test that size is mapped to image_config.aspect_ratio.""" + result = self.config.map_openai_params( + image_edit_optional_params={"size": "1024x1024"}, + model=self.model, + drop_params=False, + ) + + assert "image_config" in result + assert result["image_config"]["aspect_ratio"] == "1:1" + + def test_map_openai_params_quality(self): + """Test that quality is mapped to image_config.image_size.""" + result = self.config.map_openai_params( + image_edit_optional_params={"quality": "high"}, + 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 both size and quality are mapped correctly.""" + result = self.config.map_openai_params( + image_edit_optional_params={"size": "1792x1024", "quality": "hd"}, + model=self.model, + drop_params=False, + ) + + assert result["image_config"]["aspect_ratio"] == "16:9" + assert result["image_config"]["image_size"] == "4K" + + def test_map_openai_params_n_passthrough(self): + """Test that n parameter is passed through directly.""" + result = self.config.map_openai_params( + image_edit_optional_params={"n": 2}, + model=self.model, + drop_params=False, + ) + + assert result["n"] == 2 + + def test_map_openai_params_unknown_quality_ignored(self): + """Test that unknown quality values produce no image_size mapping.""" + result = self.config.map_openai_params( + image_edit_optional_params={"quality": "unknown_value"}, + model=self.model, + drop_params=False, + ) + + assert "image_config" not in result + + # Size-to-aspect-ratio mapping tests + + def test_map_size_to_aspect_ratio_square(self): + """Test mapping square sizes to 1:1 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 correct aspect ratios.""" + assert self.config._map_size_to_aspect_ratio("1536x1024") == "3:2" + assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" + + def test_map_size_to_aspect_ratio_portrait(self): + """Test mapping portrait sizes to correct aspect ratios.""" + assert self.config._map_size_to_aspect_ratio("1024x1536") == "2:3" + assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" + + def test_map_size_to_aspect_ratio_unknown_defaults_to_1_1(self): + """Test that unknown size defaults to 1:1.""" + assert self.config._map_size_to_aspect_ratio("999x999") == "1:1" + + # Quality-to-image-size mapping tests + + def test_map_quality_to_image_size(self): + """Test quality to image size mappings.""" + 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" + assert self.config._map_quality_to_image_size("medium") == "2K" + assert self.config._map_quality_to_image_size("high") == "4K" + assert self.config._map_quality_to_image_size("hd") == "4K" + + def test_map_quality_to_image_size_unknown_returns_none(self): + """Test that unknown quality returns None.""" + assert self.config._map_quality_to_image_size("unknown") is None + + # URL tests + + def test_get_complete_url_default(self): + """Test that default URL is OpenRouter chat completions endpoint.""" + result = self.config.get_complete_url( + model=self.model, + api_base=None, + litellm_params={}, + ) + + assert result == "https://openrouter.ai/api/v1/chat/completions" + + def test_get_complete_url_with_custom_base(self): + """Test that custom api_base gets /chat/completions appended.""" + result = self.config.get_complete_url( + model=self.model, + api_base="https://custom.openrouter.ai/api/v1", + litellm_params={}, + ) + + assert result == "https://custom.openrouter.ai/api/v1/chat/completions" + + def test_get_complete_url_with_complete_base(self): + """Test that api_base already ending in /chat/completions is not duplicated.""" + url = "https://custom.openrouter.ai/api/v1/chat/completions" + result = self.config.get_complete_url( + model=self.model, + api_base=url, + litellm_params={}, + ) + + assert result == url + + # Validate environment tests + + @patch("litellm.llms.openrouter.image_edit.transformation.get_secret_str") + def test_validate_environment_with_api_key(self, mock_get_secret): + """Test that validate_environment sets authorization header with provided key.""" + headers = {} + result = self.config.validate_environment( + headers=headers, + model=self.model, + api_key="test_api_key", + ) + + assert result["Authorization"] == "Bearer test_api_key" + mock_get_secret.assert_not_called() + + @patch("litellm.llms.openrouter.image_edit.transformation.get_secret_str") + def test_validate_environment_with_secret_key(self, mock_get_secret): + """Test that validate_environment falls back to secret key.""" + mock_get_secret.return_value = "secret_api_key" + headers = {} + result = self.config.validate_environment( + headers=headers, + model=self.model, + api_key=None, + ) + + assert result["Authorization"] == "Bearer secret_api_key" + + @patch("litellm.llms.openrouter.image_edit.transformation.litellm") + @patch("litellm.llms.openrouter.image_edit.transformation.get_secret_str") + def test_validate_environment_missing_api_key_raises(self, mock_get_secret, mock_litellm): + """Test that validate_environment raises ValueError when no API key is available.""" + mock_get_secret.return_value = None + mock_litellm.api_key = None + + with pytest.raises(ValueError, match="OPENROUTER_API_KEY is not set"): + self.config.validate_environment( + headers={}, + model=self.model, + api_key=None, + ) + + # Request transformation tests + + def test_transform_image_edit_request_basic(self): + """Test basic request transformation with image and prompt.""" + data, files = self.config.transform_image_edit_request( + model=self.model, + prompt="Add a sunset to this image", + image=self.sample_image_bytes, + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["model"] == self.model + assert data["modalities"] == ["image", "text"] + assert len(data["messages"]) == 1 + assert data["messages"][0]["role"] == "user" + + content = data["messages"][0]["content"] + assert len(content) == 2 + + # First content part should be the image + assert content[0]["type"] == "image_url" + assert content[0]["image_url"]["url"].startswith("data:image/png;base64,") + + # Second content part should be the text prompt + assert content[1]["type"] == "text" + assert content[1]["text"] == "Add a sunset to this image" + + # Files should be empty (JSON mode) + assert list(files) == [] + + def test_transform_image_edit_request_with_bytesio(self): + """Test request transformation with BytesIO image input.""" + image = BytesIO(self.sample_image_bytes) + data, files = self.config.transform_image_edit_request( + model=self.model, + prompt="Edit this", + image=image, + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + content = data["messages"][0]["content"] + assert content[0]["type"] == "image_url" + assert content[0]["image_url"]["url"].startswith("data:image/png;base64,") + + def test_transform_image_edit_request_with_multiple_images(self): + """Test request transformation with a list of images.""" + images = [self.sample_image_bytes, self.sample_image_bytes] + data, files = self.config.transform_image_edit_request( + model=self.model, + prompt="Combine these images", + image=images, + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + content = data["messages"][0]["content"] + # Two image parts + one text part + assert len(content) == 3 + assert content[0]["type"] == "image_url" + assert content[1]["type"] == "image_url" + assert content[2]["type"] == "text" + + def test_transform_image_edit_request_with_optional_params(self): + """Test that optional params are included in request body.""" + data, files = self.config.transform_image_edit_request( + model=self.model, + prompt="Edit this", + image=self.sample_image_bytes, + image_edit_optional_request_params={ + "image_config": {"aspect_ratio": "16:9"}, + "n": 2, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["image_config"]["aspect_ratio"] == "16:9" + assert data["n"] == 2 + + def test_transform_image_edit_request_base64_encoding(self): + """Test that image bytes are correctly base64-encoded in the request.""" + raw_bytes = b"test_image_data" + expected_b64 = base64.b64encode(raw_bytes).decode("utf-8") + + data, _ = self.config.transform_image_edit_request( + model=self.model, + prompt="Edit", + image=raw_bytes, + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + image_url = data["messages"][0]["content"][0]["image_url"]["url"] + # Extract the base64 part after the data URL prefix + b64_part = image_url.split(",", 1)[1] + assert b64_part == expected_b64 + + def test_transform_image_edit_request_no_prompt(self): + """Test request transformation with no prompt (image-only).""" + data, _ = self.config.transform_image_edit_request( + model=self.model, + prompt=None, + image=self.sample_image_bytes, + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + content = data["messages"][0]["content"] + # Only image, no text part + assert len(content) == 1 + assert content[0]["type"] == "image_url" + + # Response transformation tests + + def test_transform_image_edit_response_with_base64(self): + """Test response transformation with base64 image data.""" + response_data = { + "choices": [{ + "message": { + "content": "Here is the edited image.", + "role": "assistant", + "images": [{ + "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANS"}, + "type": "image_url" + }] + } + }], + "usage": { + "prompt_tokens": 300, + "completion_tokens": 1299, + "total_tokens": 1599, + "completion_tokens_details": {"image_tokens": 1290}, + "cost": 0.05 + }, + "model": self.model + } + + mock_response = MagicMock() + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + result = self.config.transform_image_edit_response( + model=self.model, + raw_response=mock_response, + logging_obj=self.logging_obj, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "iVBORw0KGgoAAAANS" + assert result.data[0].url is None + + def test_transform_image_edit_response_with_url(self): + """Test response transformation with URL image data.""" + response_data = { + "choices": [{ + "message": { + "content": "Edited.", + "role": "assistant", + "images": [{ + "image_url": {"url": "https://example.com/edited.png"}, + "type": "image_url" + }] + } + }], + "usage": {"prompt_tokens": 10, "total_tokens": 1310}, + "model": self.model + } + + mock_response = MagicMock() + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + result = self.config.transform_image_edit_response( + model=self.model, + raw_response=mock_response, + logging_obj=self.logging_obj, + ) + + assert len(result.data) == 1 + assert result.data[0].url == "https://example.com/edited.png" + assert result.data[0].b64_json is None + + def test_transform_image_edit_response_usage_and_cost(self): + """Test that usage and cost are correctly extracted from response.""" + response_data = { + "choices": [{ + "message": { + "content": "Edited.", + "role": "assistant", + "images": [{ + "image_url": {"url": "data:image/png;base64,abc123"}, + "type": "image_url" + }] + } + }], + "usage": { + "prompt_tokens": 300, + "completion_tokens": 1299, + "total_tokens": 1599, + "completion_tokens_details": {"image_tokens": 1290}, + "prompt_tokens_details": {"image_tokens": 258}, + "cost": 0.05, + "cost_details": {"input_cost": 0.01, "output_cost": 0.04} + }, + "model": self.model + } + + mock_response = MagicMock() + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + result = self.config.transform_image_edit_response( + model=self.model, + raw_response=mock_response, + logging_obj=self.logging_obj, + ) + + # Check usage + assert result.usage is not None + assert result.usage.input_tokens == 300 + assert result.usage.output_tokens == 1290 + assert result.usage.total_tokens == 1599 + assert result.usage.input_tokens_details.image_tokens == 258 + assert result.usage.input_tokens_details.text_tokens == 42 + + # Check cost + assert result._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.05 + + # Check cost details + assert result._hidden_params["response_cost_details"]["input_cost"] == 0.01 + assert result._hidden_params["response_cost_details"]["output_cost"] == 0.04 + + # Check model + assert result._hidden_params["model"] == self.model + + def test_transform_image_edit_response_multiple_images(self): + """Test response transformation with multiple output images.""" + response_data = { + "choices": [{ + "message": { + "content": "Here are your edits.", + "role": "assistant", + "images": [ + { + "image_url": {"url": "data:image/png;base64,img1data"}, + "type": "image_url" + }, + { + "image_url": {"url": "data:image/png;base64,img2data"}, + "type": "image_url" + } + ] + } + }], + "usage": {"prompt_tokens": 300, "total_tokens": 2600}, + "model": self.model + } + + mock_response = MagicMock() + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + + result = self.config.transform_image_edit_response( + model=self.model, + raw_response=mock_response, + logging_obj=self.logging_obj, + ) + + assert len(result.data) == 2 + assert result.data[0].b64_json == "img1data" + assert result.data[1].b64_json == "img2data" + + def test_transform_image_edit_response_json_error(self): + """Test that invalid JSON response raises OpenRouterException.""" + mock_response = MagicMock() + mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0) + mock_response.status_code = 500 + mock_response.headers = {} + + with pytest.raises(OpenRouterException) as exc_info: + self.config.transform_image_edit_response( + model=self.model, + raw_response=mock_response, + logging_obj=self.logging_obj, + ) + + assert "Error parsing OpenRouter response" in str(exc_info.value) + assert exc_info.value.status_code == 500 + + 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 error.status_code == 400 + + # Read image bytes tests + + def test_read_image_bytes_from_bytes(self): + """Test reading bytes directly.""" + result = self.config._read_image_bytes(b"raw_bytes") + assert result == b"raw_bytes" + + def test_read_image_bytes_from_bytesio(self): + """Test reading bytes from BytesIO.""" + bio = BytesIO(b"bytesio_data") + bio.seek(5) # Move position to test seek reset + result = self.config._read_image_bytes(bio) + assert result == b"bytesio_data" + assert bio.tell() == 5 # Position should be restored + + def test_read_image_bytes_unsupported_type(self): + """Test that unsupported image type raises ValueError.""" + with pytest.raises(ValueError, match="Unsupported image type"): + self.config._read_image_bytes("not_an_image") # type: ignore