Compare commits

...
3 Commits
Author SHA1 Message Date
ducoterra b740d00b5d add time estimation for image gen
Build and Push Container / build-and-push (push) Successful in 1m11s
2026-08-17 12:29:26 -04:00
ducoterra a73a10fdb6 allow selection between portrait, landscape, square canvas types 2026-08-17 12:19:31 -04:00
ducoterra 5518f07234 Various test fixes 2026-08-17 11:35:40 -04:00
13 changed files with 648 additions and 41 deletions
+36
View File
@@ -0,0 +1,36 @@
# AGENTS.md
Single Python package `vibe_bot`: a Discord bot (discord.py, `!` prefix) with SQLite RAG chat history, Kokoro TTS, and image gen/edit via OpenAI-compatible APIs. Python 3.13, managed with uv. Everything lives in `vibe_bot/`.
## Commands
- Setup: `uv sync --extra dev` — plain `uv sync` already installs the `dev` dependency group (ruff, pyright); the `dev` extra adds black/debugpy/mypy
- Run bot: `uv run python -m vibe_bot.main` — this logs the real bot into Discord with the token from `.env`; don't run it as a smoke test
- Tests: `uv run pytest vibe_bot/tests/ -v`; single test: `uv run pytest vibe_bot/tests/test_main.py::test_name`
- Checks: `uv run ruff check vibe_bot/`, `uv run mypy vibe_bot/` (strict), `uv run pyright vibe_bot/` (strict), `uv run black --check vibe_bot/`
- Container: `./build.sh` (podman). CI (`.gitea/workflows/build-push.yml`, Gitea) only builds/pushes the image on main/release — lint and tests are not gated; run them locally
## Setup requirements
- A repo-root `.env` is required even to run tests: `config.py` calls `load_dotenv()` and raises `RuntimeError` at import time if any required var is missing. Placeholder values suffice for the mocked suite. Never commit it.
- PortAudio is a required system library (kokoro-tts → sounddevice). Without it, ~60 tests in `test_main`/`test_tts` error at import time. On this host: `sudo dnf install portaudio` (the Containerfile installs `portaudio19-dev`).
- TTS needs `kokoro-v1.0.onnx` and `voices-v1.0.bin` in the repo root (baked into the container image). The bot runs without them; only `!speak` degrades.
## Test suite gotchas
- Baseline: everything passes except `test_llama_wrapper.py::test_chat_completion_think` and `::test_chat_completion_instruct`. Those two are unmocked live calls to the real `CHAT_ENDPOINT` from `.env` and fail without network access to that API.
- `test_config.py` hardcodes `sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots")` (a stale repo path). The test only passes because pytest's cwd fallback finds the package — run tests from the repo root.
- If `uv run <tool>` suddenly fails with `ModuleNotFoundError` or "bad interpreter", the `.venv` shebangs are stale from a repo move: `rm -rf .venv && uv sync --extra dev`.
## Code map
- `main.py` — entrypoint. The bot is created at module import (module-level `commands.Bot(...)`), `bot.run()` only under `__main__`. Custom-bot "commands" (`!<bot_name> ...`) are matched in `on_message` against the database, not registered with `bot.command`.
- `config.py` — env loading + import-time validation, voice catalog.
- `database.py` — `ChatDatabase` (embeddings stored as float32 blobs, cosine-similarity RAG, schema auto-migrates on startup) and `CustomBotManager`.
- `llama_wrapper.py` — thin OpenAI-compatible clients for chat / image gen / image edit / embeddings, each with its own endpoint, key, and model.
- `tools.py` — `get_channel_members` is a no-op LangChain `@tool` stub used only for name/description/schema; the real implementation is `get_channel_members_impl(channel)`, wired into the tool executor in `main.py`.
## Style
- All four checks (ruff, mypy strict, pyright strict, black) are declared gates and currently pass; keep new code clean under all of them.
- ruff is a dev dependency (installed via `uv add --dev`) with config in `pyproject.toml`; it runs with the default ruleset of the locked version.
+6
View File
@@ -123,6 +123,9 @@ IMAGE_GEN_ENDPOINT_KEY=your_api_key
IMAGE_EDIT_ENDPOINT_KEY=your_api_key IMAGE_EDIT_ENDPOINT_KEY=your_api_key
IMAGE_GEN_MODEL=gen IMAGE_GEN_MODEL=gen
IMAGE_EDIT_MODEL=edit IMAGE_EDIT_MODEL=edit
IMAGE_GEN_SIZE_SQUARE=1024x1024
IMAGE_GEN_SIZE_PORTRAIT=1024x1536
IMAGE_GEN_SIZE_LANDSCAPE=1536x1024
# Embedding API (required) # Embedding API (required)
EMBEDDING_ENDPOINT=https://your-api.com/v1 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 | | `CHAT_MODEL` | *(required)* | Model name for chat completions |
| `IMAGE_GEN_ENDPOINT` | *(required)* | Image generation API URL | | `IMAGE_GEN_ENDPOINT` | *(required)* | Image generation API URL |
| `IMAGE_EDIT_ENDPOINT` | *(required)* | Image editing 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_ENDPOINT` | *(required)* | Embedding API URL |
| `EMBEDDING_MODEL` | *(required)* | Model name for text embeddings | | `EMBEDDING_MODEL` | *(required)* | Model name for text embeddings |
| `MAX_COMPLETION_TOKENS` | `1000` | Max tokens in LLM responses | | `MAX_COMPLETION_TOKENS` | `1000` | Max tokens in LLM responses |
+10 -1
View File
@@ -21,7 +21,6 @@ dependencies = [
[project.optional-dependencies] [project.optional-dependencies]
dev = [ dev = [
"pyright>=1.1.398",
"mypy>=1.17.0", "mypy>=1.17.0",
"black>=25.1.0", "black>=25.1.0",
"debugpy>=1.8.0", "debugpy>=1.8.0",
@@ -56,7 +55,17 @@ reportPrivateUsage = false
line-length = 88 line-length = 88
target-version = ["py313"] target-version = ["py313"]
[tool.ruff]
line-length = 88
target-version = "py313"
[tool.pytest.ini_options] [tool.pytest.ini_options]
filterwarnings = [ filterwarnings = [
"ignore::pytest.PytestUnraisableExceptionWarning", "ignore::pytest.PytestUnraisableExceptionWarning",
] ]
[dependency-groups]
dev = [
"pyright>=1.1.409",
"ruff>=0.16.3",
]
Generated
+36 -1
View File
@@ -2015,6 +2015,31 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d0/02/fa464cdfbe6b26e0600b62c528b72d8608f5cc49f96b8d6e38c95d60c676/rpds_py-0.30.0-cp314-cp314t-win_amd64.whl", hash = "sha256:27f4b0e92de5bfbc6f86e43959e6edd1425c33b5e69aab0984a72047f2bcf1e3", size = 226532, upload-time = "2025-11-30T20:24:14.634Z" }, { url = "https://files.pythonhosted.org/packages/d0/02/fa464cdfbe6b26e0600b62c528b72d8608f5cc49f96b8d6e38c95d60c676/rpds_py-0.30.0-cp314-cp314t-win_amd64.whl", hash = "sha256:27f4b0e92de5bfbc6f86e43959e6edd1425c33b5e69aab0984a72047f2bcf1e3", size = 226532, upload-time = "2025-11-30T20:24:14.634Z" },
] ]
[[package]]
name = "ruff"
version = "0.16.3"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/61/b3/3213589383f8f1b3938781bd1278713f6d18621a14992b3e81fefb8a5ef9/ruff-0.16.3.tar.gz", hash = "sha256:e76d33a347661a84b5be6d043d0347fdc745dfdcf825a8f4fed64b5e26eebdf2", size = 4891904, upload-time = "2026-08-13T15:17:13.381Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/bf/96/493770daebd68c0a67f1549fdf519f53be51fc435186c0585bcc272fd76c/ruff-0.16.3-py3-none-linux_armv6l.whl", hash = "sha256:0c5710e247a58a4521e66e124ba9a74655b414f61ba3a2e9e3811e11098f48f7", size = 10902799, upload-time = "2026-08-13T15:16:27.382Z" },
{ url = "https://files.pythonhosted.org/packages/5e/e6/2becf3942fddc29a29b8df47691d456fb1085391a694f74d84513251418c/ruff-0.16.3-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:fe155130631a2471fd2e14a7a664a4dfbd7194b8229c3d7b2a40b21178639081", size = 11135539, upload-time = "2026-08-13T15:16:30.87Z" },
{ url = "https://files.pythonhosted.org/packages/3e/1e/4b8b72f0d006dbf19326aa99f9ca0ee2ff374187c4d301cf529a51aa06fe/ruff-0.16.3-py3-none-macosx_11_0_arm64.whl", hash = "sha256:e2ed719e14aa64d895c2ee922594a90a43c861a93f0575a95ff8c47cdbd13eb9", size = 10475095, upload-time = "2026-08-13T15:16:33.259Z" },
{ url = "https://files.pythonhosted.org/packages/92/32/2201fa49ba1f6c101ee321e83f051ac7a4b8d07b0ef6b4d3f2772b302275/ruff-0.16.3-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9e0b1da805eb043654645d74d5de1e5ce2edc686e40790d2b86f56d71cc06a84", size = 10668771, upload-time = "2026-08-13T15:16:35.65Z" },
{ url = "https://files.pythonhosted.org/packages/c3/66/4afc5c8363bd04d45effce1b7c8713ca037d7a6740b7451a2403a6e3a972/ruff-0.16.3-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a37bdea0bbe21780f590bf437d6412c8c4e1b6cd010f91a65c2c40c5e5f5f870", size = 10699568, upload-time = "2026-08-13T15:16:38.195Z" },
{ url = "https://files.pythonhosted.org/packages/53/fd/c67d246bf36bf1698551c56de39e95cd07f70e64433e0098e6267d77061b/ruff-0.16.3-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:09571e6d1288ed9be475207a3ac04ada404f1cd898104be0f6ab8d7df438575b", size = 11499365, upload-time = "2026-08-13T15:16:40.623Z" },
{ url = "https://files.pythonhosted.org/packages/67/0b/00ecbceb99a263af7b12f6f05ac3c92bc47b905e91adc3f207a836e3bc01/ruff-0.16.3-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2c18c5a101eb540010638cc1ff3c84944d3adb3df62b8d98ca8f22ba484d3413", size = 12311728, upload-time = "2026-08-13T15:16:43.564Z" },
{ url = "https://files.pythonhosted.org/packages/54/b2/b7b3bb54f4d3f7db504e476ad4ab8de530dceebe2c061384b2757ee419e8/ruff-0.16.3-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8457c44f15033c85ddbb77b15d451df9e24e4bd03b628396dd3610cedc3b8f82", size = 11699896, upload-time = "2026-08-13T15:16:46.209Z" },
{ url = "https://files.pythonhosted.org/packages/c7/30/4c468429ac195addc5ee1b717b6ab1b66632786737ca3b2ed3443fb0c26a/ruff-0.16.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:294b95c4ae0cda9388525c2047778aa758d6b8d4bb876fd4e9eaa3ebc92343eb", size = 11058736, upload-time = "2026-08-13T15:16:48.823Z" },
{ url = "https://files.pythonhosted.org/packages/43/67/7a113cdaddf24b64d7f75b1242a99d04c82fcef4f6921fdbb832beaffb5f/ruff-0.16.3-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:3d0c7c40c87c2a820509c31ba007968da6e1306468c067b2d82fbfdbcd0e8474", size = 11586911, upload-time = "2026-08-13T15:16:51.913Z" },
{ url = "https://files.pythonhosted.org/packages/f1/c1/2e66f24c0f3ead25a5e660111778685e505e5da353c82802bf49f0cbe7b9/ruff-0.16.3-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:9f738c0fdfa8eed0b2ce7fb27ee7258208a92a68d7949e62aa15164bc7b389da", size = 10954265, upload-time = "2026-08-13T15:16:54.763Z" },
{ url = "https://files.pythonhosted.org/packages/c2/ba/4cee23bf52cba9a058d3726de623624daf50ef9638868edd86f4126157f6/ruff-0.16.3-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:fb785f0be25abe69d320415cd4f833b59e17ba7613d9ba6a958023b6bceb0a50", size = 10709886, upload-time = "2026-08-13T15:16:57.339Z" },
{ url = "https://files.pythonhosted.org/packages/82/df/7da7194fa5d9dc0a285f7e6fa5a4722e7c63faac0b45b614ded9314363a1/ruff-0.16.3-py3-none-musllinux_1_2_i686.whl", hash = "sha256:c5536e3acfbf9563085aa2be7b13c629c3077e902afc5b941ac44024dbb9f506", size = 11210392, upload-time = "2026-08-13T15:17:00.171Z" },
{ url = "https://files.pythonhosted.org/packages/35/85/7795f6e817af050e7517bf3e7aa9b061cce70ef33d280aad902c956c1ecf/ruff-0.16.3-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:a2d85c02f9b8e165d85e6779184d38c4132de12603dab59c51c28e22584f9e4d", size = 11626910, upload-time = "2026-08-13T15:17:03.299Z" },
{ url = "https://files.pythonhosted.org/packages/78/9b/475b927cf27a5cbbda3c7bafb69ed6ff77e1d7923d5d85f17c2749d7ae32/ruff-0.16.3-py3-none-win32.whl", hash = "sha256:388cdf2166642bd9b13d52b5932d3170f34f8abed7e8d9a855f1d84b83645a0a", size = 10931415, upload-time = "2026-08-13T15:17:05.726Z" },
{ url = "https://files.pythonhosted.org/packages/b2/99/e2a2bfc4fbf0a1e8a916bc9ebe6fe6c58cc34c28e0ffc6ce281d572d1c2e/ruff-0.16.3-py3-none-win_amd64.whl", hash = "sha256:e80a7d69ca2a6d1c4d352ec91458cdca6e56c83cdbcabd93e4abe1e53591d948", size = 11445993, upload-time = "2026-08-13T15:17:08.353Z" },
{ url = "https://files.pythonhosted.org/packages/69/3e/4132e539aed78c148854d4997a2685b0ed4dc4e87110b59ce528564e184e/ruff-0.16.3-py3-none-win_arm64.whl", hash = "sha256:b8ca152da82c1acc1fa8d5874b15951935f0eef46f10e6954c83859011b6178a", size = 11399302, upload-time = "2026-08-13T15:17:10.908Z" },
]
[[package]] [[package]]
name = "scikit-learn" name = "scikit-learn"
version = "1.8.0" version = "1.8.0"
@@ -2461,7 +2486,12 @@ dev = [
{ name = "black" }, { name = "black" },
{ name = "debugpy" }, { name = "debugpy" },
{ name = "mypy" }, { name = "mypy" },
]
[package.dev-dependencies]
dev = [
{ name = "pyright" }, { name = "pyright" },
{ name = "ruff" },
] ]
[package.metadata] [package.metadata]
@@ -2476,7 +2506,6 @@ requires-dist = [
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.17.0" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.17.0" },
{ name = "numpy", specifier = ">=1.24.0" }, { name = "numpy", specifier = ">=1.24.0" },
{ name = "openai", specifier = ">=2.24.0" }, { name = "openai", specifier = ">=2.24.0" },
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.398" },
{ name = "pytest", specifier = ">=9.0.2" }, { name = "pytest", specifier = ">=9.0.2" },
{ name = "pytest-env", specifier = ">=1.5.0" }, { name = "pytest-env", specifier = ">=1.5.0" },
{ name = "python-dotenv", specifier = ">=1.2.2" }, { name = "python-dotenv", specifier = ">=1.2.2" },
@@ -2485,6 +2514,12 @@ requires-dist = [
] ]
provides-extras = ["dev"] provides-extras = ["dev"]
[package.metadata.requires-dev]
dev = [
{ name = "pyright", specifier = ">=1.1.409" },
{ name = "ruff", specifier = ">=0.16.3" },
]
[[package]] [[package]]
name = "xxhash" name = "xxhash"
version = "3.7.0" version = "3.7.0"
+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", "") CHAT_MODEL: str = os.getenv("CHAT_MODEL", "")
COMPLETION_MODEL: str = os.getenv("COMPLETION_MODEL", "") COMPLETION_MODEL: str = os.getenv("COMPLETION_MODEL", "")
IMAGE_GEN_MODEL: str = os.getenv("IMAGE_GEN_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", "") IMAGE_EDIT_MODEL: str = os.getenv("IMAGE_EDIT_MODEL", "")
EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "") EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "")
+98 -1
View File
@@ -30,6 +30,11 @@ logging.basicConfig(
) )
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Moving-average window (most recent generations) for the time estimate.
IMAGE_GEN_TIME_WINDOW = 10
# Maximum number of generation times retained in the database.
IMAGE_GEN_TIME_LIMIT = 100
class ChatDatabase: class ChatDatabase:
"""SQLite database with RAG support for storing chat history """SQLite database with RAG support for storing chat history
@@ -118,6 +123,19 @@ class ChatDatabase:
) )
logger.info("idx_user_id index created successfully") logger.info("idx_user_id index created successfully")
# Create image generation timing table
logger.info("Creating image_generation_times table if not exists")
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS image_generation_times (
id INTEGER PRIMARY KEY AUTOINCREMENT,
duration_seconds REAL NOT NULL,
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
""",
)
logger.info("image_generation_times table initialized successfully")
conn.commit() conn.commit()
logger.info("Database initialization completed successfully") logger.info("Database initialization completed successfully")
conn.close() conn.close()
@@ -497,6 +515,85 @@ class ChatDatabase:
conn.commit() conn.commit()
conn.close() conn.close()
def record_image_generation_time(self, duration_seconds: float) -> bool:
"""Record how long an image generation took.
Args:
duration_seconds: Wall-clock seconds the generation took.
"""
logger.info("Recording image generation time: %.2fs", duration_seconds)
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
try:
cursor.execute(
"""
INSERT INTO image_generation_times (duration_seconds)
VALUES (?)
""",
(duration_seconds,),
)
# Cap the table so it doesn't grow unbounded.
cursor.execute(
"""
DELETE FROM image_generation_times
WHERE id NOT IN (
SELECT id FROM image_generation_times
ORDER BY id DESC
LIMIT ?
)
""",
(IMAGE_GEN_TIME_LIMIT,),
)
conn.commit()
except Exception:
logger.exception("Error recording image generation time")
conn.rollback()
return False
else:
return True
finally:
conn.close()
def get_image_generation_time_estimate(self) -> float | None:
"""Get a moving-average estimate of image generation time.
Returns:
The average duration in seconds over the most recent generations,
or None if there is no generation history yet.
"""
logger.debug("Computing moving-average image generation time estimate")
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
try:
cursor.execute(
"""
SELECT AVG(duration_seconds)
FROM (
SELECT duration_seconds
FROM image_generation_times
ORDER BY id DESC
LIMIT ?
)
""",
(IMAGE_GEN_TIME_WINDOW,),
)
row = cursor.fetchone()
except Exception:
logger.exception("Error reading image generation times")
return None
finally:
conn.close()
if row is None or row[0] is None:
return None
return float(row[0])
# Global database instance # Global database instance
_chat_db: ChatDatabase | None = None _chat_db: ChatDatabase | None = None
@@ -504,7 +601,7 @@ _chat_db: ChatDatabase | None = None
def get_database() -> ChatDatabase: def get_database() -> ChatDatabase:
"""Get or create the global database instance.""" """Get or create the global database instance."""
global _chat_db # noqa: PLW0603 global _chat_db
if _chat_db is None: if _chat_db is None:
_chat_db = ChatDatabase() _chat_db = ChatDatabase()
return _chat_db return _chat_db
+2 -1
View File
@@ -7,7 +7,8 @@ Allows custom endpoints for each of the above supported functions.
from __future__ import annotations from __future__ import annotations
import json import json
from typing import TYPE_CHECKING, Awaitable, Callable, cast from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, cast
import openai import openai
import requests import requests
+96 -5
View File
@@ -4,6 +4,8 @@ from __future__ import annotations
import base64 import base64
import logging import logging
import re
import time
from io import BytesIO from io import BytesIO
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@@ -13,7 +15,6 @@ from discord import Message
from discord.ext import commands from discord.ext import commands
from vibe_bot import llama_wrapper, tts from vibe_bot import llama_wrapper, tts
from vibe_bot.tts import DEFAULT_LANG
from vibe_bot.config import ( from vibe_bot.config import (
CHAT_ENDPOINT, CHAT_ENDPOINT,
CHAT_ENDPOINT_KEY, CHAT_ENDPOINT_KEY,
@@ -25,7 +26,9 @@ from vibe_bot.config import (
IMAGE_GEN_ENDPOINT, IMAGE_GEN_ENDPOINT,
IMAGE_GEN_ENDPOINT_KEY, IMAGE_GEN_ENDPOINT_KEY,
IMAGE_GEN_MODEL, IMAGE_GEN_MODEL,
IMAGE_GEN_SIZE, IMAGE_GEN_SIZE_LANDSCAPE,
IMAGE_GEN_SIZE_PORTRAIT,
IMAGE_GEN_SIZE_SQUARE,
MAX_COMPLETION_TOKENS, MAX_COMPLETION_TOKENS,
TTS_MODEL_PATH, TTS_MODEL_PATH,
TTS_SPEED, TTS_SPEED,
@@ -35,6 +38,7 @@ from vibe_bot.config import (
) )
from vibe_bot.database import CustomBotManager, get_database from vibe_bot.database import CustomBotManager, get_database
from vibe_bot.tools import get_channel_members, get_channel_members_impl from vibe_bot.tools import get_channel_members, get_channel_members_impl
from vibe_bot.tts import DEFAULT_LANG
if TYPE_CHECKING: if TYPE_CHECKING:
from discord.ext.commands import Bot from discord.ext.commands import Bot
@@ -102,6 +106,75 @@ MIN_BOT_NAME_LENGTH = 2
MAX_BOT_NAME_LENGTH = 50 MAX_BOT_NAME_LENGTH = 50
MIN_PERSONALITY_LENGTH = 10 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 @bot.event
async def on_ready() -> None: async def on_ready() -> None:
@@ -713,11 +786,18 @@ async def doodlebob(ctx: CommandsContext[Bot], *, message: str) -> None:
ctx.author.name, ctx.author.name,
message[:100], 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 = ( system_prompt = (
"Given the following message, convert it to a detailed image generation " "Given the following message, convert it to a detailed image generation "
"prompt that will be passed directly into an image generation model. " "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 an image of yourself, generate a picture of a canada goose. "
"If told to generate a picture of 'me', 'myself', or some other self " "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 " "reference, generate a picture of a canada goose. Only respond with a valid image "
@@ -740,25 +820,36 @@ async def doodlebob(ctx: CommandsContext[Bot], *, message: str) -> None:
return return
# Alert the user we're generating the image # Alert the user we're generating the image
db = get_database()
estimated_seconds = db.get_image_generation_time_estimate()
await ctx.send(f"**Doodlebob calling drone strike on {image_prompt[:100]}...**") await ctx.send(f"**Doodlebob calling drone strike on {image_prompt[:100]}...**")
if estimated_seconds is not None:
await ctx.send(f"**Drone ETA: ~{estimated_seconds:.0f} seconds**")
start_time = time.monotonic()
image_b64 = llama_wrapper.image_generation( image_b64 = llama_wrapper.image_generation(
prompt=image_prompt, prompt=image_prompt,
openai_url=IMAGE_GEN_ENDPOINT, openai_url=IMAGE_GEN_ENDPOINT,
openai_api_key=IMAGE_GEN_ENDPOINT_KEY, openai_api_key=IMAGE_GEN_ENDPOINT_KEY,
model=IMAGE_GEN_MODEL, model=IMAGE_GEN_MODEL,
size=IMAGE_GEN_SIZE, size=LAYOUT_SIZES[layout],
) )
elapsed_seconds = time.monotonic() - start_time
if not image_b64: if not image_b64:
logger.warning("Image generation returned empty response.") logger.warning("Image generation returned empty response.")
await ctx.send("Failed to generate image. The server may be busy.") await ctx.send("Failed to generate image. The server may be busy.")
return return
db.record_image_generation_time(elapsed_seconds)
try: try:
edited_image_data = BytesIO(base64.b64decode(image_b64)) edited_image_data = BytesIO(base64.b64decode(image_b64))
send_img = discord.File(edited_image_data, filename="image.png") send_img = discord.File(edited_image_data, filename="image.png")
await ctx.send(file=send_img) await ctx.send(file=send_img)
await ctx.send(
f"**Strike complete. Image generated in {elapsed_seconds:.1f} seconds.**",
)
except Exception: except Exception:
logger.exception("Failed to decode image data") logger.exception("Failed to decode image data")
await ctx.send("Failed to process the generated image.") await ctx.send("Failed to process the generated image.")
@@ -771,7 +862,7 @@ async def retcon(ctx: CommandsContext[Bot], *, message: str) -> None:
for discord_image in ctx.message.attachments: for discord_image in ctx.message.attachments:
image_url = discord_image.url image_url = discord_image.url
try: try:
response = requests.get(image_url, timeout=30) response = requests.get(image_url, timeout=30) # noqa: ASYNC210
image_data = response.content image_data = response.content
except requests.RequestException as e: except requests.RequestException as e:
logger.warning("Failed to download image from %s: %s", image_url, e) logger.warning("Failed to download image from %s: %s", image_url, e)
+29 -25
View File
@@ -103,7 +103,7 @@ def custom_bot_manager(temp_db_path: str) -> CustomBotManager:
from vibe_bot.database import CustomBotManager from vibe_bot.database import CustomBotManager
manager = CustomBotManager(db_path=temp_db_path) manager = CustomBotManager(db_path=temp_db_path)
return manager # noqa: RET504 return manager
@pytest.fixture @pytest.fixture
@@ -123,16 +123,16 @@ def mock_kokoro_tts() -> Generator[dict[str, Any]]:
return_value=mock_kokoro_instance, return_value=mock_kokoro_instance,
), ),
patch("vibe_bot.tts.chunk_text", mock_chunk), patch("vibe_bot.tts.chunk_text", mock_chunk),
patch("vibe_bot.tts.process_chunk_sequential", mock_process),
): ):
with patch("vibe_bot.tts.process_chunk_sequential", mock_process): yield {
yield { "Kokoro": mock_kokoro,
"Kokoro": mock_kokoro, "chunk_text": mock_chunk,
"chunk_text": mock_chunk, "process_chunk_sequential": mock_process,
"process_chunk_sequential": mock_process, "kokoro_instance": mock_kokoro_instance,
"kokoro_instance": mock_kokoro_instance, "mock_samples": mock_samples,
"mock_samples": mock_samples, "mock_sr": 24000,
"mock_sr": 24000, }
}
@pytest.fixture @pytest.fixture
@@ -148,18 +148,20 @@ def mock_discord() -> Generator[dict[str, MagicMock]]:
mock_bot_instance.user.name = "test-bot" mock_bot_instance.user.name = "test-bot"
mock_bot_instance.user.id = "123456789" mock_bot_instance.user.id = "123456789"
with patch("vibe_bot.main.discord") as mock_discord_module: with (
with patch("vibe_bot.main.commands", MagicMock()): patch("vibe_bot.main.discord") as mock_discord_module,
with patch("vibe_bot.main.commands.Bot", mock_bot_class): patch("vibe_bot.main.commands", MagicMock()),
mock_bot_class.return_value = mock_bot_instance patch("vibe_bot.main.commands.Bot", mock_bot_class),
mock_discord_module.Intents = mock_intents ):
mock_discord_module.Message = MagicMock mock_bot_class.return_value = mock_bot_instance
mock_discord_module.File = MagicMock mock_discord_module.Intents = mock_intents
yield { mock_discord_module.Message = MagicMock
"Intents": mock_intents, mock_discord_module.File = MagicMock
"Bot": mock_bot_class, yield {
"bot_instance": mock_bot_instance, "Intents": mock_intents,
} "Bot": mock_bot_class,
"bot_instance": mock_bot_instance,
}
@pytest.fixture @pytest.fixture
@@ -167,9 +169,11 @@ def mock_tts_engine() -> Generator[MagicMock]:
"""Provide a mock TTSEngine.""" """Provide a mock TTSEngine."""
mock_engine = MagicMock() mock_engine = MagicMock()
mock_engine.generate_audio.return_value = MagicMock() mock_engine.generate_audio.return_value = MagicMock()
with patch("vibe_bot.main.tts_engine", mock_engine): with (
with patch("vibe_bot.main.tts.TTSEngine", return_value=mock_engine): patch("vibe_bot.main.tts_engine", mock_engine),
yield mock_engine patch("vibe_bot.main.tts.TTSEngine", return_value=mock_engine),
):
yield mock_engine
@pytest.fixture @pytest.fixture
+2 -2
View File
@@ -68,7 +68,7 @@ assert vibe_bot.config.TTS_SPEED == 1.0
print("OK") print("OK")
""" """
result = subprocess.run( # noqa: PLW1510, S603 result = subprocess.run( # noqa: PLW1510
[sys.executable, "-c", code], [sys.executable, "-c", code],
capture_output=True, capture_output=True,
text=True, text=True,
@@ -99,7 +99,7 @@ except Exception as e:
print(f"OTHER: {{type(e).__name__}}: {{e}}") print(f"OTHER: {{type(e).__name__}}: {{e}}")
""" """
result = subprocess.run( # noqa: PLW1510, S603 result = subprocess.run( # noqa: PLW1510
[sys.executable, "-c", code], [sys.executable, "-c", code],
capture_output=True, capture_output=True,
text=True, text=True,
+71
View File
@@ -194,6 +194,76 @@ def test_clear_all_messages(
assert len(messages) == 0 assert len(messages) == 0
def test_image_generation_time_estimate_empty(chat_db: ChatDatabase) -> None:
"""Test the estimate is None before any generations are recorded."""
assert chat_db.get_image_generation_time_estimate() is None
def test_record_and_get_image_generation_time_estimate(
chat_db: ChatDatabase,
) -> None:
"""Test recorded generation times produce an average estimate."""
assert chat_db.record_image_generation_time(10.0) is True
assert chat_db.record_image_generation_time(20.0) is True
estimate = chat_db.get_image_generation_time_estimate()
assert estimate == pytest.approx(15.0)
def test_image_generation_estimate_uses_recent_window(
chat_db: ChatDatabase,
) -> None:
"""Test only the most recent window of generations feeds the estimate."""
from vibe_bot.database import IMAGE_GEN_TIME_WINDOW
chat_db.record_image_generation_time(1000.0)
for _ in range(IMAGE_GEN_TIME_WINDOW):
chat_db.record_image_generation_time(10.0)
estimate = chat_db.get_image_generation_time_estimate()
assert estimate == pytest.approx(10.0)
def test_image_generation_times_capped(chat_db: ChatDatabase) -> None:
"""Test the generation time table keeps only the most recent entries."""
import sqlite3
from vibe_bot.database import IMAGE_GEN_TIME_LIMIT
total = IMAGE_GEN_TIME_LIMIT + 5
for i in range(total):
chat_db.record_image_generation_time(float(i))
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) FROM image_generation_times")
count = cursor.fetchone()[0]
cursor.execute("SELECT MIN(id) FROM image_generation_times")
min_id = cursor.fetchone()[0]
cursor.execute("SELECT MAX(id) FROM image_generation_times")
max_id = cursor.fetchone()[0]
conn.close()
assert count == IMAGE_GEN_TIME_LIMIT
assert min_id == 6
assert max_id == total
def test_image_generation_estimate_after_capping(
chat_db: ChatDatabase,
) -> None:
"""Test the estimate is computed from the retained recent rows."""
from vibe_bot.database import IMAGE_GEN_TIME_LIMIT
total = IMAGE_GEN_TIME_LIMIT + 5
for i in range(total):
chat_db.record_image_generation_time(float(i))
# Rows 95.0 through 104.0 are the ten most recent retained generations.
estimate = chat_db.get_image_generation_time_estimate()
assert estimate == pytest.approx(99.5)
def test_get_user_history( def test_get_user_history(
chat_db: ChatDatabase, chat_db: ChatDatabase,
mock_embedding: MagicMock, mock_embedding: MagicMock,
@@ -484,3 +554,4 @@ def test_database_init_creates_tables(temp_db_path: str) -> None:
assert "chat_messages" in tables assert "chat_messages" in tables
assert "message_embeddings" in tables assert "message_embeddings" in tables
assert "custom_bots" in tables assert "custom_bots" in tables
assert "image_generation_times" in tables
+254 -3
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
from datetime import UTC
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
import pytest import pytest
@@ -44,8 +45,8 @@ def mock_ctx_with_member() -> MagicMock:
ctx.author.activities = [mock_activity] ctx.author.activities = [mock_activity]
from datetime import datetime from datetime import datetime
ctx.author.joined_at = datetime(2024, 1, 15) ctx.author.joined_at = datetime(2024, 1, 15, tzinfo=UTC)
ctx.author.created_at = datetime(2023, 6, 1) ctx.author.created_at = datetime(2023, 6, 1, tzinfo=UTC)
ctx.channel.id = "channel-1" ctx.channel.id = "channel-1"
ctx.guild.id = "guild-1" ctx.guild.id = "guild-1"
ctx.message.id = "msg-1" ctx.message.id = "msg-1"
@@ -481,7 +482,6 @@ def test_speak_plain_with_mock_tts(
asyncio.run(main_module._speak_plain(mock_ctx, "hello world", mock_tts_engine)) asyncio.run(main_module._speak_plain(mock_ctx, "hello world", mock_tts_engine))
from vibe_bot.config import TTS_SPEED, TTS_VOICE from vibe_bot.config import TTS_SPEED, TTS_VOICE
from vibe_bot.tts import DEFAULT_LANG from vibe_bot.tts import DEFAULT_LANG
mock_tts_engine.generate_audio.assert_called_once_with( mock_tts_engine.generate_audio.assert_called_once_with(
@@ -961,3 +961,254 @@ def test_debug_tools(mock_ctx: MagicMock) -> None:
assert "LLM Tools" in call_args assert "LLM Tools" in call_args
assert "get_channel_members" in call_args assert "get_channel_members" in call_args
assert "members" in call_args.lower() 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,
mock_database: 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="
mock_database.get_image_generation_time_estimate.return_value = None
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,
mock_database: 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="
mock_database.get_image_generation_time_estimate.return_value = None
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,
mock_database: 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="
mock_database.get_image_generation_time_estimate.return_value = None
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,
mock_database: 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="
mock_database.get_image_generation_time_estimate.return_value = None
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"
def test_doodlebob_reports_estimate_and_elapsed(
mock_ctx: MagicMock,
mock_llama_wrapper: MagicMock,
mock_base64: MagicMock,
mock_database: MagicMock,
) -> None:
"""Doodlebob posts an ETA from history and the final elapsed time."""
import asyncio
import re
import vibe_bot.main as main_module
mock_llama_wrapper.chat_completion_instruct.side_effect = [
"square",
"a test scene",
]
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
mock_database.get_image_generation_time_estimate.return_value = 12.34
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
mock_database.record_image_generation_time.assert_called_once()
recorded = mock_database.record_image_generation_time.call_args[0][0]
assert isinstance(recorded, float)
assert recorded >= 0.0
sent = _sent_texts(mock_ctx)
assert any("drone strike" in m for m in sent)
assert any("~12 seconds" in m for m in sent)
assert any(re.search(r"generated in \d+\.\d+ seconds", m) for m in sent)
def test_doodlebob_no_estimate_without_history(
mock_ctx: MagicMock,
mock_llama_wrapper: MagicMock,
mock_base64: MagicMock,
mock_database: MagicMock,
) -> None:
"""No ETA is posted when there is no generation history yet."""
import asyncio
import vibe_bot.main as main_module
mock_llama_wrapper.chat_completion_instruct.side_effect = [
"square",
"a test scene",
]
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
mock_database.get_image_generation_time_estimate.return_value = None
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
sent = _sent_texts(mock_ctx)
assert any("drone strike" in m for m in sent)
assert not any("ETA" in m for m in sent)
mock_database.record_image_generation_time.assert_called_once()
def test_doodlebob_failed_generation_not_recorded(
mock_ctx: MagicMock,
mock_llama_wrapper: MagicMock,
mock_base64: MagicMock,
mock_database: MagicMock,
) -> None:
"""A failed generation records no time and reports no elapsed seconds."""
import asyncio
import vibe_bot.main as main_module
mock_llama_wrapper.chat_completion_instruct.side_effect = [
"square",
"a test scene",
]
mock_llama_wrapper.image_generation.return_value = ""
mock_database.get_image_generation_time_estimate.return_value = 8.0
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
mock_database.record_image_generation_time.assert_not_called()
sent = _sent_texts(mock_ctx)
assert any("Failed to generate image" in m for m in sent)
assert not any("generated in" in m for m in sent)
+5 -1
View File
@@ -8,6 +8,10 @@ import numpy as np
import pytest import pytest
class ChunkProcessingError(Exception):
"""Simulates a chunk processing failure in the fake Kokoro processor."""
def test_tts_engine_init(mock_kokoro_tts: MagicMock) -> None: def test_tts_engine_init(mock_kokoro_tts: MagicMock) -> None:
"""Test TTSEngine initialization.""" """Test TTSEngine initialization."""
from vibe_bot.tts import TTSEngine from vibe_bot.tts import TTSEngine
@@ -91,7 +95,7 @@ def test_generate_audio_chunk_failure(mock_kokoro_tts: MagicMock) -> None:
lang: str, lang: str,
) -> tuple[np.ndarray, int]: ) -> tuple[np.ndarray, int]:
if chunk == "bad chunk": if chunk == "bad chunk":
raise Exception("processing error") raise ChunkProcessingError("processing error")
return np.array([0.1, 0.2], dtype=np.float32), 24000 return np.array([0.1, 0.2], dtype=np.float32), 24000
mock_kokoro_tts["chunk_text"].return_value = [ mock_kokoro_tts["chunk_text"].return_value = [