diff --git a/README.md b/README.md index 07298df..1508592 100644 --- a/README.md +++ b/README.md @@ -140,9 +140,9 @@ Generated files are written through the `AssetStore` protocol to the references are stored in Postgres. An S3-compatible store is needed before deploying multiple application instances. -## OpenAI Image Playground +## Image Playground -To enable real GPT Image 2 generation while leaving Etsy and Printful mocked: +To enable real image generation while leaving Etsy and Printful mocked: ```text INTEGRATIONS_ENABLED=true @@ -150,6 +150,8 @@ ETSY_MODE=mock PRINTFUL_MODE=mock OPENAI_MODE=real OPENAI_API_KEY=... +GEMINI_API_KEY=... +IMAGE_GENERATION_MODEL=gpt-image-2 OPENAI_IMAGE_MODEL=gpt-image-2 OPENAI_IMAGE_QUALITY=medium OPENAI_IMAGE_WIDTH=1024 @@ -157,9 +159,20 @@ OPENAI_IMAGE_HEIGHT=1024 OPENAI_IMAGE_MAX_COST_USD=0.50 ``` -The Design Studio stores every prompt and outcome in Postgres. A request is rejected -before generation when its configured output-cost estimate exceeds the per-image -limit. Dollar amounts are estimates because the generation response reports token +The Design Studio can generate with OpenAI `gpt-image-2` or Google Gemini image +models: `gemini-3.1-flash-image`, `gemini-3-pro-image`, and +`gemini-2.5-flash-image`. `GOOGLE_API_KEY` is also accepted and takes precedence +over `GEMINI_API_KEY`. Gemini sizes are limited to documented API sizes and are +rejected before generation when unsupported. For now, Gemini REST generation is +limited to 1024x1024 because the live API rejects the optional sizing config +shown in the current docs. + +Create a Gemini key from Google AI Studio: +https://aistudio.google.com/app/apikey + +The Design Studio stores every prompt and outcome in Postgres. A request is +rejected before generation when its configured output-cost estimate exceeds the +per-image limit. Dollar amounts are estimates because provider responses report usage rather than the final billed amount. ## Trend Research And Reports diff --git a/src/ecommerce_agent/api/posting_routes.py b/src/ecommerce_agent/api/posting_routes.py index 9a4c8ef..c6d1a9c 100644 --- a/src/ecommerce_agent/api/posting_routes.py +++ b/src/ecommerce_agent/api/posting_routes.py @@ -2,7 +2,7 @@ from typing import Annotated, Any, cast import httpx -from fastapi import APIRouter, Depends, Request, status +from fastapi import APIRouter, Depends, HTTPException, Request, status from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload @@ -36,6 +36,7 @@ from ecommerce_agent.pipeline.errors import ConflictError, NotFoundError from ecommerce_agent.services import ServiceContainer from ecommerce_agent.services.etsy import EtsyOutcomeUnknown +from ecommerce_agent.services.image_models import gemini_image_size, provider_for_model from ecommerce_agent.services.printful import PrintfulError router = APIRouter(prefix="/posting", tags=["posting"]) @@ -291,6 +292,11 @@ async def generate_lifestyle_image( "The Printful mockup could not be downloaded for AI image generation." ) from error content_type = response.headers.get("content-type", "image/jpeg").split(";", 1)[0] + settings = request.app.state.settings + model = body.model or settings.image_generation_model + width = body.width or settings.openai_image_width + height = body.height or settings.openai_image_height + _validate_image_model(model, width=width, height=height, settings=settings) generated = await services.image.generate( ImageGenerationRequest( prompt=( @@ -299,8 +305,9 @@ async def generate_lifestyle_image( f"possible. {body.prompt}" ), product_concept_id=f"posting-{item.id}", - width=request.app.state.settings.openai_image_width, - height=request.app.state.settings.openai_image_height, + model=model, + width=width, + height=height, references=( ImageReference( filename="printful-reference.jpg", @@ -320,7 +327,12 @@ async def generate_lifestyle_image( approved=False, ai_drift_acknowledged=False, alt_text=item.image_alt_text, - metadata_={"warning": "AI image may drift from the exact product design."}, + metadata_={ + "warning": "AI image may drift from the exact product design.", + "model": generated.model, + "provider": generated.provider, + "generation_metadata": generated.metadata, + }, ) session.add(image) await session.commit() @@ -328,6 +340,29 @@ async def generate_lifestyle_image( return _image_view(image) +def _validate_image_model(model: str, *, width: int, height: int, settings: Any) -> None: + provider_name = provider_for_model(model, openai_model=settings.openai_image_model) + if provider_name is None: + raise HTTPException(status_code=422, detail=f"Unsupported image model: {model}.") + if provider_name == "gemini" and gemini_image_size(model, width, height) is None: + raise HTTPException( + status_code=422, + detail=( + f"{model} does not support {width}x{height}. " + "Choose one of the documented Gemini image sizes." + ), + ) + if settings.openai_mode == "real": + if provider_name == "openai" and settings.openai_api_key is None: + raise HTTPException( + status_code=422, detail="OpenAI image generation is not configured." + ) + if provider_name == "gemini" and not (settings.gemini_api_key or settings.google_api_key): + raise HTTPException( + status_code=422, detail="Gemini image generation is not configured." + ) + + @router.patch("/images/{image_id}") async def update_listing_image( image_id: uuid.UUID, diff --git a/src/ecommerce_agent/api/routes.py b/src/ecommerce_agent/api/routes.py index 948b534..8bb2e58 100644 --- a/src/ecommerce_agent/api/routes.py +++ b/src/ecommerce_agent/api/routes.py @@ -15,6 +15,7 @@ DecisionRequest, ImageGenerationCreateRequest, ImageGenerationUpdateRequest, + ProductWorkflowRequest, PublishRequest, RunResponse, ) @@ -44,6 +45,12 @@ from ecommerce_agent.pipeline.events import log_event from ecommerce_agent.pipeline.transitions import advance, require_status from ecommerce_agent.services import ServiceContainer +from ecommerce_agent.services.image_models import ( + gemini_image_size, + image_model_options, + is_gemini_image_model, + provider_for_model, +) from ecommerce_agent.services.openai_images import ( SUPPORTED_OUTPUT_FORMATS, OpenAIImageError, @@ -56,7 +63,7 @@ ServicesDep = Annotated[ServiceContainer, Depends(get_services)] MAX_REFERENCE_BYTES = 10 * 1024 * 1024 MAX_REFERENCE_PIXELS = 25_000_000 -MAX_IMAGE_PIXELS = 4_194_304 +MAX_IMAGE_PIXELS = 18_000_000 REFERENCE_CONTENT_TYPES = { "PNG": "image/png", "JPEG": "image/jpeg", @@ -68,6 +75,7 @@ class ParsedImageRequest: prompt: str title: str + model: str | None width: int | None height: int | None output_format: str @@ -80,6 +88,37 @@ async def health() -> dict[str, str]: return {"status": "ok"} +@router.get("/image-models") +async def list_image_models(request: Request) -> list[dict[str, Any]]: + settings = request.app.state.settings + mock_mode = settings.openai_mode != "real" + options = image_model_options( + openai_model=settings.openai_image_model, + openai_available=mock_mode or settings.openai_api_key is not None, + gemini_available=mock_mode + or settings.gemini_api_key is not None + or settings.google_api_key is not None, + ) + return [ + { + "model": option.model, + "label": option.label, + "provider": option.provider, + "available": option.available, + "sizes": [ + { + "width": size.width, + "height": size.height, + "aspect_ratio": size.aspect_ratio, + "image_size": size.image_size, + } + for size in option.sizes + ], + } + for option in options + ] + + @router.post("/image-generations", status_code=status.HTTP_201_CREATED) async def create_image_generation( request: Request, @@ -124,6 +163,7 @@ async def edit_image_generation( parsed = ParsedImageRequest( prompt=parsed.prompt, title=parent.title, + model=parsed.model, width=parsed.width, height=parsed.height, output_format=parsed.output_format, @@ -148,25 +188,30 @@ async def _execute_image_generation( services: ServiceContainer, ) -> dict[str, Any]: settings = request.app.state.settings + model = parsed.model or settings.image_generation_model width = parsed.width or settings.openai_image_width height = parsed.height or settings.openai_image_height + _validate_image_model(model, width=width, height=height, settings=settings) is_real = settings.openai_mode == "real" + provider_name = provider_for_model(model, openai_model=settings.openai_image_model) try: - output_estimate = ( - estimate_image_output_cost( - model=settings.openai_image_model, + if not is_real: + output_estimate = Decimal("0") + reference_reserve = Decimal("0") + elif provider_name == "gemini": + size = gemini_image_size(model, width, height) + output_estimate = size.estimated_cost_usd if size is not None else Decimal("0") + reference_reserve = Decimal("0") + else: + output_estimate = estimate_image_output_cost( + model=model, quality=settings.openai_image_quality, width=width, height=height, ) - if is_real - else Decimal("0") - ) - reference_reserve = ( - Decimal(str(settings.openai_image_reference_cost_reserve_usd)) * len(parsed.references) - if is_real - else Decimal("0") - ) + reference_reserve = Decimal( + str(settings.openai_image_reference_cost_reserve_usd) + ) * len(parsed.references) estimated_cost = output_estimate + reference_reserve except OpenAIImageError: estimated_cost = Decimal("0") @@ -178,9 +223,9 @@ async def _execute_image_generation( status=ImageGenerationStatus.PENDING, width=width, height=height, - provider="openai" if is_real else "mock", - model=settings.openai_image_model if is_real else "mock-image-v1", - quality=settings.openai_image_quality, + provider=provider_name if is_real and provider_name is not None else "mock", + model=model if is_real else "mock-image-v1", + quality="standard" if provider_name == "gemini" else settings.openai_image_quality, output_format=parsed.output_format, background_removal_friendly=parsed.background_removal_friendly, estimated_cost_usd=estimated_cost, @@ -193,6 +238,7 @@ async def _execute_image_generation( ImageGenerationRequest( prompt=generation.prompt, product_concept_id=str(generation.id), + model=model, width=generation.width, height=generation.height, references=parsed.references, @@ -235,6 +281,7 @@ async def _parse_image_request(request: Request, *, max_references: int) -> Pars content_type = request.headers.get("content-type", "") prompt_value: Any title_value: Any + model_value: Any width_value: Any height_value: Any output_format_value: Any @@ -244,6 +291,7 @@ async def _parse_image_request(request: Request, *, max_references: int) -> Pars form = await request.form() prompt_value = form.get("prompt") title_value = form.get("title", "Untitled design") + model_value = form.get("model") width_value = form.get("width") height_value = form.get("height") output_format_value = form.get("output_format", "png") @@ -251,7 +299,7 @@ async def _parse_image_request(request: Request, *, max_references: int) -> Pars uploads = [ value for key, value in form.multi_items() - if key == "reference_images" and isinstance(value, UploadFile) + if key == "reference_images" and isinstance(value, UploadFile) and value.filename ] if len(uploads) > max_references: raise HTTPException( @@ -268,6 +316,7 @@ async def _parse_image_request(request: Request, *, max_references: int) -> Pars ) from error prompt_value = body.prompt title_value = body.title + model_value = body.model width_value = body.width height_value = body.height output_format_value = body.output_format @@ -299,6 +348,7 @@ async def _parse_image_request(request: Request, *, max_references: int) -> Pars return ParsedImageRequest( prompt=prompt, title=title, + model=str(model_value).strip() if model_value else None, width=width, height=height, output_format=output_format, @@ -369,11 +419,34 @@ def _parse_dimension(value: Any, label: str) -> int | None: dimension = int(value) except (TypeError, ValueError) as error: raise HTTPException(status_code=422, detail=f"Invalid image {label}.") from error - if dimension <= 0 or dimension > 4096: + if dimension <= 0 or dimension > 5504: raise HTTPException(status_code=422, detail=f"Invalid image {label}.") return dimension +def _validate_image_model(model: str, *, width: int, height: int, settings: Any) -> None: + provider_name = provider_for_model(model, openai_model=settings.openai_image_model) + if provider_name is None: + raise HTTPException(status_code=422, detail=f"Unsupported image model: {model}.") + if settings.openai_mode == "real": + if provider_name == "openai" and settings.openai_api_key is None: + raise HTTPException( + status_code=422, detail="OpenAI image generation is not configured." + ) + if provider_name == "gemini" and not (settings.gemini_api_key or settings.google_api_key): + raise HTTPException( + status_code=422, detail="Gemini image generation is not configured." + ) + if is_gemini_image_model(model) and gemini_image_size(model, width, height) is None: + raise HTTPException( + status_code=422, + detail=( + f"{model} does not support {width}x{height}. " + "Choose one of the documented Gemini image sizes." + ), + ) + + def _extension(output_format: str) -> str: return "jpg" if output_format == "jpeg" else output_format @@ -454,14 +527,33 @@ async def start_research(session: SessionDep, services: ServicesDep) -> RunRespo response_model=RunResponse, status_code=status.HTTP_201_CREATED, ) -async def start_product(product_id: uuid.UUID, session: SessionDep) -> RunResponse: +async def start_product( + product_id: uuid.UUID, + request: Request, + session: SessionDep, + body: ProductWorkflowRequest | None = None, +) -> RunResponse: product = await session.get(Product, product_id) if product is None: raise NotFoundError("Product concept not found.") + image_model = body.image_model if body is not None else None + if image_model: + settings = request.app.state.settings + _validate_image_model( + image_model, + width=settings.openai_image_width, + height=settings.openai_image_height, + settings=settings, + ) require_status(product, ProductStatus.CONCEPT) advance(product, ProductStatus.GENERATING_DESIGN) await log_event(session, "product", product.id, "start", "Product workflow started.") - await enqueue(session, "generate_design", product.id) + await enqueue( + session, + "generate_design", + product.id, + {"image_model": image_model} if image_model else None, + ) await session.commit() return RunResponse(id=product.id, subject_type="product", status=product.status.value) diff --git a/src/ecommerce_agent/api/schemas.py b/src/ecommerce_agent/api/schemas.py index 54790ef..cdcf7a4 100644 --- a/src/ecommerce_agent/api/schemas.py +++ b/src/ecommerce_agent/api/schemas.py @@ -42,8 +42,9 @@ class GoogleAnalyticsPropertyRequest(BaseModel): class ImageGenerationCreateRequest(BaseModel): prompt: str = Field(min_length=1, max_length=32000) title: str = Field(default="Untitled design", max_length=255) - width: int | None = Field(default=None, gt=0, le=4096) - height: int | None = Field(default=None, gt=0, le=4096) + model: str | None = Field(default=None, max_length=100) + width: int | None = Field(default=None, gt=0, le=5504) + height: int | None = Field(default=None, gt=0, le=5504) output_format: str = "png" background_removal_friendly: bool = True @@ -52,6 +53,10 @@ class ImageGenerationUpdateRequest(BaseModel): title: str = Field(min_length=1, max_length=255) +class ProductWorkflowRequest(BaseModel): + image_model: str | None = Field(default=None, max_length=100) + + class PrintfulCatalogRefreshRequest(BaseModel): category_ids: list[int] | None = None @@ -199,6 +204,9 @@ class ListingImageOrderRequest(BaseModel): class LifestyleImageRequest(BaseModel): prompt: str = Field(min_length=1, max_length=8000) + model: str | None = Field(default=None, max_length=100) + width: int | None = Field(default=None, gt=0, le=5504) + height: int | None = Field(default=None, gt=0, le=5504) class PostingSubmitRequest(BaseModel): diff --git a/src/ecommerce_agent/config.py b/src/ecommerce_agent/config.py index f58b4f4..b10b09c 100644 --- a/src/ecommerce_agent/config.py +++ b/src/ecommerce_agent/config.py @@ -57,12 +57,15 @@ class Settings(BaseSettings): s3_secret_access_key: SecretStr | None = None openai_api_key: SecretStr | None = None openai_text_model: str = "gpt-5.5" + image_generation_model: str = "gpt-image-2" openai_image_model: str = "gpt-image-2" openai_image_quality: Literal["low", "medium", "high"] = "medium" openai_image_width: int = Field(default=1024, gt=0) openai_image_height: int = Field(default=1024, gt=0) openai_image_max_cost_usd: float = Field(default=0.50, gt=0) openai_image_reference_cost_reserve_usd: float = Field(default=0.06, ge=0) + gemini_api_key: SecretStr | None = None + google_api_key: SecretStr | None = None research_trend_limit: int = Field(default=8, ge=1, le=25) research_poll_seconds: int = Field(default=5, ge=1, le=60) research_max_tool_calls: int = Field(default=12, ge=1, le=50) @@ -134,6 +137,11 @@ def validate_provider_modes(self) -> "Settings": ) if self.openai_mode == "real" and self.openai_api_key is None: raise ValueError("OPENAI_API_KEY is required when OPENAI_MODE=real.") + if self.openai_mode == "real" and self.image_generation_model.startswith("gemini-"): + if self.gemini_api_key is None and self.google_api_key is None: + raise ValueError( + "GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini image models." + ) google_values = ( self.google_client_id, self.google_client_secret, diff --git a/src/ecommerce_agent/domain/dtos.py b/src/ecommerce_agent/domain/dtos.py index 50a541d..a0275d7 100644 --- a/src/ecommerce_agent/domain/dtos.py +++ b/src/ecommerce_agent/domain/dtos.py @@ -34,6 +34,7 @@ def total(self) -> int: class ImageGenerationRequest: prompt: str product_concept_id: str + model: str | None = None revision: int = 1 width: int = 1024 height: int = 1024 diff --git a/src/ecommerce_agent/jobs/worker.py b/src/ecommerce_agent/jobs/worker.py index 87ac4ef..a32c707 100644 --- a/src/ecommerce_agent/jobs/worker.py +++ b/src/ecommerce_agent/jobs/worker.py @@ -1,7 +1,9 @@ import asyncio +import inspect import logging -from collections.abc import Callable +from collections.abc import Awaitable, Callable from datetime import UTC, datetime, timedelta +from typing import cast from sqlalchemy import and_, or_, select from sqlalchemy.ext.asyncio import AsyncSession @@ -64,11 +66,17 @@ async def process_one( if job is None: return False job_id, job_type, subject_id, attempts = job.id, job.job_type, job.subject_id, job.attempts + payload = job.payload handler = JOB_HANDLERS.get(job_type) try: if handler is None: raise ValueError(f"Unknown job type: {job_type}") - await handler(session, services, subject_id) + if len(inspect.signature(handler).parameters) == 4: + await cast(Callable[..., Awaitable[object]], handler)( + session, services, subject_id, payload + ) + else: + await handler(session, services, subject_id) except DeferJob as deferred: await session.rollback() job = await session.get(Job, job_id) diff --git a/src/ecommerce_agent/pipeline/steps.py b/src/ecommerce_agent/pipeline/steps.py index 551d3bc..007f4f3 100644 --- a/src/ecommerce_agent/pipeline/steps.py +++ b/src/ecommerce_agent/pipeline/steps.py @@ -30,7 +30,10 @@ async def _get_product(session: AsyncSession, product_id: uuid.UUID) -> Product: async def generate_design( - session: AsyncSession, services: ServiceContainer, product_id: uuid.UUID + session: AsyncSession, + services: ServiceContainer, + product_id: uuid.UUID, + payload: dict[str, Any] | None = None, ) -> Product: """Create the next design revision: brief, generated image, and review handoff. @@ -59,6 +62,7 @@ async def generate_design( ImageGenerationRequest( prompt=brief["prompt"], product_concept_id=str(product.id), + model=(payload or {}).get("image_model"), revision=revision, ) ) diff --git a/src/ecommerce_agent/services/factory.py b/src/ecommerce_agent/services/factory.py index a0d57e9..e9e4d8a 100644 --- a/src/ecommerce_agent/services/factory.py +++ b/src/ecommerce_agent/services/factory.py @@ -9,6 +9,8 @@ SignedLocalAssetStore, ) from ecommerce_agent.services.etsy import EtsyClient +from ecommerce_agent.services.gemini_images import GeminiImageProvider +from ecommerce_agent.services.image_router import ImageGenerationRouter from ecommerce_agent.services.interfaces import ( CopyGenerationProvider, EtsyGateway, @@ -54,12 +56,14 @@ class ServiceContainer: def build_service_container(settings: Settings | None = None) -> ServiceContainer: settings = settings or get_settings() assets = LocalAssetStore(settings.asset_root) + image: ImageGenerationProvider if settings.openai_mode == "real": if not settings.integrations_enabled or settings.openai_api_key is None: raise RuntimeError( "Real OpenAI construction requires enabled integrations and an API key." ) - image: ImageGenerationProvider = OpenAIImageProvider( + providers: dict[str, ImageGenerationProvider] = {} + providers["openai"] = OpenAIImageProvider( api_key=settings.openai_api_key.get_secret_value(), asset_store=assets, model=settings.openai_image_model, @@ -69,6 +73,19 @@ def build_service_container(settings: Settings | None = None) -> ServiceContaine str(settings.openai_image_reference_cost_reserve_usd) ), ) + gemini_key = settings.google_api_key or settings.gemini_api_key + if gemini_key is not None: + providers["gemini"] = GeminiImageProvider( + api_key=gemini_key.get_secret_value(), + asset_store=assets, + default_model=settings.image_generation_model, + max_cost_usd=Decimal(str(settings.openai_image_max_cost_usd)), + ) + image = ImageGenerationRouter( + default_model=settings.image_generation_model, + openai_model=settings.openai_image_model, + providers=providers, + ) copy: CopyGenerationProvider = OpenAICopyProvider( api_key=settings.openai_api_key.get_secret_value(), model=settings.openai_text_model, diff --git a/src/ecommerce_agent/services/gemini_images.py b/src/ecommerce_agent/services/gemini_images.py new file mode 100644 index 0000000..c363506 --- /dev/null +++ b/src/ecommerce_agent/services/gemini_images.py @@ -0,0 +1,208 @@ +import asyncio +import base64 +import binascii +import io +import re +from collections.abc import Awaitable, Callable +from decimal import Decimal +from typing import Any, cast + +import httpx +from PIL import Image, UnidentifiedImageError + +from ecommerce_agent.domain.dtos import GeneratedAsset, ImageGenerationRequest +from ecommerce_agent.services.image_models import ( + GEMINI_IMAGE_MODELS, + gemini_image_size, +) +from ecommerce_agent.services.interfaces import AssetStore +from ecommerce_agent.services.openai_images import ( + OUTPUT_CONTENT_TYPES, + SUPPORTED_OUTPUT_FORMATS, + OpenAIImageError, + OpenAIImageOutcomeUnknown, + build_effective_prompt, +) + +GEMINI_GENERATE_URL = "https://generativelanguage.googleapis.com/v1/models/{model}:generateContent" + + +class GeminiImageProvider: + def __init__( + self, + *, + api_key: str, + asset_store: AssetStore, + default_model: str, + max_cost_usd: Decimal, + client: httpx.AsyncClient | None = None, + max_attempts: int = 3, + retry_delay_seconds: float = 0.5, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ): + self._api_key = api_key + self._asset_store = asset_store + self._default_model = default_model + self._max_cost_usd = max_cost_usd + self._client = client + self._max_attempts = max_attempts + self._retry_delay_seconds = retry_delay_seconds + self._sleep = sleep + + async def generate(self, request: ImageGenerationRequest) -> GeneratedAsset: + model = request.model or self._default_model + if model not in GEMINI_IMAGE_MODELS: + raise OpenAIImageError(f"Unsupported Gemini image model: {model}.") + if request.output_format not in SUPPORTED_OUTPUT_FORMATS: + raise OpenAIImageError(f"Unsupported output format: {request.output_format}.") + size = gemini_image_size(model, request.width, request.height) + if size is None: + raise OpenAIImageError( + f"{model} does not support {request.width}x{request.height}. " + "Choose one of the documented Gemini image sizes." + ) + if size.estimated_cost_usd > self._max_cost_usd: + raise OpenAIImageError( + f"Estimated image cost ${size.estimated_cost_usd} exceeds the configured " + f"${self._max_cost_usd} per-image limit." + ) + + response = await self._post_with_safe_retries( + model, + json=_request_payload(request), + ) + image_bytes, usage, mime_type = _parse_response(response) + image_bytes = _convert_image( + image_bytes, + expected_width=request.width, + expected_height=request.height, + output_format=request.output_format, + ) + safe_id = re.sub(r"[^A-Za-z0-9_.-]", "-", request.product_concept_id) + extension = "jpg" if request.output_format == "jpeg" else request.output_format + key = f"image-generations/{safe_id}/revision-{request.revision}.{extension}" + uri, checksum = await self._asset_store.put_bytes( + key=key, + content=image_bytes, + content_type=OUTPUT_CONTENT_TYPES[request.output_format], + ) + return GeneratedAsset( + uri=uri, + checksum=checksum, + width=request.width, + height=request.height, + provider="gemini", + model=model, + metadata={ + "aspect_ratio": size.aspect_ratio, + "image_size": size.image_size, + "output_format": request.output_format, + "background_removal_friendly": request.background_removal_friendly, + "reference_count": len(request.references), + "source_mime_type": mime_type, + "usage": usage, + "estimated_cost_usd": str(size.estimated_cost_usd), + "request_id": response.headers.get("x-request-id"), + }, + ) + + async def _post_with_safe_retries(self, model: str, *, json: dict[str, Any]) -> httpx.Response: + owns_client = self._client is None + client = self._client or httpx.AsyncClient(timeout=httpx.Timeout(300.0)) + try: + for attempt in range(1, self._max_attempts + 1): + try: + response = await client.post( + GEMINI_GENERATE_URL.format(model=model), + headers={ + "x-goog-api-key": self._api_key, + "Content-Type": "application/json", + }, + json=json, + ) + except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as error: + if attempt == self._max_attempts: + raise OpenAIImageError( + "Gemini image generation could not establish a connection." + ) from error + await self._sleep(self._retry_delay_seconds * (2 ** (attempt - 1))) + continue + except (httpx.ReadTimeout, httpx.WriteTimeout) as error: + raise OpenAIImageOutcomeUnknown( + "Gemini image generation timed out after the request began; " + "the remote outcome is unknown and was not retried." + ) from error + + if response.status_code == 429 and attempt < self._max_attempts: + await self._sleep(self._retry_delay_seconds * (2 ** (attempt - 1))) + continue + if response.is_error: + raise OpenAIImageError( + f"Gemini image generation failed with HTTP {response.status_code}." + ) + return response + finally: + if owns_client: + await client.aclose() + raise OpenAIImageError("Gemini image generation failed after safe retries.") + + +def _request_payload(request: ImageGenerationRequest) -> dict[str, Any]: + parts: list[dict[str, Any]] = [{"text": build_effective_prompt(request)}] + parts.extend( + { + "inlineData": { + "mimeType": reference.content_type, + "data": base64.b64encode(reference.content).decode(), + } + } + for reference in request.references + ) + # ponytail: Gemini REST rejects the documented sizing config today; only expose square sizes. + return {"contents": [{"parts": parts}]} + + +def _parse_response(response: httpx.Response) -> tuple[bytes, dict[str, Any], str | None]: + try: + payload = cast(dict[str, Any], response.json()) + usage_value = payload.get("usageMetadata", payload.get("usage", {})) + usage = cast(dict[str, Any], usage_value) if isinstance(usage_value, dict) else {} + candidates = cast(list[dict[str, Any]], payload["candidates"]) + parts = cast(list[dict[str, Any]], candidates[0]["content"]["parts"]) + for part in parts: + inline = part.get("inlineData") or part.get("inline_data") + if isinstance(inline, dict) and isinstance(inline.get("data"), str): + return ( + base64.b64decode(inline["data"], validate=True), + usage, + cast(str | None, inline.get("mimeType") or inline.get("mime_type")), + ) + except (KeyError, IndexError, TypeError, ValueError, binascii.Error) as error: + raise OpenAIImageError("Gemini returned an invalid image response.") from error + raise OpenAIImageError("Gemini returned no image data.") + + +def _convert_image( + image_bytes: bytes, + *, + expected_width: int, + expected_height: int, + output_format: str, +) -> bytes: + try: + with Image.open(io.BytesIO(image_bytes)) as image: + image.load() + if image.size != (expected_width, expected_height): + raise OpenAIImageError( + f"Gemini returned an unexpected image size: {image.size[0]}x{image.size[1]}." + ) + output_image = image.convert("RGB") if output_format == "jpeg" else image + output = io.BytesIO() + output_image.save( + output, format={"png": "PNG", "jpeg": "JPEG", "webp": "WEBP"}[output_format] + ) + return output.getvalue() + except OpenAIImageError: + raise + except (UnidentifiedImageError, OSError) as error: + raise OpenAIImageError("Gemini returned invalid image data.") from error diff --git a/src/ecommerce_agent/services/image_models.py b/src/ecommerce_agent/services/image_models.py new file mode 100644 index 0000000..c5cdb39 --- /dev/null +++ b/src/ecommerce_agent/services/image_models.py @@ -0,0 +1,99 @@ +from dataclasses import dataclass +from decimal import Decimal +from typing import Literal + +ImageProviderName = Literal["openai", "gemini"] + +GEMINI_31_FLASH_IMAGE = "gemini-3.1-flash-image" +GEMINI_3_PRO_IMAGE = "gemini-3-pro-image" +GEMINI_25_FLASH_IMAGE = "gemini-2.5-flash-image" +GEMINI_IMAGE_MODELS = { + GEMINI_31_FLASH_IMAGE, + GEMINI_3_PRO_IMAGE, + GEMINI_25_FLASH_IMAGE, +} + + +@dataclass(frozen=True) +class GeminiImageSize: + width: int + height: int + aspect_ratio: str + image_size: str | None + estimated_cost_usd: Decimal + + +@dataclass(frozen=True) +class ImageModelOption: + model: str + label: str + provider: ImageProviderName + available: bool + sizes: tuple[GeminiImageSize, ...] = () + + +GEMINI_IMAGE_SIZES: dict[str, tuple[GeminiImageSize, ...]] = { + GEMINI_31_FLASH_IMAGE: (GeminiImageSize(1024, 1024, "1:1", "1K", Decimal("0.067")),), + GEMINI_3_PRO_IMAGE: (GeminiImageSize(1024, 1024, "1:1", "1K", Decimal("0.134")),), + GEMINI_25_FLASH_IMAGE: (GeminiImageSize(1024, 1024, "1:1", None, Decimal("0.039")),), +} + + +def is_gemini_image_model(model: str) -> bool: + return model in GEMINI_IMAGE_MODELS + + +def gemini_image_size(model: str, width: int, height: int) -> GeminiImageSize | None: + return next( + ( + size + for size in GEMINI_IMAGE_SIZES.get(model, ()) + if size.width == width and size.height == height + ), + None, + ) + + +def image_model_options( + *, + openai_model: str, + openai_available: bool, + gemini_available: bool, +) -> tuple[ImageModelOption, ...]: + return ( + ImageModelOption( + model=openai_model, + label=f"OpenAI {openai_model}", + provider="openai", + available=openai_available, + ), + ImageModelOption( + model=GEMINI_31_FLASH_IMAGE, + label="Nano Banana 2", + provider="gemini", + available=gemini_available, + sizes=GEMINI_IMAGE_SIZES[GEMINI_31_FLASH_IMAGE], + ), + ImageModelOption( + model=GEMINI_3_PRO_IMAGE, + label="Nano Banana Pro", + provider="gemini", + available=gemini_available, + sizes=GEMINI_IMAGE_SIZES[GEMINI_3_PRO_IMAGE], + ), + ImageModelOption( + model=GEMINI_25_FLASH_IMAGE, + label="Nano Banana", + provider="gemini", + available=gemini_available, + sizes=GEMINI_IMAGE_SIZES[GEMINI_25_FLASH_IMAGE], + ), + ) + + +def provider_for_model(model: str, *, openai_model: str) -> ImageProviderName | None: + if model == openai_model: + return "openai" + if is_gemini_image_model(model): + return "gemini" + return None diff --git a/src/ecommerce_agent/services/image_router.py b/src/ecommerce_agent/services/image_router.py new file mode 100644 index 0000000..723cac1 --- /dev/null +++ b/src/ecommerce_agent/services/image_router.py @@ -0,0 +1,27 @@ +from dataclasses import replace + +from ecommerce_agent.domain.dtos import GeneratedAsset, ImageGenerationRequest +from ecommerce_agent.services.image_models import provider_for_model +from ecommerce_agent.services.interfaces import ImageGenerationProvider +from ecommerce_agent.services.openai_images import OpenAIImageError + + +class ImageGenerationRouter: + def __init__( + self, + *, + default_model: str, + openai_model: str, + providers: dict[str, ImageGenerationProvider], + ): + self._default_model = default_model + self._openai_model = openai_model + self._providers = providers + + async def generate(self, request: ImageGenerationRequest) -> GeneratedAsset: + model = request.model or self._default_model + provider_name = provider_for_model(model, openai_model=self._openai_model) + if provider_name is None or provider_name not in self._providers: + raise OpenAIImageError(f"Image model is not configured: {model}.") + request = request if request.model else replace(request, model=model) + return await self._providers[provider_name].generate(request) diff --git a/src/ecommerce_agent/services/openai_images.py b/src/ecommerce_agent/services/openai_images.py index 37b3905..397de4c 100644 --- a/src/ecommerce_agent/services/openai_images.py +++ b/src/ecommerce_agent/services/openai_images.py @@ -103,8 +103,9 @@ def __init__( async def generate(self, request: ImageGenerationRequest) -> GeneratedAsset: if request.output_format not in SUPPORTED_OUTPUT_FORMATS: raise OpenAIImageError(f"Unsupported output format: {request.output_format}.") + model = request.model or self._model output_cost = estimate_image_output_cost( - model=self._model, + model=model, quality=self._quality, width=request.width, height=request.height, @@ -122,7 +123,7 @@ async def generate(self, request: ImageGenerationRequest) -> GeneratedAsset: response = await self._post_with_safe_retries( OPENAI_IMAGE_EDITS_URL, data={ - "model": self._model, + "model": model, "prompt": effective_prompt, "n": "1", "size": f"{request.width}x{request.height}", @@ -141,7 +142,7 @@ async def generate(self, request: ImageGenerationRequest) -> GeneratedAsset: response = await self._post_with_safe_retries( OPENAI_IMAGE_GENERATIONS_URL, json={ - "model": self._model, + "model": model, "prompt": effective_prompt, "n": 1, "size": f"{request.width}x{request.height}", @@ -171,7 +172,7 @@ async def generate(self, request: ImageGenerationRequest) -> GeneratedAsset: width=request.width, height=request.height, provider="openai", - model=self._model, + model=model, metadata={ "quality": self._quality, "output_format": request.output_format, diff --git a/src/ecommerce_agent/web/routes.py b/src/ecommerce_agent/web/routes.py index 245ccd7..3b1b5e1 100644 --- a/src/ecommerce_agent/web/routes.py +++ b/src/ecommerce_agent/web/routes.py @@ -12,6 +12,8 @@ from fastapi.templating import Jinja2Templates from markupsafe import Markup +from ecommerce_agent.services.image_models import image_model_options + _TEMPLATES_DIR = Path(__file__).parent / "templates" templates = Jinja2Templates(directory=str(_TEMPLATES_DIR)) @@ -88,8 +90,8 @@ def asset_url(uri: str | None) -> str: {"label": "Agent Runs · Logs", "slug": "agent-runs", "icon": "pulse"}, ] -# Topbar model dropdowns are cosmetic selectors (the app uses a fixed -# server-side model). Options mirror the prototype's WS_DATA.models. +# Text model selection remains cosmetic. Image model selection is submitted by +# the Design Studio composer. TEXT_MODELS = [ "gpt-5", "gpt-5-mini", @@ -98,14 +100,6 @@ def asset_url(uri: str | None) -> str: "gemini-2.5-pro", "llama-4-70b", ] -IMAGE_MODELS = [ - "gpt-image-2", - "gpt-image-2-mini", - "dall-e-3", - "flux-1.1-pro", - "imagen-4", - "sd-3.5-large", -] def shell_context( @@ -121,6 +115,28 @@ def shell_context( """Build the template context shared by every screen (sidebar + topbar).""" settings = request.app.state.settings cap = getattr(settings, "image_generation_cap_usd", None) or getattr(settings, "cap_usd", None) + options = image_model_options( + openai_model=settings.openai_image_model, + openai_available=settings.openai_mode != "real" or settings.openai_api_key is not None, + gemini_available=settings.openai_mode != "real" + or settings.gemini_api_key is not None + or settings.google_api_key is not None, + ) + available = [option for option in options if option.available] or list(options) + image_model = settings.image_generation_model + if image_model not in {option.model for option in available}: + image_model = available[0].model + image_models = [ + { + "value": option.model, + "label": option.label, + "provider": option.provider, + } + for option in available + ] + image_model_label = next( + option["label"] for option in image_models if option["value"] == image_model + ) ctx: dict[str, Any] = { "request": request, "app_name": "WORKSHOP", @@ -131,9 +147,10 @@ def shell_context( "primary": primary, "badges": badges or {}, "text_models": TEXT_MODELS, - "image_models": IMAGE_MODELS, + "image_models": image_models, "text_model": TEXT_MODELS[0], - "image_model": IMAGE_MODELS[0], + "image_model": image_model, + "image_model_label": image_model_label, "integrations": "LIVE", "cap_usd": f"{float(cap):.2f}" if cap is not None else "0.50", } diff --git a/src/ecommerce_agent/web/static/app.js b/src/ecommerce_agent/web/static/app.js index 8934125..4ea7845 100644 --- a/src/ecommerce_agent/web/static/app.js +++ b/src/ecommerce_agent/web/static/app.js @@ -57,9 +57,15 @@ } }); var valueEl = wrap2.querySelector(".ws-dd-value"); - if (valueEl) valueEl.textContent = value; + if (valueEl) valueEl.textContent = opt.getAttribute("data-label") || value; var key = wrap2.getAttribute("data-dd"); try { if (key) localStorage.setItem("ws-model-" + key, value); } catch (err) {} + if (key === "image") { + document.querySelectorAll("[data-image-model]").forEach(function (input) { + input.value = value; + syncImageSize(input.closest("[data-aspect-form]"), value); + }); + } closeAllMenus(null); return; } @@ -67,6 +73,23 @@ if (!e.target.closest(".ws-dd-menu")) closeAllMenus(null); }); + function syncImageSize(form, model) { + if (!form) return; + var gemini = model.indexOf("gemini-") === 0; + form.querySelectorAll(".ws-aspect input[name='aspect']").forEach(function (radio) { + var disabled = gemini && radio.value !== "1:1"; + radio.disabled = disabled; + radio.closest(".ws-seg-cell").classList.toggle("is-disabled", disabled); + }); + if (gemini) { + var square = form.querySelector(".ws-aspect input[value='1:1']"); + if (square && !square.checked) { + square.checked = true; + square.dispatchEvent(new Event("change", { bubbles: true })); + } + } + } + // ── Master/detail selection highlight ───────────────────────────────────── document.addEventListener("click", function (e) { var card = e.target.closest("[data-detail]"); @@ -145,6 +168,9 @@ var menu = wrap.querySelector(".ws-dd-menu"); if (menu) menu.hidden = true; }); + document.querySelectorAll("[data-image-model]").forEach(function (input) { + syncImageSize(input.closest("[data-aspect-form]"), input.value); + }); syncToggleLabels(); } if (document.readyState !== "loading") restore(); diff --git a/src/ecommerce_agent/web/static/b1.css b/src/ecommerce_agent/web/static/b1.css index 024f0eb..26757a4 100644 --- a/src/ecommerce_agent/web/static/b1.css +++ b/src/ecommerce_agent/web/static/b1.css @@ -567,3 +567,7 @@ textarea.ws-input { line-height: 1.5; font-size: 13px; resize: vertical; } * { transition: none !important; } .ws-card:hover, .ws-thumb:hover, .ws-primary:hover { transform: none; } } +.ws-seg-cell.is-disabled { + opacity: 0.38; + cursor: not-allowed; +} diff --git a/src/ecommerce_agent/web/templates/design_studio.html b/src/ecommerce_agent/web/templates/design_studio.html index 78028b5..46e7a6e 100644 --- a/src/ecommerce_agent/web/templates/design_studio.html +++ b/src/ecommerce_agent/web/templates/design_studio.html @@ -7,6 +7,7 @@ {{ chip("v1 → generate", "live") }}
+
{{ field_label("Title") }}
{{ field_label("Prompt", "preview available") }}
@@ -19,12 +20,11 @@
- {{ field_label("Aspect ratio", "1024 × 1280") }} + {{ field_label("Aspect ratio", "1024 × 1024") }}
- - - + +
diff --git a/src/ecommerce_agent/web/templates/partials/_macros.html b/src/ecommerce_agent/web/templates/partials/_macros.html index da15196..51f760f 100644 --- a/src/ecommerce_agent/web/templates/partials/_macros.html +++ b/src/ecommerce_agent/web/templates/partials/_macros.html @@ -42,21 +42,25 @@ {%- endmacro %} -{% macro model_dropdown(label, value, options, key) -%} +{% macro model_dropdown(label, value, options, key, value_label="") -%}
diff --git a/src/ecommerce_agent/web/templates/partials/_topbar.html b/src/ecommerce_agent/web/templates/partials/_topbar.html index 80f5bd2..ea06191 100644 --- a/src/ecommerce_agent/web/templates/partials/_topbar.html +++ b/src/ecommerce_agent/web/templates/partials/_topbar.html @@ -8,7 +8,7 @@

{{ title }}

{{ chip(integrations, "live") }} {{ model_dropdown("TXT", text_model, text_models, "text") }} - {{ model_dropdown("IMG", image_model, image_models, "image") }} + {{ model_dropdown("IMG", image_model, image_models, "image", image_model_label) }} {{ mode_toggle() }} {{ chip("cap $" ~ cap_usd) }}
diff --git a/tests/smoke/test_api_and_dashboard.py b/tests/smoke/test_api_and_dashboard.py index 538a390..0f3597a 100644 --- a/tests/smoke/test_api_and_dashboard.py +++ b/tests/smoke/test_api_and_dashboard.py @@ -318,12 +318,14 @@ async def _get_session(): async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: response = await client.post( "/api/image-generations", - json={ + data={ "title": "Hat patch", "prompt": "A clean geometric hat patch concept", - "width": 1024, - "height": 1280, + "model": "gpt-image-2", + "width": "1024", + "height": "1280", }, + files={"reference_images": ("", b"", "application/octet-stream")}, ) assert response.status_code == 201 generation = response.json() @@ -349,6 +351,7 @@ async def _get_session(): f"/api/image-generations/{generation['id']}/edits", data={ "prompt": "Make the shape blue", + "model": "gpt-image-2", "width": "1536", "height": "768", "output_format": "webp", @@ -713,9 +716,16 @@ async def drain() -> None: assert len(ready) == 1 product_id = ready[0]["id"] - response = await client.post(f"/api/concepts/{product_id}/workflows") + response = await client.post( + f"/api/concepts/{product_id}/workflows", + json={"image_model": "gpt-image-2"}, + ) assert response.status_code == 201 assert response.json()["status"] == "generating_design" + async with db_sessions() as session: + job = await session.scalar(select(Job).where(Job.subject_id == uuid.UUID(product_id))) + assert job is not None + assert job.payload == {"image_model": "gpt-image-2"} await drain() queue = (await client.get("/api/approval-queue")).json() diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index e4ad686..f1fb39c 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -4,8 +4,8 @@ from ecommerce_agent.config import Settings from ecommerce_agent.services import build_service_container from ecommerce_agent.services.assets import S3AssetStore, SignedLocalAssetStore +from ecommerce_agent.services.image_router import ImageGenerationRouter from ecommerce_agent.services.mocks import MockEtsyGateway, MockPrintfulGateway -from ecommerce_agent.services.openai_images import OpenAIImageProvider from ecommerce_agent.services.printful import DisabledPrintfulGateway, PrintfulClient @@ -38,6 +38,17 @@ def test_real_openai_requires_api_key() -> None: ) +def test_gemini_default_image_model_requires_api_key() -> None: + with pytest.raises(ValidationError, match="GEMINI_API_KEY"): + Settings( + _env_file=None, + integrations_enabled=True, + openai_mode="real", + openai_api_key="test-key", + image_generation_model="gemini-2.5-flash-image", + ) + + def test_real_etsy_requires_credentials() -> None: with pytest.raises(ValidationError, match="ETSY_API_KEY"): Settings( @@ -71,7 +82,7 @@ def test_factory_can_use_real_openai_with_mock_etsy_and_disabled_printful(tmp_pa ) services = build_service_container(settings) - assert isinstance(services.image, OpenAIImageProvider) + assert isinstance(services.image, ImageGenerationRouter) assert isinstance(services.etsy, MockEtsyGateway) assert isinstance(services.printful, DisabledPrintfulGateway) diff --git a/tests/unit/test_gemini_images.py b/tests/unit/test_gemini_images.py new file mode 100644 index 0000000..96b5aa5 --- /dev/null +++ b/tests/unit/test_gemini_images.py @@ -0,0 +1,159 @@ +import base64 +import io +from decimal import Decimal +from pathlib import Path +from typing import Any + +import httpx +import pytest +from PIL import Image + +from ecommerce_agent.domain.dtos import ImageGenerationRequest, ImageReference +from ecommerce_agent.services.assets import LocalAssetStore +from ecommerce_agent.services.gemini_images import GEMINI_GENERATE_URL, GeminiImageProvider +from ecommerce_agent.services.openai_images import OpenAIImageError, OpenAIImageOutcomeUnknown + + +def image_bytes(output_format: str = "png", width: int = 1024, height: int = 1024) -> bytes: + output = io.BytesIO() + Image.new("RGB", (width, height), "yellow").save( + output, format={"png": "PNG", "jpeg": "JPEG", "webp": "WEBP"}[output_format] + ) + return output.getvalue() + + +def success_response( + request: httpx.Request, *, width: int = 1024, height: int = 1024 +) -> httpx.Response: + encoded = base64.b64encode(image_bytes("png", width, height)).decode() + return httpx.Response( + 200, + headers={"x-request-id": "gemini_req"}, + json={ + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": encoded, + } + } + ] + } + } + ], + "usageMetadata": {"totalTokenCount": 1290}, + }, + request=request, + ) + + +def provider(tmp_path: Path, handler, *, max_attempts: int = 3) -> GeminiImageProvider: + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + return GeminiImageProvider( + api_key="gemini-key", + asset_store=LocalAssetStore(tmp_path), + default_model="gemini-2.5-flash-image", + max_cost_usd=Decimal("0.50"), + client=client, + max_attempts=max_attempts, + retry_delay_seconds=0, + ) + + +async def test_generates_with_reference_payload_and_converts_format(tmp_path) -> None: + captured: dict[str, Any] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["api_key"] = request.headers["x-goog-api-key"] + captured["payload"] = request.read().decode() + return success_response(request) + + result = await provider(tmp_path, handler).generate( + ImageGenerationRequest( + prompt="Put Reference 1 on a product poster", + product_concept_id="gemini-id", + model="gemini-2.5-flash-image", + references=(ImageReference("logo.png", "image/png", b"logo-bytes"),), + output_format="webp", + background_removal_friendly=False, + ) + ) + + assert captured["url"] == GEMINI_GENERATE_URL.format(model="gemini-2.5-flash-image") + assert captured["api_key"] == "gemini-key" + assert '"inlineData"' in captured["payload"] + assert "generationConfig" not in captured["payload"] + assert result.provider == "gemini" + assert result.model == "gemini-2.5-flash-image" + assert result.uri.endswith(".webp") + assert result.metadata["usage"]["totalTokenCount"] == 1290 + with Image.open(tmp_path / "image-generations/gemini-id/revision-1.webp") as image: + assert image.format == "WEBP" + + +async def test_gemini_31_uses_minimal_payload(tmp_path) -> None: + captured: dict[str, Any] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["payload"] = request.read().decode() + return success_response(request) + + await provider(tmp_path, handler).generate( + ImageGenerationRequest( + prompt="Poster", + product_concept_id="id", + model="gemini-3.1-flash-image", + ) + ) + + assert "generationConfig" not in captured["payload"] + + +async def test_unsupported_size_is_rejected_before_network_call(tmp_path) -> None: + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return success_response(request) + + with pytest.raises(OpenAIImageError, match="does not support"): + await provider(tmp_path, handler).generate( + ImageGenerationRequest( + prompt="Bad size", + product_concept_id="id", + model="gemini-2.5-flash-image", + width=1024, + height=1280, + ) + ) + assert calls == 0 + + +async def test_invalid_response_is_rejected(tmp_path) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"candidates": []}, request=request) + + with pytest.raises(OpenAIImageError, match="invalid|no image"): + await provider(tmp_path, handler).generate( + ImageGenerationRequest(prompt="Invalid", product_concept_id="id") + ) + + +async def test_ambiguous_timeout_is_not_retried(tmp_path) -> None: + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + raise httpx.ReadTimeout("timed out", request=request) + + with pytest.raises(OpenAIImageOutcomeUnknown, match="not retried"): + await provider(tmp_path, handler).generate( + ImageGenerationRequest(prompt="Timeout", product_concept_id="id") + ) + assert calls == 1 diff --git a/tests/web/test_design_studio.py b/tests/web/test_design_studio.py index 1d4d0f2..0a122ec 100644 --- a/tests/web/test_design_studio.py +++ b/tests/web/test_design_studio.py @@ -68,3 +68,14 @@ async def test_empty_state(web) -> None: resp = await web.client.get("/design-studio") assert resp.status_code == 200 assert "No designs yet" in resp.text + + +async def test_image_model_picker_controls_generation_form(web) -> None: + resp = await web.client.get("/design-studio") + assert resp.status_code == 200 + assert 'name="model" value="gpt-image-2" data-image-model' in resp.text + assert 'data-value="gemini-3.1-flash-image"' in resp.text + assert "Nano Banana 2" in resp.text + assert 'data-value="gemini-3-pro-image"' in resp.text + assert "Nano Banana Pro" in resp.text + assert 'data-value="gemini-2.5-flash-image"' in resp.text