From c86d59542cdf0ab44fd13f8ce1695a3654a33bce Mon Sep 17 00:00:00 2001 From: Abi Raja Date: Thu, 26 Feb 2026 12:54:18 -0500 Subject: [PATCH] Support per-image aspect ratios in generate_images --- backend/agent/providers/openai.py | 3 + backend/agent/tools/definitions.py | 27 ++- backend/agent/tools/runtime.py | 109 ++++++++- backend/agent/tools/summaries.py | 16 +- backend/image_generation/aspect_ratios.py | 51 ++++ backend/image_generation/generation.py | 21 +- backend/prompts/system_prompt.py | 2 +- backend/tests/test_agent_tools.py | 22 ++ backend/tests/test_batching.py | 217 +++++++++++++++++- .../src/components/agent/AgentActivity.tsx | 40 +++- 10 files changed, 476 insertions(+), 32 deletions(-) create mode 100644 backend/image_generation/aspect_ratios.py diff --git a/backend/agent/providers/openai.py b/backend/agent/providers/openai.py index 28c6d2cfb..6b3213505 100644 --- a/backend/agent/providers/openai.py +++ b/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 diff --git a/backend/agent/tools/definitions.py b/backend/agent/tools/definitions.py index d29f4c79a..99963ab9a 100644 --- a/backend/agent/tools/definitions.py +++ b/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"], } diff --git a/backend/agent/tools/runtime.py b/backend/agent/tools/runtime.py index b16805075..8310eec56 100644 --- a/backend/agent/tools/runtime.py +++ b/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) diff --git a/backend/agent/tools/summaries.py b/backend/agent/tools/summaries.py index c87ba88ed..24ec7c33d 100644 --- a/backend/agent/tools/summaries.py +++ b/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": diff --git a/backend/image_generation/aspect_ratios.py b/backend/image_generation/aspect_ratios.py new file mode 100644 index 000000000..5c85a1baf --- /dev/null +++ b/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" diff --git a/backend/image_generation/generation.py b/backend/image_generation/generation.py index c412043d0..ffc17ff7e 100644 --- a/backend/image_generation/generation.py +++ b/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, diff --git a/backend/prompts/system_prompt.py b/backend/prompts/system_prompt.py index 0fa1d9b94..f8cfb8282 100644 --- a/backend/prompts/system_prompt.py +++ b/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. diff --git a/backend/tests/test_agent_tools.py b/backend/tests/test_agent_tools.py index e85e1fa53..3e2375565 100644 --- a/backend/tests/test_agent_tools.py +++ b/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"] diff --git a/backend/tests/test_batching.py b/backend/tests/test_batching.py index 492905e60..8f5148e45 100644 --- a/backend/tests/test_batching.py +++ b/backend/tests/test_batching.py @@ -3,6 +3,7 @@ import pytest from image_generation import generation +from image_generation.aspect_ratios import DEFAULT_ASPECT_RATIO from agent.tools.runtime import AgentToolRuntime from agent.tools.types import ToolCall from agent.state import AgentFileState @@ -16,11 +17,13 @@ async def test_process_tasks_batches_replicate_calls( concurrent = 0 max_concurrent = 0 + seen_aspect_ratios: list[str] = [] - async def tracking_generate(prompt: str, api_key: str) -> str: + async def tracking_generate(prompt: str, api_key: str, aspect_ratio: str) -> str: nonlocal concurrent, max_concurrent concurrent += 1 max_concurrent = max(max_concurrent, concurrent) + seen_aspect_ratios.append(aspect_ratio) await asyncio.sleep(0.01) concurrent -= 1 return f"url-for-{prompt}" @@ -33,6 +36,218 @@ async def tracking_generate(prompt: str, api_key: str) -> str: assert len(results) == 7 assert results == [f"url-for-prompt-{i}" for i in range(7)] assert max_concurrent <= 3 + assert seen_aspect_ratios == [DEFAULT_ASPECT_RATIO] * 7 + + +@pytest.mark.asyncio +async def test_process_tasks_ignores_aspect_ratio_for_dalle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + async def tracking_dalle(prompt: str, api_key: str, base_url: str | None) -> str: + captured["prompt"] = prompt + captured["api_key"] = api_key + captured["base_url"] = base_url + return "url-for-dalle" + + monkeypatch.setattr(generation, "generate_image_dalle", tracking_dalle) + + results = await generation.process_tasks( + ["prompt-1"], + "openai-key", + None, + "dalle3", + aspect_ratio="9:16", + ) + + assert results == ["url-for-dalle"] + assert captured["prompt"] == "prompt-1" + assert captured["api_key"] == "openai-key" + assert captured["base_url"] is None + + +@pytest.mark.asyncio +async def test_generate_images_uses_default_aspect_ratio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("agent.tools.runtime.REPLICATE_API_KEY", "fake-key") + calls: list[dict[str, object]] = [] + + async def fake_process_tasks( + prompts: list[str], + api_key: str, + base_url: str | None, + model: str, + aspect_ratio: str = DEFAULT_ASPECT_RATIO, + ) -> list[str]: + calls.append( + { + "prompts": prompts, + "api_key": api_key, + "base_url": base_url, + "model": model, + "aspect_ratio": aspect_ratio, + } + ) + return [f"https://example.com/{prompt}.png" for prompt in prompts] + + monkeypatch.setattr("agent.tools.runtime.process_tasks", fake_process_tasks) + + runtime = AgentToolRuntime( + file_state=AgentFileState(), + should_generate_images=True, + openai_api_key=None, + openai_base_url=None, + ) + result = await runtime.execute( + ToolCall( + id="test", + name="generate_images", + arguments={"prompts": [{"prompt": "test prompt"}]}, + ) + ) + + assert result.ok + assert len(calls) == 1 + assert calls[0]["aspect_ratio"] == DEFAULT_ASPECT_RATIO + assert result.result["images"][0]["aspect_ratio"] == DEFAULT_ASPECT_RATIO + + +@pytest.mark.asyncio +async def test_generate_images_supports_per_prompt_aspect_ratios( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("agent.tools.runtime.REPLICATE_API_KEY", "fake-key") + calls: list[dict[str, object]] = [] + + async def fake_process_tasks( + prompts: list[str], + api_key: str, + base_url: str | None, + model: str, + aspect_ratio: str = DEFAULT_ASPECT_RATIO, + ) -> list[str]: + calls.append( + { + "prompts": prompts, + "aspect_ratio": aspect_ratio, + } + ) + return [ + f"https://example.com/{aspect_ratio}/{index}.png" + for index, _prompt in enumerate(prompts) + ] + + monkeypatch.setattr("agent.tools.runtime.process_tasks", fake_process_tasks) + + runtime = AgentToolRuntime( + file_state=AgentFileState(), + should_generate_images=True, + openai_api_key=None, + openai_base_url=None, + ) + result = await runtime.execute( + ToolCall( + id="test", + name="generate_images", + arguments={ + "prompts": [ + {"prompt": "square prompt"}, + {"prompt": "wide prompt", "aspect_ratio": "16:9"}, + {"prompt": "tall prompt", "aspect_ratio": "9:16"}, + ] + }, + ) + ) + + assert result.ok + assert [call["aspect_ratio"] for call in calls] == [ + DEFAULT_ASPECT_RATIO, + "16:9", + "9:16", + ] + assert [item["aspect_ratio"] for item in result.result["images"]] == [ + DEFAULT_ASPECT_RATIO, + "16:9", + "9:16", + ] + + +@pytest.mark.asyncio +async def test_generate_images_rejects_batch_level_aspect_ratio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("agent.tools.runtime.REPLICATE_API_KEY", "fake-key") + + runtime = AgentToolRuntime( + file_state=AgentFileState(), + should_generate_images=True, + openai_api_key=None, + openai_base_url=None, + ) + result = await runtime.execute( + ToolCall( + id="test", + name="generate_images", + arguments={"prompts": [{"prompt": "test prompt"}], "aspect_ratio": "16:9"}, + ) + ) + + assert not result.ok + assert result.summary["error"] == "Invalid aspect_ratio" + + +@pytest.mark.asyncio +async def test_generate_images_rejects_invalid_per_prompt_aspect_ratio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("agent.tools.runtime.REPLICATE_API_KEY", "fake-key") + + runtime = AgentToolRuntime( + file_state=AgentFileState(), + should_generate_images=True, + openai_api_key=None, + openai_base_url=None, + ) + result = await runtime.execute( + ToolCall( + id="test", + name="generate_images", + arguments={ + "prompts": [ + {"prompt": "test prompt", "aspect_ratio": "7:5"}, + ] + }, + ) + ) + + assert not result.ok + assert result.summary["error"] == "Invalid aspect_ratio" + + +@pytest.mark.asyncio +async def test_generate_images_rejects_non_object_prompt_items( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("agent.tools.runtime.REPLICATE_API_KEY", "fake-key") + + runtime = AgentToolRuntime( + file_state=AgentFileState(), + should_generate_images=True, + openai_api_key=None, + openai_base_url=None, + ) + result = await runtime.execute( + ToolCall( + id="test", + name="generate_images", + arguments={"prompts": ["legacy string prompt"]}, + ) + ) + + assert not result.ok + assert result.summary["error"] == "Invalid prompts payload" @pytest.mark.asyncio diff --git a/frontend/src/components/agent/AgentActivity.tsx b/frontend/src/components/agent/AgentActivity.tsx index d01af32d2..bb3f849da 100644 --- a/frontend/src/components/agent/AgentActivity.tsx +++ b/frontend/src/components/agent/AgentActivity.tsx @@ -245,11 +245,32 @@ function renderToolDetails(event: AgentEvent, variantCode?: string) { {/* While running: show prompts with dividers */} {event.status === "running" && input?.prompts && Array.isArray(input.prompts) && (
- {input.prompts.map((prompt: string, index: number) => ( -
- {prompt} -
- ))} + {input.prompts.map((promptEntry: unknown, index: number) => { + const promptText = + typeof promptEntry === "object" && + promptEntry !== null && + "prompt" in promptEntry && + typeof (promptEntry as { prompt?: unknown }).prompt === "string" + ? (promptEntry as { prompt: string }).prompt + : ""; + const promptAspectRatio = + typeof promptEntry === "object" && + promptEntry !== null && + "aspect_ratio" in promptEntry && + typeof (promptEntry as { aspect_ratio?: unknown }).aspect_ratio === "string" + ? (promptEntry as { aspect_ratio: string }).aspect_ratio + : null; + return ( +
+
{promptText}
+ {promptAspectRatio && ( +
+ {promptAspectRatio} +
+ )} +
+ ); + })}
)} {/* After complete: 50/50 image left, prompt right */} @@ -271,8 +292,13 @@ function renderToolDetails(event: AgentEvent, variantCode?: string) { )} -
- {item.prompt} +
+
{item.prompt}
+ {item.aspect_ratio && ( +
+ {item.aspect_ratio} +
+ )}
))}