From feeacf720b736841029123f94bc4d2c885e549fb Mon Sep 17 00:00:00 2001 From: Aniket Dixit Date: Wed, 9 Sep 2026 17:22:32 +0530 Subject: [PATCH 1/2] image aspect ratio --- tee_gateway/controllers/chat_controller.py | 12 ++ tee_gateway/image_generation.py | 33 +++++- tee_gateway/model_registry.py | 111 +++++++++++++++++- .../models/create_chat_completion_request.py | 3 + tee_gateway/openapi/openapi.yaml | 8 ++ tee_gateway/test/test_image_generation.py | 64 ++++++++-- 6 files changed, 217 insertions(+), 14 deletions(-) diff --git a/tee_gateway/controllers/chat_controller.py b/tee_gateway/controllers/chat_controller.py index bda0c7a..7231851 100644 --- a/tee_gateway/controllers/chat_controller.py +++ b/tee_gateway/controllers/chat_controller.py @@ -35,6 +35,7 @@ canonical_user_content, ) from tee_gateway.image_generation import ( + _aspect_ratio_params, create_image_generation_response, create_image_generation_streaming_response, ) @@ -295,6 +296,10 @@ def _create_non_streaming_response( ), ) + if cfg.image_output and chat_request.aspect_ratio: + aspect = _aspect_ratio_params(cfg, chat_request.aspect_ratio) + model = model.bind(image_config={"aspect_ratio": aspect["aspect_ratio"]}) + # Bind user tools and/or the native web search tool if requested. if tools_list: model = model.bind_tools(tools_list) @@ -473,6 +478,10 @@ def _create_streaming_response( ), ) + if image_output_model and chat_request.aspect_ratio: + aspect = _aspect_ratio_params(cfg, chat_request.aspect_ratio) + model = model.bind(image_config={"aspect_ratio": aspect["aspect_ratio"]}) + # Bind user tools and/or the native web search tool if requested. if tools_list: model = model.bind_tools(tools_list) @@ -956,6 +965,8 @@ def _chat_request_to_dict(chat_request: CreateChatCompletionRequest) -> dict: d["response_format"] = _normalize_response_format(chat_request.response_format) if chat_request.web_search: d["web_search"] = True + if chat_request.aspect_ratio: + d["aspect_ratio"] = chat_request.aspect_ratio return d @@ -979,6 +990,7 @@ def _parse_chat_request(chat_request_dict: dict) -> CreateChatCompletionRequest: tool_choice=chat_request_dict.get("tool_choice"), user=chat_request_dict.get("user"), web_search=chat_request_dict.get("web_search", False), + aspect_ratio=chat_request_dict.get("aspect_ratio"), ) diff --git a/tee_gateway/image_generation.py b/tee_gateway/image_generation.py index fa66e5d..8aee5e1 100644 --- a/tee_gateway/image_generation.py +++ b/tee_gateway/image_generation.py @@ -290,7 +290,11 @@ def _build_reference_uploads( def _build_generations_payload( - cfg: Any, prompt: str, count: int, refs: Optional[List[str]] + cfg: Any, + prompt: str, + count: int, + refs: Optional[List[str]], + aspect_ratio: Optional[str] = None, ) -> dict[str, Any]: """Build the JSON body for a ``/images/generations`` request. @@ -305,16 +309,37 @@ def _build_generations_payload( payload["n"] = count if cfg.image_extra_params: payload.update(cfg.image_extra_params) + payload.update(_aspect_ratio_params(cfg, aspect_ratio)) if refs: payload["image"] = refs[0] if len(refs) == 1 else refs return payload +def _aspect_ratio_params(cfg: Any, aspect_ratio: Optional[str]) -> dict[str, str]: + """Translate the public ratio into this provider's request shape.""" + if aspect_ratio is None: + return {} + if not isinstance(aspect_ratio, str): + raise ValueError("aspect_ratio must be a string") + ratio = aspect_ratio.strip() + if not ratio or ratio == "auto": + return {} + supported = cfg.image_aspect_ratios or {} + value = supported.get(ratio) + if value is None: + choices = ", ".join(supported) or "none" + raise ValueError( + f"Unsupported aspect_ratio {ratio!r} for this model; supported: {choices}" + ) + return {cfg.image_aspect_ratio_param: value} + + def generate_images( model: str, prompt: str, n: int = 1, reference_images: Optional[List[str]] = None, + aspect_ratio: Optional[str] = None, ) -> tuple[list[str], int]: """Generate images via a provider's OpenAI-compatible images endpoint. @@ -382,6 +407,7 @@ def generate_images( form["response_format"] = cfg.image_response_format if cfg.image_extra_params: form.update({k: str(v) for k, v in cfg.image_extra_params.items()}) + form.update(_aspect_ratio_params(cfg, aspect_ratio)) resp = client.post(edit_endpoint, data=form, files=uploads) else: # No edit endpoint (or nothing uploadable): JSON generations. Inline @@ -390,7 +416,9 @@ def generate_images( json_refs = refs if not cfg.image_edit_endpoint else None resp = client.post( _IMAGE_GENERATION_PATH, - json=_build_generations_payload(cfg, prompt, count, json_refs), + json=_build_generations_payload( + cfg, prompt, count, json_refs, aspect_ratio=aspect_ratio + ), ) _raise_for_status_with_detail(resp) data = resp.json().get("data", []) or [] @@ -497,6 +525,7 @@ def _run_image_generation( prompt, n=chat_request.n or 1, reference_images=reference_images, + aspect_ratio=chat_request.aspect_ratio, ) timestamp = int(time.time()) diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index 4380f79..1d85859 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -63,6 +63,11 @@ class ModelConfig: # Static extra params merged verbatim into the request payload (e.g. size, # watermark). Keyed by field name; values must be JSON-serializable. image_extra_params: Optional[Mapping[str, Any]] = None + # Optional aspect-ratio choices exposed to callers. Keys are normalized + # public ``aspect_ratio`` values and values are provider-specific values. + # Omitted/"auto" requests send no ratio override. + image_aspect_ratios: Optional[Mapping[str, str]] = None + image_aspect_ratio_param: str = "aspect_ratio" # USD per image-modality output token, for ``image_output`` models (Gemini # "nano banana"). These providers bill image output at a higher rate than # text/thinking output: image tokens at this rate, text + thinking tokens at @@ -117,6 +122,61 @@ class ModelConfig: "stream": False, } +# 2K-ish explicit dimensions for ModelArk models. The bare ``size: 2K`` config +# remains the auto path, where ModelArk chooses the shape from the prompt or +# reference image. +_BYTEDANCE_2K_ASPECT_SIZES: dict[str, str] = { + "1:1": "2048x2048", + "16:9": "2560x1440", + "9:16": "1440x2560", + "4:3": "2368x1776", + "3:4": "1776x2368", + "3:2": "2496x1664", + "2:3": "1664x2496", + "21:9": "3136x1344", +} + +_GEMINI_IMAGE_ASPECT_RATIOS: dict[str, str] = { + ratio: ratio + for ratio in ( + "1:1", + "1:4", + "1:8", + "2:3", + "3:2", + "3:4", + "4:1", + "4:3", + "4:5", + "5:4", + "8:1", + "9:16", + "16:9", + "21:9", + ) +} + +_GROK_IMAGE_ASPECT_RATIOS: dict[str, str] = { + ratio: ratio + for ratio in ( + "1:1", + "16:9", + "9:16", + "4:3", + "3:4", + "3:2", + "2:3", + "2:1", + "1:2", + "19.5:9", + "9:19.5", + "20:9", + "9:20", + "21:9", + "5:2", + ) +} + @unique class SupportedModel(Enum): @@ -240,9 +300,9 @@ class SupportedModel(Enum): # live on OpenAI's separate ``/images/edits`` endpoint, which takes the # reference images as multipart file uploads rather than a JSON ``image`` # field — so reference turns are routed there via ``image_edit_endpoint`` - # (up to 10 references per request). Size/quality are pinned so the flat - # per-image price stays predictable. Billed at a flat $0.05 per generated - # image; token prices unused. + # (up to 10 references per request). Quality remains pinned to medium while + # size defaults to auto unless the caller selects one of the three shapes. + # Billed at a flat $0.05 per generated image; token prices unused. GPT_IMAGE_2 = ModelConfig( provider="openai", api_name="gpt-image-2", @@ -253,7 +313,13 @@ class SupportedModel(Enum): image_response_format=None, image_supports_reference=True, image_edit_endpoint="/images/edits", - image_extra_params={"size": "1024x1024", "quality": "medium"}, + image_extra_params={"quality": "medium"}, + image_aspect_ratios={ + "1:1": "1024x1024", + "3:2": "1536x1024", + "2:3": "1024x1536", + }, + image_aspect_ratio_param="size", ) # ── Anthropic ─────────────────────────────────────────────────────── @@ -392,6 +458,21 @@ class SupportedModel(Enum): output_price_usd=Decimal("0.0000015"), image_output=True, image_output_price_usd=Decimal("0.00003"), + image_aspect_ratios={ + ratio: ratio + for ratio in ( + "1:1", + "2:3", + "3:2", + "3:4", + "4:3", + "4:5", + "5:4", + "9:16", + "16:9", + "21:9", + ) + }, ) # Native image generation ("nano banana 2"), the latest Gemini image model. # Google bills output at two rates: text/thinking at $3/MTok and images at @@ -404,6 +485,7 @@ class SupportedModel(Enum): output_price_usd=Decimal("0.000003"), image_output=True, image_output_price_usd=Decimal("0.00006"), + image_aspect_ratios=_GEMINI_IMAGE_ASPECT_RATIOS, ) GEMINI_3_5_FLASH = ModelConfig( provider="google", @@ -510,6 +592,7 @@ class SupportedModel(Enum): image_generation=True, per_image_price_usd=Decimal("0.02"), image_response_format="url", + image_aspect_ratios=_GROK_IMAGE_ASPECT_RATIOS, ) # Grok Imagine Image 2.0 — xAI's newer, higher-quality image model # (released ~2026-08-11), offered alongside grok-imagine-image rather than @@ -527,6 +610,7 @@ class SupportedModel(Enum): image_generation=True, per_image_price_usd=Decimal("0.04"), image_response_format="url", + image_aspect_ratios=_GROK_IMAGE_ASPECT_RATIOS, ) # ── ByteDance (BytePlus ModelArk, OpenAI-compatible) ──────────────── @@ -579,6 +663,8 @@ class SupportedModel(Enum): image_generation=True, per_image_price_usd=Decimal("0.03"), image_supports_reference=True, + image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES, + image_aspect_ratio_param="size", ) # Seedream 5.0 Lite image generation via a ModelArk deployment endpoint. # Seedream 5.0 Lite image generation via a ModelArk deployment endpoint @@ -595,6 +681,8 @@ class SupportedModel(Enum): image_send_n=False, image_supports_reference=True, image_extra_params=_BYTEDANCE_EP_IMAGE_PARAMS, + image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES, + image_aspect_ratio_param="size", ) # Seedance 4.5 image generation via a ModelArk deployment endpoint. # Returns hosted URLs (fetched and inlined by the gateway) and takes the @@ -610,6 +698,8 @@ class SupportedModel(Enum): image_send_n=False, image_supports_reference=True, image_extra_params=_BYTEDANCE_EP_IMAGE_PARAMS, + image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES, + image_aspect_ratio_param="size", ) # Seedance 5.0 image generation via a ModelArk deployment endpoint. # Returns hosted URLs (fetched and inlined by the gateway). Unlike the @@ -628,6 +718,8 @@ class SupportedModel(Enum): image_send_n=False, image_supports_reference=True, image_extra_params=_SEEDANCE_5_IMAGE_PARAMS, + image_aspect_ratios=_BYTEDANCE_2K_ASPECT_SIZES, + image_aspect_ratio_param="size", ) # ── OpenRouter (OpenAI-compatible) ────────────────────────────────── @@ -674,7 +766,16 @@ class SupportedModel(Enum): per_image_price_usd=Decimal("0.015"), image_response_format=None, image_send_n=False, - image_extra_params={"size": "1280x1280"}, + image_aspect_ratios={ + "1:1": "1280x1280", + "3:2": "1568x1056", + "2:3": "1056x1568", + "4:3": "1472x1088", + "3:4": "1088x1472", + "16:9": "1728x960", + "9:16": "960x1728", + }, + image_aspect_ratio_param="size", ) # ── Legacy models (not in current SDK — retained for older SDK versions) ── diff --git a/tee_gateway/models/create_chat_completion_request.py b/tee_gateway/models/create_chat_completion_request.py index 3d3300b..2aa36cc 100644 --- a/tee_gateway/models/create_chat_completion_request.py +++ b/tee_gateway/models/create_chat_completion_request.py @@ -34,6 +34,7 @@ def __init__( function_call=None, functions=None, web_search=False, + aspect_ratio=None, ): self.messages = messages self.model = model @@ -66,6 +67,7 @@ def __init__( self.function_call = function_call self.functions = functions self.web_search = web_search + self.aspect_ratio = aspect_ratio @classmethod def from_dict(cls, dikt) -> "CreateChatCompletionRequest": @@ -103,5 +105,6 @@ def from_dict(cls, dikt) -> "CreateChatCompletionRequest": "function_call", "functions", "web_search", + "aspect_ratio", } return cls(**{k: v for k, v in dikt.items() if k in known}) diff --git a/tee_gateway/openapi/openapi.yaml b/tee_gateway/openapi/openapi.yaml index 5426630..2edc141 100644 --- a/tee_gateway/openapi/openapi.yaml +++ b/tee_gateway/openapi/openapi.yaml @@ -2884,6 +2884,14 @@ components: type: array model: $ref: "#/components/schemas/CreateChatCompletionRequest_model" + aspect_ratio: + description: | + Optional output aspect ratio for image-generation models. Omit this + field (or pass `auto`) to let the provider choose automatically. + Supported values depend on the selected model. + nullable: true + title: aspect_ratio + type: string web_search: default: false deprecated: true diff --git a/tee_gateway/test/test_image_generation.py b/tee_gateway/test/test_image_generation.py index 6b8ac5a..0ac0f2d 100644 --- a/tee_gateway/test/test_image_generation.py +++ b/tee_gateway/test/test_image_generation.py @@ -108,15 +108,15 @@ def test_zai_glm_image_uses_documented_payload_and_fetches_url(self): payload = kwargs["json"] self.assertEqual(payload["model"], "glm-image") self.assertEqual(payload["prompt"], "a poster") - self.assertEqual(payload["size"], "1280x1280") + self.assertNotIn("size", payload) self.assertNotIn("n", payload) self.assertNotIn("response_format", payload) - def test_openai_gpt_image_omits_response_format_and_pins_size_quality(self): + def test_openai_gpt_image_defaults_size_to_auto_and_pins_quality(self): # gpt-image models always return base64 and reject `response_format`, so - # the field must be omitted; size/quality are pinned for predictable - # billing. The shared openai_http_client is reused (base_url ends in /v1, - # so the request lands on OpenAI's /v1/images/generations). + # the field must be omitted. Quality stays pinned while size is absent + # on the automatic path. The shared openai_http_client is reused + # (base_url ends in /v1, so this lands on /v1/images/generations). client = MagicMock() client.post.return_value = _mock_response([{"b64_json": "aGVsbG8="}]) with patch.object(llm_backend, "openai_http_client", client): @@ -129,10 +129,46 @@ def test_openai_gpt_image_omits_response_format_and_pins_size_quality(self): self.assertEqual(payload["model"], get_model_config(GPT_IMAGE).api_name) self.assertEqual(payload["prompt"], "a red cube") self.assertEqual(payload["n"], 1) - self.assertEqual(payload["size"], "1024x1024") + self.assertNotIn("size", payload) self.assertEqual(payload["quality"], "medium") self.assertNotIn("response_format", payload) + def test_gpt_image_translates_aspect_ratio_to_supported_size(self): + client = MagicMock() + client.post.return_value = _mock_response([{"b64_json": "aGVsbG8="}]) + with patch.object(llm_backend, "openai_http_client", client): + generate_images(GPT_IMAGE, "a landscape", aspect_ratio="3:2") + + payload = client.post.call_args.kwargs["json"] + self.assertEqual(payload["size"], "1536x1024") + self.assertEqual(payload["quality"], "medium") + + def test_grok_forwards_aspect_ratio_and_auto_omits_it(self): + client = MagicMock() + client.post.return_value = _mock_response([]) + with patch.object(llm_backend, "xai_http_client", client): + generate_images(GROK_IMAGE, "a banner", aspect_ratio="21:9") + selected = client.post.call_args.kwargs["json"] + generate_images(GROK_IMAGE, "surprise me", aspect_ratio="auto") + automatic = client.post.call_args.kwargs["json"] + + self.assertEqual(selected["aspect_ratio"], "21:9") + self.assertNotIn("aspect_ratio", automatic) + + def test_invalid_aspect_ratio_is_rejected_before_provider_call(self): + client = MagicMock() + with patch.object(llm_backend, "openai_http_client", client): + with self.assertRaisesRegex(ValueError, "Unsupported aspect_ratio"): + generate_images(GPT_IMAGE, "a banner", aspect_ratio="16:9") + client.post.assert_not_called() + + def test_gemini_aspect_ratio_uses_image_config_shape(self): + cfg = get_model_config("gemini-3.1-flash-image") + self.assertEqual( + image_generation._aspect_ratio_params(cfg, "16:9"), + {"aspect_ratio": "16:9"}, + ) + def test_seedance_uses_url_format_and_extra_params(self): client = MagicMock() client.post.return_value = _mock_response([{"url": "https://cdn/img.jpg"}]) @@ -283,7 +319,7 @@ def test_gpt_image_edits_uploads_references_as_multipart(self): self.assertEqual(form["model"], get_model_config(GPT_IMAGE).api_name) self.assertEqual(form["prompt"], "add the logo to the photo") self.assertEqual(form["n"], "1") - self.assertEqual(form["size"], "1024x1024") + self.assertNotIn("size", form) self.assertEqual(form["quality"], "medium") self.assertNotIn("response_format", form) # Both references are uploaded under the repeated image[] field, decoded @@ -296,6 +332,20 @@ def test_gpt_image_edits_uploads_references_as_multipart(self): self.assertEqual(uploads[1][1][0], "image_1.jpg") self.assertEqual(uploads[1][1][1], b"DEF") + def test_gpt_image_edit_translates_aspect_ratio_to_form_size(self): + client = MagicMock() + client.post.return_value = _mock_response([{"b64_json": "aGVsbG8="}]) + refs = ["data:image/png;base64,QUJD"] + with patch.object(llm_backend, "openai_http_client", client): + generate_images( + GPT_IMAGE, + "make it portrait", + reference_images=refs, + aspect_ratio="2:3", + ) + + self.assertEqual(client.post.call_args.kwargs["data"]["size"], "1024x1536") + def test_gpt_image_without_references_uses_generations(self): # No references -> plain text-to-image on the JSON generations endpoint. client = MagicMock() From 6ccd7147e75ae5b8a5de77701c0bc21af1ee45f5 Mon Sep 17 00:00:00 2001 From: Aniket Dixit Date: Wed, 9 Sep 2026 18:35:10 +0530 Subject: [PATCH 2/2] fixes --- tee_gateway/controllers/chat_controller.py | 35 ++++-- tee_gateway/image_generation.py | 31 +++++- tee_gateway/model_registry.py | 37 +++++-- tee_gateway/test/test_chat_controller.py | 123 +++++++++++++++++++++ tee_gateway/test/test_image_generation.py | 103 ++++++++++++++++- 5 files changed, 307 insertions(+), 22 deletions(-) diff --git a/tee_gateway/controllers/chat_controller.py b/tee_gateway/controllers/chat_controller.py index 7231851..a2c1296 100644 --- a/tee_gateway/controllers/chat_controller.py +++ b/tee_gateway/controllers/chat_controller.py @@ -35,9 +35,10 @@ canonical_user_content, ) from tee_gateway.image_generation import ( - _aspect_ratio_params, + aspect_ratio_params, create_image_generation_response, create_image_generation_streaming_response, + validate_aspect_ratio, ) from tee_gateway.model_registry import get_model_config from tee_gateway.moderation import ( @@ -103,6 +104,13 @@ def create_chat_completion(body): except AttachmentValidationError as e: return {"error": "Invalid attachment", "message": str(e)}, 400 + # A ratio the target model can't produce is a client error, so it is caught + # here rather than surfacing from the provider call as a 500. + try: + validate_aspect_ratio(chat_request.model, chat_request.aspect_ratio) + except ValueError as e: + return {"error": "Invalid aspect_ratio", "message": str(e)}, 400 + # Score the newest user turn of image requests before any provider work # (text chat is not moderated — see moderation.should_moderate_model). # Fail-open: an unavailable moderation endpoint yields an unchecked @@ -296,14 +304,20 @@ def _create_non_streaming_response( ), ) - if cfg.image_output and chat_request.aspect_ratio: - aspect = _aspect_ratio_params(cfg, chat_request.aspect_ratio) - model = model.bind(image_config={"aspect_ratio": aspect["aspect_ratio"]}) - # Bind user tools and/or the native web search tool if requested. if tools_list: model = model.bind_tools(tools_list) + # Aspect ratio for Gemini's inline-image models rides in Gemini's + # image_config (these models leave image_aspect_ratio_param at its + # default, which is ImageConfig's own field name). Bound after + # bind_tools, which re-binds the base model and drops kwargs bound + # before it. Empty params mean the automatic path: bind nothing. + if cfg.image_output: + image_config = aspect_ratio_params(cfg, chat_request.aspect_ratio) + if image_config: + model = model.bind(image_config=image_config) + # Bind response_format if provided (json_object or json_schema). # Anthropic does not support response_format via bind(); use # with_structured_output() for json_schema instead (json_object has no @@ -478,14 +492,17 @@ def _create_streaming_response( ), ) - if image_output_model and chat_request.aspect_ratio: - aspect = _aspect_ratio_params(cfg, chat_request.aspect_ratio) - model = model.bind(image_config={"aspect_ratio": aspect["aspect_ratio"]}) - # Bind user tools and/or the native web search tool if requested. if tools_list: model = model.bind_tools(tools_list) + # Aspect ratio for the inline-image models (see the non-streaming path + # for why this is bound after bind_tools). + if image_output_model: + image_config = aspect_ratio_params(cfg, chat_request.aspect_ratio) + if image_config: + model = model.bind(image_config=image_config) + # Bind response_format if provided (json_object or json_schema). # Anthropic does not support response_format via bind(); use # with_structured_output() for json_schema instead (json_object has no diff --git a/tee_gateway/image_generation.py b/tee_gateway/image_generation.py index 8aee5e1..70d7c3e 100644 --- a/tee_gateway/image_generation.py +++ b/tee_gateway/image_generation.py @@ -309,14 +309,19 @@ def _build_generations_payload( payload["n"] = count if cfg.image_extra_params: payload.update(cfg.image_extra_params) - payload.update(_aspect_ratio_params(cfg, aspect_ratio)) + payload.update(aspect_ratio_params(cfg, aspect_ratio)) if refs: payload["image"] = refs[0] if len(refs) == 1 else refs return payload -def _aspect_ratio_params(cfg: Any, aspect_ratio: Optional[str]) -> dict[str, str]: - """Translate the public ratio into this provider's request shape.""" +def aspect_ratio_params(cfg: Any, aspect_ratio: Optional[str]) -> dict[str, str]: + """Translate the public ratio into this provider's request shape. + + Returns an empty dict for the automatic path (field omitted, or ``auto``), + so callers must treat "no params" as "let the provider choose" rather than + indexing the result. + """ if aspect_ratio is None: return {} if not isinstance(aspect_ratio, str): @@ -334,6 +339,24 @@ def _aspect_ratio_params(cfg: Any, aspect_ratio: Optional[str]) -> dict[str, str return {cfg.image_aspect_ratio_param: value} +def validate_aspect_ratio(model: str, aspect_ratio: Optional[str]) -> None: + """Reject an unsupported ratio before any provider work is done. + + Raises ``ValueError`` so the controller can answer 400 rather than letting + the same error surface from deep inside a generation as a 500. Models that + cannot shape their output (every text model) ignore the field instead of + failing the request — like the deprecated ``web_search`` flag. + """ + if aspect_ratio is None: + return + try: + cfg = get_model_config(model) + except ValueError: + return # unknown model: reported by the normal request path + if cfg.image_generation or cfg.image_output: + aspect_ratio_params(cfg, aspect_ratio) + + def generate_images( model: str, prompt: str, @@ -407,7 +430,7 @@ def generate_images( form["response_format"] = cfg.image_response_format if cfg.image_extra_params: form.update({k: str(v) for k, v in cfg.image_extra_params.items()}) - form.update(_aspect_ratio_params(cfg, aspect_ratio)) + form.update(aspect_ratio_params(cfg, aspect_ratio)) resp = client.post(edit_endpoint, data=form, files=uploads) else: # No edit endpoint (or nothing uploadable): JSON generations. Inline diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index 1d85859..d72895d 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -65,8 +65,13 @@ class ModelConfig: image_extra_params: Optional[Mapping[str, Any]] = None # Optional aspect-ratio choices exposed to callers. Keys are normalized # public ``aspect_ratio`` values and values are provider-specific values. - # Omitted/"auto" requests send no ratio override. + # Omitted/"auto" requests send no ratio override. Required on any image + # model: without it, callers get no shape choice at all. image_aspect_ratios: Optional[Mapping[str, str]] = None + # Request field the ratio is sent under — a provider taking explicit pixel + # dimensions uses "size". ``image_output`` (inline-image chat) models must + # leave this at the default: it doubles as the field name inside Gemini's + # ImageConfig, which the chat controller binds the ratio through. image_aspect_ratio_param: str = "aspect_ratio" # USD per image-modality output token, for ``image_output`` models (Gemini # "nano banana"). These providers bill image output at a higher rate than @@ -136,6 +141,28 @@ class ModelConfig: "21:9": "3136x1344", } +# gpt-image-2 takes explicit ``WIDTHxHEIGHT`` sizes rather than ratios: it +# accepts any size whose edges are multiples of 16, whose long:short ratio is +# at most 3:1, whose longest edge is <= 3840px, and whose pixel count is +# between 655,360 and 8,294,400. OpenAI's three documented presets +# (1024x1024, 1536x1024, 1024x1536) are kept verbatim; the rest are the +# exact-ratio sizes closest to them in pixel count, so every shape stays +# inside the ~1.0-1.6MP band the flat per-image price is set for (output +# tokens, and so our real cost, scale with pixel count). Omitting the field +# leaves size at OpenAI's ``auto``. +_GPT_IMAGE_ASPECT_SIZES: dict[str, str] = { + "1:1": "1024x1024", + "3:2": "1536x1024", + "2:3": "1024x1536", + "4:3": "1408x1056", + "3:4": "1056x1408", + "5:4": "1280x1024", + "4:5": "1024x1280", + "16:9": "1536x864", + "9:16": "864x1536", + "21:9": "1792x768", +} + _GEMINI_IMAGE_ASPECT_RATIOS: dict[str, str] = { ratio: ratio for ratio in ( @@ -301,7 +328,7 @@ class SupportedModel(Enum): # reference images as multipart file uploads rather than a JSON ``image`` # field — so reference turns are routed there via ``image_edit_endpoint`` # (up to 10 references per request). Quality remains pinned to medium while - # size defaults to auto unless the caller selects one of the three shapes. + # size defaults to auto unless the caller selects a shape. # Billed at a flat $0.05 per generated image; token prices unused. GPT_IMAGE_2 = ModelConfig( provider="openai", @@ -314,11 +341,7 @@ class SupportedModel(Enum): image_supports_reference=True, image_edit_endpoint="/images/edits", image_extra_params={"quality": "medium"}, - image_aspect_ratios={ - "1:1": "1024x1024", - "3:2": "1536x1024", - "2:3": "1024x1536", - }, + image_aspect_ratios=_GPT_IMAGE_ASPECT_SIZES, image_aspect_ratio_param="size", ) diff --git a/tee_gateway/test/test_chat_controller.py b/tee_gateway/test/test_chat_controller.py index 43f9190..435d8c5 100644 --- a/tee_gateway/test/test_chat_controller.py +++ b/tee_gateway/test/test_chat_controller.py @@ -1,8 +1,10 @@ import unittest +from unittest.mock import Mock, patch import connexion from flask import json +from tee_gateway.controllers.chat_controller import create_chat_completion from tee_gateway.encoder import JSONEncoder from tee_gateway.test import BaseTestCase @@ -70,6 +72,127 @@ def test_pdf_file_part_passes_schema_validation(self): self.assertNotEqual(400, resp.status_code, resp.data.decode("utf-8")) +class TestAspectRatioRequests(unittest.TestCase): + """``aspect_ratio`` handling on /v1/chat/completions. + + A ratio the model can't produce is a client error (400) rather than a + provider failure surfaced as a 500, and the automatic path ("auto", or the + field omitted) must bind nothing at all. + """ + + def _request(self, **extra): + body = { + "model": "gemini-3.1-flash-image", + "messages": [{"role": "user", "content": "a red cube"}], + "stream": False, + } + body.update(extra) + return body + + @patch("tee_gateway.controllers.chat_controller.connexion") + def test_unsupported_ratio_is_a_400(self, mock_connexion): + mock_connexion.request.is_json = True + mock_connexion.request.get_json.return_value = self._request(aspect_ratio="7:3") + + result, status = create_chat_completion(None) + + self.assertEqual(400, status) + self.assertEqual("Invalid aspect_ratio", result["error"]) + self.assertIn("supported: ", result["message"]) + + @patch("tee_gateway.controllers.chat_controller.get_tee_keys") + @patch("tee_gateway.controllers.chat_controller.get_chat_model_cached") + @patch("tee_gateway.controllers.chat_controller.connexion") + def test_ratio_is_bound_as_gemini_image_config( + self, mock_connexion, mock_get_model, mock_get_tee_keys + ): + mock_connexion.request.is_json = True + mock_connexion.request.get_json.return_value = self._request( + aspect_ratio="16:9" + ) + model = _mock_image_model() + mock_get_model.return_value = model + mock_get_tee_keys.return_value = _mock_tee_keys() + + create_chat_completion(None) + + model.bind.assert_called_once_with(image_config={"aspect_ratio": "16:9"}) + + @patch("tee_gateway.controllers.chat_controller.get_tee_keys") + @patch("tee_gateway.controllers.chat_controller.get_chat_model_cached") + @patch("tee_gateway.controllers.chat_controller.connexion") + def test_ratio_survives_bound_tools( + self, mock_connexion, mock_get_model, mock_get_tee_keys + ): + # bind_tools() re-binds the base model, so a kwarg bound before it is + # dropped: the image_config has to be bound onto the tools-bound model. + mock_connexion.request.is_json = True + mock_connexion.request.get_json.return_value = self._request( + aspect_ratio="16:9", + tools=[ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + ) + model = _mock_image_model() + tools_bound = _mock_image_model() + model.bind_tools.return_value = tools_bound + mock_get_model.return_value = model + mock_get_tee_keys.return_value = _mock_tee_keys() + + create_chat_completion(None) + + model.bind.assert_not_called() + tools_bound.bind.assert_called_once_with(image_config={"aspect_ratio": "16:9"}) + + @patch("tee_gateway.controllers.chat_controller.get_tee_keys") + @patch("tee_gateway.controllers.chat_controller.get_chat_model_cached") + @patch("tee_gateway.controllers.chat_controller.connexion") + def test_auto_binds_nothing( + self, mock_connexion, mock_get_model, mock_get_tee_keys + ): + mock_connexion.request.is_json = True + mock_connexion.request.get_json.return_value = self._request( + aspect_ratio="auto" + ) + model = _mock_image_model() + mock_get_model.return_value = model + mock_get_tee_keys.return_value = _mock_tee_keys() + + result = create_chat_completion(None) + + model.bind.assert_not_called() + self.assertIn("choices", result) + + +def _mock_image_model() -> Mock: + """A LangChain chat model stand-in that returns one inline image.""" + response = Mock() + response.content = [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,QUJD"}} + ] + response.tool_calls = [] + response.usage_metadata = { + "input_tokens": 5, + "output_tokens": 3, + "total_tokens": 8, + } + model = Mock() + model.invoke.return_value = response + model.bind.return_value = model + model.bind_tools.return_value = model + return model + + +def _mock_tee_keys() -> Mock: + keys = Mock() + keys.sign_data.return_value = "bW9ja3NpZ25hdHVyZQ==" + keys.get_tee_id.return_value = "abcdef01" * 8 + return keys + + class TestChatController(BaseTestCase): """ChatController integration test stubs""" diff --git a/tee_gateway/test/test_image_generation.py b/tee_gateway/test/test_image_generation.py index 0ac0f2d..4dcdd0e 100644 --- a/tee_gateway/test/test_image_generation.py +++ b/tee_gateway/test/test_image_generation.py @@ -159,13 +159,13 @@ def test_invalid_aspect_ratio_is_rejected_before_provider_call(self): client = MagicMock() with patch.object(llm_backend, "openai_http_client", client): with self.assertRaisesRegex(ValueError, "Unsupported aspect_ratio"): - generate_images(GPT_IMAGE, "a banner", aspect_ratio="16:9") + generate_images(GPT_IMAGE, "a banner", aspect_ratio="7:3") client.post.assert_not_called() def test_gemini_aspect_ratio_uses_image_config_shape(self): cfg = get_model_config("gemini-3.1-flash-image") self.assertEqual( - image_generation._aspect_ratio_params(cfg, "16:9"), + image_generation.aspect_ratio_params(cfg, "16:9"), {"aspect_ratio": "16:9"}, ) @@ -836,6 +836,105 @@ def test_malformed_image_parts_are_ignored(self): self.assertEqual(refs, []) +class TestAspectRatioSupport(unittest.TestCase): + """The public ratio -> provider-request translation, and the size tables. + + The tables are checked against each provider's documented size rules so a + new entry that a provider would reject fails here rather than in + production, where it would 400 the whole generation. + """ + + def test_auto_and_unset_take_the_provider_default(self): + cfg = get_model_config(GPT_IMAGE) + for value in (None, "auto", "", " "): + with self.subTest(value=value): + self.assertEqual(image_generation.aspect_ratio_params(cfg, value), {}) + + def test_ratio_is_trimmed_before_lookup(self): + cfg = get_model_config(GPT_IMAGE) + self.assertEqual( + image_generation.aspect_ratio_params(cfg, " 16:9 "), {"size": "1536x864"} + ) + + def test_validate_rejects_unsupported_ratio_for_an_image_model(self): + with self.assertRaisesRegex(ValueError, "Unsupported aspect_ratio"): + image_generation.validate_aspect_ratio(GPT_IMAGE, "7:3") + + def test_validate_ignores_the_field_for_models_that_cannot_use_it(self): + # Text models have no output shape to set; the field is a no-op there + # rather than a rejection (old clients may send it unconditionally). + image_generation.validate_aspect_ratio("gpt-4.1", "16:9") + # An unknown model is reported by the normal request path, not here. + image_generation.validate_aspect_ratio("not-a-model", "16:9") + + def test_gpt_image_sizes_satisfy_openai_size_constraints(self): + # https://developers.openai.com/api/docs/guides/image-generation: + # edges must be multiples of 16, the long:short ratio at most 3:1, + # neither edge over 3840px, and total pixels within [655360, 8294400]. + sizes = get_model_config(GPT_IMAGE).image_aspect_ratios or {} + self.assertIn("16:9", sizes) + for ratio, size in sizes.items(): + with self.subTest(ratio=ratio): + width, height = (int(part) for part in size.split("x")) + self.assertEqual((width % 16, height % 16), (0, 0)) + self.assertLessEqual(max(width, height), 3840) + self.assertLessEqual(max(width, height) / min(width, height), 3.0) + self.assertGreaterEqual(width * height, 655_360) + self.assertLessEqual(width * height, 8_294_400) + # Pixel count stays in the band the flat per-image price + # assumes (gpt-image-2 bills by output tokens, which scale + # with area), and the size really is the advertised shape. + self.assertLessEqual(width * height, 1_600_000) + self._assert_matches_ratio(ratio, width, height) + + def test_bytedance_sizes_satisfy_modelark_size_constraints(self): + # Every ModelArk image model in the registry shares one size table, so + # it has to satisfy the tightest documented window across them: + # pixels >= 2560x1440 (Seedream 4.5 / 5.0 lite) and <= 2048x2048x1.1025 + # (Seedream 5.0 pro), with the ratio inside [1/16, 16]. + sizes = get_model_config(SEEDREAM).image_aspect_ratios or {} + self.assertEqual(sizes, get_model_config(SEEDANCE_5).image_aspect_ratios) + for ratio, size in sizes.items(): + with self.subTest(ratio=ratio): + width, height = (int(part) for part in size.split("x")) + self.assertGreaterEqual(width * height, 3_686_400) + self.assertLessEqual(width * height, 4_624_220) + self.assertLessEqual(max(width, height) / min(width, height), 16) + self._assert_matches_ratio(ratio, width, height) + + def test_glm_sizes_are_zai_recommended_resolutions(self): + # Z.ai documents these seven resolutions for glm-image and requires + # both edges to be multiples of 32, within 512-2048px. + sizes = get_model_config(GLM_IMAGE).image_aspect_ratios or {} + self.assertEqual( + set(sizes.values()), + { + "1280x1280", + "1568x1056", + "1056x1568", + "1472x1088", + "1088x1472", + "1728x960", + "960x1728", + }, + ) + for ratio, size in sizes.items(): + with self.subTest(ratio=ratio): + width, height = (int(part) for part in size.split("x")) + self.assertEqual((width % 32, height % 32), (0, 0)) + self.assertGreaterEqual(min(width, height), 512) + self.assertLessEqual(max(width, height), 2048) + # Z.ai's own recommendations only approximate their labels. + self._assert_matches_ratio(ratio, width, height, tolerance=0.03) + + def _assert_matches_ratio(self, ratio, width, height, tolerance=0.001): + """Assert WxH is the shape its public ``aspect_ratio`` label claims.""" + left, right = (float(part) for part in ratio.split(":")) + self.assertAlmostEqual( + width / height, left / right, delta=(left / right) * tolerance + ) + + class TestPerImageBilling(unittest.TestCase): """Flat per-image pricing, independent of token usage."""