services/vibe-bot
2
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity

allow selection between portrait, landscape, square canvas types

Browse Source
This commit is contained in:
ducoterra
2026-08-17 12:19:31 -04:00
parent 5518f07234
commit a73a10fdb6
4 changed files with 247 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+6
View File
@@ -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
+3 -1
View File
@@ -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
+82 -3
View File
@@ -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
+156
View File
@@ -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"