allow selection between portrait, landscape, square canvas types

vibe_bot/main.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import base64
 import logging
+import re
 from io import BytesIO
 from typing import TYPE_CHECKING
 
@@ -24,7 +25,9 @@ from vibe_bot.config import (
     IMAGE_GEN_ENDPOINT,
     IMAGE_GEN_ENDPOINT_KEY,
     IMAGE_GEN_MODEL,
-    IMAGE_GEN_SIZE,
+    IMAGE_GEN_SIZE_LANDSCAPE,
+    IMAGE_GEN_SIZE_PORTRAIT,
+    IMAGE_GEN_SIZE_SQUARE,
     MAX_COMPLETION_TOKENS,
     TTS_MODEL_PATH,
     TTS_SPEED,
@@ -102,6 +105,75 @@ MIN_BOT_NAME_LENGTH = 2
 MAX_BOT_NAME_LENGTH = 50
 MIN_PERSONALITY_LENGTH = 10
 
+# Image layout (canvas orientation) selection for doodlebob.
+DEFAULT_IMAGE_LAYOUT = "square"
+VALID_IMAGE_LAYOUTS = ("portrait", "landscape", "square")
+LAYOUT_SIZES: dict[str, str] = {
+    "portrait": IMAGE_GEN_SIZE_PORTRAIT,
+    "landscape": IMAGE_GEN_SIZE_LANDSCAPE,
+    "square": IMAGE_GEN_SIZE_SQUARE,
+}
+
+IMAGE_LAYOUT_SYSTEM_PROMPT = (
+    "You decide the aspect ratio (layout) of an image that will be generated "
+    "from a user's request. Choose exactly ONE layout from these three options:\n"
+    "- portrait: a tall, vertical image (taller than wide). Use for subjects that "
+    "are taller than they are wide, such as a single standing person or animal, "
+    "a full-body character, a tall building, a skyscraper, a tree, a rocket, or "
+    "any vertical composition.\n"
+    "- landscape: a wide, horizontal image (wider than tall). Use for scenes that "
+    "are wider than they are tall, such as wide landscapes, panoramas, cityscapes, "
+    "seas and horizons, battle or group scenes spread out horizontally, or any "
+    "horizontal composition.\n"
+    "- square: an image that is as wide as it is tall. Use for balanced subjects, "
+    "close-ups, faces, single objects, logos, emblems, or whenever no strong tall "
+    "or wide orientation is implied.\n"
+    "Rules:\n"
+    "- Base your choice ONLY on the orientation the content implies.\n"
+    "- Respond with ONLY the single word portrait, landscape, or square.\n"
+    "- Do NOT include any other text, punctuation, explanation, or reasoning.\n"
+)
+
+
+def parse_image_layout(response: str) -> str:
+    """Parse an LLM response into a valid image layout.
+
+    Args:
+        response: The raw LLM response text.
+
+    Returns:
+        One of "portrait", "landscape", or "square". Falls back to "square"
+        when the response is empty or does not contain a valid layout.
+
+    """
+    text = response.strip().lower()
+    for layout in VALID_IMAGE_LAYOUTS:
+        if re.search(rf"\b{layout}\b", text):
+            return layout
+    return DEFAULT_IMAGE_LAYOUT
+
+
+def select_image_layout(user_message: str) -> str:
+    """Ask the LLM to pick an image layout for the given content.
+
+    Args:
+        user_message: The user's original image request.
+
+    Returns:
+        One of "portrait", "landscape", or "square". Falls back to "square"
+        when the LLM returns an empty or malformed response.
+
+    """
+    response = llama_wrapper.chat_completion_instruct(
+        system_prompt=IMAGE_LAYOUT_SYSTEM_PROMPT,
+        user_prompt=user_message,
+        openai_url=CHAT_ENDPOINT,
+        openai_api_key=CHAT_ENDPOINT_KEY,
+        model=CHAT_MODEL,
+        max_tokens=MAX_COMPLETION_TOKENS,
+    )
+    return parse_image_layout(response)
+
 
 @bot.event
 async def on_ready() -> None:
@@ -713,11 +785,18 @@ async def doodlebob(ctx: CommandsContext[Bot], *, message: str) -> None:
         ctx.author.name,
         message[:100],
     )
-    await ctx.send(f"**Doodlebob erasing {message[:100]}...**")
+    await ctx.send("**Doodlebob shopping for a canvas...**")
+
+    # Let the LLM pick the canvas orientation based on the content.
+    layout = select_image_layout(message)
+    logger.info("Doodlebob selected layout %r for %s", layout, ctx.author.name)
+    await ctx.send(f"**Doodlebob selected {layout}**")
 
     system_prompt = (
         "Given the following message, convert it to a detailed image generation "
         "prompt that will be passed directly into an image generation model. "
+        f"The final image will use a {layout} canvas, so compose the scene to "
+        f"fit that orientation. "
         "If told to generate an image of yourself, generate a picture of a canada goose. "
         "If told to generate a picture of 'me', 'myself', or some other self "
         "reference, generate a picture of a canada goose. Only respond with a valid image "
@@ -747,7 +826,7 @@ async def doodlebob(ctx: CommandsContext[Bot], *, message: str) -> None:
         openai_url=IMAGE_GEN_ENDPOINT,
         openai_api_key=IMAGE_GEN_ENDPOINT_KEY,
         model=IMAGE_GEN_MODEL,
-        size=IMAGE_GEN_SIZE,
+        size=LAYOUT_SIZES[layout],
     )
 
     if not image_b64: