allow selection between portrait, landscape, square canvas types

README.md

@@ -123,6 +123,9 @@ IMAGE_GEN_ENDPOINT_KEY=your_api_key
 IMAGE_EDIT_ENDPOINT_KEY=your_api_key
 IMAGE_GEN_MODEL=gen
 IMAGE_EDIT_MODEL=edit
+IMAGE_GEN_SIZE_SQUARE=1024x1024
+IMAGE_GEN_SIZE_PORTRAIT=1024x1536
+IMAGE_GEN_SIZE_LANDSCAPE=1536x1024
 
 # Embedding API (required)
 EMBEDDING_ENDPOINT=https://your-api.com/v1
@@ -269,6 +272,9 @@ uv run black --check vibe_bot/
 | `CHAT_MODEL`            | *(required)*       | Model name for chat completions       |
 | `IMAGE_GEN_ENDPOINT`    | *(required)*       | Image generation API URL              |
 | `IMAGE_EDIT_ENDPOINT`   | *(required)*       | Image editing API URL                 |
+| `IMAGE_GEN_SIZE_SQUARE`   | `1024x1024`        | Square canvas size for image generation     |
+| `IMAGE_GEN_SIZE_PORTRAIT` | `1024x1536`        | Portrait (tall) canvas size for image generation |
+| `IMAGE_GEN_SIZE_LANDSCAPE`| `1536x1024`        | Landscape (wide) canvas size for image generation |
 | `EMBEDDING_ENDPOINT`    | *(required)*       | Embedding API URL                     |
 | `EMBEDDING_MODEL`       | *(required)*       | Model name for text embeddings        |
 | `MAX_COMPLETION_TOKENS` | `1000`             | Max tokens in LLM responses           |

vibe_bot/config.py

@@ -38,7 +38,9 @@ EMBEDDING_ENDPOINT_KEY: str = os.getenv("EMBEDDING_ENDPOINT_KEY", "placeholder")
 CHAT_MODEL: str = os.getenv("CHAT_MODEL", "")
 COMPLETION_MODEL: str = os.getenv("COMPLETION_MODEL", "")
 IMAGE_GEN_MODEL: str = os.getenv("IMAGE_GEN_MODEL", "")
-IMAGE_GEN_SIZE: str = os.getenv("IMAGE_GEN_SIZE", "1024x1024")
+IMAGE_GEN_SIZE_SQUARE: str = os.getenv("IMAGE_GEN_SIZE_SQUARE", "1024x1024")
+IMAGE_GEN_SIZE_PORTRAIT: str = os.getenv("IMAGE_GEN_SIZE_PORTRAIT", "1024x1536")
+IMAGE_GEN_SIZE_LANDSCAPE: str = os.getenv("IMAGE_GEN_SIZE_LANDSCAPE", "1536x1024")
 IMAGE_EDIT_MODEL: str = os.getenv("IMAGE_EDIT_MODEL", "")
 EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "")
 

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:

vibe_bot/tests/test_main.py

@@ -961,3 +961,159 @@ def test_debug_tools(mock_ctx: MagicMock) -> None:
     assert "LLM Tools" in call_args
     assert "get_channel_members" in call_args
     assert "members" in call_args.lower()
+
+
+# ---------------------------------------------------------------------------
+# Image layout selection (doodlebob)
+# ---------------------------------------------------------------------------
+
+LAYOUT_TEST_SIZES = {
+    "portrait": "1024x1536",
+    "landscape": "1536x1024",
+    "square": "1024x1024",
+}
+
+
+def _sent_texts(mock_ctx: MagicMock) -> list[str]:
+    """Collect positional (text) arguments sent via ctx.send."""
+    return [c.args[0] for c in mock_ctx.send.call_args_list if c.args]
+
+
+@pytest.mark.parametrize(
+    ("response", "expected"),
+    [
+        ("portrait", "portrait"),
+        ("landscape", "landscape"),
+        ("square", "square"),
+        ("  Portrait  ", "portrait"),
+        ("LANDSCAPE", "landscape"),
+        ("square.", "square"),
+        ("I would use portrait.", "portrait"),
+        ("This scene is best as landscape.", "landscape"),
+        ("A square composition works here.", "square"),
+    ],
+)
+def test_parse_image_layout_valid(response: str, expected: str) -> None:
+    """Test that valid LLM layout responses parse to the right layout."""
+    import vibe_bot.main as main_module
+
+    assert main_module.parse_image_layout(response) == expected
+
+
+@pytest.mark.parametrize(
+    "response",
+    [
+        "",
+        "   ",
+        "banana",
+        "1024x1024",
+        "tall and wide",
+        "squareness",  # word boundary should prevent a match on "square"
+    ],
+)
+def test_parse_image_layout_defaults_to_square(response: str) -> None:
+    """Test that empty or malformed responses fall back to square."""
+    import vibe_bot.main as main_module
+
+    assert main_module.parse_image_layout(response) == "square"
+
+
+def test_doodlebob_selects_portrait(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test doodlebob picks portrait and passes the portrait size."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "portrait",  # layout selection
+        "a tall portrait of a lighthouse",  # prompt rewrite
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="a tall lighthouse"))
+
+    # Layout selection is the first instruct call, using the layout prompt.
+    layout_call = mock_llama_wrapper.chat_completion_instruct.call_args_list[0]
+    assert layout_call.kwargs["system_prompt"] == main_module.IMAGE_LAYOUT_SYSTEM_PROMPT
+    assert layout_call.kwargs["user_prompt"] == "a tall lighthouse"
+
+    mock_llama_wrapper.image_generation.assert_called_once()
+    assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1536"
+
+    sent = _sent_texts(mock_ctx)
+    assert any("shopping for a canvas" in m for m in sent)
+    assert any("selected portrait" in m for m in sent)
+    assert any("drone strike" in m for m in sent)
+
+
+def test_doodlebob_selects_landscape(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test doodlebob picks landscape and passes the landscape size."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "landscape",
+        "a wide panoramic coastline",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="wide coastline"))
+
+    assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1536x1024"
+
+
+def test_doodlebob_malformed_layout_defaults_square(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test a malformed layout response falls back to the square size."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "banana",  # malformed layout
+        "a balanced composition",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="a logo"))
+
+    assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1024"
+    sent = _sent_texts(mock_ctx)
+    assert any("selected square" in m for m in sent)
+
+
+def test_doodlebob_empty_layout_defaults_square(
+    mock_ctx: MagicMock,
+    mock_llama_wrapper: MagicMock,
+    mock_base64: MagicMock,
+) -> None:
+    """Test an empty layout response (LLM failure) falls back to square."""
+    import asyncio
+
+    import vibe_bot.main as main_module
+
+    mock_llama_wrapper.chat_completion_instruct.side_effect = [
+        "",  # empty layout response
+        "a balanced composition",
+    ]
+    mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
+
+    with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
+        asyncio.run(main_module.doodlebob(mock_ctx, message="a logo"))
+
+    assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1024"