Merge pull request #22403 from shivaaang/fix/openrouter-image-edit
feat(openrouter): add image edit support for OpenRouter models
This commit is contained in:
Cesar Garcia
GitHub
commit
8a1e915bb6
@ -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)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="openrouter" label="OpenRouter">
|
||||
|
||||
#### 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)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
@ -398,6 +439,34 @@ curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
|
||||
-F "size=1024x1024"
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="openrouter" label="OpenRouter">
|
||||
|
||||
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 <YOUR-LITELLM-KEY>" \
|
||||
-F "model=openrouter-image-edit" \
|
||||
-F "image=@original_image.png" \
|
||||
-F "prompt=Make the sky a vibrant purple sunset" \
|
||||
-F "size=1024x1024"
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
|
||||
@ -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)
|
||||
```
|
||||
|
||||
11
litellm/llms/openrouter/image_edit/__init__.py
Normal file
11
litellm/llms/openrouter/image_edit/__init__.py
Normal file
@ -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()
|
||||
367
litellm/llms/openrouter/image_edit/transformation.py
Normal file
367
litellm/llms/openrouter/image_edit/transformation.py
Normal file
@ -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.")
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
Loading…
Reference in New Issue
Block a user