Add image generation tool support

Enable the Responses API image_generation tool with configurable quality,
size, and background options. Handle streaming events including partial
images, and render generated images inline in the chat.

Closes #42

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: chriscarrollsmith

routers/chat.py

@@ -25,8 +25,10 @@
     ResponseMcpCallInProgressEvent, ResponseMcpCallArgumentsDeltaEvent,
     ResponseWebSearchCallInProgressEvent, ResponseWebSearchCallSearchingEvent,
     ResponseWebSearchCallCompletedEvent,
+    ResponseImageGenCallInProgressEvent, ResponseImageGenCallGeneratingEvent,
+    ResponseImageGenCallCompletedEvent, ResponseImageGenCallPartialImageEvent,
 )
-from openai.types.responses.response_output_item import McpApprovalRequest
+from openai.types.responses.response_output_item import McpApprovalRequest, ImageGenerationCall
 from openai.types.responses import ResponseFunctionToolCall, ResponseComputerToolCall
 from openai.types.responses.response_code_interpreter_tool_call import ResponseCodeInterpreterToolCall
 from openai._types import NOT_GIVEN
@@ -183,6 +185,18 @@ async def event_generator() -> AsyncGenerator[str, None]:
             tools.append(ws_tool)
         if "computer_use" in enabled_tools:
             tools.append(build_computer_tool())
+        if "image_generation" in enabled_tools:
+            ig_tool: Dict[str, Any] = {"type": "image_generation"}
+            ig_quality = os.getenv("IMAGE_GENERATION_QUALITY", "auto").strip()
+            if ig_quality in {"low", "medium", "high"}:
+                ig_tool["quality"] = ig_quality
+            ig_size = os.getenv("IMAGE_GENERATION_SIZE", "auto").strip()
+            if ig_size in {"1024x1024", "1536x1024", "1024x1536"}:
+                ig_tool["size"] = ig_size
+            ig_background = os.getenv("IMAGE_GENERATION_BACKGROUND", "auto").strip()
+            if ig_background in {"opaque", "transparent"}:
+                ig_tool["background"] = ig_background
+            tools.append(ig_tool)
 
         try:
             stream = await client.responses.create(  # type: ignore[call-overload]
@@ -227,7 +241,8 @@ async def iterate_stream(s, response_id: str = "") -> AsyncGenerator[str, None]:
                                 ResponseMcpCallInProgressEvent() | \
                                 ResponseFunctionCallArgumentsDoneEvent() | \
                                 ResponseWebSearchCallInProgressEvent() | \
-                                ResponseWebSearchCallCompletedEvent():
+                                ResponseWebSearchCallCompletedEvent() | \
+                                ResponseImageGenCallCompletedEvent():
                                 # Don't need to handle "in progress" or intermediate "done" events
                                 # (though long-running code interpreter interpreting might warrant handling)
                                 continue
@@ -351,6 +366,28 @@ async def iterate_stream(s, response_id: str = "") -> AsyncGenerator[str, None]:
                                         )
                                     )
 
+                            case ResponseImageGenCallInProgressEvent() | ResponseImageGenCallGeneratingEvent():
+                                current_item_id = event.item_id
+                                yield sse_format(
+                                    "toolCallCreated",
+                                    templates.get_template('components/assistant-step.html').render(
+                                        step_type='toolCall',
+                                        step_id=event.item_id,
+                                        content="Generating image..."
+                                    )
+                                )
+
+                            case ResponseImageGenCallPartialImageEvent():
+                                # Display partial image as it streams in
+                                img_html = (
+                                    f'<div class="imageOutput">'
+                                    f'<img src="data:image/png;base64,{event.partial_image_b64}" '
+                                    f'alt="Partial image (generating...)" '
+                                    f'onclick="openImagePreview(this.src)" style="cursor:pointer" />'
+                                    f'</div>'
+                                )
+                                yield sse_format("imageOutput", img_html)
+
                             case ResponseContentPartAddedEvent():
                                 # This event indicates the start of annotations; skip creating a new assistantMessage
                                 continue
@@ -504,6 +541,18 @@ async def iterate_stream(s, response_id: str = "") -> AsyncGenerator[str, None]:
                                         async for out in iterate_stream(next_stream, response_id):
                                             yield out
 
+                                elif isinstance(event.item, ImageGenerationCall):
+                                    current_item_id = event.item.id
+                                    if event.item.result:
+                                        img_html = (
+                                            f'<div class="imageOutput">'
+                                            f'<img src="data:image/png;base64,{event.item.result}" '
+                                            f'alt="Generated image" '
+                                            f'onclick="openImagePreview(this.src)" style="cursor:pointer" />'
+                                            f'</div>'
+                                        )
+                                        yield sse_format("imageOutput", img_html)
+
                                 elif isinstance(event.item, ResponseComputerToolCall):
                                     current_item_id = event.item.id
                                     call_id = event.item.call_id

routers/setup.py

@@ -88,6 +88,11 @@ async def read_setup(
     web_search_region = os.getenv("WEB_SEARCH_LOCATION_REGION", "")
     web_search_timezone = os.getenv("WEB_SEARCH_LOCATION_TIMEZONE", "")
 
+    # Image generation config
+    image_generation_quality = os.getenv("IMAGE_GENERATION_QUALITY", "auto")
+    image_generation_size = os.getenv("IMAGE_GENERATION_SIZE", "auto")
+    image_generation_background = os.getenv("IMAGE_GENERATION_BACKGROUND", "auto")
+
     if not openai_api_key:
         setup_message = "OpenAI API key is missing."
 
@@ -141,6 +146,9 @@ async def read_setup(
             "web_search_city": web_search_city,
             "web_search_region": web_search_region,
             "web_search_timezone": web_search_timezone,
+            "image_generation_quality": image_generation_quality,
+            "image_generation_size": image_generation_size,
+            "image_generation_background": image_generation_background,
         }
     )
 
@@ -235,6 +243,9 @@ async def save_app_config(
     web_search_city: Optional[str] = Form(default=None),
     web_search_region: Optional[str] = Form(default=None),
     web_search_timezone: Optional[str] = Form(default=None),
+    image_generation_quality: Optional[str] = Form(default=None),
+    image_generation_size: Optional[str] = Form(default=None),
+    image_generation_background: Optional[str] = Form(default=None),
 ) -> RedirectResponse:
     status = "success"
     message_text = ""
@@ -340,6 +351,21 @@ def get_or_empty(items: List[str], i: int) -> str:
             update_env_file("WEB_SEARCH_LOCATION_TIMEZONE", (web_search_timezone or "").strip())
             status = "success"
             message_text = "Web search configuration saved."
+        elif action == "save_image_generation_config":
+            quality = (image_generation_quality or "auto").strip()
+            if quality not in {"auto", "low", "medium", "high"}:
+                quality = "auto"
+            size = (image_generation_size or "auto").strip()
+            if size not in {"auto", "1024x1024", "1536x1024", "1024x1536"}:
+                size = "auto"
+            background = (image_generation_background or "auto").strip()
+            if background not in {"auto", "opaque", "transparent"}:
+                background = "auto"
+            update_env_file("IMAGE_GENERATION_QUALITY", quality)
+            update_env_file("IMAGE_GENERATION_SIZE", size)
+            update_env_file("IMAGE_GENERATION_BACKGROUND", background)
+            status = "success"
+            message_text = "Image generation configuration saved."
         else:
             # Standard app config save
             if model is None or instructions is None:

templates/setup.html

@@ -95,6 +95,10 @@ <h3>Model &amp; Instructions</h3>
                       <input type="checkbox" name="tool_types" value="computer_use" {% if "computer_use" in current_tools %}checked{% endif %}>
                       Computer Use
                     </label>
+                    <label>
+                      <input type="checkbox" name="tool_types" value="image_generation" {% if "image_generation" in current_tools %}checked{% endif %}>
+                      Image Generation
+                    </label>
                   </fieldset>
 
                   {# ── Display ── #}
@@ -165,6 +169,50 @@ <h3>Web Search</h3>
 
 
 
+              {# ── Image Generation Configuration ── #}
+              {% if "image_generation" in current_tools %}
+                <div class="setupSection">
+                  <h3>Image Generation</h3>
+                  <p class="setupMessage">Configure image generation quality, size, and background options.</p>
+                  <form action="{{ url_for('save_app_config') }}" method="POST">
+                    <input type="hidden" name="action" value="save_image_generation_config">
+
+                    <div>
+                      <label for="image-gen-quality">Quality:</label>
+                      <select name="image_generation_quality" id="image-gen-quality" class="input">
+                        <option value="auto" {% if image_generation_quality == "auto" %}selected{% endif %}>Auto (default)</option>
+                        <option value="low" {% if image_generation_quality == "low" %}selected{% endif %}>Low (fast, lower cost)</option>
+                        <option value="medium" {% if image_generation_quality == "medium" %}selected{% endif %}>Medium</option>
+                        <option value="high" {% if image_generation_quality == "high" %}selected{% endif %}>High</option>
+                      </select>
+                    </div>
+
+                    <div>
+                      <label for="image-gen-size">Size:</label>
+                      <select name="image_generation_size" id="image-gen-size" class="input">
+                        <option value="auto" {% if image_generation_size == "auto" %}selected{% endif %}>Auto (default)</option>
+                        <option value="1024x1024" {% if image_generation_size == "1024x1024" %}selected{% endif %}>1024x1024 (square)</option>
+                        <option value="1536x1024" {% if image_generation_size == "1536x1024" %}selected{% endif %}>1536x1024 (landscape)</option>
+                        <option value="1024x1536" {% if image_generation_size == "1024x1536" %}selected{% endif %}>1024x1536 (portrait)</option>
+                      </select>
+                    </div>
+
+                    <div>
+                      <label for="image-gen-background">Background:</label>
+                      <select name="image_generation_background" id="image-gen-background" class="input">
+                        <option value="auto" {% if image_generation_background == "auto" %}selected{% endif %}>Auto (default)</option>
+                        <option value="opaque" {% if image_generation_background == "opaque" %}selected{% endif %}>Opaque</option>
+                        <option value="transparent" {% if image_generation_background == "transparent" %}selected{% endif %}>Transparent</option>
+                      </select>
+                    </div>
+
+                    <div class="centered-button-container">
+                      <button type="submit" class="button">Save image generation config</button>
+                    </div>
+                  </form>
+                </div>
+              {% endif %}
+
               {# ── Custom Functions & MCP Servers ── #}
               {% if "function" in current_tools or "mcp" in current_tools %}
                 <form action="{{ url_for('save_app_config') }}" method="POST" id="tool-config-form">

tests/conftest.py

@@ -332,7 +332,25 @@ def env_computer_use_only(app_server: int) -> Iterator[None]:
         yield
 
 
+@pytest.fixture
+def env_image_generation_only(app_server: int) -> Iterator[None]:
+    with _dotenv({**_BASE_ENV, "ENABLED_TOOLS": "image_generation"}):
+        yield
+
+
+@pytest.fixture
+def env_image_generation_with_config(app_server: int) -> Iterator[None]:
+    with _dotenv({
+        **_BASE_ENV,
+        "ENABLED_TOOLS": "image_generation",
+        "IMAGE_GENERATION_QUALITY": "high",
+        "IMAGE_GENERATION_SIZE": "1024x1536",
+        "IMAGE_GENERATION_BACKGROUND": "transparent",
+    }):
+        yield
+
+
 @pytest.fixture
 def env_all_tools(app_server: int) -> Iterator[None]:
-    with _dotenv({**_BASE_ENV, "ENABLED_TOOLS": "file_search,function,mcp,web_search,computer_use"}):
+    with _dotenv({**_BASE_ENV, "ENABLED_TOOLS": "file_search,function,mcp,web_search,computer_use,image_generation"}):
         yield

tests/test_setup_page_rendering.py

@@ -551,6 +551,85 @@ def test_no_other_tool_sections(
         expect(mcp_rows).not_to_be_visible()
 
 
+class TestImageGenerationTool:
+    """Tests for image_generation tool conditional rendering."""
+
+    def test_image_generation_checkbox_visible(
+        self, page: Page, base_url: str, app_server, env_api_key_no_tools
+    ):
+        """Image generation checkbox should be visible in the tools fieldset."""
+        page.goto(f"{base_url}/setup/")
+
+        image_gen_cb = page.locator('input[value="image_generation"]')
+        expect(image_gen_cb).to_be_visible()
+
+    def test_image_generation_checkbox_checked_when_enabled(
+        self, page: Page, base_url: str, app_server, env_image_generation_only
+    ):
+        """The image_generation checkbox should be checked when enabled."""
+        page.goto(f"{base_url}/setup/")
+
+        image_gen_cb = page.locator('input[value="image_generation"]')
+        expect(image_gen_cb).to_be_checked()
+
+    def test_shows_image_generation_config_section(
+        self, page: Page, base_url: str, app_server, env_image_generation_only
+    ):
+        """When image_generation is enabled, the config section should be visible."""
+        page.goto(f"{base_url}/setup/")
+
+        heading = page.locator('h3:has-text("Image Generation")')
+        expect(heading).to_be_visible()
+
+        quality = page.locator("#image-gen-quality")
+        expect(quality).to_be_visible()
+
+        size = page.locator("#image-gen-size")
+        expect(size).to_be_visible()
+
+        background = page.locator("#image-gen-background")
+        expect(background).to_be_visible()
+
+    def test_hides_image_generation_config_when_not_enabled(
+        self, page: Page, base_url: str, app_server, env_api_key_no_tools
+    ):
+        """Image generation config section should NOT be visible when not enabled."""
+        page.goto(f"{base_url}/setup/")
+
+        quality = page.locator("#image-gen-quality")
+        expect(quality).not_to_be_visible()
+
+    def test_populates_saved_config_values(
+        self, page: Page, base_url: str, app_server, env_image_generation_with_config
+    ):
+        """Saved image generation config values should be populated in the form."""
+        page.goto(f"{base_url}/setup/")
+
+        quality = page.locator("#image-gen-quality")
+        expect(quality).to_have_value("high")
+
+        size = page.locator("#image-gen-size")
+        expect(size).to_have_value("1024x1536")
+
+        background = page.locator("#image-gen-background")
+        expect(background).to_have_value("transparent")
+
+    def test_no_other_tool_sections(
+        self, page: Page, base_url: str, app_server, env_image_generation_only
+    ):
+        """File search, function, and MCP sections should NOT be visible."""
+        page.goto(f"{base_url}/setup/")
+
+        upload_form = page.locator("#upload-form")
+        expect(upload_form).not_to_be_visible()
+
+        registry_rows = page.locator("#registry-rows")
+        expect(registry_rows).not_to_be_visible()
+
+        mcp_rows = page.locator("#mcp-rows")
+        expect(mcp_rows).not_to_be_visible()
+
+
 class TestAllToolsEnabled:
     """Tests for when all three conditional tools are enabled."""
 
@@ -585,6 +664,15 @@ def test_shows_web_search_config(
         context_size = page.locator("#web-search-context-size")
         expect(context_size).to_be_visible()
 
+    def test_shows_image_generation_config(
+        self, page: Page, base_url: str, app_server, env_all_tools
+    ):
+        """Image generation config section should be visible when all tools enabled."""
+        page.goto(f"{base_url}/setup/")
+
+        quality = page.locator("#image-gen-quality")
+        expect(quality).to_be_visible()
+
     def test_no_computer_use_config(
         self, page: Page, base_url: str, app_server, env_all_tools
     ):
@@ -614,3 +702,6 @@ def test_all_tool_checkboxes_checked(
 
         computer_use_cb = page.locator('input[value="computer_use"]')
         expect(computer_use_cb).to_be_checked()
+
+        image_gen_cb = page.locator('input[value="image_generation"]')
+        expect(image_gen_cb).to_be_checked()