V0

backend/app/api/routes.py

@@ -2,7 +2,9 @@
 
 from fastapi import APIRouter
 
+from app.controllers.image import ImageController
 from app.controllers.messages import MessagesController
+from app.services.image import ImageService
 from app.services.messages import MessagesService
 
 log = logging.getLogger(__name__)
@@ -31,3 +33,18 @@ def get_messages_controller_router():
     tags=["messages"],
     prefix="/api/messages",
 )
+
+
+### Image Generation
+
+
+def get_image_controller_router():
+    service = ImageService()
+    return ImageController(service=service).router
+
+
+router.include_router(
+    get_image_controller_router(),
+    tags=["image"],
+    prefix="/api/generate-image",
+)

backend/app/controllers/image.py

@@ -0,0 +1,44 @@
+import logging
+
+from fastapi import APIRouter, HTTPException
+
+from app.models.image import ImageGenerationRequest, ImageGenerationResponse
+from app.services.image import ImageService
+
+log = logging.getLogger(__name__)
+
+
+class ImageController:
+    def __init__(self, service: ImageService):
+        self.router = APIRouter()
+        self.service = service
+        self.setup_routes()
+
+    def setup_routes(self):
+        router = self.router
+
+        @router.post(
+            "",
+            response_model=ImageGenerationResponse,
+        )
+        async def generate_image(
+            input: ImageGenerationRequest,
+        ) -> ImageGenerationResponse:
+            log.info(f"Generating image with prompt: {input.prompt}")
+            try:
+                response: ImageGenerationResponse = await self.service.generate_image(
+                    input=input
+                )
+                log.info("Image generation completed successfully")
+                return response
+            except ValueError as e:
+                log.error(f"Validation error: {e}")
+                raise HTTPException(status_code=400, detail=str(e))
+            except RuntimeError as e:
+                log.error(f"Service error: {e}")
+                raise HTTPException(status_code=500, detail=str(e))
+            except Exception as e:
+                log.error(f"Unexpected error: {e}")
+                raise HTTPException(
+                    status_code=500, detail="An unexpected error occurred"
+                )

backend/app/models/image.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class ImageGenerationRequest(BaseModel):
+    prompt: str = Field(description="The text prompt describing the image to generate.")
+    image_data: Optional[str] = Field(
+        default=None,
+        description="Optional base64 encoded image data to use as input for image generation.",
+    )
+
+
+class ImageGenerationResponse(BaseModel):
+    image_data: str = Field(description="Base64 encoded generated image data.")
+    text_response: Optional[str] = Field(
+        default=None, description="Any text response from the model."
+    )

backend/app/services/image.py

@@ -0,0 +1,87 @@
+import base64
+import logging
+from io import BytesIO
+
+from google import genai
+from PIL import Image
+
+from app.models.image import ImageGenerationRequest, ImageGenerationResponse
+
+log = logging.getLogger(__name__)
+
+
+class ImageService:
+    def __init__(self):
+        self.client = genai.Client()
+        self.model = "gemini-2.5-flash-image"
+
+    async def generate_image(
+        self, input: ImageGenerationRequest
+    ) -> ImageGenerationResponse:
+        """
+        Generate an image using Google's Imagen API.
+
+        Args:
+            input: The image generation request containing prompt and optional input image
+
+        Returns:
+            ImageGenerationResponse containing the generated image data
+        """
+        log.info(f"Generating image with prompt: {input.prompt}")
+
+        prompt = f"Generate an whiteboard drawingimage based on the following prompt and the reference image: {input.prompt}"
+
+        # Prepare reference image if provided
+        reference_image = None
+        if input.image_data:
+            try:
+                image_bytes = base64.b64decode(input.image_data)
+                reference_image = Image.open(BytesIO(image_bytes))
+                log.info("Added reference image to request")
+            except Exception as e:
+                log.error(f"Error decoding input image: {e}")
+                raise ValueError(f"Invalid image data: {e}")
+
+        # Call Gemini API with image generation model
+        try:
+            # Build the contents list - prompt is required, reference image is optional
+            contents = [prompt]
+            if reference_image:
+                contents.append(reference_image)
+
+            response = self.client.models.generate_content(
+                model=self.model, contents=contents
+            )
+
+            # Parse the response - can contain text and/or image parts
+            if not response.candidates or len(response.candidates) == 0:
+                log.error("No candidates received from Gemini API")
+                raise ValueError("No response generated by the model")
+
+            generated_image_data = None
+            text_response = None
+
+            # Iterate through response parts to extract text and image
+            for part in response.candidates[0].content.parts:
+                if part.text is not None:
+                    text_response = part.text
+                    log.info(f"Received text response: {text_response[:100]}...")
+                elif part.inline_data is not None:
+                    # Convert inline data to base64
+                    generated_image_data = base64.b64encode(
+                        part.inline_data.data
+                    ).decode("utf-8")
+                    log.info("Generated image received and encoded")
+
+            # Ensure we got at least an image
+            if not generated_image_data:
+                log.error("No image data received from Gemini API")
+                raise ValueError("No image generated by the model")
+
+            return ImageGenerationResponse(
+                image_data=generated_image_data, text_response=text_response
+            )
+
+        except Exception as e:
+            log.error(f"Error calling Gemini API: {e}")
+            raise RuntimeError(f"Failed to generate image: {e}")