mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
fix(openrouter): send image generation to the dedicated /api/v1/images endpoint
litellm.image_generation and the proxy /v1/images/generations route posted
openrouter/ image models to OpenRouter /chat/completions. Image-only models
such as openai/gpt-image-1-mini get a 404 there that says to use
/api/v1/images instead, so the call fails and no spend is recorded
The image generation transform now posts {model, prompt, size, quality, n}
plus provider-specific fields to {api_base}/images and reads data[].b64_json.
usage.cost from that response becomes the response cost through the existing
llm_provider-x-litellm-response-cost path. Image+text models such as
google/gemini-2.5-flash-image are served by the same endpoint, so every
openrouter/ image generation call goes there
Older configs keep working: an api_base ending in /chat/completions is
rewritten to /images, and image_config.aspect_ratio and image_config.image_size
map to aspect_ratio and resolution. extra_headers no longer ends up in the JSON
body. Chat image output through completion(modalities=[...]) and image edit are
unchanged
Co-authored-by: Joly0 <13993216+Joly0@users.noreply.github.com>
Co-authored-by: Robinnnnn <12162433+Robinnnnn@users.noreply.github.com>
This commit is contained in:
tinysolver
parent
8efb4a21f6
commit
f545af5b7d
3 changed files with 387 additions and 788 deletions
|
|
@ -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,<base64_data>
|
||||
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."""
|
||||
|
|
|
|||
|
|
@ -2009,7 +2009,7 @@
|
|||
"messages": true,
|
||||
"responses": true,
|
||||
"embeddings": true,
|
||||
"image_generations": false,
|
||||
"image_generations": true,
|
||||
"audio_transcriptions": false,
|
||||
"audio_speech": false,
|
||||
"moderations": false,
|
||||
|
|
|
|||
|
|
@ -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"<html>bad gateway</html>"))
|
||||
|
||||
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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue