Add aspect ratio and image size options to image generation; implement chat-based image generation handling

Co-authored-by: Copilot <copilot@github.com>
2026-04-27 19:28:04 +02:00
parent 17ae8d9477
commit 58c2cb4490
6 changed files with 211 additions and 18 deletions
@@ -149,6 +149,105 @@ async def test_generate_image_upstream_error(client):
    assert resp.status_code == 502


+# --- Chat-based image generation (FLUX, GPT-5 Image Mini) ---
+
+FAKE_IMAGE_CHAT_FLUX = {
+    "id": "gen-img-chat-1",
+    "model": "black-forest-labs/flux.2-klein-4b",
+    "choices": [{
+        "message": {
+            "role": "assistant",
+            "content": "Here is your generated image.",
+            "images": [{
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,abc123"},
+            }],
+        }
+    }],
+}
+
+FAKE_IMAGE_CHAT_GPT5 = {
+    "id": "gen-img-chat-2",
+    "model": "openai/gpt-5-image-mini",
+    "choices": [{
+        "message": {
+            "role": "assistant",
+            "content": "Generated image.",
+            "images": [{
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,xyz789"},
+            }],
+        }
+    }],
+}
+
+
+async def test_generate_image_chat_flux(client):
+    token = await _user_token(client)
+    with patch("backend.app.routers.generate.openrouter.generate_image_chat", new_callable=AsyncMock, return_value=FAKE_IMAGE_CHAT_FLUX):
+        resp = await client.post(
+            "/generate/image",
+            json={"model": "black-forest-labs/flux.2-klein-4b",
+                  "prompt": "A sunset"},
+            headers={"Authorization": f"Bearer {token}"},
+        )
+    assert resp.status_code == 200
+    data = resp.json()
+    assert data["id"] == "gen-img-chat-1"
+    assert len(data["images"]) == 1
+    assert data["images"][0]["url"] == "data:image/png;base64,abc123"
+
+
+async def test_generate_image_chat_gpt5_image_mini(client):
+    token = await _user_token(client)
+    with patch("backend.app.routers.generate.openrouter.generate_image_chat", new_callable=AsyncMock, return_value=FAKE_IMAGE_CHAT_GPT5):
+        resp = await client.post(
+            "/generate/image",
+            json={"model": "openai/gpt-5-image-mini", "prompt": "A cat"},
+            headers={"Authorization": f"Bearer {token}"},
+        )
+    assert resp.status_code == 200
+    data = resp.json()
+    assert data["model"] == "openai/gpt-5-image-mini"
+    assert len(data["images"]) == 1
+
+
+async def test_generate_image_chat_with_image_config(client):
+    token = await _user_token(client)
+    mock = AsyncMock(return_value=FAKE_IMAGE_CHAT_FLUX)
+    with patch("backend.app.routers.generate.openrouter.generate_image_chat", mock):
+        await client.post(
+            "/generate/image",
+            json={
+                "model": "black-forest-labs/flux.2-klein-4b",
+                "prompt": "A landscape",
+                "aspect_ratio": "16:9",
+                "image_size": "2K",
+            },
+            headers={"Authorization": f"Bearer {token}"},
+        )
+    call_kwargs = mock.call_args.kwargs
+    assert call_kwargs["image_config"]["aspect_ratio"] == "16:9"
+    assert call_kwargs["image_config"]["image_size"] == "2K"
+    assert call_kwargs["modalities"] == ["image"]
+
+
+async def test_generate_image_chat_unauthenticated(client):
+    resp = await client.post("/generate/image", json={"model": "flux.2-klein-4b", "prompt": "Hi"})
+    assert resp.status_code == 401
+
+
+async def test_generate_image_chat_upstream_error(client):
+    token = await _user_token(client)
+    with patch("backend.app.routers.generate.openrouter.generate_image_chat", new_callable=AsyncMock, side_effect=Exception("timeout")):
+        resp = await client.post(
+            "/generate/image",
+            json={"model": "black-forest-labs/flux.2-klein-4b", "prompt": "Hi"},
+            headers={"Authorization": f"Bearer {token}"},
+        )
+    assert resp.status_code == 502
+
+
 # ---------------------------------------------------------------------------
 # POST /generate/video
 # ---------------------------------------------------------------------------