feat: add async image generation (agenerate_image)

eric-tramel · claude · eric-tramel · commit a87f98b6afba · 2026-02-12T20:56:48.000-05:00
Add agenerate_image(), _agenerate_image_chat_completion(), and _agenerate_image_diffusion() async methods mirroring the sync generate_image() added in #317. The chat completion path uses acompletion(), the diffusion path uses router.aimage_generation(). Includes 5 new tests covering both paths, error cases, and usage tracking. Also fixes F821 lint errors for type annotations. Co-Authored-By: Remi <noreply@anthropic.com>
diff --git a/packages/data-designer-engine/src/data_designer/engine/models/facade.py b/packages/data-designer-engine/src/data_designer/engine/models/facade.py
@@ -729,3 +729,169 @@ async def agenerate(
             )
 
         return output_obj, messages
+
+    @acatch_llm_exceptions
+    async def agenerate_image(
+        self,
+        prompt: str,
+        multi_modal_context: list[dict[str, Any]] | None = None,
+        skip_usage_tracking: bool = False,
+        **kwargs: Any,
+    ) -> list[str]:
+        """Async version of generate_image. Generate image(s) and return base64-encoded data.
+
+        Automatically detects the appropriate API based on model name:
+        - Diffusion models (DALL-E, Stable Diffusion, Imagen, etc.) → image_generation API
+        - All other models → chat/completions API (default)
+
+        Both paths return base64-encoded image data. If the API returns multiple images,
+        all are returned in the list.
+
+        Args:
+            prompt: The prompt for image generation
+            multi_modal_context: Optional list of image contexts for multi-modal generation.
+                Only used with autoregressive models via chat completions API.
+            skip_usage_tracking: Whether to skip usage tracking
+            **kwargs: Additional arguments to pass to the model (including n=number of images)
+
+        Returns:
+            List of base64-encoded image strings (without data URI prefix)
+
+        Raises:
+            ImageGenerationError: If image generation fails or returns invalid data
+        """
+        logger.debug(
+            f"Generating image with model {self.model_name!r}...",
+            extra={"model": self.model_name, "prompt": prompt},
+        )
+
+        # Auto-detect API type based on model name
+        if is_image_diffusion_model(self.model_name):
+            images = await self._agenerate_image_diffusion(prompt, skip_usage_tracking, **kwargs)
+        else:
+            images = await self._agenerate_image_chat_completion(
+                prompt, multi_modal_context, skip_usage_tracking, **kwargs
+            )
+
+        # Track image usage
+        if not skip_usage_tracking and len(images) > 0:
+            self._usage_stats.extend(image_usage=ImageUsageStats(total_images=len(images)))
+
+        return images
+
+    async def _agenerate_image_chat_completion(
+        self,
+        prompt: str,
+        multi_modal_context: list[dict[str, Any]] | None = None,
+        skip_usage_tracking: bool = False,
+        **kwargs: Any,
+    ) -> list[str]:
+        """Async version of _generate_image_chat_completion.
+
+        Generate image(s) using autoregressive model via chat completions API.
+
+        Args:
+            prompt: The prompt for image generation
+            multi_modal_context: Optional list of image contexts for multi-modal generation
+            skip_usage_tracking: Whether to skip usage tracking
+            **kwargs: Additional arguments to pass to the model
+
+        Returns:
+            List of base64-encoded image strings
+        """
+        messages = prompt_to_messages(user_prompt=prompt, multi_modal_context=multi_modal_context)
+
+        response = None
+        try:
+            response = await self.acompletion(
+                messages=messages,
+                skip_usage_tracking=skip_usage_tracking,
+                **kwargs,
+            )
+
+            logger.debug(
+                f"Received image(s) from autoregressive model {self.model_name!r}",
+                extra={"model": self.model_name, "response": response},
+            )
+
+            # Validate response structure
+            if not response.choices or len(response.choices) == 0:
+                raise ImageGenerationError("Image generation response missing choices")
+
+            message = response.choices[0].message
+            images = []
+
+            # Extract base64 from images attribute (primary path)
+            if hasattr(message, "images") and message.images:
+                for image in message.images:
+                    # Handle different response formats
+                    if isinstance(image, dict) and "image_url" in image:
+                        image_url = image["image_url"]
+
+                        if isinstance(image_url, dict) and "url" in image_url:
+                            if (b64 := _try_extract_base64(image_url["url"])) is not None:
+                                images.append(b64)
+                        elif isinstance(image_url, str):
+                            if (b64 := _try_extract_base64(image_url)) is not None:
+                                images.append(b64)
+                    # Fallback: treat as base64 string
+                    elif isinstance(image, str):
+                        if (b64 := _try_extract_base64(image)) is not None:
+                            images.append(b64)
+
+            # Fallback: check content field if it looks like image data
+            if not images:
+                content = message.content or ""
+                if content and (content.startswith("data:image/") or is_base64_image(content)):
+                    if (b64 := _try_extract_base64(content)) is not None:
+                        images.append(b64)
+
+            if not images:
+                raise ImageGenerationError("No image data found in image generation response")
+
+            return images
+
+        except Exception:
+            raise
+
+    async def _agenerate_image_diffusion(
+        self, prompt: str, skip_usage_tracking: bool = False, **kwargs: Any
+    ) -> list[str]:
+        """Async version of _generate_image_diffusion.
+
+        Generate image(s) using diffusion model via image_generation API.
+
+        Always returns base64. If the API returns URLs instead of inline base64,
+        the images are downloaded and converted automatically.
+
+        Returns:
+            List of base64-encoded image strings
+        """
+        kwargs = self.consolidate_kwargs(**kwargs)
+
+        response = None
+
+        try:
+            response = await self._router.aimage_generation(prompt=prompt, model=self.model_name, **kwargs)
+
+            logger.debug(
+                f"Received {len(response.data)} image(s) from diffusion model {self.model_name!r}",
+                extra={"model": self.model_name, "response": response},
+            )
+
+            # Validate response
+            if not response.data or len(response.data) == 0:
+                raise ImageGenerationError("Image generation returned no data")
+
+            images = [b64 for img in response.data if (b64 := _try_extract_base64(img)) is not None]
+
+            if not images:
+                raise ImageGenerationError("No image data could be extracted from response")
+
+            return images
+
+        except Exception:
+            raise
+        finally:
+            if not skip_usage_tracking and response is not None:
+                self._track_token_usage_from_image_diffusion(response)
diff --git a/packages/data-designer-engine/tests/engine/models/test_facade.py b/packages/data-designer-engine/tests/engine/models/test_facade.py
@@ -18,6 +18,7 @@
 
 if TYPE_CHECKING:
     import litellm
+    from litellm.types.utils import EmbeddingResponse, ModelResponse
 
 
 def mock_oai_response_object(response_text: str) -> StubResponse:
@@ -1403,3 +1404,108 @@ async def test_agenerate_success(
     # Trace should contain at least the user prompt and the assistant response
     assert any(msg.role == "user" for msg in trace)
     assert any(msg.role == "assistant" and msg.content == "parsed output" for msg in trace)
+
+
+# =============================================================================
+# Async image generation tests
+# =============================================================================
+
+
+@patch("data_designer.engine.models.facade.CustomRouter.aimage_generation", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_agenerate_image_diffusion_success(
+    mock_aimage_generation: AsyncMock,
+    stub_model_facade: ModelFacade,
+) -> None:
+    """Test async image generation via diffusion API."""
+    mock_response = litellm.types.utils.ImageResponse(
+        data=[
+            litellm.types.utils.ImageObject(b64_json="image1_base64"),
+            litellm.types.utils.ImageObject(b64_json="image2_base64"),
+        ]
+    )
+    mock_aimage_generation.return_value = mock_response
+
+    with patch("data_designer.engine.models.facade.is_image_diffusion_model", return_value=True):
+        images = await stub_model_facade.agenerate_image(prompt="test prompt")
+
+    assert len(images) == 2
+    assert images == ["image1_base64", "image2_base64"]
+    assert mock_aimage_generation.call_count == 1
+    # Verify image usage was tracked
+    assert stub_model_facade.usage_stats.image_usage.total_images == 2
+
+
+@patch.object(ModelFacade, "acompletion", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_agenerate_image_chat_completion_success(
+    mock_acompletion: AsyncMock,
+    stub_model_facade: ModelFacade,
+) -> None:
+    """Test async image generation via chat completion API."""
+    mock_message = litellm.types.utils.Message(
+        role="assistant",
+        content="",
+        images=[
+            litellm.types.utils.ImageURLListItem(
+                type="image_url", image_url={"url": "data:image/png;base64,image1"}, index=0
+            ),
+        ],
+    )
+    mock_response = litellm.types.utils.ModelResponse(choices=[litellm.types.utils.Choices(message=mock_message)])
+    mock_acompletion.return_value = mock_response
+
+    with patch("data_designer.engine.models.facade.is_image_diffusion_model", return_value=False):
+        images = await stub_model_facade.agenerate_image(prompt="test prompt")
+
+    assert len(images) == 1
+    assert images == ["image1"]
+    assert mock_acompletion.call_count == 1
+    assert stub_model_facade.usage_stats.image_usage.total_images == 1
+
+
+@patch("data_designer.engine.models.facade.CustomRouter.aimage_generation", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_agenerate_image_diffusion_no_data(
+    mock_aimage_generation: AsyncMock,
+    stub_model_facade: ModelFacade,
+) -> None:
+    """Test async image generation raises error when diffusion API returns no data."""
+    mock_response = litellm.types.utils.ImageResponse(data=[])
+    mock_aimage_generation.return_value = mock_response
+
+    with patch("data_designer.engine.models.facade.is_image_diffusion_model", return_value=True):
+        with pytest.raises(ImageGenerationError, match="Image generation returned no data"):
+            await stub_model_facade.agenerate_image(prompt="test prompt")
+
+
+@patch.object(ModelFacade, "acompletion", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_agenerate_image_chat_completion_no_choices(
+    mock_acompletion: AsyncMock,
+    stub_model_facade: ModelFacade,
+) -> None:
+    """Test async image generation raises error when response has no choices."""
+    mock_response = litellm.types.utils.ModelResponse(choices=[])
+    mock_acompletion.return_value = mock_response
+
+    with patch("data_designer.engine.models.facade.is_image_diffusion_model", return_value=False):
+        with pytest.raises(ImageGenerationError, match="Image generation response missing choices"):
+            await stub_model_facade.agenerate_image(prompt="test prompt")
+
+
+@patch("data_designer.engine.models.facade.CustomRouter.aimage_generation", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_agenerate_image_skip_usage_tracking(
+    mock_aimage_generation: AsyncMock,
+    stub_model_facade: ModelFacade,
+) -> None:
+    """Test that async image generation respects skip_usage_tracking flag."""
+    mock_response = litellm.types.utils.ImageResponse(data=[litellm.types.utils.ImageObject(b64_json="image1_base64")])
+    mock_aimage_generation.return_value = mock_response
+
+    with patch("data_designer.engine.models.facade.is_image_diffusion_model", return_value=True):
+        images = await stub_model_facade.agenerate_image(prompt="test prompt", skip_usage_tracking=True)
+
+    assert len(images) == 1
+    assert stub_model_facade.usage_stats.image_usage.total_images == 0
