Support per-image aspect ratios in generate_images

backend/agent/providers/openai.py

@@ -97,6 +97,9 @@ def transform(node: Dict[str, Any], in_object_property: bool = False) -> None:
 
         if in_object_property and node_type is not None:
             node["type"] = _nullable_type(node_type)
+            enum_values = node.get("enum")
+            if isinstance(enum_values, list) and None not in enum_values:
+                node["enum"] = [*enum_values, None]
 
     transform(schema_copy, in_object_property=False)
     return schema_copy

backend/agent/tools/definitions.py

@@ -1,6 +1,7 @@
 from typing import Any, Dict, List
 
 from agent.tools.types import CanonicalToolDefinition
+from image_generation.aspect_ratios import SUPPORTED_ASPECT_RATIOS
 
 
 def _create_schema() -> Dict[str, Any]:
@@ -57,16 +58,32 @@ def _edit_schema() -> Dict[str, Any]:
 
 
 def _image_schema() -> Dict[str, Any]:
+    prompt_item_schema: Dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "prompt": {
+                "type": "string",
+                "description": "Prompt describing a single image to generate.",
+            },
+            "aspect_ratio": {
+                "type": "string",
+                "description": (
+                    "Optional image aspect ratio for this prompt. Choose one that "
+                    "fits where the image will be used."
+                ),
+                "enum": list(SUPPORTED_ASPECT_RATIOS),
+            },
+        },
+        "required": ["prompt"],
+        "additionalProperties": False,
+    }
     return {
         "type": "object",
         "properties": {
             "prompts": {
                 "type": "array",
-                "items": {
-                    "type": "string",
-                    "description": "Prompt describing a single image to generate.",
-                },
-            }
+                "items": prompt_item_schema,
+            },
         },
         "required": ["prompts"],
     }

backend/agent/tools/runtime.py

@@ -1,9 +1,15 @@
 # pyright: reportUnknownVariableType=false
 import asyncio
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, Union, cast
 
 from codegen.utils import extract_html_content
 from config import REPLICATE_API_KEY
+from image_generation.aspect_ratios import (
+    AspectRatio,
+    DEFAULT_ASPECT_RATIO,
+    SUPPORTED_ASPECT_RATIOS,
+    is_supported_aspect_ratio,
+)
 from image_generation.generation import process_tasks
 from image_generation.replicate import remove_background
 
@@ -12,6 +18,11 @@
 from agent.tools.summaries import summarize_text
 
 
+class ImagePromptRequest(TypedDict):
+    prompt: str
+    aspect_ratio: AspectRatio
+
+
 class AgentToolRuntime:
     def __init__(
         self,
@@ -195,14 +206,67 @@ async def _generate_images(self, args: Dict[str, Any]) -> ToolExecutionResult:
                 summary={"error": "Missing prompts"},
             )
 
-        cleaned = [prompt.strip() for prompt in prompts if isinstance(prompt, str)]
-        unique_prompts = list(dict.fromkeys([p for p in cleaned if p]))
-        if not unique_prompts:
+        if "aspect_ratio" in args:
+            return ToolExecutionResult(
+                ok=False,
+                result={
+                    "error": (
+                        "Top-level aspect_ratio is not supported. "
+                        "Set aspect_ratio per prompt item in prompts[]."
+                    )
+                },
+                summary={"error": "Invalid aspect_ratio"},
+            )
+
+        requests: list[ImagePromptRequest] = []
+        for item in prompts:
+            if not isinstance(item, dict):
+                return ToolExecutionResult(
+                    ok=False,
+                    result={
+                        "error": (
+                            "Each prompts[] entry must be an object with "
+                            "'prompt' and optional 'aspect_ratio'."
+                        )
+                    },
+                    summary={"error": "Invalid prompts payload"},
+                )
+
+            prompt = ensure_str(item.get("prompt")).strip()
+            if not prompt:
+                return ToolExecutionResult(
+                    ok=False,
+                    result={"error": "Each prompts[] entry requires a non-empty prompt"},
+                    summary={"error": "Missing prompt"},
+                )
+
+            raw_aspect_ratio = item.get("aspect_ratio", DEFAULT_ASPECT_RATIO)
+            if not is_supported_aspect_ratio(raw_aspect_ratio):
+                return ToolExecutionResult(
+                    ok=False,
+                    result={
+                        "error": (
+                            f"Unsupported aspect_ratio: {raw_aspect_ratio}. "
+                            f"Supported values: {', '.join(SUPPORTED_ASPECT_RATIOS)}"
+                        )
+                    },
+                    summary={"error": "Invalid aspect_ratio"},
+                )
+            requests.append(
+                {
+                    "prompt": prompt,
+                    "aspect_ratio": cast(AspectRatio, raw_aspect_ratio),
+                }
+            )
+
+        if not requests:
             return ToolExecutionResult(
                 ok=False,
                 result={"error": "No valid prompts provided"},
                 summary={"error": "No valid prompts"},
             )
+
+        model: Literal["dalle3", "flux"]
         if REPLICATE_API_KEY:
             model = "flux"
             api_key = REPLICATE_API_KEY
@@ -218,19 +282,42 @@ async def _generate_images(self, args: Dict[str, Any]) -> ToolExecutionResult:
             api_key = self.openai_api_key
             base_url = self.openai_base_url
 
-        generated = await process_tasks(unique_prompts, api_key, base_url, model)  # type: ignore
-        merged_results = {
-            prompt: url for prompt, url in zip(unique_prompts, generated)
-        }
+        grouped_requests: dict[AspectRatio, list[tuple[int, str]]] = {}
+        for index, request in enumerate(requests):
+            grouped_requests.setdefault(request["aspect_ratio"], []).append(
+                (index, request["prompt"])
+            )
+
+        grouped_items = list(grouped_requests.items())
+        grouped_results = await asyncio.gather(
+            *[
+                process_tasks(
+                    prompts=[prompt for _, prompt in entries],
+                    api_key=api_key,
+                    base_url=base_url,
+                    model=model,
+                    aspect_ratio=aspect_ratio,
+                )
+                for aspect_ratio, entries in grouped_items
+            ]
+        )
+
+        generated: list[str | None] = [None] * len(requests)
+        for (aspect_ratio, entries), urls in zip(grouped_items, grouped_results):
+            _ = aspect_ratio
+            for (index, _prompt), url in zip(entries, urls):
+                generated[index] = url
+
         summary_items = [
             {
-                "prompt": prompt,
+                "prompt": request["prompt"],
                 "url": url,
+                "aspect_ratio": request["aspect_ratio"],
                 "status": "ok" if url else "error",
             }
-            for prompt, url in merged_results.items()
+            for request, url in zip(requests, generated)
         ]
-        result = {"images": merged_results}
+        result = {"images": summary_items}
         summary = {"images": summary_items}
         return ToolExecutionResult(ok=True, result=result, summary=summary)
 

backend/agent/tools/summaries.py

@@ -49,9 +49,21 @@ def summarize_tool_input(tool_call: ToolCall, file_state: AgentFileState) -> Dic
     if tool_call.name == "generate_images":
         prompts = args.get("prompts") or []
         if isinstance(prompts, list):
+            summarized_prompts: list[dict[str, str]] = []
+            for entry in prompts:
+                if not isinstance(entry, dict):
+                    continue
+                prompt = ensure_str(entry.get("prompt")).strip()
+                if not prompt:
+                    continue
+                summarized_entry: dict[str, str] = {"prompt": prompt}
+                aspect_ratio = entry.get("aspect_ratio")
+                if isinstance(aspect_ratio, str) and aspect_ratio.strip():
+                    summarized_entry["aspect_ratio"] = aspect_ratio
+                summarized_prompts.append(summarized_entry)
             return {
-                "count": len(prompts),
-                "prompts": [ensure_str(p) for p in prompts],
+                "count": len(summarized_prompts),
+                "prompts": summarized_prompts,
             }
 
     if tool_call.name == "remove_background":

backend/image_generation/aspect_ratios.py

@@ -0,0 +1,51 @@
+from typing import Literal, cast
+
+
+AspectRatio = Literal[
+    "1:1",
+    "16:9",
+    "9:16",
+    "3:2",
+    "2:3",
+    "4:3",
+    "3:4",
+    "5:4",
+    "4:5",
+    "21:9",
+    "9:21",
+]
+
+SUPPORTED_ASPECT_RATIOS: tuple[AspectRatio, ...] = (
+    "1:1",
+    "16:9",
+    "9:16",
+    "3:2",
+    "2:3",
+    "4:3",
+    "3:4",
+    "5:4",
+    "4:5",
+    "21:9",
+    "9:21",
+)
+DEFAULT_ASPECT_RATIO: AspectRatio = "1:1"
+
+DalleImageSize = Literal["1024x1024", "1792x1024", "1024x1792"]
+
+
+def is_supported_aspect_ratio(value: object) -> bool:
+    return isinstance(value, str) and value in SUPPORTED_ASPECT_RATIOS
+
+
+def normalize_aspect_ratio(value: object) -> AspectRatio:
+    if is_supported_aspect_ratio(value):
+        return cast(AspectRatio, value)
+    return DEFAULT_ASPECT_RATIO
+
+
+def aspect_ratio_to_dalle_size(aspect_ratio: AspectRatio) -> DalleImageSize:
+    if aspect_ratio == "1:1":
+        return "1024x1024"
+
+    width, height = (int(part) for part in aspect_ratio.split(":"))
+    return "1792x1024" if width > height else "1024x1792"

backend/image_generation/generation.py

@@ -4,6 +4,10 @@
 
 from openai import AsyncOpenAI
 
+from image_generation.aspect_ratios import (
+    AspectRatio,
+    DEFAULT_ASPECT_RATIO,
+)
 from image_generation.replicate import call_replicate
 
 
@@ -15,16 +19,17 @@ async def process_tasks(
     api_key: str,
     base_url: str | None,
     model: Literal["dalle3", "flux"],
+    aspect_ratio: AspectRatio = DEFAULT_ASPECT_RATIO,
 ) -> List[Union[str, None]]:
     start_time = time.time()
     if model == "dalle3":
         tasks = [generate_image_dalle(prompt, api_key, base_url) for prompt in prompts]
         results = await asyncio.gather(*tasks, return_exceptions=True)
     else:
-        results: list[str | BaseException] = []
+        results: list[str | None | BaseException] = []
         for i in range(0, len(prompts), REPLICATE_BATCH_SIZE):
             batch = prompts[i : i + REPLICATE_BATCH_SIZE]
-            tasks = [generate_image_replicate(p, api_key) for p in batch]
+            tasks = [generate_image_replicate(p, api_key, aspect_ratio) for p in batch]
             results.extend(await asyncio.gather(*tasks, return_exceptions=True))
     end_time = time.time()
     generation_time = end_time - start_time
@@ -42,7 +47,9 @@ async def process_tasks(
 
 
 async def generate_image_dalle(
-    prompt: str, api_key: str, base_url: str | None
+    prompt: str,
+    api_key: str,
+    base_url: str | None,
 ) -> Union[str, None]:
     client = AsyncOpenAI(api_key=api_key, base_url=base_url)
     res = await client.images.generate(
@@ -59,12 +66,16 @@ async def generate_image_dalle(
     return res.data[0].url
 
 
-async def generate_image_replicate(prompt: str, api_key: str) -> str:
+async def generate_image_replicate(
+    prompt: str,
+    api_key: str,
+    aspect_ratio: AspectRatio = DEFAULT_ASPECT_RATIO,
+) -> str:
     # We use Flux 2 Klein
     return await call_replicate(
         {
             "prompt": prompt,
-            "aspect_ratio": "1:1",
+            "aspect_ratio": aspect_ratio,
             "output_format": "png",
         },
         api_key,

backend/prompts/system_prompt.py

@@ -15,7 +15,7 @@
 - For a brand new app, call create_file exactly once with the full HTML.
 - For updates, call edit_file using exact string replacements. Do NOT regenerate the entire file.
 - Do not output raw HTML in chat. Any code changes must go through tools.
-- When available, use generate_images to create image URLs from prompts (you may pass multiple prompts). The image generation AI is not capable of generating images with a transparent background.
+- When available, use generate_images to create image URLs from prompts (you may pass multiple prompts). Pass `prompts` as objects like `{ "prompt": "...", "aspect_ratio": "16:9" }`; `aspect_ratio` is optional and applies per image (default: 1:1). Supported values: 1:1, 16:9, 9:16, 3:2, 2:3, 4:3, 3:4, 5:4, 4:5, 21:9, 9:21. Some models may ignore this parameter. The image generation AI is not capable of generating images with a transparent background.
 - Use remove_background to remove backgrounds from provided image URLs when needed (you may pass multiple image URLs).
 - Use retrieve_option to fetch the full HTML for a specific option (1-based option_number) when a user references another option.
 

backend/tests/test_agent_tools.py

@@ -1,4 +1,6 @@
 from agent.tools import canonical_tool_definitions
+from agent.providers.openai import serialize_openai_tools
+from image_generation.aspect_ratios import SUPPORTED_ASPECT_RATIOS
 
 
 def test_canonical_tool_definitions_include_generate_images_when_enabled() -> None:
@@ -9,3 +11,23 @@ def test_canonical_tool_definitions_include_generate_images_when_enabled() -> No
 def test_canonical_tool_definitions_exclude_generate_images_when_disabled() -> None:
     tool_names = [tool.name for tool in canonical_tool_definitions(False)]
     assert "generate_images" not in tool_names
+
+
+def test_generate_images_schema_exposes_supported_aspect_ratios() -> None:
+    tools = canonical_tool_definitions(True)
+    generate_images_tool = next(tool for tool in tools if tool.name == "generate_images")
+    schema = generate_images_tool.parameters
+    prompt_item = schema["properties"]["prompts"]["items"]
+    assert prompt_item["type"] == "object"
+    assert prompt_item["properties"]["aspect_ratio"]["enum"] == list(SUPPORTED_ASPECT_RATIOS)
+
+
+def test_openai_generate_images_schema_is_compatible() -> None:
+    tools = canonical_tool_definitions(True)
+    serialized_tools = serialize_openai_tools(tools)
+    generate_images_tool = next(
+        tool for tool in serialized_tools if tool["name"] == "generate_images"
+    )
+    prompt_item = generate_images_tool["parameters"]["properties"]["prompts"]["items"]
+    assert "oneOf" not in prompt_item
+    assert None in prompt_item["properties"]["aspect_ratio"]["enum"]