From a73a10fdb67768f8ad929803ae1af5c3537e2908 Mon Sep 17 00:00:00 2001 From: ducoterra Date: Mon, 17 Aug 2026 12:19:31 -0400 Subject: [PATCH] allow selection between portrait, landscape, square canvas types --- README.md | 6 ++ vibe_bot/config.py | 4 +- vibe_bot/main.py | 85 +++++++++++++++++++- vibe_bot/tests/test_main.py | 156 ++++++++++++++++++++++++++++++++++++ 4 files changed, 247 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 3b8918d..45431f0 100644 --- a/README.md +++ b/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 | diff --git a/vibe_bot/config.py b/vibe_bot/config.py index 6dde2be..7afc6f5 100644 --- a/vibe_bot/config.py +++ b/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", "") diff --git a/vibe_bot/main.py b/vibe_bot/main.py index d2c3035..9a5a0ed 100644 --- a/vibe_bot/main.py +++ b/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: diff --git a/vibe_bot/tests/test_main.py b/vibe_bot/tests/test_main.py index c250380..9c28fa3 100644 --- a/vibe_bot/tests/test_main.py +++ b/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"