hanzoai_agent/sdk/python/hanzo_agents/vision.py

"""
Image Generation Module

Handles image generation across multiple providers (LLM, OpenRouter).
Keeps provider-specific implementation details separate from the main agent code.

Supported Providers:
- LLM: DALL-E, Azure DALL-E, Bedrock Stable Diffusion, etc.
- OpenRouter: Gemini image generation, etc.
"""

from typing import Any, Optional
from hanzo_agents.logger import log_error


async def generate_image_llm(
    prompt: str,
    model: str,
    size: str,
    quality: str,
    style: Optional[str],
    response_format: str,
    **kwargs,
) -> Any:
    """
    Generate image using LLM's image generation API.

    This function uses LLM's `aimage_generation()` which supports:
    - OpenAI DALL-E (dall-e-3, dall-e-2)
    - Azure DALL-E
    - AWS Bedrock Stable Diffusion
    - And other LLM-supported image generation models

    Args:
        prompt: Text prompt for image generation
        model: Model to use (e.g., "dall-e-3", "azure/dall-e-3")
        size: Image size (e.g., "1024x1024", "1792x1024")
        quality: Image quality ("standard", "hd")
        style: Image style ("vivid", "natural") - DALL-E 3 only
        response_format: Response format ("url", "b64_json")
        **kwargs: Additional LLM parameters

    Returns:
        MultimodalResponse with generated image(s)

    Raises:
        ImportError: If llm is not installed
        Exception: If image generation fails
    """
    try:
        import llm
    except ImportError:
        raise ImportError(
            "llm is not installed. Please install it with `pip install llm`."
        )

    # Prepare image generation parameters
    image_params = {
        "prompt": prompt,
        "model": model,
        "size": size,
        "quality": quality,
        "response_format": response_format,
        **kwargs,
    }

    # Add style parameter only for DALL-E 3
    if style and "dall-e-3" in model:
        image_params["style"] = style

    try:
        # Use LLM's image generation function
        response = await llm.aimage_generation(**image_params)

        # Import multimodal response detection
        from hanzo_agents.multimodal_response import detect_multimodal_response

        # Detect and wrap multimodal content
        return detect_multimodal_response(response)

    except Exception as e:
        log_error(f"LLM image generation failed: {e}")
        raise


async def generate_image_openrouter(
    prompt: str,
    model: str,
    size: str,
    quality: str,
    style: Optional[str],
    response_format: str,
    **kwargs,
) -> Any:
    """
    Generate image using OpenRouter's chat completions API.

    OpenRouter uses modalities to enable image generation through
    the standard chat completions endpoint. This is different from
    LLM's dedicated image generation API.

    Supported models:
    - google/gemini-2.5-flash-image-preview
    - And other OpenRouter models with image generation capabilities

    Args:
        prompt: Text prompt for image generation
        model: OpenRouter model (must start with "openrouter/")
        size: Image size (may not be used by all OpenRouter models)
        quality: Image quality (may not be used by all OpenRouter models)
        style: Image style (may not be used by all OpenRouter models)
        response_format: Response format (may not be used by all OpenRouter models)
        **kwargs: Additional OpenRouter-specific parameters (e.g., image_config)

    Returns:
        MultimodalResponse with generated image(s)

    Raises:
        ImportError: If llm is not installed
        Exception: If image generation fails

    Note:
        OpenRouter-specific parameters like `image_config` should be passed via kwargs.
        Example: image_config={"aspect_ratio": "16:9"}
    """
    try:
        import llm
    except ImportError:
        raise ImportError(
            "llm is not installed. Please install it with `pip install llm`."
        )

    from hanzo_agents.multimodal_response import ImageOutput, MultimodalResponse

    # Build messages for OpenRouter chat completions
    messages = [{"role": "user", "content": prompt}]

    # Prepare parameters for OpenRouter
    # OpenRouter uses chat completions with modalities parameter
    completion_params = {
        "model": model,
        "messages": messages,
        "modalities": ["image", "text"],
        **kwargs,  # Pass through any additional kwargs (e.g., image_config)
    }

    try:
        # Use LLM's completion function (OpenRouter uses chat API)
        response = await llm.acompletion(**completion_params)

        # Extract images from OpenRouter response
        # OpenRouter returns images in choices[0].message.images
        images = []
        text_content = ""

        if hasattr(response, "choices") and len(response.choices) > 0:
            message = response.choices[0].message

            # Extract text content
            if hasattr(message, "content") and message.content:
                text_content = message.content

            # Extract images
            if hasattr(message, "images") and message.images:
                for img_data in message.images:
                    # OpenRouter images have structure: {"type": "image_url", "image_url": {"url": "data:..."}}
                    if hasattr(img_data, "image_url"):
                        image_url = (
                            img_data.image_url.url
                            if hasattr(img_data.image_url, "url")
                            else None
                        )
                    elif isinstance(img_data, dict) and "image_url" in img_data:
                        image_url = img_data["image_url"].get("url")
                    else:
                        image_url = None

                    if image_url:
                        images.append(
                            ImageOutput(
                                url=image_url,
                                b64_json=None,
                                revised_prompt=None,
                            )
                        )

        # Create MultimodalResponse
        return MultimodalResponse(
            text=text_content or prompt,
            audio=None,
            images=images,
            files=[],
            raw_response=response,
        )

    except Exception as e:
        log_error(f"OpenRouter image generation failed: {e}")
        raise