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

complete restructure

Browse Source
This commit is contained in:
ducoterra
2026-08-19 13:16:43 -04:00
parent d7b6f28cbd
commit f87e1d51ef
60 changed files with 8176 additions and 5143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitea/workflows/build-push.yml
+38
View File
@@ -4,13 +4,51 @@ on:
push: push:
branches: branches:
- main - main
pull_request:
release: release:
types: types:
- published - published
jobs: jobs:
test:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Setup Python
uses: actions/setup-python@v5
with:
python-version: "3.13"
- name: Setup uv
uses: astral-sh/setup-uv@v5
- name: Install PortAudio
run: sudo apt-get update && sudo apt-get install -y --no-install-recommends portaudio19-dev
- name: Install dependencies
run: uv sync
- name: Ruff
run: uv run ruff check vibe_bot/
- name: Mypy
run: uv run mypy vibe_bot/
- name: Pyright
run: uv run pyright vibe_bot/
- name: Black
run: uv run black --check vibe_bot/
- name: Test
run: uv run pytest vibe_bot/tests/ -m "not live" -v
build-and-push: build-and-push:
runs-on: ubuntu-latest runs-on: ubuntu-latest
needs: [test]
if: github.event_name != 'pull_request'
steps: steps:
- name: Checkout - name: Checkout
uses: actions/checkout@v4 uses: actions/checkout@v4
.gitignore
+2
View File
@@ -14,3 +14,5 @@ wheels/
.mypy_cache/ .mypy_cache/
.pytest_cache/ .pytest_cache/
.agent/
AGENTS.md
+18 -12
View File
@@ -1,34 +1,40 @@
# AGENTS.md # 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/`. 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/` (plus the standalone benchmark `scripts/bench_rag.py`, which is never imported by the package).
## Commands ## Commands
- Setup: `uv sync --extra dev` — plain `uv sync` already installs the `dev` dependency group (ruff, pyright); the `dev` extra adds black/debugpy/mypy - Setup: `uv sync` — installs the project plus the `dev` dependency group (ruff, mypy, pyright, black, pytest); there are no extras
- 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 - 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` - Tests: `uv run pytest vibe_bot/tests/ -v` — hermetic by default (`addopts = "-m 'not live'"`); `uv run pytest -m live` runs the network-dependent live tests; single test: `uv run pytest vibe_bot/tests/test_app.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/` - 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 - RAG benchmark: `uv run python scripts/bench_rag.py` — hermetic (embeddings monkeypatched, temp DBs); prints p95 latency of `get_conversation_context` at 1k and 5k rows
- Container: `./build.sh` (podman). CI (`.gitea/workflows/build-push.yml`, Gitea) runs a `test` job (all four checks + hermetic pytest) on pushes/PRs/releases, and `build-and-push` (image build/push) only on main/release, gated on `test`
## Setup requirements ## 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. - 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`). - PortAudio is a required system library (kokoro-tts → sounddevice). Without it, the test suite errors at import time (every module loads `conftest.py`, which imports `vibe_bot.app` → `vibe_bot.tts` → `sounddevice`). 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. - 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 ## 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. - The network-dependent tests are marked `live` (in `test_llm_client.py`): unmocked calls to the real `CHAT_ENDPOINT` from `.env`. They are deselected by default and only run via `uv run pytest -m live`, which needs 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`.
- 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 ## 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`. - `main.py` — entrypoint: `validate_config()`, `configure_logging()`, then `build_bot(app).run(token)`.
- `app.py` — composition root: `configure_logging()` (the only `basicConfig`), `@dataclass App` (flat fields: db/manager/registry/tts/chat/image/speech/conversation/bot_cache), `create_app()`, `build_bot(app)` (intents; event handlers; `commands.register_all(bot, app)`). Custom-bot "commands" (`!<bot_name> ...`) are matched in `on_message` against `app.bot_cache`, not registered with `bot.command`.
- `commands/` — command groups registered by `register_all`: `custom_bots` (custom-bot, list-custom-bots, delete-custom-bot), `speech` (speak, voices), `images` (doodlebob, retcon), `conversation` (talkforme), `admin` (lobotomize, debug, history), `chat` (docstring only — custom-bot chat flows through `on_message`); `_state.py` holds the shared App.
- `services/` — the LLM-backed logic: `ChatService`, `ImageService`, `SpeechService`, `ConversationService`; the Discord command handlers are thin wrappers.
- `db/` — `connection` (shared connect), `schema` (init + migrations log at INFO; backfills the `norm` column on existing DBs), `messages` (`ChatDatabase`: embeddings as float32 blobs with a stored L2 `norm`, cleanup, RAG context), `bots` (`CustomBotManager`), `search` (single-JOIN RAG retrieval scored by one matrix multiply over the stored norms), `timing` (image-gen ETA stats), `vectors` (vector math). `database.py` is a thin facade re-exporting `ChatDatabase`/`CustomBotManager`.
- `llm/` — `chat.py` (completion clients: instruct / with-history / with-tools), `images.py` (gen + edit clients), `registry.py` (tool registry: schema + dispatch). `llm_client.py` is the public facade over `llm/` plus the embedding HTTP plumbing. (`llama_wrapper.py` is gone.)
- `config.py` — env loading + import-time validation, voice catalog. - `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`. - `prompts.py` — LLM prompt constants (image layout/prompt/verify, length hint, `build_system_prompt`).
- `llama_wrapper.py` — thin OpenAI-compatible clients for chat / image gen / image edit / embeddings, each with its own endpoint, key, and model. - `textutil.py` — `split_message`, `get_user_info`.
- `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`. - `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 `app.py`.
- `tts.py` — `TTSEngine` (Kokoro wrapper; `AudioResult.partial` signals failed chunks); voice/speed defaults are aliases of `config.TTS_VOICE`/`TTS_SPEED` (single source: config).
## Style ## Style
Containerfile
+4
View File
@@ -15,4 +15,8 @@ ENV UV_NO_DEV=1
WORKDIR /app WORKDIR /app
RUN uv sync --locked RUN uv sync --locked
# Run as a non-root user: /app holds the relative default DB (chat_history.db), /db is the mounted volume
RUN useradd -m bot && chown -R bot:bot /app && mkdir -p /db && chown bot:bot /db
USER bot
CMD uv run python -m vibe_bot.main CMD uv run python -m vibe_bot.main
MANUAL_TESTS.md
+299
View File
@@ -0,0 +1,299 @@
# Manual Testing Guide (Discord)
How to manually verify each bot command from a Discord server. The automated suite
(`uv run pytest vibe_bot/tests/ -v`) covers the internals; this guide is for
end-to-end behavior against a real Discord server and the real LLM/image/embedding
APIs. Every command below makes real API calls, so expect API usage and latency.
## Before You Start
1. A working `.env` at the repo root: `DISCORD_TOKEN`, the chat/image/edit/embedding
endpoints with their keys and model names (see the Configuration section of
`README.md`).
2. The bot running and online in your test server, either:
```bash
uv run python -m vibe_bot.main
```
or the container image (`./build.sh`, then run with `--env-file .env`).
Watch the console output while testing — the bot logs each command trigger.
3. In the Discord Developer Portal, the bot must have **Message Content Intent**,
**Server Members Intent**, and **Presence Intent** enabled (the bot requests all
three at startup and will refuse to log in without them).
4. For `!speak`: `kokoro-v1.0.onnx` and `voices-v1.0.bin` in the project root
(baked into the container image). Without them the bot still runs; only
`!speak` degrades (it replies "TTS engine not initialized..."). `!voices` is a
static list and works either way.
5. Use a dedicated test channel. Start every clean pass by wiping stored memory:
```
!lobotomize
```
Expected: "All conversation history and memory has been cleared."
(This deletes **all** chat messages and embeddings for **all** users and bots.
It does **not** delete the custom bots themselves or the image ETA stats.)
6. Create the two standard test bots used throughout this guide. Use lowercase
names: incoming chat text is lowercased before being matched against the stored
bot name, so a lowercase name matches whatever capitalization you type.
```
!custom-bot alfred you are a proper british butler who speaks in short, polite sentences
!custom-bot jarvis you are a dry, sarcastic AI assistant
```
## Command Reference
### `!custom-bot <name> <personality>` — create a custom bot
**Happy path**
- Send: `!custom-bot alfred you are a proper british butler`
- Expect: `Custom bot **'alfred'** has been created with personality: *you are a proper british butler*`
followed by `You can now use this bot with: `!alfred <your message>``
- `!list-custom-bots` now includes alfred, and `!alfred <msg>` works immediately
(the bot-name cache is invalidated on create).
**Replacing an existing name** — bot names are a global namespace, not per-user.
- Send the same command again with a different personality.
- Expect: `Custom bot **'alfred'** already existed and has been **replaced** with personality: *...*`
**Validation** (names 2–50 chars, personality 10–1000 chars):
| Send | Expected reply |
| ---- | -------------- |
| `!custom-bot a you are a test bot` (1-char name) | `Invalid bot name. Name must be between 2 and 50 characters.` |
| a 51+-character name | same |
| `!custom-bot bob short` (personality < 10 chars) | `Invalid personality. Description must be at least 10 characters.` |
| a 1001+-character personality | `Personality too long. Max 1000 characters.` |
### `!list-custom-bots` — list custom bots
- Empty state (fresh database, or after deleting every bot):
`No custom bots have been created yet. Use `!custom-bot <name> <personality>` to create one.`
- After creating bots: `Available Custom Bots:` followed by one `* <name>` line per bot.
### `!delete-custom-bot <name>` — delete a custom bot (owner only)
- As the creator: `!delete-custom-bot jarvis`
Expect: `Custom bot 'jarvis' has been deleted.`
Afterwards `!jarvis hi` produces no response and `!list-custom-bots` omits it
(the bot-name cache is invalidated on delete).
- Unknown name: `!delete-custom-bot nosuchbot`
Expect: `Custom bot 'nosuchbot' not found.`
- Ownership: have a second Discord account create `tempbot`, then run
`!delete-custom-bot tempbot` from a third account.
Expect: `You can only delete your own custom bots.` (the bot is **not** deleted).
The creator's delete should then succeed.
### Chatting with a custom bot — `!<bot_name> <message>`
This is not a registered command; messages starting with `!<known bot name> ` are
matched in `on_message` and dispatched to the chat service.
**Happy path**
- Send: `!alfred Could you fetch me some tea?`
- Expect, in order:
1. `alfred is searching its databanks for Could you fetch me some tea?...`
2. `alfred response`
3. The reply, in character, under 2–3 sentences.
**Long replies** — a reply longer than 1000 characters arrives as multiple
consecutive messages (chunked at 1000 chars).
**Missing the trailing text** — `!alfred` alone (no message after the name)
matches nothing: the bot stays silent and no command error is raised.
**Memory / RAG**
1. `!alfred My favorite color is teal. Remember that.`
2. Later — different topic, ideally a different channel or day:
`!alfred What is my favorite color?`
3. Expect: it answers teal. RAG prepends your 10 most recent messages plus up to 5
semantically similar stored messages (similarity threshold 0.7) to the prompt.
4. Cross-bot check (shared memory is by design — history is not siloed per bot):
`!jarvis What is this user's favorite color?` should also know.
**Tool call**
- Send: `!alfred Who is in this channel?`
- The bot may call the `get_channel_members` tool; expect
`alfred is looking at the channel members...` before a reply that names server
members. (Not guaranteed every time — the LLM decides whether to call it.)
**Failure path** — stop the LLM API (or point `CHAT_ENDPOINT` at a dead URL and
restart the bot), then send `!alfred hi`.
Expect: `An error occurred while processing your request.` and a traceback in the
bot console.
### `!speak <text> [--voice <name>]` — plain text-to-speech
- `!speak hello world`
Expect: `Generating speech...` then a `speech.mp3` attachment. Play it.
- Without `--voice`, the configured `TTS_VOICE` (default `af_sarah`) is used.
- `!speak hello world --voice af_bella`
Expect: same flow, different voice. The `--voice` flag is only recognized at the
**very end** of the message; `!speak the --voice flag is fun` speaks the entire
string verbatim.
- Unknown voice: `!speak hi --voice af_nobody`
Expect: `Unknown voice 'af_nobody'. Use `!voices` to see available voices.`
- No text: `!speak --voice af_bella`
Expect: `Please provide text to speak.`
- Over 5000 characters of text:
Expect: `Text too long to speak. Max 5000 characters.`
- Cooldown: 3 uses per 30 seconds per user. A 4th immediate call:
`You're using that too quickly, try again in Ns.`
- Missing TTS model files (run without the onnx/bin files):
`TTS engine not initialized. Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.`
### `!speak <bot_name> <text> [--voice <name>]` — custom bot responds and speaks
Requires an existing custom bot; the first word is compared **exactly**
(case-sensitive) against stored bot names.
- `!speak alfred what time is it`
Expect, in order:
1. `**alfred** is thinking...`
2. `**alfred**: <in-character reply>` (the LLM reply, not your text)
3. `Generating speech for **alfred**...`
4. `speech.mp3` containing the **bot's reply** spoken aloud.
The exchange is also stored in history (visible later via `!history alfred`).
- `!speak alfred what time is it --voice am_puck` — same flow, am_puck voice.
- `!speak alfred` (no text after the name):
`Please provide text for the bot to respond to.`
- `!speak nonexistent hello`: `Custom bot 'nonexistent' not found.`
### `!voices` — list TTS voices
- No arguments. Expect the static catalog grouped by category:
en-US female (`af_alloy` … `af_sky`), en-US male (`am_adam` … `am_puck`),
en-GB (`bf_*`, `bm_*`), fr-FR (`ff_siwis`), Italian (`if_sara`, `im_nicola`),
Japanese (`jf_*`, `jm_kumo`), Mandarin (`zf_*`, `zm_*`), ending with
`Use `!speak <text> --voice <voice_name>` to choose a voice.`
- Spot-check one voice per language with `!speak <text in that language> --voice <name>`
and listen for the expected accent.
### `!doodlebob <prompt>` — generate an image
- `!doodlebob a cat sitting on the moon`
Expect, in order:
1. `**Doodlebob shopping for a canvas...**`
2. `**Doodlebob selected <layout>**` — the LLM picks `square`, `portrait`, or
`landscape` from the prompt (a tall subject should yield portrait, a wide
scene landscape). The canvas size follows the layout.
3. `**Doodlebob calling drone strike on <first 100 chars of the final prompt>...**`
4. From the **second** generation onward: `**Drone ETA: ~N seconds**` (moving
average of recorded generation times).
5. An `image.png` attachment that matches the prompt.
6. `**Strike complete. Image generated in N.N seconds.**`
- Prompt over 2000 characters: `Prompt too long. Max 2000 characters.`
- Cooldown: 1 per 60 seconds per user.
- Image API failure: `Failed to generate image. The server may be busy.`
### `!retcon <prompt>` — edit an attached image
- Attach an image **on the same message** as the command (drag it into the
channel): `!retcon make it look like it is on fire [image]`
Expect: `**Rewriting history to match make it look like it is on fire...**`
then an edited `image.png` attachment.
- No attachment: `!retcon make it sunny`
Expect: `Please attach an image to edit.`
(Only real attachments count; pasted external URLs are ignored, and only
Discord CDN hosts are downloaded, capped at 8 MB.)
- Prompt over 2000 characters: `Prompt too long. Max 2000 characters.`
- Image edit API failure: `Failed to edit the image.`
### `!talkforme <bot1> <bot2> <n> <topic>` — bot-vs-bot conversation
Requires two existing custom bots.
- `!talkforme alfred jarvis 4 the meaning of life`
Expect: `alfred is going to talk to jarvis about "the meaning of life" for 4 replies.`
then alternating messages headed `## alfred` / `## jarvis`, each reply in
character. Total chat messages is **n + 1** (bot1's opener plus n loop
replies); n is capped at 20, so at most 21 messages.
- Too few parts: `!talkforme alfred jarvis`
Expect: `Usage: !talkforme bot1 bot2 <number> <topic>`
- Unknown bot: `!talkforme alfred nosuchbot 4 tea`
Expect: `nosuchbot is not a real bot...`
- Non-integer count: `!talkforme alfred jarvis four tea`
Expect: `Message limit must be an integer.`
- Topic over 500 characters: `Topic too long. Max 500 characters.`
- Cooldown: 1 per 30 seconds per user.
### `!history <bot_name>` — view a bot's chat history
- After chatting with alfred: `!history alfred`
Expect: `Chat History for **alfred**:` followed by the latest up-to-20
exchanges (oldest of the 20 first), each pair as:
```text
<your message>
---
alfred: <its reply>
```
Turns from `!speak alfred ...` are included too (they are stored under the
bot name).
- Unknown bot: `!history nosuchbot`
Expect: `Custom bot 'nosuchbot' not found.`
- A bot with no stored messages: `No chat history found for **alfred**.`
- After `!lobotomize`: same "No chat history found" reply.
### `!debug [subcommand]` — debug menu
- `!debug` — prints the menu listing `members`, `whoami`, `tools`.
- `!debug members` — `Members in this channel (N total):` followed by an
alphabetized roster of **server** members (Discord has no per-channel
membership; the channel only identifies the guild), each entry showing display
name, nickname, global name, and `[status]`. In a DM (no guild):
`No members found in this channel.`
- `!debug whoami` — a block about your own account: Global Name (if set),
Nickname (if set), Top Role (if not @everyone), Activities (if any, custom
status excluded), Joined date, Username, User ID, Account Created date. Verify
each value against your account.
- `!debug tools` — the `get_channel_members` tool's name, description, and
parameter JSON schema.
- `!debug banana` — `Unknown debug sub-command: `banana`` plus a pointer to `!debug`.
### `!lobotomize` — wipe all stored memory
- Send `!lobotomize`.
Expect: `All conversation history and memory has been cleared.`
- Verify:
- `!history alfred` → `No chat history found for **alfred**.`
- The RAG recall test above no longer remembers the fact.
- `!list-custom-bots` is **unchanged** (bots survive; only
`chat_messages` and `message_embeddings` are cleared).
## Suggested End-to-End Pass
A full manual pass in order:
1. `!lobotomize` — clean slate
2. `!custom-bot alfred ...` and `!custom-bot jarvis ...`
3. `!list-custom-bots` — both listed
4. `!alfred ...` — a few turns; tell it a fact to remember later
5. `!jarvis What did I tell you about ...?` — cross-bot memory
6. `!speak hello world --voice af_bella`, then `!speak alfred tell me a joke`
7. `!voices`
8. `!doodlebob ...` twice (second run should show the Drone ETA)
9. `!retcon ... [image]`
10. `!talkforme alfred jarvis 4 ...`
11. `!history alfred` — contains the steps 4–6 exchanges
12. `!debug`, `!debug members`, `!debug whoami`, `!debug tools`
13. `!delete-custom-bot jarvis` — then confirm `!jarvis hi` is silent
14. `!lobotomize` — confirm `!history alfred` is empty
## Covered Only by the Automated Suite
Not practical to verify by hand (see `uv run pytest vibe_bot/tests/ -v`):
- embedding similarity thresholds and vector math
- database migrations and the `MAX_HISTORY_MESSAGES` cleanup (needs 1000+ messages)
- message chunking edge cases
- RAG benchmark behavior (`uv run python scripts/bench_rag.py`)
README.md
+168 -64
View File
@@ -9,7 +9,9 @@ A Discord bot that stores long-term chat history using SQLite with RAG (Retrieva
- [Text-to-Speech](#text-to-speech) - [Text-to-Speech](#text-to-speech)
- [Image Commands](#image-commands) - [Image Commands](#image-commands)
- [Bot Conversations](#bot-conversations) - [Bot Conversations](#bot-conversations)
- [Chat History](#chat-history) - [Admin & Debug](#admin--debug)
- [Design Notes](#design-notes)
- [Limits & Cooldowns](#limits--cooldowns)
- [Features](#features) - [Features](#features)
- [Setup](#setup) - [Setup](#setup)
- [Prerequisites](#prerequisites) - [Prerequisites](#prerequisites)
@@ -47,24 +49,24 @@ Once you create a custom bot, interact with it by prefixing your message with th
**Example:** **Example:**
1. Create a bot: `!custom-bot alfred you are a proper british butler` 1. Create a bot: `!custom-bot alfred you are a proper british butler`
2. Use the bot: `alfred Could you fetch me some tea?` 2. Use the bot: `!alfred Could you fetch me some tea?`
3. The bot will respond in character as a British butler 3. The bot will respond in character as a British butler
### Text-to-Speech ### Text-to-Speech
| Command | Description | Example Usage | | Command | Description | Example Usage |
| ------------------------------------ | ----------------------------------------------------- | ------------------------------------------ | | ------------------------------------ | ----------------------------------------------------- | ----------------------------------------- |
| `!speak <text>` | Convert text to speech (MP3 attachment) | `!speak hello world` | | `!speak <text>` | Convert text to speech (MP3 attachment) | `!speak hello world` |
| `!speak <text> --voice <voice>` | Convert text to speech with a specific voice | `!speak hello world --voice af_bella` | | `!speak <text> --voice <voice>` | Convert text to speech with a specific voice | `!speak hello world --voice af_bella` |
| `!speak <bot_name> <text>` | Have a custom bot respond and speak | `!speak alfred what time is it` | | `!speak <bot_name> <text>` | Have a custom bot respond and speak | `!speak alfred what time is it` |
| `!speak <bot_name> <text> --voice` | Have a custom bot respond and speak with a voice | `!speak alfred what time is it --voice am_puck` | | `!speak <bot_name> <text> --voice` | Have a custom bot respond and speak with a voice | `!speak alfred what time is it --voice am_puck` |
| `!voices` | List all available TTS voices by category | `!voices` | | `!voices` | List all available TTS voices by category | `!voices` |
### Image Commands ### Image Commands
| Command | Description | Example Usage | | Command | Description | Example Usage |
| ------------ | ------------------------------------ | ------------------------------------------ | | ------------ | ------------------------------------ | ----------------------------------------- |
| `!doodlebob` | Generate an image from a text prompt | `!doodlebob a cat sitting on the moon` | | `!doodlebob` | Generate an image from a text prompt | `!doodlebob a cat sitting on the moon` |
| `!retcon` | Edit an attached image with text | `!retcon <image attachment> Make it sunny` | | `!retcon` | Edit an attached image with text | `!retcon <image attachment> Make it sunny` |
### Bot Conversations ### Bot Conversations
@@ -73,11 +75,53 @@ Once you create a custom bot, interact with it by prefixing your message with th
| -------------------------------------- | ------------------------------------------- | ------------------------------------------------ | | -------------------------------------- | ------------------------------------------- | ------------------------------------------------ |
| `!talkforme <bot1> <bot2> <n> <topic>` | Have two bots discuss a topic for n replies | `!talkforme alfred jarvis 4 the meaning of life` | | `!talkforme <bot1> <bot2> <n> <topic>` | Have two bots discuss a topic for n replies | `!talkforme alfred jarvis 4 the meaning of life` |
### Chat History ### Admin & Debug
| Command | Description | Example Usage | | Command | Description | Example Usage |
| --------------------- | ------------------------------------- | ----------------- | | --------------------- | ------------------------------------------------------- | --------------------- |
| `!history <bot_name>` | View the chat history of a custom bot | `!history alfred` | | `!lobotomize` | Clear all conversation history and memory for all bots | `!lobotomize` |
| `!debug` | Show the debug menu | `!debug` |
| `!debug members` | List all members in the current channel | `!debug members` |
| `!debug whoami` | Show all information the bot has about you | `!debug whoami` |
| `!debug tools` | Show the LLM's available tools | `!debug tools` |
| `!history <bot_name>` | View the chat history of a custom bot | `!history alfred` |
## Design Notes
These semantics are intentional:
- **Shared memory**: message history is not siloed per bot. Any user's history can inform
any bot's RAG context — the semantic (embedding) search runs over the whole message
store, so a fact one user told bot A can resurface in a later conversation with bot B.
- **Shared bot namespace**: bot names are global, not per-user. Creating a bot with a
name that already exists **replaces** it (you will be told it was replaced); only
deletion is restricted to the bot's creator.
- **Any member may use any bot**: custom-bot chat (`!<bot_name> <message>`) is open to
every member of the server, regardless of who created the bot.
## Limits & Cooldowns
**Input limits** (over-limit input is rejected with a friendly message before any
LLM/DB work):
| Input | Limit |
| ------------------------- | ------------------------------------------------ |
| Custom bot name | 2–50 characters |
| Custom bot personality | 10–1000 characters |
| `!speak` text | 5000 characters (after `--voice` parsing) |
| `!doodlebob` / `!retcon` prompt | 2000 characters |
| `!talkforme` topic | 500 characters |
| `!talkforme` reply count | Capped at 20 replies per invocation |
| LLM responses | `MAX_COMPLETION_TOKENS` tokens (default 1000) |
| `!retcon` downloads | Discord CDN hosts only, 8 MB per image |
**Cooldowns** (per user; a friendly "try again in Ns" message is sent when hit):
| Command | Cooldown |
| ------------- | ----------- |
| `!doodlebob` | 1 / 60 s |
| `!talkforme` | 1 / 120 s |
| `!speak` | 3 / 30 s |
## Features ## Features
@@ -110,11 +154,8 @@ DISCORD_TOKEN=your_discord_bot_token
# Chat/Completion API (required) # Chat/Completion API (required)
CHAT_ENDPOINT=https://your-api.com/v1 CHAT_ENDPOINT=https://your-api.com/v1
COMPLETION_ENDPOINT=https://your-api.com/v1
CHAT_ENDPOINT_KEY=your_api_key CHAT_ENDPOINT_KEY=your_api_key
COMPLETION_ENDPOINT_KEY=your_api_key
CHAT_MODEL=your_model_name CHAT_MODEL=your_model_name
COMPLETION_MODEL=your_model_name
# Image Generation (required) # Image Generation (required)
IMAGE_GEN_ENDPOINT=https://your-api.com/v1 IMAGE_GEN_ENDPOINT=https://your-api.com/v1
@@ -169,16 +210,16 @@ uv run python -m vibe_bot.main
### Database Structure ### Database Structure
The system uses SQLite with three tables: The system uses SQLite with four tables:
1. **chat_messages**: Stores message metadata 1. **chat_messages**: Stores message metadata
- `message_id`, `user_id`, `username`, `content`, `bot_name`, `timestamp`, `channel_id`, `guild_id` - `message_id`, `user_id`, `username`, `content`, `bot_name`, `role`, `timestamp`, `channel_id`, `guild_id`
2. **message_embeddings**: Stores vector embeddings for RAG 2. **message_embeddings**: Stores vector embeddings for RAG
- `message_id` (PK), `embedding` (binary blob of float32 values) - `message_id` (PK), `embedding` (binary blob of float32 values)
3. **custom_bots**: Stores custom bot configurations 3. **custom_bots**: Stores custom bot configurations
- `bot_name` (PK), `system_prompt`, `created_by`, `created_at`, `is_active` - `bot_name` (PK), `system_prompt`, `created_by`, `created_at`, `is_active`
4. **image_generation_times**: Stores per-generation durations for the ETA estimate
- `id` (PK), `duration_seconds`, `generated_at`
### RAG Process ### RAG Process
@@ -195,26 +236,76 @@ The system uses SQLite with three tables:
```text ```text
vibe_discord_bots/ vibe_discord_bots/
├── vibe_bot/ ├── vibe_bot/
│ ├── __init__.py # Package marker │ ├── __init__.py # Package marker
│ ├── main.py # Main bot application (commands, event handlers) │ ├── app.py # Composition root: App, services, bot event handlers
│ ├── config.py # Environment variable loading and validation │ ├── commands/
│ ├── database.py # SQLite database with RAG + CustomBotManager │ │ ├── __init__.py # register_all: wires every command group
│ ├── llama_wrapper.py # OpenAI-compatible API wrappers (chat, images, embeddings) │ │ ├── _state.py # Shared App holder for command handlers
│ ├── tts.py # Kokoro TTS engine │ │ ├── admin.py # !lobotomize, !debug, !history
│ │ ├── chat.py # Custom-bot chat group (dispatched via on_message)
│ │ ├── conversation.py # !talkforme
│ │ ├── custom_bots.py # !custom-bot, !list-custom-bots, !delete-custom-bot
│ │ ├── images.py # !doodlebob, !retcon
│ │ └── speech.py # !speak, !voices
│ ├── config.py # Environment variable loading and validation
│ ├── database.py # Facade over the db/ package
│ ├── db/
│ │ ├── __init__.py
│ │ ├── bots.py # CustomBotManager (create/read/list/delete)
│ │ ├── connection.py # Shared SQLite connect() helper
│ │ ├── messages.py # ChatDatabase: persistence, cleanup, RAG context
│ │ ├── schema.py # Table creation and migrations
│ │ ├── search.py # Similarity search and history lookups
│ │ ├── timing.py # Image-generation timing statistics
│ │ └── vectors.py # Embedding vector helpers
│ ├── llm/
│ │ ├── __init__.py
│ │ ├── chat.py # Chat completion clients (instruct/tools/history)
│ │ ├── images.py # Image gen/edit clients
│ │ └── registry.py # Tool registry (schema + dispatch)
│ ├── llm_client.py # Public LLM/image/embedding client facade
│ ├── main.py # Entrypoint
│ ├── prompts.py # LLM prompt constants
│ ├── services/
│ │ ├── __init__.py
│ │ ├── chat_service.py # Custom-bot chat turn (context -> LLM -> persist)
│ │ ├── conversation_service.py # Bot-vs-bot conversation runner
│ │ ├── image_service.py # !doodlebob / !retcon logic
│ │ └── speech_service.py # !speak logic
│ ├── textutil.py # Message splitting and user-info helpers
│ ├── tools.py # LLM tool stubs and channel-member helper
│ ├── tts.py # Kokoro TTS engine
│ └── tests/ │ └── tests/
│ ├── conftest.py # Shared test fixtures │ ├── __init__.py
│ ├── test_main.py # Bot command tests │ ├── _helpers.py # Shared test helpers
│ ├── test_config.py # Config loading tests │ ├── conftest.py # Shared test fixtures
│ ├── test_database.py # Database + CustomBotManager tests │ ├── test_app.py # Composition root and bot handler tests
│ ├── test_llama_wrapper.py # API wrapper tests │ ├── test_commands.py # Discord command tests
│ └── test_tts.py # TTS engine tests │ ├── test_config.py # Config loading tests
├── pyproject.toml # Project dependencies (uv) │ ├── test_database.py # Database + CustomBotManager tests
├── uv.lock # Locked dependency versions │ ├── test_docs.py # README sync tests (commands + file tree)
├── .env # Environment variables │ ├── test_llm_client.py # LLM client tests
├── kokoro-v1.0.onnx # Kokoro TTS model │ ├── test_logging.py # No-content logging tests
├── voices-v1.0.bin # Kokoro voice definitions │ ├── test_prompts.py # Prompt constant tests
├── Containerfile # Podman/Docker build file │ ├── test_services.py # Service-layer tests
└── README.md # This file │ ├── test_textutil.py # Text utility tests
│ ├── test_tools.py # Tool stub tests
│ └── test_tts.py # TTS engine tests
├── .gitea/
│ └── workflows/
│ └── build-push.yml # CI: build and push the container image
├── AGENTS.md # Agent instructions
├── Containerfile # Podman/Docker build file
├── MANUAL_TESTS.md # Manual Discord testing guide
├── README.md # This file
├── build.sh # Podman build wrapper
├── kokoro-v1.0.onnx # Kokoro TTS model
├── pyproject.toml # Project metadata and dependencies (uv)
├── scripts/
│ └── bench_rag.py # RAG p95 latency benchmark (standalone)
├── uv.lock # Locked dependency versions
├── vibe-bot.container # systemd unit for running the bot container
└── voices-v1.0.bin # Kokoro voice definitions
``` ```
## Building ## Building
@@ -233,7 +324,7 @@ uv run python -m vibe_bot.main
```bash ```bash
# Build the container image # Build the container image
podman build -t vibe-bot:latest . podman build -t localhost/vibe-bot:latest -f Containerfile .
# Run with environment file # Run with environment file
podman run --env-file .env localhost/vibe-bot:latest podman run --env-file .env localhost/vibe-bot:latest
@@ -241,12 +332,19 @@ podman run --env-file .env localhost/vibe-bot:latest
## Testing ## Testing
Run the full test suite: Run the full test suite (hermetic by default; network-dependent `live` tests are
deselected):
```bash ```bash
uv run pytest vibe_bot/tests/ -v uv run pytest vibe_bot/tests/ -v
``` ```
Run the network-dependent live tests explicitly:
```bash
uv run pytest -m live
```
Run linters: Run linters:
```bash ```bash
@@ -265,24 +363,30 @@ uv run black --check vibe_bot/
## Configuration ## Configuration
| Variable | Default | Description | | Variable | Default | Description |
| ----------------------- | ------------------ | ------------------------------------- | | -------------------------- | ------------------ | --------------------------------------------- |
| `DISCORD_TOKEN` | *(required)* | Discord bot authentication token | | `DISCORD_TOKEN` | *(required)* | Discord bot authentication token |
| `CHAT_ENDPOINT` | *(required)* | OpenAI-compatible chat API URL | | `CHAT_ENDPOINT` | *(required)* | OpenAI-compatible chat API URL |
| `CHAT_MODEL` | *(required)* | Model name for chat completions | | `CHAT_ENDPOINT_KEY` | `placeholder` | API key for the chat endpoint |
| `IMAGE_GEN_ENDPOINT` | *(required)* | Image generation API URL | | `CHAT_MODEL` | *(required)* | Model name for chat completions |
| `IMAGE_EDIT_ENDPOINT` | *(required)* | Image editing API URL | | `IMAGE_GEN_ENDPOINT` | *(required)* | Image generation API URL |
| `IMAGE_GEN_SIZE_SQUARE` | `1024x1024` | Square canvas size for image generation | | `IMAGE_GEN_ENDPOINT_KEY` | `placeholder` | API key for the image generation endpoint |
| `IMAGE_GEN_SIZE_PORTRAIT` | `1024x1536` | Portrait (tall) canvas size for image generation | | `IMAGE_GEN_MODEL` | *(required)* | Model name for image generation |
| `IMAGE_GEN_SIZE_LANDSCAPE`| `1536x1024` | Landscape (wide) canvas size for image generation | | `IMAGE_GEN_SIZE_SQUARE` | `1024x1024` | Square canvas size for image generation |
| `EMBEDDING_ENDPOINT` | *(required)* | Embedding API URL | | `IMAGE_GEN_SIZE_PORTRAIT` | `1024x1536` | Portrait (tall) canvas size for image generation |
| `EMBEDDING_MODEL` | *(required)* | Model name for text embeddings | | `IMAGE_GEN_SIZE_LANDSCAPE` | `1536x1024` | Landscape (wide) canvas size for image generation |
| `MAX_COMPLETION_TOKENS` | `1000` | Max tokens in LLM responses | | `IMAGE_EDIT_ENDPOINT` | *(required)* | Image editing API URL |
| `MAX_HISTORY_MESSAGES` | `1000` | Max messages kept in the database | | `IMAGE_EDIT_ENDPOINT_KEY` | `placeholder` | API key for the image editing endpoint |
| `SIMILARITY_THRESHOLD` | `0.7` | Min cosine similarity for RAG context | | `IMAGE_EDIT_MODEL` | *(required)* | Model name for image editing |
| `TOP_K_RESULTS` | `5` | Number of similar messages retrieved | | `EMBEDDING_ENDPOINT` | *(required)* | Embedding API URL |
| `TTS_MODEL_PATH` | `kokoro-v1.0.onnx` | Path to Kokoro ONNX model file | | `EMBEDDING_ENDPOINT_KEY` | `placeholder` | API key for the embedding endpoint |
| `TTS_VOICES_PATH` | `voices-v1.0.bin` | Path to Kokoro voices binary file | | `EMBEDDING_MODEL` | *(required)* | Model name for text embeddings |
| `TTS_VOICE` | `af_sarah` | Default voice for TTS | | `MAX_COMPLETION_TOKENS` | `1000` | Max tokens in LLM responses |
| `TTS_SPEED` | `1.0` | Speech speed multiplier | | `DB_PATH` | `chat_history.db` | SQLite database file path |
| `DB_PATH` | `chat_history.db` | SQLite database file path | | `MAX_HISTORY_MESSAGES` | `1000` | Max messages kept in the database |
| `SIMILARITY_THRESHOLD` | `0.7` | Min cosine similarity for RAG context |
| `TOP_K_RESULTS` | `5` | Number of similar messages retrieved |
| `TTS_MODEL_PATH` | `kokoro-v1.0.onnx` | Path to Kokoro ONNX model file |
| `TTS_VOICES_PATH` | `voices-v1.0.bin` | Path to Kokoro voices binary file |
| `TTS_VOICE` | `af_sarah` | Default voice for TTS |
| `TTS_SPEED` | `1.0` | Speech speed multiplier |
REMEDIATION_PLAN.md
-754
View File
@@ -1,754 +0,0 @@
# REMEDIATION_PLAN.md — Vibe Bot "Gold Quality" Roadmap
Date of audit: 2026-08-17
Status: Approved for execution (maintainer confirmed scope decisions in §1)
This document is written for a coding agent. Work through it top to bottom:
§1 (scope), then §2 (the audit, for reference), then §3 (the phases, in order).
## 0. Agent Protocol
### Commands (run from the repo root)
```bash
uv sync --extra dev # setup (plain `uv sync` already includes ruff+pyright)
uv run pytest vibe_bot/tests/ -v # tests (must run from repo root)
uv run ruff check vibe_bot/ # lint
uv run mypy vibe_bot/ # strict types
uv run pyright vibe_bot/ # strict types
uv run black --check vibe_bot/ # format check
```
- `uv run python -m vibe_bot.main` logs into the real Discord server. **Never run it as a smoke test.**
- A repo-root `.env` with placeholder values is required for imports/tests (see AGENTS.md).
- If `uv run <tool>` fails with `ModuleNotFoundError`/bad interpreter: `rm -rf .venv && uv sync --extra dev`.
- Baseline at audit time: 156/156 tests pass; all four static gates pass. Two tests
(`test_llama_wrapper.py::test_chat_completion_think`, `::test_chat_completion_instruct`)
make **live network calls** and only pass when `CHAT_ENDPOINT` is reachable.
### Per-task rules
1. Work **one task at a time**, in the order given within a phase.
2. After each task, run the full gate set (pytest + ruff + mypy + pyright + black). All must pass before starting the next task.
3. Add/adjust tests as specified in each task. Tests are the safety net — do not delete existing assertions to make a gate pass.
4. Check off the task's `[ ]` box in this file when its "Done when" criteria are met.
5. Do not commit unless the maintainer asks. (Suggested: one commit per task, message `P1-T1: <summary>`.)
6. Line numbers cited in §2/§3 are from the audit date and will drift as you edit — treat them as orientation, then locate by symbol name.
### Severity key
Critical = bot harms itself/availability · High = data corruption or dead-end UX · Medium = maintainability/cost · Low = polish.
## 1. Scope Decisions (confirmed with maintainer — do not revisit)
| Topic | Decision | Consequence for this plan |
|---|---|---|
| Multi-user privacy | Private, trusted group. **Shared chat history is a feature.** Anyone may overwrite any bot; anyone may wipe history. | Findings 1.1, 1.2, 5.1 (authorization) are **accepted by design**. Only scoped residuals remain: fix the misleading "It may already exist" message, document shared-memory semantics in README, add cost-abuse cooldowns. |
| Structural change | **Full restructure approved** (§3 Phase 2). | `main.py` is split into `commands/` + `services/` + `prompts.py`; `llama_wrapper` renamed. |
| TTS model files in git (325MB `.onnx` + 26MB `.bin`) | **Leave as-is.** | Finding 6.3 dropped; no LFS/build-download work. |
## 2. PART I — THE AUDIT REPORT
### State of the Codebase
~5,900 lines across 8 source files + 6 test files. Hygiene is good (strict mypy+pyright clean,
ruff/black clean, 156 green tests, env-based secrets, fully parameterized SQL, solid docstrings).
The core problems are structural: an async runtime doing almost everything synchronously on the
event loop, a data-retention logic bug, unguarded failure paths, and a 1,276-line god module
mixing command handlers, LLM orchestration, and 3 KB of prompt constants.
Tally: 2 Critical · 6 High · 9 Medium · 8 Low (some Highs downgraded to accepted-by-design in §1).
### Findings
#### Dimension 1: Code Quality & Correctness
- **1.1** Sev: High (→ accepted, scoped). Custom-bot namespace is global; `INSERT OR REPLACE`
(`database.py` `create_custom_bot`) lets any user silently overwrite another's bot; user is told
"It may already exist" when it was actually overwritten.
→ Residual: message must say "updated/replaced".
- **1.2** Sev: High (→ accepted). `search_similar_messages` is unscoped; RAG context can include
other users' conversations. Intended shared-memory behavior.
→ Residual: document explicitly in README; do not "fix".
- **1.3** Sev: High. `_cleanup_old_messages` (`database.py`): the second DELETE's subquery runs
**after** the first delete mutated `chat_messages`, so it removes embeddings for the next-oldest
*live* messages and never the just-deleted ones. Unbounded orphan growth in `message_embeddings`
+ lost RAG recall.
- **1.4** Sev: Medium. Magic string `WHERE username != 'vibe-bot'` (`database.py` in
`search_similar_messages` and `get_user_history`). Silently breaks if the bot's Discord
username differs; bot rows leak into user search results.
- **1.5** Sev: Medium. `llama_wrapper.embedding()`: `resp.json()` sits **outside** the try block;
a non-JSON 2xx body raises `JSONDecodeError` out of `get_conversation_context`, which is called
outside `handle_chat`'s try → unhandled `CommandError`, user left hanging.
- **1.6** Sev: High. Doodlebob/retcon failure holes: `image_generation` catches only
`openai.APIConnectionError` (4xx/5xx propagate, no user feedback); `!retcon` has no
attachment-count check (empty image list → API error), and empty `image_b64` decodes to `b""`
and an empty file is uploaded with a success message.
- **1.7** Sev: Low. `talkforme`: int limit validated **after** the announcement message; the 20-cap
never announced; first reply sent unchunked (can exceed Discord's 2000-char limit and abort the
conversation).
- **1.8** Sev: Low. `!speak` `--voice` parse uses `message.rsplit("--voice ", 1)`; text containing
the flag mid-string is corrupted.
- **1.9** Sev: Low. `get_recent_messages` and `get_custom_bot` annotate `datetime` but SQLite
returns `str` (no row factory).
- **1.10** Sev: Low. `handle_chat`/`_speak_with_bot` ignore `add_message`'s `False` return —
silent persistence failure.
#### Dimension 2: Architecture & Structure
- **2.1** Sev: **Critical.** Synchronous work on the asyncio event loop. `chat_completion_with_tools`
is `async def` but wraps a **sync** OpenAI client; `doodlebob`/`retcon`/`talkforme`/`speak`/
`handle_chat` call sync LLM (60s timeout), image (300s timeout), embedding, TTS (CPU), and
`requests.get` directly inside async handlers. One slow `!doodlebob` can miss Discord heartbeats
and drop the whole bot offline for every user.
- **2.2** Sev: High. `main.py` is a 1,276-line god module: command handlers + LLM orchestration +
prompt constants + chunking + TTS wiring; `bot = commands.Bot(...)` at import time makes the
module un-importable without full env.
- **2.3** Sev: Medium. `config.py` raises `RuntimeError` at **import** if env vars are missing;
`logging.basicConfig` is called in three modules (`config.py`, `database.py`, `main.py`) —
first wins, rest dead.
- **2.4** Sev: Medium. `on_message` constructs a `CustomBotManager` (connection + DDL) and runs a
full `list_custom_bots()` SELECT on **every message**, including non-`!` chatter.
- **2.5** Sev: Medium. Inconsistent lifetimes: `get_database()` singleton vs per-command
`CustomBotManager()`; `ChatDatabase.client` (OpenAI) is created but **never used** (dead code —
embeddings go through raw `requests`); `llama_wrapper` creates a new `openai.OpenAI` per call
(no connection reuse).
- **2.6** Sev: Low. Dead config: `COMPLETION_ENDPOINT`/`_KEY`/`_MODEL` are required at import but
used nowhere; `EMBEDDING_DIMENSION` unused.
- **2.7** Sev: Medium. Tool-calling scaffolding (tool-def dict + executor + notifier closures)
copy-pasted between `handle_chat` and `_speak_with_bot`; adding a tool means editing 3+ places.
#### Dimension 3: Readability & Maintainability
- **3.1** Sev: Medium. Three near-identical chat functions in `llama_wrapper.py`
(`chat_completion` / `chat_completion_instruct` / `chat_completion_with_history`) differing only
by `seed=-1` and message shape; `chat_completion` has zero production callers.
- **3.2** Sev: Low. Misleading names: `llama_wrapper` (no Llama inside); `get_channel_members_impl`
actually returns **guild-wide** members (Discord has no per-channel membership) and its docstring
overpromises; `bot_name` loop var in `on_message` shadows the module `bot`.
- **3.3** Sev: Medium. The 1900-char chunking while-loop is copy-pasted **five** times in
`main.py` (debug members/whoami/tools, voices, history); two 1000-char mid-word splitters
(handle_chat, talkforme) can break words and multi-codepoint emoji.
- **3.4** Sev: Medium. Log noise + PII: INFO-level logs per cursor operation in `database.py`;
`logger.info("Chat prompts: %s", prompts)` logs full conversation + user profile every message.
- **3.5** Sev: Low. 11 `type: ignore` suppressions; most hide that `tool_call.function` is
`Function | None` and should be narrowed with a guard.
#### Dimension 4: Documentation
- **4.1** Sev: Medium. README stale: file tree missing `tools.py`, `build.sh`, `vibe-bot.container`,
`.gitea/`, `AGENTS.md`; lists `.env` as a repo file; `!debug` and `!lobotomize` undocumented;
RAG description omits the shared-memory (cross-user) behavior.
- **4.2** Sev: Medium. `test_config.py` hardcodes
`sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots")` (two places); passes only
via cwd fallback.
- **4.3** Sev: Low. Command reference exists in three places (README, docstrings, in-chat `!debug`)
with no single source of truth.
#### Dimension 5: Robustness & Error Handling
- **5.1** Sev: High (→ accepted, scoped). No command authorization or rate limiting anywhere.
Accepted for a trusted group; residual: per-user cooldowns on paid commands
(`!doodlebob`, `!talkforme`, `!speak`) as a cost-abuse guard.
- **5.2** Sev: Medium. No input size bounds: personality (min only), `!speak` text, `!doodlebob`
prompt, `!talkforme` topic. Unbounded personality is also a prompt-injection/cost surface
(concatenated verbatim into system prompts).
- **5.3** Sev: Medium. `!retcon` fetches attachment URLs with no allowlist, no size cap, blocking
`requests.get`.
- **5.4** Sev: Medium. No SQLite `WAL`/`busy_timeout` — required once 2.1 is fixed with threads,
or concurrent writes fail with "database is locked".
- **5.5** Sev: Low. Swallowed failures with silent degradation: TTS skips failing chunks (partial
audio, no user notice); `get_channel_members_impl` returns an error *string* the LLM treats as
real roster data.
#### Dimension 6: Performance & Security
- **6.1** Sev: Medium. RAG retrieval is O(N): loads **all** embeddings (up to 1000 × 8 KB) into
Python per query, plus one extra `SELECT` per candidate for the `_response` join (N+1).
- **6.2** Sev: Medium. Redundant embedding API calls: ~3 per turn, and the **bot response**
embedding is stored but never retrieved (search excludes bot rows). ~⅓ of embedding spend wasted.
- **6.3** Sev: Medium (→ **dropped** per §1). 351 MB of model binaries committed to git.
- **6.4** Sev: Medium. Dependency hygiene: `mypy` (two conflicting versions: main ≥2.1.0, dev
extra ≥1.17.0), `pytest`, `pytest-env`, `types-requests` are **runtime** deps shipped into the
container; `[tool.uv] required-environments` pins the lock to linux/x86_64.
- **6.5** Sev: Low. CI builds/pushes with **no lint/test gate** and only on main/release (no PR
signal); container runs as root, no healthcheck.
- **6.6** Security positives (no action): no hardcoded secrets (env only, `.env*` gitignored,
registry creds via Gitea secrets); all SQL parameterized; no XSS surface. Residuals covered by
1.1/1.2 (accepted), 3.4 (PII logging), 5.3 (fetch allowlist).
## 3. PART II — THE REMEDIATION ROADMAP
Phase order is by dependency and risk: Phase 1 removes the ways the bot harms itself
(behavior-preserving fixes), Phase 2 is the restructure on the stabilized codebase,
Phase 3 makes the dev loop and CI trustworthy, Phase 4 optimizes the final shapes.
---
### PHASE 1 — Stability & Foundation
**Goal:** The bot can no longer freeze globally, corrupt retention data, or strand users in
unhandled error states. Zero structural changes; behavior-preserving fixes behind the existing
test suite.
#### P1-T1 — SQLite concurrency preparation (finding 5.4)
Prerequisite for P1-T2; land it first.
- File: `vibe_bot/database.py` → `ChatDatabase._initialize_database` (and
`CustomBotManager._initialize_custom_bots_table` if you want both tables covered).
- After connecting, execute:
```sql
PRAGMA journal_mode=WAL;
PRAGMA busy_timeout=5000;
```
(WAL is persistent; busy_timeout is per-connection, so set it wherever connections are opened —
introduce a small `_connect()` helper returning a configured connection and use it in every
method instead of raw `sqlite3.connect`.)
- Tests: none required (pragma-only), but keep the suite green.
- Done when: [ ] all `sqlite3.connect` calls in `database.py` go through the helper; gates green.
#### P1-T2 — Unblock the event loop (finding 2.1 — THE critical fix)
- Files: `vibe_bot/main.py`, `vibe_bot/llama_wrapper.py`.
- Strategy: **`asyncio.to_thread` first** (lowest risk, behavior-preserving). True async clients
come in Phase 2.
1. In every async command/handler, wrap blocking calls:
- `handle_chat`: `db.get_conversation_context(...)`, `llama_wrapper.chat_completion_with_tools(...)`
(it is `async def` but its inner client calls are sync — either keep `await`ing it after
converting its inner `.create()` calls to run in a thread, or simplest: change
`chat_completion_with_tools` to a plain sync function and `await asyncio.to_thread(...)` it),
`db.add_message(...)` (×2).
- `doodlebob`: `select_image_layout`, `chat_completion_instruct` (×2: prompt + verify),
`image_generation`.
- `retcon`: the `requests.get` download (remove the `# noqa: ASYNC210` once async via
`to_thread`) and `llama_wrapper.image_edit`.
- `talkforme`: every `llama_wrapper.chat_completion_with_history` call.
- `_speak_with_bot` / `_speak_plain`: `db.get_conversation_context`,
`chat_completion_with_tools`, `db.add_message`, `engine.generate_audio` (CPU-bound TTS).
2. Keep all function signatures otherwise identical.
- Tests:
- New test: while a mocked 0.5s LLM call is in flight inside `handle_chat` (or `doodlebob`),
a concurrent `asyncio.create_task` sleep completes — i.e. prove the loop is not blocked
(assert with `asyncio.wait` / timing: the concurrent task's elapsed time < the mocked call
time by a wide margin, e.g. concurrent task finishes while the call is still running).
- Done when: [ ] no direct blocking call (openai sync client, `requests`, sqlite, TTS) executes
on the event loop thread; new concurrency test passes; gates green.
#### P1-T3 — Fix `_cleanup_old_messages` (finding 1.3)
- File: `vibe_bot/database.py` → `ChatDatabase._cleanup_old_messages`.
- Rewrite: first `SELECT id, message_id FROM chat_messages ORDER BY timestamp ASC LIMIT ?`
(where `? = count - MAX_HISTORY_MESSAGES`); derive the affected message_ids **including**
their `_response` companions in Python (`mid` and `f"{mid}_response"`); then
`DELETE FROM chat_messages WHERE id IN (...)` and
`DELETE FROM message_embeddings WHERE message_id IN (...)` — both from the captured set,
in the same transaction (already the case: caller commits).
- Tests (in `test_database.py`):
- Insert `MAX_HISTORY_MESSAGES + 3` user+response rows (with embeddings). Call the internal
cleanup path via `add_message`. Assert: exactly 3 oldest user rows gone; their embeddings
gone; their `_response` embeddings gone; the **next** oldest surviving user row still has
its embedding; `SELECT COUNT(*) FROM message_embeddings` == rows that have a live message.
- Done when: [ ] regression test passes; no orphaned embeddings possible after cleanup; gates green.
#### P1-T4 — Guard `embedding()` JSON parsing (finding 1.5)
- File: `vibe_bot/llama_wrapper.py` → `embedding`.
- Move `data = resp.json()` inside the `try` and catch `ValueError` (JSONDecodeError's base)
alongside `requests.RequestException` → return `[]`.
- Test: mock `requests.post` returning a 200 with body `<html>rate limited</html>` → assert
`embedding(...) == []` and no exception.
- Done when: [ ] non-JSON 2xx can't escape `embedding()`; test passes; gates green.
#### P1-T5 — Close doodlebob/retcon failure holes (finding 1.6, part of 5.3)
- Files: `vibe_bot/llama_wrapper.py` (`image_generation`), `vibe_bot/main.py` (`doodlebob`, `retcon`).
- 1. In `image_generation` and `image_edit`, catch `openai.OpenAIError` (covers
`APIStatusError`, `APIConnectionError`, `APITimeoutError`) → return `""` (log the error).
Callers already handle `""`.
- 2. `retcon`: before calling `image_edit`, if `image_data_list` is empty → send
"Please attach an image to edit." and return. After `image_edit`, if result is `""` → send
"Failed to edit the image." and return. Wrap `base64.b64decode` in try/except
(`binascii.Error`) with a user-facing failure message.
- 3. `retcon` download: restrict to Discord CDN hosts (`discordcdn.com`, `discordapp.com`,
`discord.com`, `discord.media`) — reject anything else with a warning log; cap download at
8 MB (stream with `stream=True`, abort if `Content-Length` or accumulated bytes exceed cap).
- 4. `doodlebob`: after `select_image_layout`/prompt/verify, the existing `== ""` checks remain;
no extra try needed once wrappers swallow `OpenAIError`.
- Tests:
- retcon with zero attachments → friendly message, `image_edit` not called.
- retcon with `image_edit` returning `""` → failure message, no file sent.
- retcon with non-Discord attachment URL → not downloaded.
- `image_generation` with mocked `openai.APIStatusError` → returns `""`.
- Done when: [ ] all four tests pass; no `!retcon`/`!doodlebob` path ends without a user message
on API failure; gates green.
#### P1-T6 — Replace the `'vibe-bot'` magic string with a role column (finding 1.4)
- File: `vibe_bot/database.py`.
- 1. In `_initialize_database`, after the existing `bot_name` migration pattern, add: if
`role` column missing → `ALTER TABLE chat_messages ADD COLUMN role TEXT`.
Backfill in the same migration:
`UPDATE chat_messages SET role = 'assistant' WHERE message_id LIKE '%_response' AND role IS NULL;`
`UPDATE chat_messages SET role = 'user' WHERE role IS NULL;`
- 2. `add_message`: accept `role: str = "user"` and store it (callers in `main.py` pass
`role="assistant"` for the bot-response rows).
- 3. `search_similar_messages`: `WHERE cm.username != 'vibe-bot'` → `WHERE cm.role = 'user'`.
`get_user_history`: same replacement.
- 4. `main.py`: update both `add_message` call sites (handle_chat, _speak_with_bot) to pass role.
- Tests:
- Fresh DB: add user + response rows, assert `get_user_history` excludes responses and
`search_similar_messages` only matches user rows.
- Migration: create a legacy row (no role, `message_id` ending `_response`) by inserting
directly via sqlite, run `ChatDatabase()`, assert role backfilled to `assistant`.
- Done when: [ ] no `vibe-bot` string remains in `database.py` (grep); tests pass; gates green.
#### P1-T7 — Fix `talkforme` ordering and chunking (finding 1.7)
- File: `vibe_bot/main.py` → `talkforme`.
- 1. Move `int(limit)` parsing/validation **before** the "is going to talk" announcement.
- 2. Announce the effective cap: `for {min(limit, talk_limit)} replies`.
- 3. Send the first reply through the same 1000-char loop used for subsequent replies
(extract a local helper now; the global `split_message` arrives in Phase 2).
- Tests:
- non-integer limit → usage/error message, **no** announcement message sent first.
- first reply > 1000 chars → sent in chunks (assert `ctx.send` call count).
- Done when: [ ] both tests pass; gates green.
#### P1-T8 — Input size bounds (finding 5.2)
- File: `vibe_bot/main.py` (constants near `MIN_BOT_NAME_LENGTH`).
- Add and enforce with friendly rejection messages:
- `MAX_PERSONALITY_LENGTH = 1000` in `custom_bot`.
- `MAX_SPEAK_LENGTH = 5000` in `speak` (check the text part after `--voice` parsing).
- `MAX_IMAGE_PROMPT_LENGTH = 2000` in `doodlebob` and `retcon`.
- `MAX_TOPIC_LENGTH = 500` in `talkforme`.
- Tests: one per bound — over-limit input → rejection message, no LLM/DB call.
- Done when: [ ] four tests pass; gates green.
#### P1-T9 — Cost-abuse cooldowns (scoped finding 5.1)
- File: `vibe_bot/main.py`.
- Add `@commands.cooldown(rate, per, type=commands.BucketType.user)` (pick sane values, e.g.
doodlebob 1/60s, talkforme 1/120s, speak 3/30s) to the three paid commands.
- Handle `commands.MaxConcurrency`/`CooldownRetry` gracefully: register a
`on_command_error` handler that sends "You're using that too quickly, try again in Ns."
for `commands.CooldownRaise`.
- Test: mock a second immediate invocation → cooldown message, LLM not called.
- Done when: [ ] test passes; gates green.
#### P1-T10 — Honest overwrite message (scoped finding 1.1)
- File: `vibe_bot/database.py` (`create_custom_bot`), `vibe_bot/main.py` (`custom_bot`).
- Make `create_custom_bot` distinguish: check existence first (`SELECT 1 ... WHERE bot_name=?`)
→ return `"created" | "replaced" | False`. `custom_bot` command: on `"replaced"` send
"Custom bot **'X'** already existed and has been **replaced**." (Shared namespace is intended.)
- Test: create same name twice → second call reports replaced; both succeed.
- Done when: [ ] test passes; no "It may already exist" string remains; gates green.
#### P1-T11 — Stop silent TTS partial audio (scoped finding 5.5)
- Files: `vibe_bot/tts.py`, `vibe_bot/main.py`.
- `generate_audio`: track failed-chunk count; if any chunk failed, log a warning and either
(a) attach a note — simplest: return the audio and let the caller detect via a new
`partial: bool` attribute on a small result dataclass, or (b) raise `ValueError` listing the
failed chunk count and let the command send "Some audio chunks failed; audio may be incomplete."
Choose (b) if the maintainer prefers fail-fast — **default to (a)**: keep the audio, and have
the command send a short warning line when partial.
- Test: mock `process_chunk_sequential` to fail on chunk 1 of 2 → audio still produced, warning
sent (or ValueError raised, matching the chosen behavior).
- Done when: [ ] test passes; gates green.
#### Phase 1 verification & success criteria
Run the full gate set, then:
- [ ] Concurrency proof: a slow (mocked ≥60s) LLM call no longer blocks a concurrent task
(P1-T2 test) — the bot survives a slow `!doodlebob` globally.
- [ ] No orphaned embeddings and correct retention after cleanup (P1-T3 test).
- [ ] No command path ends without a user-facing message on API failure (mock 4xx/5xx tests:
doodlebob, retcon, handle_chat).
- [ ] Full suite green: `uv run pytest vibe_bot/tests/ -v` (156 original + new tests).
- [ ] `ruff` / `mypy` / `pyright` / `black --check` all pass.
---
### PHASE 2 — Structural Integrity (full restructure)
**Goal:** Kill the god module and duplicated orchestration; commands become thin adapters over
services; lifetimes explicit; true async clients. Do this **after Phase 1 gates are green**.
Order matters: utilities and clients first (P2-T1/T2), services next (P2-T3…T5), command split
after (P2-T6), DB layer last (P2-T7) because services depend on it.
#### P2-T1 — `textutil.split_message` (finding 3.3 — do first, Phase 1 code already uses it indirectly)
- New file: `vibe_bot/textutil.py`.
- `def split_message(text: str, limit: int = 1900) -> list[str]`:
- Split on newlines first (never break a line mid-text when avoidable); if a single line
exceeds `limit`, split it **codepoint-safely** (Python `str` slicing is already codepoint-safe,
but never split inside a grapheme cluster — use `regex` package's `\X` only if you add the dep;
otherwise plain slicing at `limit` is acceptable and matches today's behavior — document the
choice).
- Preserve existing behavior for the 5 call sites; replace all five 1900-char loops in
`main.py` (debug members/whoami/tools, voices, history) and both 1000-char splitters
(`handle_chat`, `talkforme` → `limit=1000`).
- Tests: exact-limit text; text with no newlines; emoji at the split boundary (assert no lone
surrogate/combining char corruption by round-tripping `"".join(chunks) == original`);
multi-paragraph text respects newline preference.
- Done when: [ ] `main.py` contains zero hand-rolled chunking loops; `"".join(split_message(t)) == t`
property test passes for a corpus including emoji; gates green.
#### P2-T2 — `llm_client.py`: one core, shared async clients (findings 2.1-completion, 3.1, 2.5, 3.2, 3.5)
- New file: `vibe_bot/llm_client.py` (delete `llama_wrapper.py`; update all imports).
- 1. **Shared clients:** module-level lazy singletons `get_chat_client()`,
`get_image_gen_client()`, `get_image_edit_client()`, `get_embedding_http_session()` returning
`openai.AsyncOpenAI(base_url=..., api_key=..., max_retries=0 where currently set)` instances
built once. Delete the per-call `openai.OpenAI(...)` constructions.
- Also delete `ChatDatabase.client` (dead code) and its `openai` import in `database.py`;
drop `db.client.close()` from the `chat_db` test fixture.
- 2. **One core function:**
`async def chat_complete(messages: list[ChatCompletionMessageParam], *, model: str, max_tokens: int, seed: int | None = None, tools: ... | None = None, tool_executor: ... | None = None, tool_call_notifier: ... | None = None, max_tool_rounds: int = 5) -> str`
— contains the (now async) tool loop currently in `chat_completion_with_tools`.
Replace the `# type: ignore[union-attr]` cluster with a real guard:
`if tool_call.function is None: continue`.
- Keep thin compat wrappers *only if* it reduces churn: `chat_completion_instruct(system_prompt, user_prompt, ...)` and `chat_completion_with_history(system_prompt, prompts, ...)` become 3-line adapters over `chat_complete`. Delete `chat_completion` (test-only caller — port its test to `chat_complete`).
- 3. **`image_generation` / `image_edit`:** async (`await client.images.generate/edit`), catch
`openai.OpenAIError` → `""` (carries P1-T5 behavior).
- 4. **`embedding`:** async via `httpx.AsyncClient` (add `httpx` to deps) **or** keep `requests`
inside `asyncio.to_thread`. Prefer `httpx.AsyncClient` with a shared session; keep the
OpenAI-style and Ollama-style response handling and P1-T4's JSON guard.
- 5. `ToolRegistry` (finding 2.7): small class in `llm_client.py` (or `tools.py`):
- `register(name, description, args_schema, impl)` — seed it with `get_channel_members`
(schema) → `get_channel_members_impl` (impl).
- `to_openai_tools() -> list[dict]` and `async execute(name, args, channel) -> str`.
- The notifier ("is looking at the channel members...") becomes a registry-level callback.
- Tests: port all `test_llama_wrapper.py` tests to `test_llm_client.py` (rename module + async
awaits); add: client singleton identity test (same object across calls); tool registry
execute/dispatch test incl. unknown-tool path; guard test for `function is None` tool call.
- Done when: [ ] `llama_wrapper.py` deleted; grep shows zero `openai.OpenAI(` per-call
constructions; zero `type: ignore` in the new client (except the two `import-untyped` in
`tts.py`); gates green.
#### P2-T3 — `prompts.py` (part of finding 2.2)
- New file: `vibe_bot/prompts.py` — move all prompt constants from `main.py`:
`IMAGE_LAYOUT_SYSTEM_PROMPT`, `IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE`,
`IMAGE_PROMPT_VERIFY_SYSTEM_PROMPT`, plus the "Keep your responses under 2-3 sentences."
suffix and "User Information:" assembly as a helper
`build_system_prompt(personality, user_info) -> str` (currently duplicated string in
`handle_chat` and `_speak_with_bot`).
- Move the layout parsing (`parse_image_layout`, `VALID_IMAGE_LAYOUTS`, `DEFAULT_IMAGE_LAYOUT`,
`LAYOUT_SIZES`) here too.
- Tests: existing prompt-related tests in `test_main.py` re-point at `vibe_bot.prompts`;
add a test that `build_system_prompt` output contains personality + user info + length hint.
- Done when: [ ] `main.py` (pre-split) contains no prompt text; gates green.
#### P2-T4 — `services/` (findings 2.2, 2.4, 2.5, 6.2)
New package `vibe_bot/services/` with a dataclass `BotContext` (or pass explicit args):
- `services/chat_service.py` — `ChatService.handle(ctx, bot_name, message, system_prompt, response_prefix)`:
RAG context → prompts → `chat_complete` with tools (via `ToolRegistry`) → persist
(user row `role="user"`, response row `role="assistant"`; **log a warning when
`add_message` returns False** — finding 1.10) → chunked reply via `split_message`.
**Do not embed bot responses** (finding 6.2): only user messages get embeddings — remove the
response-row embedding by not calling the embedding path for assistant rows
(`add_message(role="assistant", embed=False)` parameter, or a separate non-embedding insert).
- `services/image_service.py` — `ImageService.generate(ctx, message)`:
layout select (`max_tokens=2` — Phase 4 item, do it here since we're touching it anyway) →
prompt engineering → verify → ETA message → generate → record time → send file.
`ImageService.edit(ctx, message)`: the P1-T5-hardened retcon flow.
- `services/speech_service.py` — `SpeechService.speak(ctx, message)`: voice-flag parse
(replace `rsplit` with a **trailing-anchored regex**: `r"^(?P<text>.*)\s+--voice\s+(?P<voice>\S+)$"`
— fixes finding 1.8), voice validation, bot-vs-plain dispatch, TTS via `asyncio.to_thread`,
language lookup via a precomputed `VOICE_LANGUAGES: dict[str, str]` built from `VOICES_LIST`
at import (replaces the per-call list scan).
- `services/conversation_service.py` — `ConversationService.run(ctx, bot1, bot2, limit, topic)`:
the Phase-1-fixed talkforme loop, using `split_message(1000)`.
- Each service: constructor takes the llm client(s), `ChatDatabase`, `CustomBotManager`,
`TTSEngine | None`, `ToolRegistry` (dependency injection; no module-level singletons inside).
- Tests: port `test_main.py` command tests to service-level tests where the logic moved;
keep a thin command-layer test per command proving wiring (mock service, assert called with
parsed args). Add: voice-regex tests (flag mid-text is preserved; trailing flag extracted;
missing flag → None).
- Done when: [ ] four service modules exist; no service imports `discord` except for
`ctx`/`channel` type hints under `TYPE_CHECKING`; gates green.
#### P2-T5 — `app.py` + composition root (findings 2.2, 2.3, 2.4, 2.5)
- New file: `vibe_bot/app.py`:
- `configure_logging()` — the **only** `logging.basicConfig` call in the codebase (remove the
other two from `config.py` and `database.py`).
- `create_app() -> App` dataclass: builds `ChatDatabase`, `CustomBotManager` (singleton —
no more per-message construction), LLM clients, `ToolRegistry`, `TTSEngine | None`
(keep the tolerant init-or-None behavior), and the four services.
- Bot construction moves here: `build_bot(app) -> commands.Bot` (intents, `on_ready`,
command registration, `on_message`, `on_command_error`).
- `on_message`: skip bots; **if not `message.content.startswith("!"): return`** before any
DB/manager touch (finding 2.4); then look up the bot-name cache.
- Bot-name cache: `app.custom_bots: dict[str, tuple[prompt, creator]]` rebuilt from
`list_custom_bots()`; invalidated by `custom_bot`/`delete_custom_bot` commands (call
`app.invalidate_bot_cache()`).
- `vibe_bot/main.py` shrinks to: `def main() -> None: validate_config(); configure_logging(); app = create_app(); build_bot(app).run(DISCORD_TOKEN)` + `if __name__ == "__main__": main()`.
- `config.py`: split loading from validation — module only **loads**; add
`def validate_config() -> None` containing today's `RuntimeError` checks (called from
`main()` only, so the package becomes importable without full env).
- Tests: `test_main.py` reorganized: command tests now call `build_bot(app)`-registered
functions or the service layer; the `mock_discord` fixture is replaced by direct service
tests (Phase 4 cleanup of leftovers). Import-time behavior test: importing `vibe_bot.config`
with empty env must NOT raise (only `validate_config()` raises).
- Done when: [ ] `main.py` < 40 lines; grep shows exactly one `logging.basicConfig` in
`vibe_bot/`; no `CustomBotManager()` construction outside `app.py`; gates green.
#### P2-T6 — `commands/` package (completes finding 2.2, 2.6)
- New package `vibe_bot/commands/` — one module per group, each exposing
`def register(bot: commands.Bot, app: App) -> None`:
- `custom_bots.py`: `custom_bot`, `list_custom_bots`, `delete_custom_bot` (use
`app.bot_cache`; keep P1-T10's replaced/created semantics).
- `chat.py`: nothing registered (custom-bot chat flows through `on_message` →
`app.services.chat`) — keep the `on_message` dispatch in `app.py` calling
`app.services.chat.handle(...)`.
- `speech.py`: `speak`, `voices` (via `SpeechService`).
- `images.py`: `doodlebob`, `retcon` (via `ImageService`; keep P1-T9 cooldowns).
- `conversation.py`: `talkforme` (via `ConversationService`).
- `admin.py`: `lobotomize`, `debug`, `history`.
- `config.py`: delete dead `COMPLETION_ENDPOINT`/`COMPLETION_ENDPOINT_KEY`/`COMPLETION_MODEL`
and `EMBEDDING_DIMENSION` (finding 2.6); update README env section in Phase 3.
- `tools.py`: fix `get_channel_members_impl` docstring → "guild members" (finding 3.2);
keep the `@tool` stub as the schema source for the registry.
- Tests: each command module gets a wiring test (mock service + mock ctx); delete now-dead
fixtures (`mock_env_vars`, `mock_discord` if unused) from `conftest.py`.
- Done when: [ ] no module in `vibe_bot/` exceeds ~400 lines; every command traceable to a
single service method; gates green.
#### P2-T7 — Database layer cleanup (findings 1.9, 6.1, 2.5)
- File: `vibe_bot/database.py`.
- 1. **Single-JOIN RAG fetch** (6.1): in `search_similar_messages`, replace
"fetch all embeddings + per-candidate SELECT" with one query:
```sql
SELECT cm.message_id, cm.content,
(SELECT content FROM chat_messages r WHERE r.message_id = cm.message_id || '_response') AS response,
me.embedding
FROM chat_messages cm
JOIN message_embeddings me ON cm.message_id = me.message_id
WHERE cm.role = 'user'
```
(correlated subquery avoids a LEFT JOIN fan-out; if a real LEFT JOIN is cleaner, use
`LEFT JOIN chat_messages r ON r.message_id = cm.message_id || '_response'`).
Vectorize the cosine loop: stack blobs into one `np.ndarray` of shape `(n, dim)` and compute
similarities with one matrix op (normalize query vector once; keep the per-vector norm
recompute for now — norms get precomputed in Phase 4).
- 2. **Honest types** (1.9): register a sqlite `datetime` converter
(`sqlite3.register_adapter`/`detect_types` or a row factory mapping the `timestamp` column)
so `get_recent_messages`/`get_custom_bot` really return `datetime`; adjust tests that compare
strings.
- 3. `add_message`: add `embed: bool = True` (assistant rows skip the embedding call — finding 6.2
wiring); keep `role` parameter from P1-T6.
- 4. Delete the dead `OpenAI` import/client (if P2-T2 didn't already).
- Tests: update RAG tests for the JOIN (assert identical results vs old behavior on a seeded DB:
same top-k, same ordering); property: one query executed per search (count with
`sqlite3` tracing or a wrapper).
- Done when: [ ] `search_similar_messages` issues exactly 1 SELECT per call; type annotations
match runtime values; gates green.
#### Phase 2 verification & success criteria
- [ ] No module > ~400 lines (`wc -l vibe_bot/**/*.py`).
- [ ] `main.py` no longer imports anything except `app`.
- [ ] Every command testable via mock services without discord module mocking.
- [ ] Behavior unchanged: original 156 tests (ported) + new tests all green.
- [ ] Strict gates green with **net-fewer** `type: ignore`s.
- [ ] grep: zero `logging.basicConfig` outside `app.py`; zero `openai.OpenAI(` per-call
constructions; zero `'vibe-bot'` in `database.py`; zero `llama_wrapper` references.
---
### PHASE 3 — Developer Experience & Documentation
**Goal:** Hermetic network-free test suite; CI that gates merges; docs that match reality.
Items 3.1–3.4 can start as soon as Phase 2's P2-T6 lands (they're independent of P2-T7).
#### P3-T1 — Hermetic test suite (AGENTS.md baseline retirement)
- Files: `vibe_bot/tests/test_llm_client.py` (or `test_llama_wrapper.py` if renamed later).
- Mark the two live tests `@pytest.mark.live` and add to `pyproject.toml`:
`[tool.pytest.ini_options] addopts = "-m 'not live'"` (keep `filterwarnings` as-is).
Document in the Testing section: `uv run pytest -m live` runs them.
- Done when: [ ] `rm -rf .venv && uv sync --extra dev && uv run pytest vibe_bot/tests/ -v`
passes on a machine with placeholder `.env` and **no network** to `CHAT_ENDPOINT`
(verify by pointing `CHAT_ENDPOINT` at an unroutable address for the test run).
#### P3-T2 — Test hygiene (finding 4.2)
- Files: `vibe_bot/tests/test_config.py`, `conftest.py`, `test_llm_client.py`.
- Delete both hardcoded `sys.path.insert(0, "/var/home/ducoterra/...")` lines; if `test_config.py`
then fails on import, add `pythonpath = ["."]` under `[tool.pytest.ini_options]` instead.
- Delete dead fixtures (`mock_env_vars` if still unused) and `TEMPDIR`.
- Fix `test_bot_intents_set` to actually assert intents (or rename to match what it asserts).
- Done when: [ ] grep for `/var/home` in `vibe_bot/tests/` returns nothing; suite green.
#### P3-T3 — Coverage for new surface
- Add tests (in the matching module):
- `split_message`: property test `"".join(split_message(t, L)) == t` over a corpus including
emoji, combining marks, and code spans; exact-multiple-of-limit input.
- Voice-flag regex: trailing flag, flag mid-text (must be preserved), no flag, flag with
missing value.
- Cleanup regression (P1-T3) — already required, verify present.
- `role` migration backfill (P1-T6) — verify present.
- Cooldowns (P1-T9): immediate re-invocation → friendly message.
- Service error paths: `ChatService` with failing `chat_complete` → user sees the error
message and no rows persisted; `ImageService.edit` with rejected URL.
- Done when: [ ] all listed tests exist and pass.
#### P3-T4 — Dependency hygiene (finding 6.4)
- File: `pyproject.toml`.
- Move to `[dependency-groups] dev` (and delete the now-redundant `[project.optional-dependencies] dev`):
`mypy` (single version pin — drop the conflicting second one), `pytest`, `pytest-env`,
`black`, `debugpy`; keep `pyright`, `ruff` where they are.
- Add `httpx` to runtime deps (P2-T2).
- Widen or remove `[tool.uv] required-environments` (recommend: remove, accept multi-platform lock).
- `uv lock` after changes; verify `uv sync --extra dev` and the Containerfile's `uv sync --locked`
both work (`./build.sh` dry-run optional — do **not** push images).
- Done when: [ ] `uv run python -c "import mypy"` fails in a prod-only sync
(`uv sync --no-dev`); container build still passes `uv sync --locked`; gates green.
#### P3-T5 — CI gates (finding 6.5)
- File: `.gitea/workflows/build-push.yml` (or a new `ci.yml`).
- New job `test` on `pull_request` (and `push` to main):
```yaml
- uses: actions/setup-python@v5 (3.13)
- name: Install uv && uv sync --extra dev
- run: uv run ruff check vibe_bot/
- run: uv run mypy vibe_bot/
- run: uv run pyright vibe_bot/
- run: uv run black --check vibe_bot/
- run: uv run pytest vibe_bot/tests/ -v # hermetic per P3-T1
```
(Adapt to the Gitea runner flavor already in use; note: portaudio system dep must be present
on CI runners — `sudo dnf install portaudio` per AGENTS.md, or install in the job; if the
runner is apt-based: `apt-get install -y portaudio19-dev`.)
- Make the build/push job depend on `test` passing.
- Done when: [ ] a PR with a deliberate type error cannot pass CI (verify locally by simulating
the job steps); merge is blocked on red `test`.
#### P3-T6 — Container hardening (finding 6.5)
- File: `Containerfile`.
- Add a non-root user: `RUN useradd -m bot && chown -R bot /app` … `USER bot`; ensure
`DB_PATH` default dir is writable (`vibe-bot.container` mounts `/db` — `mkdir -p /db && chown bot /db`).
- Optional: `HEALTHCHECK` — a Discord bot has no HTTP endpoint; use a lightweight liveness probe
(e.g. a small asyncio task writing a heartbeat file, checked by `test -f`) **only if** the
maintainer wants it; otherwise skip and note why.
- Done when: [ ] `./build.sh` succeeds; `podman run --rm localhost/vibe-bot:latest uv run python -c "print(1)"`
runs as non-root (verify with `podman run ... id`).
#### P3-T7 — Logging discipline (finding 3.4)
- Files: `vibe_bot/database.py`, `vibe_bot/services/*`.
- Per-cursor-operation `logger.info` → `logger.debug` (keep init/migration at INFO).
- `logger.info("Chat prompts: %s", prompts)` → log metadata only:
`logger.info("chat: bot=%s user=%s context_msgs=%d", ...)`.
- Done when: [ ] no log line can contain message content or user profile data
(spot-check by grepping for `prompts`, `content`, `get_user_info` in `logger.` calls).
#### P3-T8 — Documentation (findings 4.1, 4.3)
- Files: `README.md`, `AGENTS.md`.
- README:
- Regenerate the file tree from the actual post-Phase-2 layout.
- Document **all** commands including `!debug` and `!lobotomize`.
- Add a "Design notes" subsection stating the intentional semantics: shared memory (any
user's history can inform any bot's RAG context), shared bot namespace (creating an existing
name replaces it), any member may use any bot.
- Document input limits and cooldowns (P1-T8/P1-T9 values).
- Remove `COMPLETION_*` from the env table (P2-T6).
- AGENTS.md: update code map to new layout; retire the fixed gotchas (hardcoded sys.path,
two live tests → now `-m live`); keep the portaudio/.env notes.
- Done when: [ ] README tree matches `git ls-files` (ignore dotfiles/binaries); every `@bot.command`
name appears in the README; AGENTS.md gotcha list only contains still-true items.
#### Phase 3 verification & success criteria
- [ ] Fresh-clone hermetic run green with no network (P3-T1 command).
- [ ] A failing lint/type/test on a PR is blocked in CI (P3-T5).
- [ ] Container runs non-root (P3-T6).
- [ ] README/AGENTS.md accurate (P3-T8).
- [ ] Full gate set green.
---
### PHASE 4 — Optimization & Polish
**Goal:** Flatten RAG cost at scale, remove per-call allocations, pick off residual nits.
Only now — optimizations land in the final shapes.
#### P4-T1 — Precomputed vector norms (completes finding 6.1)
- File: `vibe_bot/database.py`.
- Store the L2 norm alongside the embedding (new column `norm REAL` on
`message_embeddings`, backfilled on migration via numpy over existing blobs) so
`search_similar_messages` computes `query_norm * (Q @ V.T) / norms` with one matrix multiply
and no per-vector renormalization.
- Benchmark script (add `scripts/bench_rag.py`, not committed to the package): seed N messages
(1k/5k), measure p95 of `get_conversation_context`; record before/after numbers in this file
next to the task.
- Done when: [ ] p95 at 5k ≤ p95 at 1k × 1.2 (flat-ish); suite green. (If `sqlite-vec` is
considered, only then — default is to stay on pure SQLite.)
#### P4-T2 — Client reuse verification
- Profile/trace a full `!doodlebob` + chat turn: assert exactly one HTTP session per endpoint
(httpx/AsyncOpenAI client identity stable; no new `AsyncOpenAI(...)` after startup).
- Done when: [ ] verified (add a unit test asserting `get_chat_client() is get_chat_client()`);
no per-call construction in traces.
#### P4-T3 — Remaining caches & token economy
- Confirm `VOICE_LANGUAGES` dict lookup (P2-T4) and `app.bot_cache` (P2-T5) are actually used on
the hot path (add asserts/measurements in tests if trivial).
- Verify layout-selection call uses `max_tokens=2` and verify-call a bounded small value
(already set in P2-T4 — confirm and measure end-to-end `doodlebob` latency before/after;
record the number next to this task).
- Done when: [ ] latency numbers recorded; no regressions in suite.
#### P4-T4 — Polish sweep
- [ ] `!retcon` output size: match source aspect ratio instead of hardcoded `768x768`
(revisit the original "keep generation time down" rationale with the configurable
`IMAGE_GEN_SIZE_*` knobs; keep 768 if the endpoint is slow — document the decision in
`llm_client.py` docstring).
- [ ] Remove `DEFAULT_VOICE`/`DEFAULT_SPEED` duplication between `tts.py` and `config.py`
(single source in `config.py`; `tts.py` defaults reference it).
- [ ] Delete any leftover `# noqa: ASYNC210` (should be gone after P1-T2/P2-T4).
- [ ] Remove the two `type: ignore[import-untyped]` in `tts.py` if `kokoro-tts`/`soundfile`
stubs now exist; otherwise keep and add a comment that stubs are upstream-missing.
- [ ] Sweep for dead code: `get_recent_messages` (has callers? if not, delete or keep with a
docstring noting it's a utility), `deactivate_custom_bot` (unused? wire into delete or
delete), unused test fixtures.
- Done when: [ ] all sub-items checked; `ruff` clean (no unused imports warnings suppressed).
#### Phase 4 verification & success criteria
- [ ] RAG p95 flat 1k→5k (P4-T1 numbers recorded).
- [ ] Embedding API calls per chat turn == 2 (user + query) — assert via mock call count in a
service test.
- [ ] `!doodlebob` latency improved vs Phase 1 baseline (numbers recorded).
- [ ] Full gate set green; no module > ~400 lines; zero dead fixtures/imports.
---
## 4. Final Acceptance (end of Phase 4)
- [ ] `uv run pytest vibe_bot/tests/ -v` green, hermetic (no network).
- [ ] `uv run ruff check vibe_bot/` / `mypy` / `pyright` / `black --check` all green.
- [ ] CI green on a PR (lint+test job gates merge).
- [ ] Container builds and runs as non-root.
- [ ] README + AGENTS.md match the code.
- [ ] No Critical or High finding from §2 is open (accepted-by-design items 1.1/1.2/5.1 have
their scoped residuals complete and are documented in the README).
pyproject.toml
+11 -16
View File
@@ -8,29 +8,13 @@ dependencies = [
"discord>=2.3.2", "discord>=2.3.2",
"openai>=2.24.0", "openai>=2.24.0",
"requests>=2.32.5", "requests>=2.32.5",
"types-requests>=2.32.4.20260107",
"numpy>=1.24.0", "numpy>=1.24.0",
"pytest>=9.0.2",
"python-dotenv>=1.2.2", "python-dotenv>=1.2.2",
"pytest-env>=1.5.0",
"kokoro-tts>=2.3.1", "kokoro-tts>=2.3.1",
"mypy>=2.1.0",
"langchain-openai>=0.4.0", "langchain-openai>=0.4.0",
"langchain-core>=0.3.0", "langchain-core>=0.3.0",
] ]
[project.optional-dependencies]
dev = [
"mypy>=1.17.0",
"black>=25.1.0",
"debugpy>=1.8.0",
]
[tool.uv]
required-environments = [
"sys_platform == 'linux' and platform_machine == 'x86_64'",
]
[tool.mypy] [tool.mypy]
strict = true strict = true
python_version = "3.13" python_version = "3.13"
@@ -60,12 +44,23 @@ line-length = 88
target-version = "py313" target-version = "py313"
[tool.pytest.ini_options] [tool.pytest.ini_options]
addopts = "-m 'not live'"
markers = [
"live: network-dependent integration tests, deselected by default",
]
filterwarnings = [ filterwarnings = [
"ignore::pytest.PytestUnraisableExceptionWarning", "ignore::pytest.PytestUnraisableExceptionWarning",
] ]
[dependency-groups] [dependency-groups]
dev = [ dev = [
"black>=25.1.0",
"debugpy>=1.8.0",
"mypy>=2.1.0",
"pyright>=1.1.409", "pyright>=1.1.409",
"pytest>=9.0.2",
"pytest-cov>=7.0.0",
"pytest-env>=1.5.0",
"ruff>=0.16.3", "ruff>=0.16.3",
"types-requests>=2.32.4.20260107",
] ]
scripts/bench_rag.py
+133
View File
@@ -0,0 +1,133 @@
"""RAG retrieval benchmark: p95 latency of get_conversation_context.
Standalone dev tool (do not import into the package). Seeds temporary
SQLite databases with 1k and 5k user/response rows using deterministic
fake embeddings (the embedding HTTP call is monkeypatched, so no network
is needed), then measures p95 of ChatDatabase.get_conversation_context
over repeated queries.
Run from the repo root:
uv run python scripts/bench_rag.py
Prints p95@1k and p95@5k in milliseconds.
"""
from __future__ import annotations
import hashlib
import sys
import tempfile
import time
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import vibe_bot.db.messages as db_messages
from vibe_bot import llm_client
from vibe_bot.database import ChatDatabase
EMBEDDING_DIM = 256
NUM_CLUSTERS = 16
TOPICS = [
"sailing",
"baking",
"gardening",
"astronomy",
"chess",
"pottery",
"mountaineering",
"photography",
"brewing",
"carpentry",
"weaving",
"falconry",
"cartography",
"metallurgy",
"botany",
"masonry",
]
SIZES = (1000, 5000)
WARMUP_QUERIES = 20
MEASURED_QUERIES = 200
USER_ID = "bench-user"
def fake_embedding(text: str, *, model: str, url: str, api_key: str) -> list[float]:
"""Deterministic cluster-structured fake embedding (no network)."""
topic = text.split(maxsplit=1)[0].lower()
seed = int.from_bytes(hashlib.sha256(topic.encode("utf-8")).digest()[:8], "big")
rng = np.random.default_rng(seed)
center = np.zeros(EMBEDDING_DIM, dtype=np.float32)
span = EMBEDDING_DIM // NUM_CLUSTERS
cluster = seed % NUM_CLUSTERS
center[cluster * span : (cluster + 1) * span] = 1.0
noise = rng.standard_normal(EMBEDDING_DIM, dtype=np.float32)
noise /= np.linalg.norm(noise)
vector = center + 0.3 * noise
return [float(x) for x in vector]
def seed_db(db_path: str, n_rows: int) -> list[str]:
"""Seed n_rows user/response pairs; return the user contents as queries."""
db = ChatDatabase(db_path=db_path)
queries: list[str] = []
for i in range(n_rows):
topic = TOPICS[i % NUM_CLUSTERS]
content = f"{topic} question number {i}"
queries.append(content)
db.add_message(
message_id=f"bench-{i}",
user_id=USER_ID,
username="bench",
content=content,
)
db.add_message(
message_id=f"bench-{i}_response",
user_id="bench-bot",
username="bench-bot",
content=f"response {i}",
role="assistant",
embed=False,
)
return queries
def measure(db_path: str, queries: list[str]) -> float:
"""Return p95 in ms of get_conversation_context over repeated queries."""
db = ChatDatabase(db_path=db_path)
for query in queries[:WARMUP_QUERIES]:
db.get_conversation_context(USER_ID, query)
latencies_ms: list[float] = []
for query in queries[:MEASURED_QUERIES]:
start = time.perf_counter()
db.get_conversation_context(USER_ID, query)
latencies_ms.append((time.perf_counter() - start) * 1000)
return float(np.percentile(latencies_ms, 95))
def main() -> None:
llm_client.embedding = fake_embedding # type: ignore[assignment]
db_messages.MAX_HISTORY_MESSAGES = 10**9
results: dict[int, float] = {}
for size in SIZES:
with tempfile.TemporaryDirectory(prefix="bench_rag_") as tmp:
db_path = str(Path(tmp) / f"bench_{size}.db")
queries = seed_db(db_path, size)
results[size] = measure(db_path, queries)
p95_1k = results[1000]
p95_5k = results[5000]
passed = p95_1k <= 10.0 and p95_5k <= 75.0
print(f"p95@1k = {p95_1k:.2f} ms")
print(f"p95@5k = {p95_5k:.2f} ms")
print(f"ratio = {p95_5k / p95_1k:.2f}x (informational)")
print(f"target: p95@1k <= 10ms, p95@5k <= 75ms -> {'PASS' if passed else 'FAIL'}")
if __name__ == "__main__":
main()
uv.lock
Generated
+112 -22
View File
@@ -5,9 +5,6 @@ resolution-markers = [
"python_full_version >= '3.15'", "python_full_version >= '3.15'",
"python_full_version < '3.15'", "python_full_version < '3.15'",
] ]
required-markers = [
"platform_machine == 'x86_64' and sys_platform == 'linux'",
]
[[package]] [[package]]
name = "aiohappyeyeballs" name = "aiohappyeyeballs"
@@ -414,6 +411,90 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/6d/c1/e419ef3723a074172b68aaa89c9f3de486ed4c2399e2dbd8113a4fdcaf9e/colorlog-6.10.1-py3-none-any.whl", hash = "sha256:2d7e8348291948af66122cff006c9f8da6255d224e7cf8e37d8de2df3bad8c9c", size = 11743, upload-time = "2025-10-16T16:14:10.512Z" }, { url = "https://files.pythonhosted.org/packages/6d/c1/e419ef3723a074172b68aaa89c9f3de486ed4c2399e2dbd8113a4fdcaf9e/colorlog-6.10.1-py3-none-any.whl", hash = "sha256:2d7e8348291948af66122cff006c9f8da6255d224e7cf8e37d8de2df3bad8c9c", size = 11743, upload-time = "2025-10-16T16:14:10.512Z" },
] ]
[[package]]
name = "coverage"
version = "7.15.4"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/be/c3/4f2195f512fb172aa425a8803a874b2baa9ba7f80ff7b6080998761fc701/coverage-7.15.4.tar.gz", hash = "sha256:0548198fff07ccf4faf469520bce1c2eceb1ce3e62891921138dec10907f9d00", size = 936952, upload-time = "2026-08-06T13:50:24.442Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f1/84/651a9310859673aaa3b3203f1aa1641ca60fcf2494683e1c9474c7172780/coverage-7.15.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c705b28feb2775dc82a25f1d473a370bc37ff93f5177f4e29ce2425f560f6921", size = 222565, upload-time = "2026-08-06T13:48:00.796Z" },
{ url = "https://files.pythonhosted.org/packages/82/f9/4dcf700137e8af550670f4d74d1b63828ce93e1e2b05e5f10710eb2ea987/coverage-7.15.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3ff205ab5e3ecc670f6a4dd19d9cbf12ede53dd41cfc1e15716ec961ea6d314e", size = 222936, upload-time = "2026-08-06T13:48:02.391Z" },
{ url = "https://files.pythonhosted.org/packages/07/4a/612ff1e780b3fbfd637486f542f84adc5503873d8b5d279dec1ffeef9414/coverage-7.15.4-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5172326e861a38b48b48befca15e0f477a26b283337a33a739c8fed229934e36", size = 253926, upload-time = "2026-08-06T13:48:04.382Z" },
{ url = "https://files.pythonhosted.org/packages/b0/04/d1cff1c2ead4708a6a79c01d3736b6a25bd38a36678398f72a8dd33dfad9/coverage-7.15.4-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:12b59c90084e3234fb11184886bf4a40f4f16a8c8f867be2e087b81f8e8868d4", size = 256523, upload-time = "2026-08-06T13:48:05.996Z" },
{ url = "https://files.pythonhosted.org/packages/b9/80/d34e13fb4b293cbdb9665838cf5522077b8ad14ef947550631a4bced36a5/coverage-7.15.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:349062d66f00b40fa2c1c222438bad25fabf755631b5d82937fe985c8008615c", size = 257759, upload-time = "2026-08-06T13:48:08.036Z" },
{ url = "https://files.pythonhosted.org/packages/0f/e7/2c5fe7636fdb0732fe0f09f308a5b066864078b7fc61f6678e8478554f2e/coverage-7.15.4-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4256ced708e598e05209bc1a8ab4074e04a51dba4c62fb45926a229af675ace7", size = 259890, upload-time = "2026-08-06T13:48:09.834Z" },
{ url = "https://files.pythonhosted.org/packages/92/28/9689f0858dfff59c2ea688938ab9fa2925631235df67126a42b6c5c70ae1/coverage-7.15.4-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d80f974b20782d9612c8b4c9beeca867074c7cf4079d1419843fa25a26428b25", size = 254121, upload-time = "2026-08-06T13:48:11.459Z" },
{ url = "https://files.pythonhosted.org/packages/f9/e2/785077c230c157243eb5aa9a26c3be260ecd02001bead54a3cada3df8e03/coverage-7.15.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2e179f19bfe1d31f8eeeaa12990194d761c4f62f0759661000bca6cd8729f40b", size = 255891, upload-time = "2026-08-06T13:48:13.209Z" },
{ url = "https://files.pythonhosted.org/packages/d4/90/e20371b17b40f912f21305c2db2f30efa3de306f7320fc916804872c85a4/coverage-7.15.4-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:8bc16bb47b7679670eceff71d78bfb7d6e5b143f6c2cd117487ec7c75e0d4b78", size = 253859, upload-time = "2026-08-06T13:48:14.736Z" },
{ url = "https://files.pythonhosted.org/packages/05/49/25371987ee459a5f67c0427fb75c74f9358e65f2c71fe75bf41c1b6c5fcb/coverage-7.15.4-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:1cd685005cd2c4200adfc14cf39a603b9320efab3f18a8f7f156d20c9cc3345f", size = 258011, upload-time = "2026-08-06T13:48:16.464Z" },
{ url = "https://files.pythonhosted.org/packages/30/6e/32e67467f6154bf4f1c4f63b05acc5097cba4237d45bbeeea446b52e8ac1/coverage-7.15.4-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:337399ad2c93b3acd2a937627dae8b3e86b66707cd3d3e856347999aadf1ef8d", size = 253676, upload-time = "2026-08-06T13:48:18.493Z" },
{ url = "https://files.pythonhosted.org/packages/03/c1/8b24192e89286399765155251f99ee9f070a9d637109018ac23d99b99f6f/coverage-7.15.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:96e257121228ec5cd2bb919276e94ac11074471bc37d68dbae0e8308cce15fff", size = 255453, upload-time = "2026-08-06T13:48:20.057Z" },
{ url = "https://files.pythonhosted.org/packages/16/6f/8b41ebdf67c87854e17c035336a90f1cfbad0c14c2a584301be6ff148718/coverage-7.15.4-cp313-cp313-win32.whl", hash = "sha256:c65a9e0dfc6143491879da4e13b5e30f8be192055de508d737fb14601edbd22c", size = 224605, upload-time = "2026-08-06T13:48:21.655Z" },
{ url = "https://files.pythonhosted.org/packages/e0/e2/2946c7f0b42b152ecb21ff1bdad72e3d301e790c0c487e4a86e8c9f69347/coverage-7.15.4-cp313-cp313-win_amd64.whl", hash = "sha256:2ff8f5e9b8f7a94f0c11c45631eee103dbcb7d63274edd12c56efe1be690b3b4", size = 225148, upload-time = "2026-08-06T13:48:23.376Z" },
{ url = "https://files.pythonhosted.org/packages/9e/83/3f4a69957f48ae7a0aba76c34743f88963d607b19e03f3f8e66f91cae0f9/coverage-7.15.4-cp313-cp313-win_arm64.whl", hash = "sha256:6e0a8a5083b096487d6cfced94cdd514d8f5db6f113610fb36c0620edb1028cf", size = 224536, upload-time = "2026-08-06T13:48:25.117Z" },
{ url = "https://files.pythonhosted.org/packages/ea/ac/748cf29eeb2d6be34a3176ce26a4f49e38085ee08e8935f05f6f26ed7e0f/coverage-7.15.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:770e9325ab5ea6d56f77e59b29ecfe0ac20b57a82a601876f90494a4dda0386f", size = 222608, upload-time = "2026-08-06T13:48:26.806Z" },
{ url = "https://files.pythonhosted.org/packages/0b/02/1abbf5c984677b0aa439cdacaccbf38d248939d8ef8fe1cc7a50d73edb77/coverage-7.15.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d12b33a3a50a1676b7784dc8d00a0c6d66a9f2add4b85a041c19b6a7e53ef23c", size = 222940, upload-time = "2026-08-06T13:48:28.432Z" },
{ url = "https://files.pythonhosted.org/packages/eb/e1/ff8f9f53d9fcf586125b55d0b1f04ec1c14955fee41e83d5814bee141bb5/coverage-7.15.4-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5669c8378ebde86f5def7a25d29586631b58acc27ffde04399f678f3dfc6e082", size = 253985, upload-time = "2026-08-06T13:48:29.995Z" },
{ url = "https://files.pythonhosted.org/packages/a1/26/595759762e514e81be1d7d01ed03444303bcd152226a6529998d253f9201/coverage-7.15.4-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:ff97a14362eef486483ed44042ca2027ea257df6ff768e62358ee0c9776925ac", size = 256492, upload-time = "2026-08-06T13:48:31.634Z" },
{ url = "https://files.pythonhosted.org/packages/24/68/b79aabac54d482be23b5fcdd4f4662bff24a78edc4ee29201726929936d5/coverage-7.15.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5a325e815318638aed1655d9c06e6d7c2d3d46c09231ce988070428a8762d734", size = 257837, upload-time = "2026-08-06T13:48:33.186Z" },
{ url = "https://files.pythonhosted.org/packages/09/0f/bf7f297885a5bf6fd71e5782404e0ff059ca09e8711ceb3a08544abde45a/coverage-7.15.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:474223409d88eb20d2d6a0d37ea60e8647a65a90cc008dc1f0410af5f64f1e0d", size = 260152, upload-time = "2026-08-06T13:48:34.75Z" },
{ url = "https://files.pythonhosted.org/packages/fd/f1/296744e854ff8368542343457414380465e9ceefb9192342feb9d3bc461d/coverage-7.15.4-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7f2f62ae3cd189dd2e13aece758c57b3eecbd27be070dbd4cbd10936049e5dbf", size = 253978, upload-time = "2026-08-06T13:48:36.434Z" },
{ url = "https://files.pythonhosted.org/packages/55/b0/bbdb2e9057493e66220a2e149ca2d301ba0e3a58a83bd6b90de9826d16f3/coverage-7.15.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:39ece820e29e0a2ba34b3ecb3be83c27e997eed8926f2ba6fe7ce7a0bda5843b", size = 255846, upload-time = "2026-08-06T13:48:38.317Z" },
{ url = "https://files.pythonhosted.org/packages/96/e4/38015b2b6d21258713bd17e76b59d033b191efb5703589cffd037dfbca20/coverage-7.15.4-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:f21b56dcace11dfe013014201f577dcd592b2a9b72182d930361b47cf6f73f25", size = 253808, upload-time = "2026-08-06T13:48:39.993Z" },
{ url = "https://files.pythonhosted.org/packages/0b/64/0d515c1e60ee6fbfd1a0e79c07cd87d388a233b7adc37758735677203808/coverage-7.15.4-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:93a3a0b662abcc10c73a47cbc72cd60f63618d6989fb2d1286e50eacd974f303", size = 258081, upload-time = "2026-08-06T13:48:41.971Z" },
{ url = "https://files.pythonhosted.org/packages/91/71/04d9e7a3642146c6351338aef4ef85ab11dbbb54744c13245caba1aad1c0/coverage-7.15.4-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:141fae2cabf5569b782c10afc4c850ce10f618c13f8db54765cba99cc839da1f", size = 253624, upload-time = "2026-08-06T13:48:43.731Z" },
{ url = "https://files.pythonhosted.org/packages/b4/a7/6c28b74c81ebff66987b0e2522ba5cffa3e90b0c33cb6a2eb264d4ee8cf1/coverage-7.15.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:81294c7e6ab30c5f74c0353b11b2fd6320e72d9bee6ac73b357caa8b916323a5", size = 255280, upload-time = "2026-08-06T13:48:45.58Z" },
{ url = "https://files.pythonhosted.org/packages/52/af/bc19996a7014b98d7bbb0f0939453c67074af65784a3aa16a789a07381fa/coverage-7.15.4-cp314-cp314-win32.whl", hash = "sha256:7bbd7d6418e0dab31a206af5203bd43ae36edb8e7fba1940b055d3e9249290d7", size = 224768, upload-time = "2026-08-06T13:48:47.525Z" },
{ url = "https://files.pythonhosted.org/packages/ee/90/219484e476d6e101ba0a444852579e05f5b75c37c611a42ed1190f73ef62/coverage-7.15.4-cp314-cp314-win_amd64.whl", hash = "sha256:f0204ed122758782970526057093f448051a39db9d810d4e344bb87a3546f425", size = 225259, upload-time = "2026-08-06T13:48:49.513Z" },
{ url = "https://files.pythonhosted.org/packages/b7/66/fa77daf4e383e5f776dac62c2409b6af81910ae6fe326bd5170dba74cc63/coverage-7.15.4-cp314-cp314-win_arm64.whl", hash = "sha256:9e71e7bc71c686a123347ae47a0de33a175e797a85bb57b791492adf4eec8ed8", size = 224684, upload-time = "2026-08-06T13:48:51.235Z" },
{ url = "https://files.pythonhosted.org/packages/58/5b/f03bf0ce362bbf3f785fa5219620d00778d4ac6fc9e407734828e9c672f6/coverage-7.15.4-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7c922735321eef3f87c280a3d39afff6b646723a2880b862cda4ac7a093b8aa8", size = 223338, upload-time = "2026-08-06T13:48:52.896Z" },
{ url = "https://files.pythonhosted.org/packages/0f/76/e77d0ae22501831cc9f92193e8a957a5caa1dd177f90a6d1d9b106242d92/coverage-7.15.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f41c17c4668a655ce96d090d8d5ffdc24ef64b5a02f9753884d08483e8a4a41a", size = 223609, upload-time = "2026-08-06T13:48:54.688Z" },
{ url = "https://files.pythonhosted.org/packages/82/1a/b1f089da8d38ac612fa2dd6dc7f4a1a7657d12f3e261d2996edd3a838d0b/coverage-7.15.4-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:46822e9b6ff1c6a72b518c162c44a8f45a61a1d609c51084bf5b16c023c5037b", size = 264970, upload-time = "2026-08-06T13:48:56.403Z" },
{ url = "https://files.pythonhosted.org/packages/bf/31/e66d98d6e9c7fcc88470f1e234eaf6b1950dc0dfbf797f7282c1c861da24/coverage-7.15.4-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3d6f4955b73b5445271379a59e3792b0d978f42d4a01e0cf7a67d9c33a3bb0a5", size = 267088, upload-time = "2026-08-06T13:48:58.41Z" },
{ url = "https://files.pythonhosted.org/packages/59/a1/ae94eb2c541add426378408379f233591e069040b1e2cdb33df9498a0682/coverage-7.15.4-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3fc9e047706fb4a9abb54f719d3aa643e80e5bb3818182c40aee01ac0f0247ba", size = 269508, upload-time = "2026-08-06T13:49:00.42Z" },
{ url = "https://files.pythonhosted.org/packages/9c/c7/88a10694a1c6a213569766aba9f25847b28155d4ac731b13226db216356d/coverage-7.15.4-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:05e491d4f3165d62d4f5c8fd48dfeabf2ae8f42cbbd484319af33ea851b78982", size = 270629, upload-time = "2026-08-06T13:49:02.234Z" },
{ url = "https://files.pythonhosted.org/packages/b3/34/d8b8232e5e55169933b59aabcef2fedfa4b9d8897361bb80fcbda146505f/coverage-7.15.4-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:226c66e80ec0598d3b9b4874123df167ccca342aca8714f77cac6829688ee09c", size = 264043, upload-time = "2026-08-06T13:49:04.102Z" },
{ url = "https://files.pythonhosted.org/packages/7e/35/58b009dbf8c471c7224716478b9fed4a7e1af15320e1ed41660978504663/coverage-7.15.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ac41cc14bebda0dbfb0628036b7f75706935c95bcc07fefe9a0f93614aa60a57", size = 266963, upload-time = "2026-08-06T13:49:05.821Z" },
{ url = "https://files.pythonhosted.org/packages/62/aa/57fbda1b42c892968273c56b6ee9dc0f1310850859230a507bc7873b1f65/coverage-7.15.4-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:8af623e5cd92080acddd02b38f2f406a2c3a0893c38950b211890361448fbf26", size = 264569, upload-time = "2026-08-06T13:49:07.706Z" },
{ url = "https://files.pythonhosted.org/packages/98/8a/360e6e7f24d477b7e889703af0afa878d15b6d4d8d2a822b2835c169a879/coverage-7.15.4-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:07545711d4f0f32852a18f18ad11f76f0109909d09e78b9008b4cfc67e829429", size = 268299, upload-time = "2026-08-06T13:49:09.587Z" },
{ url = "https://files.pythonhosted.org/packages/4e/89/6f701261aee21b6b5fa8f7872229406dc917e125069448292223bf213606/coverage-7.15.4-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:a0865421cfdc53654b342d515e5a233187590882d20b95752150e53f65460017", size = 263413, upload-time = "2026-08-06T13:49:11.604Z" },
{ url = "https://files.pythonhosted.org/packages/3f/0f/6f04036edc260ed425af83e834f627fad48941ce97b50bfe6edd8b6fa623/coverage-7.15.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:460115e32ee40566476db5048f9bec1e842c127ad8e6f8be745aad3ac9cbc839", size = 265725, upload-time = "2026-08-06T13:49:13.38Z" },
{ url = "https://files.pythonhosted.org/packages/c4/ce/d19b5d4d5c49a7bfb925fd74310fee7d28bc99520ac3367ccbc54e662518/coverage-7.15.4-cp314-cp314t-win32.whl", hash = "sha256:cbde877ef9dd7baf272b9bfef2b8a25edd45d9170fc326951dd20eb480335e85", size = 225079, upload-time = "2026-08-06T13:49:15.265Z" },
{ url = "https://files.pythonhosted.org/packages/26/bb/7aa1b3b173faee0679037ca950bbbe1247273656697994d8d13f80f8d4b4/coverage-7.15.4-cp314-cp314t-win_amd64.whl", hash = "sha256:3da9e92d1c551fd7563833e9ade686efb0c4b7363ab7681a94283958c950bf5e", size = 225911, upload-time = "2026-08-06T13:49:17.279Z" },
{ url = "https://files.pythonhosted.org/packages/81/1c/4ea9e47426d80038d9222db3c4534cb6021a74b237d3ff97ffd33b6600dd/coverage-7.15.4-cp314-cp314t-win_arm64.whl", hash = "sha256:3a54f5a0d85050c73a38f6793090ee83974531e67fe5e57a1da9bee11398aa5e", size = 225219, upload-time = "2026-08-06T13:49:19.293Z" },
{ url = "https://files.pythonhosted.org/packages/2b/c4/dc5d2ac8f9142e7ec7de66e7bf0591db29d78955a040bd915870d9c0e657/coverage-7.15.4-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:2c9872e4d9dc5d3cf616bf4b382f5a00359305a5be666a3dd0b5cdb4e49597f9", size = 222604, upload-time = "2026-08-06T13:49:21.279Z" },
{ url = "https://files.pythonhosted.org/packages/70/39/33e63df81fe2ee100897451841c821467635923e58e37c6bd4b46dd8106c/coverage-7.15.4-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:e101dbb4b9b72f0cddd8cdc8c9c5b47f456766f5e0ac82dbfb75e5c55409b78a", size = 222944, upload-time = "2026-08-06T13:49:23.187Z" },
{ url = "https://files.pythonhosted.org/packages/99/1f/ef3ffb5557febc75a0d97aa459d0266d7d741110265121cc6d8539343d44/coverage-7.15.4-cp315-cp315-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:7d1abebdb047729e852b9c77a00497dfbeb11eb3a117e037d7dbc3ac8e5f5c54", size = 254050, upload-time = "2026-08-06T13:49:25.008Z" },
{ url = "https://files.pythonhosted.org/packages/6f/f5/1f0f6f77698c3601ca0ae7431e34b24c62ca2f06fecb23b73ed1f651d2be/coverage-7.15.4-cp315-cp315-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:d28a4a899354d0ea6214cc59b4fa19eefbce1b9ff1688ab579acf49e894bd3fb", size = 256967, upload-time = "2026-08-06T13:49:26.896Z" },
{ url = "https://files.pythonhosted.org/packages/03/7a/2ed9bed79925f4367c83c77f66a89e5ca7229c288d2d19ad5f36d1ca0070/coverage-7.15.4-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ffb3c2aacea411cc7e1d27712490c11108e2de1d39019ae32915493a59a8b9ed", size = 258587, upload-time = "2026-08-06T13:49:28.692Z" },
{ url = "https://files.pythonhosted.org/packages/45/8c/fa34044f71b7cc4ecb6da9c2408770959b0591fa9b5fb6fb6bca38f94298/coverage-7.15.4-cp315-cp315-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a9447978a92f405d301123cfd39ff49895490efb769a758fe2734c7f631bf8ce", size = 260785, upload-time = "2026-08-06T13:49:30.472Z" },
{ url = "https://files.pythonhosted.org/packages/4f/54/d5727ce36b4524a7394ab9f5f1df378e1f23affcdab01037dc8655185cc7/coverage-7.15.4-cp315-cp315-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:050467a7983b8e2fe7dd41a78bb30c3e7f8c0b8cafda14b1c46f8b5e3cf2dd3c", size = 254545, upload-time = "2026-08-06T13:49:32.271Z" },
{ url = "https://files.pythonhosted.org/packages/dc/e6/6e3783e576719590194bdffb6dd6d85490801785b7c331e35a245d8cb8b5/coverage-7.15.4-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:d003b7a5708ddad5c206c79607a6b92abb6fc13c57d99d8a4468cc03a2941ced", size = 256682, upload-time = "2026-08-06T13:49:34.089Z" },
{ url = "https://files.pythonhosted.org/packages/dc/f2/bacdbde18b69ed2de424fcf64d9fb0a4913753d4f0eca8bae9daad69f4bd/coverage-7.15.4-cp315-cp315-musllinux_1_2_i686.whl", hash = "sha256:c38efe30fd74e5c19e9433f11fb1f5dc9c6522770971b7c6145bbaa413dc8800", size = 254560, upload-time = "2026-08-06T13:49:36.052Z" },
{ url = "https://files.pythonhosted.org/packages/6c/a3/1fb927196e3477c1b48831169ab58ba08f451ba87ae311ff1de68b26a616/coverage-7.15.4-cp315-cp315-musllinux_1_2_ppc64le.whl", hash = "sha256:1f4f826d70f772ab8b0c052329580d7fe8b8abd191e4ce0c8f81aec6614665d3", size = 258792, upload-time = "2026-08-06T13:49:38.01Z" },
{ url = "https://files.pythonhosted.org/packages/41/58/30d4c149c69053de0edfe325614c1d28d508f62b1783e0e4a234d2e49136/coverage-7.15.4-cp315-cp315-musllinux_1_2_riscv64.whl", hash = "sha256:4a4bf917c9953f57c957be31c1cd504e3bd2f34d4a352b9d391a3025336f6768", size = 253968, upload-time = "2026-08-06T13:49:39.934Z" },
{ url = "https://files.pythonhosted.org/packages/89/e4/77f639371b918aad30dda4051f95404b43578f7f2e2f87ba73e02ed1ff37/coverage-7.15.4-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:1c9bf40ebef178a45192c75c4964760bb261b0e6ad725da5fc4c93f674f19753", size = 255893, upload-time = "2026-08-06T13:49:41.825Z" },
{ url = "https://files.pythonhosted.org/packages/5c/62/13be29b3ddab35f14c87967a4820a05106d2a3eccb4fa4ff550bf30b75e0/coverage-7.15.4-cp315-cp315-win32.whl", hash = "sha256:43619d04c3671792d2c4706ae8bf45e265dc87bbd4078189ef8b847ea1e74be2", size = 224768, upload-time = "2026-08-06T13:49:44.08Z" },
{ url = "https://files.pythonhosted.org/packages/a1/70/af0c6be0f964af6954f6b74bc109b0dbca02824696d2520fb17fe1ab06e3/coverage-7.15.4-cp315-cp315-win_amd64.whl", hash = "sha256:be619439dbcd31a2eab10b32de9fff62c26ed4bab69dc32b8363fdaaa0882809", size = 225242, upload-time = "2026-08-06T13:49:45.899Z" },
{ url = "https://files.pythonhosted.org/packages/4f/2d/f3bd3aab899fc9efc18b53133ee68f5f98574ef480649b23e12962226387/coverage-7.15.4-cp315-cp315-win_arm64.whl", hash = "sha256:def597967dafc2e8d97c9097ea453c464e0bb8ed38f193a43070f10dc623bb6d", size = 224674, upload-time = "2026-08-06T13:49:48.322Z" },
{ url = "https://files.pythonhosted.org/packages/f5/ca/f69251cd63eabc6438321aea22148754cce758a26bde07dd490e3fe7cfc5/coverage-7.15.4-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:c7dbc748ac8a1e3e59a2b28bea47675e6e778081dbbf081bde0d75def2fcbe1d", size = 223333, upload-time = "2026-08-06T13:49:50.293Z" },
{ url = "https://files.pythonhosted.org/packages/a7/a7/037b53b2885b0d8447064432491a4d5a1014cd9f97a594d53acd0c04541a/coverage-7.15.4-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:2413074a5ecbb61a01a7888fc72db0ca324d13588c5b38bc0dd8564cdcdfea26", size = 223630, upload-time = "2026-08-06T13:49:52.637Z" },
{ url = "https://files.pythonhosted.org/packages/80/4f/152b8a4779ae90da11bb24f7467df8a59f0be48a5c52acb856325ca48289/coverage-7.15.4-cp315-cp315t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4e6f6f632b7b2f714bf7a1346e8f97b650ee71f3c298aaad42a2ab60f0f07645", size = 264489, upload-time = "2026-08-06T13:49:54.52Z" },
{ url = "https://files.pythonhosted.org/packages/10/2d/84b4b9e0e1dd6528a51920ff7031f35b789382e467a28ec6a5a578cb8812/coverage-7.15.4-cp315-cp315t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8df457da2249d3c75ca2e5e835d59c725abfe92d27fdff6cd99eed85b51d5e9a", size = 267567, upload-time = "2026-08-06T13:49:56.721Z" },
{ url = "https://files.pythonhosted.org/packages/53/fc/ba01cc25299f9f8a2c8b02d3b28c53f3543d9fbfbe4e74fa2760b48f163e/coverage-7.15.4-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:050f66a08805acb5b8a23c6d4a517b1ecf82c08e81ed0e4bd727df065e5c6624", size = 270123, upload-time = "2026-08-06T13:49:58.736Z" },
{ url = "https://files.pythonhosted.org/packages/cf/d0/db2647cbf40b14f8c308f94ff7bf89c06d564e59f396906edf50086ec788/coverage-7.15.4-cp315-cp315t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1587fb771d1ccceef708fdde1e5af8c7ed24b486b61d13a321acb7d8145390aa", size = 271107, upload-time = "2026-08-06T13:50:00.811Z" },
{ url = "https://files.pythonhosted.org/packages/70/ff/4d2d17924552c458bb4f77dd631f0e3bc92fbbdf2d2d916cd4b33bbfd5b1/coverage-7.15.4-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:8b4f1c3a69ca580f3fbd6b2046915f536d7f586874f25c1bb23add2a3c88d50f", size = 264955, upload-time = "2026-08-06T13:50:03.023Z" },
{ url = "https://files.pythonhosted.org/packages/ee/de/dc010c7a3691f396d93bbc26bfcafa1c2a3a351cd520470f15faf5795bd5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:ffb58d7eff5b7f6ecc6fa21d6288ab7f968a212cb67d682c269c09b9eba3b66f", size = 267949, upload-time = "2026-08-06T13:50:05.557Z" },
{ url = "https://files.pythonhosted.org/packages/78/ea/dc96a11375e83c045c2f7c61fb6918277cfe9401db7c0f7b1d111a84b2e5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_i686.whl", hash = "sha256:d9df165544774574ee004b953023d1bebada1894a80b1052a43d798b0f676e67", size = 264421, upload-time = "2026-08-06T13:50:07.612Z" },
{ url = "https://files.pythonhosted.org/packages/c8/86/b77131a0f9503ce461cd577076147d7a9040f0c5dda772686f729e2cc9cb/coverage-7.15.4-cp315-cp315t-musllinux_1_2_ppc64le.whl", hash = "sha256:f9de0a24a4079b53e523b5c5e2c5945ec251ab486652659955187cf255a259bc", size = 269121, upload-time = "2026-08-06T13:50:09.58Z" },
{ url = "https://files.pythonhosted.org/packages/24/24/944bc35007862955e7ebf05754e645419dcf5d7526c52735cfa2715e8ebf/coverage-7.15.4-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:150089274bdc9f940628552cb92844e0223c987f1902ab8efe9f45a2ec758d88", size = 264565, upload-time = "2026-08-06T13:50:11.722Z" },
{ url = "https://files.pythonhosted.org/packages/c7/cc/a3bb9f93e7e740659163e2ea584f8196ddcd2c456a5dbe15f6c50105fec1/coverage-7.15.4-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:a58a94fed5da6997d258e8f7668c1e195fbd04a691d781b7558f1e468f9e68bc", size = 266522, upload-time = "2026-08-06T13:50:13.786Z" },
{ url = "https://files.pythonhosted.org/packages/49/dd/e0e40f3560d878d888c580698ff5ad1179f5e1c3ac949684ef66b41a3817/coverage-7.15.4-cp315-cp315t-win32.whl", hash = "sha256:ebd5a6d8466ff30836572f3ba2cae8a5e8f85029b1c6d5e2ed338dc472a5166a", size = 225068, upload-time = "2026-08-06T13:50:15.825Z" },
{ url = "https://files.pythonhosted.org/packages/c6/7e/37732ea80eebc30e976e4cdab15c190bc42d96959a42e38ddf6f8c60468f/coverage-7.15.4-cp315-cp315t-win_amd64.whl", hash = "sha256:288bde2a2d7ab6b6c2d7252fcde8b524387f2d970bdba9658fc6f8bbcaef0f9b", size = 225895, upload-time = "2026-08-06T13:50:17.928Z" },
{ url = "https://files.pythonhosted.org/packages/c6/08/1e00f7923eaaba45fb3d51dd794125fc766304b1df264f3a9c6557bfb30e/coverage-7.15.4-cp315-cp315t-win_arm64.whl", hash = "sha256:68be5e1de60ff13c9095bbec0e5a7fa45b33b101752215b91345ea1f61c4a278", size = 225213, upload-time = "2026-08-06T13:50:19.981Z" },
{ url = "https://files.pythonhosted.org/packages/b4/d9/e70c286c979378f061d8266e279b686ab0b0b688e1fe0af864684f23a77d/coverage-7.15.4-py3-none-any.whl", hash = "sha256:964730a1e9de9c0cf11be6a1a3c79ce419c34882842abd256086ba4698705e84", size = 214332, upload-time = "2026-08-06T13:50:22.192Z" },
]
[[package]] [[package]]
name = "csvw" name = "csvw"
version = "4.0.0" version = "4.0.0"
@@ -1722,6 +1803,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" },
] ]
[[package]]
name = "pytest-cov"
version = "7.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "coverage" },
{ name = "pluggy" },
{ name = "pytest" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b1/51/a849f96e117386044471c8ec2bd6cfebacda285da9525c9106aeb28da671/pytest_cov-7.1.0.tar.gz", hash = "sha256:30674f2b5f6351aa09702a9c8c364f6a01c27aae0c1366ae8016160d1efc56b2", size = 55592, upload-time = "2026-03-21T20:11:16.284Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" },
]
[[package]] [[package]]
name = "pytest-env" name = "pytest-env"
version = "1.5.0" version = "1.5.0"
@@ -2471,53 +2566,48 @@ dependencies = [
{ name = "kokoro-tts" }, { name = "kokoro-tts" },
{ name = "langchain-core" }, { name = "langchain-core" },
{ name = "langchain-openai" }, { name = "langchain-openai" },
{ name = "mypy" },
{ name = "numpy" }, { name = "numpy" },
{ name = "openai" }, { name = "openai" },
{ name = "pytest" },
{ name = "pytest-env" },
{ name = "python-dotenv" }, { name = "python-dotenv" },
{ name = "requests" }, { name = "requests" },
{ name = "types-requests" },
]
[package.optional-dependencies]
dev = [
{ name = "black" },
{ name = "debugpy" },
{ name = "mypy" },
] ]
[package.dev-dependencies] [package.dev-dependencies]
dev = [ dev = [
{ name = "black" },
{ name = "debugpy" },
{ name = "mypy" },
{ name = "pyright" }, { name = "pyright" },
{ name = "pytest" },
{ name = "pytest-cov" },
{ name = "pytest-env" },
{ name = "ruff" }, { name = "ruff" },
{ name = "types-requests" },
] ]
[package.metadata] [package.metadata]
requires-dist = [ requires-dist = [
{ name = "black", marker = "extra == 'dev'", specifier = ">=25.1.0" },
{ name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.8.0" },
{ name = "discord", specifier = ">=2.3.2" }, { name = "discord", specifier = ">=2.3.2" },
{ name = "kokoro-tts", specifier = ">=2.3.1" }, { name = "kokoro-tts", specifier = ">=2.3.1" },
{ name = "langchain-core", specifier = ">=0.3.0" }, { name = "langchain-core", specifier = ">=0.3.0" },
{ name = "langchain-openai", specifier = ">=0.4.0" }, { name = "langchain-openai", specifier = ">=0.4.0" },
{ name = "mypy", specifier = ">=2.1.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 = "pytest", specifier = ">=9.0.2" },
{ name = "pytest-env", specifier = ">=1.5.0" },
{ name = "python-dotenv", specifier = ">=1.2.2" }, { name = "python-dotenv", specifier = ">=1.2.2" },
{ name = "requests", specifier = ">=2.32.5" }, { name = "requests", specifier = ">=2.32.5" },
{ name = "types-requests", specifier = ">=2.32.4.20260107" },
] ]
provides-extras = ["dev"]
[package.metadata.requires-dev] [package.metadata.requires-dev]
dev = [ dev = [
{ name = "black", specifier = ">=25.1.0" },
{ name = "debugpy", specifier = ">=1.8.0" },
{ name = "mypy", specifier = ">=2.1.0" },
{ name = "pyright", specifier = ">=1.1.409" }, { name = "pyright", specifier = ">=1.1.409" },
{ name = "pytest", specifier = ">=9.0.2" },
{ name = "pytest-cov", specifier = ">=7.0.0" },
{ name = "pytest-env", specifier = ">=1.5.0" },
{ name = "ruff", specifier = ">=0.16.3" }, { name = "ruff", specifier = ">=0.16.3" },
{ name = "types-requests", specifier = ">=2.32.4.20260107" },
] ]
[[package]] [[package]]
vibe_bot/app.py
+203
View File
@@ -0,0 +1,203 @@
"""Composition root: owns every singleton, the four services, and the bot."""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
import discord
from discord import Message
from discord.ext import commands
from vibe_bot import llm_client
from vibe_bot.config import TTS_MODEL_PATH, TTS_VOICES_PATH
from vibe_bot.database import ChatDatabase, CustomBotManager
from vibe_bot.llm_client import ToolRegistry
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.conversation_service import ConversationService
from vibe_bot.services.image_service import ImageService
from vibe_bot.services.speech_service import SpeechService
from vibe_bot.tts import TTSEngine
if TYPE_CHECKING:
from discord.ext.commands import Bot
from discord.ext.commands import Context as CommandsContext
logger = logging.getLogger(__name__)
def configure_logging() -> None:
"""Configure root logging (the single ``basicConfig`` in the codebase)."""
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
@dataclass
class App:
"""The composition root: every singleton and the four services."""
db: ChatDatabase
manager: CustomBotManager
registry: ToolRegistry
tts: TTSEngine | None
chat: ChatService
image: ImageService
speech: SpeechService
conversation: ConversationService
bot_cache: dict[str, tuple[str, str]]
def _build_bot_cache(manager: CustomBotManager) -> dict[str, tuple[str, str]]:
"""Snapshot the custom bots as name -> (system_prompt, creator)."""
return {
name: (system_prompt, creator)
for name, system_prompt, creator in manager.list_custom_bots()
}
def invalidate_bot_cache(app: App) -> None:
"""Rebuild ``app.bot_cache`` from the database after a change."""
app.bot_cache = _build_bot_cache(app.manager)
def create_app() -> App:
"""Build the App: singletons, tolerant TTS init, and the four services."""
db = ChatDatabase()
manager = CustomBotManager()
registry = llm_client.get_tool_registry()
engine: TTSEngine | None = None
try:
engine = TTSEngine(TTS_MODEL_PATH, TTS_VOICES_PATH)
logger.info("TTS engine initialized successfully")
except Exception:
logger.exception("Failed to initialize TTS engine")
logger.info(
"Make sure kokoro-v1.0.onnx and voices-v1.0.bin are in the project directory",
)
return App(
db=db,
manager=manager,
registry=registry,
tts=engine,
chat=ChatService(db, registry),
image=ImageService(db, discord.File),
speech=SpeechService(db, manager, engine, discord.File),
conversation=ConversationService(manager),
bot_cache=_build_bot_cache(manager),
)
# Module-level holders wired by build_bot(); the event handlers below are
# module-level so they stay importable and testable.
_app: App | None = None
_bot: commands.Bot | None = None
def _require_app() -> App:
"""The App wired by build_bot(); handlers must not run before it."""
if _app is None:
msg = "App is not initialized; build_bot(app) must be called first."
raise RuntimeError(msg)
return _app
def _require_bot() -> commands.Bot:
"""The Bot wired by build_bot(); event handlers must not run before it."""
if _bot is None:
msg = "Bot is not initialized; build_bot(app) must be called first."
raise RuntimeError(msg)
return _bot
async def on_ready() -> None:
"""Log when the bot is ready and logged in."""
bot = _require_bot()
logger.info("Bot is starting up...")
logger.info("Bot logged in as %s", bot.user)
async def on_message(message: Message) -> None:
"""Handle incoming messages for custom bot command detection."""
app = _require_app()
bot = _require_bot()
if message.author == bot.user:
return
if not message.content.startswith("!"):
return
message_content = message.content.lower()
logger.debug(
"Processing message from user %s (chars=%d)",
message.author.id,
len(message_content),
)
for bot_name, (system_prompt, _creator) in app.bot_cache.items():
if message_content.startswith(f"!{bot_name} "):
logger.info(
"Custom bot %r triggered by user %s", bot_name, message.author.id
)
user_message = message.content[len(f"!{bot_name} ") :]
logger.debug(
"Extracted user message for bot %r (chars=%d)",
bot_name,
len(user_message),
)
response_prefix = f"{bot_name} response"
logger.info("Sending request to chat service for bot %r", bot_name)
ctx = await bot.get_context(message)
await app.chat.handle(
ctx=ctx,
bot_name=bot_name,
message=user_message,
system_prompt=system_prompt,
response_prefix=response_prefix,
)
return
# If no custom bot matched, call the default event handler
await bot.process_commands(message)
async def on_command_error(
ctx: CommandsContext[Bot],
error: commands.CommandError,
) -> None:
"""Send a friendly message when a command hits its per-user cooldown."""
if isinstance(error, commands.CommandOnCooldown):
retry_after = max(1, round(error.retry_after))
await ctx.send(f"You're using that too quickly, try again in {retry_after}s.")
return
logger.exception("Unhandled command error in %s: %s", ctx.command, error)
def build_bot(app: App) -> commands.Bot:
"""Create the Discord bot with all event and command handlers attached."""
global _app, _bot
_app = app
intents = discord.Intents.default()
intents.message_content = True
intents.members = True
intents.presences = True
bot = commands.Bot(command_prefix="!", intents=intents)
_bot = bot
bot.event(on_ready)
bot.event(on_message)
bot.event(on_command_error)
# Imported here (not at module top) to break the import cycle.
from vibe_bot.commands import register_all
register_all(bot, app)
return bot
vibe_bot/commands/__init__.py
+22
View File
@@ -0,0 +1,22 @@
"""Discord command groups; ``register_all`` wires every group onto the bot."""
from __future__ import annotations
from typing import TYPE_CHECKING
from vibe_bot.commands import admin, chat, conversation, custom_bots, images, speech
if TYPE_CHECKING:
from discord.ext import commands
from vibe_bot.app import App
def register_all(bot: commands.Bot, app: App) -> None:
"""Register every command group on the bot."""
custom_bots.register(bot, app)
speech.register(bot, app)
images.register(bot, app)
conversation.register(bot, app)
admin.register(bot, app)
chat.register(bot, app)
vibe_bot/commands/_state.py
+24
View File
@@ -0,0 +1,24 @@
"""Shared App holder for the command modules, wired by each group's register."""
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from vibe_bot.app import App
_app: App | None = None
def set_app(app: App) -> None:
"""Store the App so command handlers can reach the services."""
global _app
_app = app
def require_app() -> App:
"""The App wired by build_bot(); handlers must not run before it."""
if _app is None:
msg = "App is not initialized; build_bot(app) must be called first."
raise RuntimeError(msg)
return _app
vibe_bot/commands/admin.py
+135
View File
@@ -0,0 +1,135 @@
"""Admin and debug commands: wipe history, debug menu, chat history."""
from __future__ import annotations
import asyncio
import logging
from typing import TYPE_CHECKING
from discord.ext import commands
from vibe_bot.commands._state import require_app, set_app
from vibe_bot.prompts import get_user_info
from vibe_bot.textutil import split_message
from vibe_bot.tools import get_channel_members, get_channel_members_impl
if TYPE_CHECKING:
from discord.ext.commands import Bot
from discord.ext.commands import Context as CommandsContext
from vibe_bot.app import App
logger = logging.getLogger(__name__)
def register(bot: commands.Bot, app: App) -> None:
"""Register the admin and debug commands."""
set_app(app)
bot.command(name="lobotomize")(lobotomize)
bot.command(name="debug")(debug)
bot.command(name="history")(history)
async def lobotomize(ctx: CommandsContext[Bot]) -> None:
"""Clear all conversation history and memory for all bots."""
app = require_app()
logger.info("Lobotomize command triggered by user %s", ctx.author.id)
await asyncio.to_thread(app.db.clear_all_messages)
await ctx.send("All conversation history and memory has been cleared. 🧠✨")
async def debug(
ctx: CommandsContext[Bot],
*,
subcommand: str | None = None,
) -> None:
"""Debug menu for various debugging sub-commands.
Usage: !debug <subcommand>
Available sub-commands:
- members: List all members in the current channel
- whoami: Show all information the bot has about you
- tools: Show the LLM's available tools
"""
logger.info(
"Debug command triggered by user %s with subcommand %r",
ctx.author.id,
subcommand,
)
if not subcommand:
menu = "Debug Menu:\n\n"
menu += "Available sub-commands:\n"
menu += "- `members` - List all members in the current channel\n"
menu += "- `whoami` - Show all information the bot has about you\n"
menu += "- `tools` - Show the LLM's available tools"
await ctx.send(menu)
return
if subcommand == "members":
result = get_channel_members_impl(ctx.channel)
for chunk in split_message(result):
await ctx.send(chunk)
return
if subcommand == "whoami":
user_info = get_user_info(ctx.author)
for chunk in split_message(user_info):
await ctx.send(chunk)
return
if subcommand == "tools":
tool_list = "LLM Tools:\n\n"
tool_list += f"- `{get_channel_members.name}`\n"
tool_list += f" Description: {get_channel_members.description}\n"
args_schema = get_channel_members.args_schema
if isinstance(args_schema, type):
tool_list += f" Parameters: {args_schema.model_json_schema()}"
for chunk in split_message(tool_list):
await ctx.send(chunk)
return
await ctx.send(
f"Unknown debug sub-command: `{subcommand}`\n\n"
f"Use `!debug` to see available sub-commands.",
)
async def history(ctx: CommandsContext[Bot], bot_name: str) -> None:
"""View the chat history of a custom bot.
Usage: !history <bot_name>
"""
app = require_app()
logger.info(
"History command triggered by user %s for bot %r",
ctx.author.id,
bot_name,
)
bot_info = await asyncio.to_thread(app.manager.get_custom_bot, bot_name)
if not bot_info:
await ctx.send(f"Custom bot '{bot_name}' not found.")
return
history = await asyncio.to_thread(
app.db.get_bot_history, bot_name=bot_name, limit=20
)
if not history:
await ctx.send(f"No chat history found for **{bot_name}**. ")
return
history.reverse()
formatted_history: list[str] = []
for user_msg, bot_resp in history:
formatted_history.append(user_msg)
formatted_history.append(f"{bot_name}: {bot_resp}")
header = f"Chat History for **{bot_name}**:\n\n"
full_text = header + "\n---\n".join(formatted_history)
for chunk in split_message(full_text):
await ctx.send(chunk)
vibe_bot/commands/chat.py
+19
View File
@@ -0,0 +1,19 @@
"""Custom-bot chat has no command of its own, so this group registers nothing.
A message like ``!alfred hello`` is matched against the bot-name cache in
``vibe_bot.app.on_message`` and dispatched to ``ChatService.handle``; ordinary
``!`` messages fall through to the commands registered by the other groups.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from discord.ext import commands
from vibe_bot.app import App
def register(bot: commands.Bot, app: App) -> None:
"""Register no commands: custom-bot chat flows through ``on_message``."""
vibe_bot/commands/conversation.py
+38
View File
@@ -0,0 +1,38 @@
"""Bot-vs-bot conversation commands."""
from __future__ import annotations
from typing import TYPE_CHECKING
from discord.ext import commands
from vibe_bot.commands._state import require_app, set_app
if TYPE_CHECKING:
from discord.ext.commands import Bot
from discord.ext.commands import Context as CommandsContext
from vibe_bot.app import App
MIN_TALKFORME_PARTS = 4
def register(bot: commands.Bot, app: App) -> None:
"""Register the conversation commands."""
set_app(app)
bot.command(name="talkforme")(talkforme)
@commands.cooldown(rate=1, per=30, type=commands.BucketType.user)
async def talkforme(ctx: CommandsContext[Bot], *, message: str) -> None:
"""Have two bots talk to each other about a topic.
Usage: !talkforme bot1 bot2 4 some conversation topic
"""
app = require_app()
parts = message.split(" ", maxsplit=MIN_TALKFORME_PARTS - 1)
if len(parts) < MIN_TALKFORME_PARTS:
await ctx.send("Usage: !talkforme bot1 bot2 <number> <topic>")
return
await app.conversation.run(ctx, parts[0], parts[1], parts[2], " ".join(parts[3:]))
vibe_bot/commands/custom_bots.py
+241
View File
@@ -0,0 +1,241 @@
"""Custom-bot management commands: create, list, and delete."""
from __future__ import annotations
import asyncio
import logging
from typing import TYPE_CHECKING
from discord.ext import commands
from vibe_bot.app import invalidate_bot_cache
from vibe_bot.commands._state import require_app, set_app
if TYPE_CHECKING:
from discord.ext.commands import Bot
from discord.ext.commands import Context as CommandsContext
from vibe_bot.app import App
logger = logging.getLogger(__name__)
MIN_BOT_NAME_LENGTH = 2
MAX_BOT_NAME_LENGTH = 50
MIN_PERSONALITY_LENGTH = 10
MAX_PERSONALITY_LENGTH = 1000
def register(bot: commands.Bot, app: App) -> None:
"""Register the custom-bot management commands."""
set_app(app)
bot.command(name="custom-bot")(custom_bot)
bot.command(name="list-custom-bots")(list_custom_bots)
bot.command(name="delete-custom-bot")(delete_custom_bot)
async def custom_bot(
ctx: CommandsContext[Bot],
bot_name: str,
*,
personality: str,
) -> None:
"""Create a custom bot with a name and personality.
Usage: !custom-bot <bot_name> <personality_description>
Example: !custom-bot alfred you are a proper british butler
"""
app = require_app()
logger.info(
"Custom bot command initiated by user %s: name=%r, personality_chars=%d",
ctx.author.id,
bot_name,
len(personality),
)
# Validate bot name
name_length = 0 if not bot_name else len(bot_name)
if (
not bot_name
or name_length < MIN_BOT_NAME_LENGTH
or name_length > MAX_BOT_NAME_LENGTH
):
logger.warning(
"Invalid bot name from user %s: %r (length: %d)",
ctx.author.id,
bot_name,
name_length,
)
await ctx.send("Invalid bot name. Name must be between 2 and 50 characters.")
return
logger.debug("Bot name validation passed for %r", bot_name)
# Validate personality
personality_length = 0 if not personality else len(personality)
if not personality or personality_length < MIN_PERSONALITY_LENGTH:
logger.warning(
"Invalid personality from user %s: length=%d",
ctx.author.id,
personality_length,
)
await ctx.send(
"Invalid personality. Description must be at least 10 characters.",
)
return
if personality_length > MAX_PERSONALITY_LENGTH:
logger.warning(
"Personality too long from user %s: length=%d",
ctx.author.id,
personality_length,
)
await ctx.send(
f"Personality too long. " f"Max {MAX_PERSONALITY_LENGTH} characters.",
)
return
logger.debug("Personality validation passed for bot %r", bot_name)
# Create the custom bot
logger.debug(
"Attempting to create custom bot %r for user %s",
bot_name,
ctx.author.id,
)
result = await asyncio.to_thread(
app.manager.create_custom_bot,
bot_name=bot_name,
system_prompt=personality,
created_by=str(ctx.author.id),
)
if result is False:
logger.warning(
"Failed to create custom bot %r for user %s",
bot_name,
ctx.author.id,
)
await ctx.send("Failed to create custom bot.")
return
await asyncio.to_thread(invalidate_bot_cache, app)
if result == "replaced":
logger.info(
"Replaced existing custom bot %r for user %s",
bot_name,
ctx.author.id,
)
await ctx.send(
f"Custom bot **'{bot_name}'** already existed and has been "
f"**replaced** with personality: *{personality}*",
)
else:
logger.info(
"Successfully created custom bot %r for user %s",
bot_name,
ctx.author.id,
)
await ctx.send(
f"Custom bot **'{bot_name}'** has been created "
f"with personality: *{personality}*",
)
await ctx.send(
f"\nYou can now use this bot with: " f"`!{bot_name} <your message>`",
)
async def list_custom_bots(ctx: CommandsContext[Bot]) -> None:
"""List all custom bots available in the server."""
app = require_app()
logger.info("Listing custom bots requested by user %s", ctx.author.id)
logger.debug("Fetching list of custom bots from database")
bots = await asyncio.to_thread(app.manager.list_custom_bots)
if not bots:
logger.debug("No custom bots found for user %s", ctx.author.id)
await ctx.send(
"No custom bots have been created yet. "
"Use `!custom-bot <name> <personality>` to create one.",
)
return
logger.debug(
"Found %d custom bots, displaying top 10 for user %s",
len(bots),
ctx.author.id,
)
bot_list = "Available Custom Bots:\n\n"
for name, _prompt, _creator in bots:
bot_list += f"* {name}\n"
logger.debug("Sending bot list response to user %s", ctx.author.id)
await ctx.send(bot_list)
async def delete_custom_bot(ctx: CommandsContext[Bot], bot_name: str) -> None:
"""Delete a custom bot (only the creator can delete).
Usage: !delete-custom-bot <bot_name>
"""
app = require_app()
logger.info(
"Delete custom bot command initiated by user %s: bot_name=%r",
ctx.author.id,
bot_name,
)
# Get bot info
logger.debug("Looking up custom bot %r in database", bot_name)
bot_info = await asyncio.to_thread(app.manager.get_custom_bot, bot_name)
if not bot_info:
logger.warning(
"Custom bot %r not found by user %s",
bot_name,
ctx.author.id,
)
await ctx.send(f"Custom bot '{bot_name}' not found.")
return
logger.debug(
"Custom bot %r found, owned by user %s",
bot_name,
bot_info[2],
)
# Check ownership
if bot_info[2] != str(ctx.author.id):
logger.warning(
"User %s attempted to delete bot %r they don't own",
ctx.author.id,
bot_name,
)
await ctx.send("You can only delete your own custom bots.")
return
logger.debug(
"User %s is authorized to delete bot %r",
ctx.author.id,
bot_name,
)
# Delete the bot
logger.debug("Deleting custom bot %r from database", bot_name)
success = await asyncio.to_thread(app.manager.delete_custom_bot, bot_name)
if success:
logger.info(
"Successfully deleted custom bot %r by user %s",
bot_name,
ctx.author.id,
)
await asyncio.to_thread(invalidate_bot_cache, app)
await ctx.send(f"Custom bot '{bot_name}' has been deleted.")
else:
logger.warning(
"Failed to delete custom bot %r by user %s",
bot_name,
ctx.author.id,
)
await ctx.send("Failed to delete custom bot.")
vibe_bot/commands/images.py
+35
View File
@@ -0,0 +1,35 @@
"""Image commands: generation and editing."""
from __future__ import annotations
from typing import TYPE_CHECKING
from discord.ext import commands
from vibe_bot.commands._state import require_app, set_app
if TYPE_CHECKING:
from discord.ext.commands import Bot
from discord.ext.commands import Context as CommandsContext
from vibe_bot.app import App
def register(bot: commands.Bot, app: App) -> None:
"""Register the image commands."""
set_app(app)
bot.command(name="doodlebob")(doodlebob)
bot.command(name="retcon")(retcon)
@commands.cooldown(rate=1, per=60, type=commands.BucketType.user)
async def doodlebob(ctx: CommandsContext[Bot], *, message: str) -> None:
"""Convert a message into an image using Doodlebob."""
app = require_app()
await app.image.generate(ctx, message=message)
async def retcon(ctx: CommandsContext[Bot], *, message: str) -> None:
"""Edit an attached image based on a text prompt."""
app = require_app()
await app.image.edit(ctx, message=message)
vibe_bot/commands/speech.py
+56
View File
@@ -0,0 +1,56 @@
"""Speech commands: text-to-speech and the voice catalog."""
from __future__ import annotations
from typing import TYPE_CHECKING
from discord.ext import commands
from vibe_bot.commands._state import require_app, set_app
from vibe_bot.config import VOICES_LIST
from vibe_bot.textutil import split_message
if TYPE_CHECKING:
from discord.ext.commands import Bot
from discord.ext.commands import Context as CommandsContext
from vibe_bot.app import App
def register(bot: commands.Bot, app: App) -> None:
"""Register the speech commands."""
set_app(app)
bot.command(name="speak")(speak)
bot.command(name="voices")(voices)
async def voices(ctx: CommandsContext[Bot]) -> None:
"""List all available TTS voices organized by category."""
voice_list = "Available Voices:\n\n"
for category, info in VOICES_LIST.items():
voice_list += f"{category} ({info['language']}):\n"
for v in info["voices"]:
voice_list += f"- {v}\n"
voice_list += "\n"
voice_list += "Use `!speak <text> --voice <voice_name>` to choose a voice."
for chunk in split_message(voice_list):
await ctx.send(chunk)
@commands.cooldown(rate=3, per=30, type=commands.BucketType.user)
async def speak(
ctx: CommandsContext[Bot],
*,
message: str,
) -> None:
"""Have the bot speak the given text using Kokoro TTS, or have a custom bot speak.
Usage: !speak <text> --voice <voice_name> - plain text to speech
Usage: !speak <bot_name> <text> --voice <voice_name> - have a custom bot respond and speak
Example: !speak hello world
Example: !speak hello world --voice af_bella
Example: !speak alfred what time is it --voice am_puck
"""
app = require_app()
await app.speech.speak(ctx, message=message)
vibe_bot/config.py
+33 -52
View File
@@ -7,11 +7,6 @@ import os
from dotenv import load_dotenv from dotenv import load_dotenv
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
load_dotenv() load_dotenv()
@@ -21,7 +16,6 @@ DISCORD_TOKEN: str = os.getenv("DISCORD_TOKEN", "")
# Endpoints # Endpoints
CHAT_ENDPOINT: str = os.getenv("CHAT_ENDPOINT", "") CHAT_ENDPOINT: str = os.getenv("CHAT_ENDPOINT", "")
COMPLETION_ENDPOINT: str = os.getenv("COMPLETION_ENDPOINT", "")
IMAGE_GEN_ENDPOINT: str = os.getenv("IMAGE_GEN_ENDPOINT", "") IMAGE_GEN_ENDPOINT: str = os.getenv("IMAGE_GEN_ENDPOINT", "")
IMAGE_EDIT_ENDPOINT: str = os.getenv("IMAGE_EDIT_ENDPOINT", "") IMAGE_EDIT_ENDPOINT: str = os.getenv("IMAGE_EDIT_ENDPOINT", "")
EMBEDDING_ENDPOINT: str = os.getenv("EMBEDDING_ENDPOINT", "") EMBEDDING_ENDPOINT: str = os.getenv("EMBEDDING_ENDPOINT", "")
@@ -29,14 +23,12 @@ MAX_COMPLETION_TOKENS: int = int(os.getenv("MAX_COMPLETION_TOKENS", "1000"))
# API Keys # API Keys
CHAT_ENDPOINT_KEY: str = os.getenv("CHAT_ENDPOINT_KEY", "placeholder") CHAT_ENDPOINT_KEY: str = os.getenv("CHAT_ENDPOINT_KEY", "placeholder")
COMPLETION_ENDPOINT_KEY: str = os.getenv("COMPLETION_ENDPOINT_KEY", "placeholder")
IMAGE_GEN_ENDPOINT_KEY: str = os.getenv("IMAGE_GEN_ENDPOINT_KEY", "placeholder") IMAGE_GEN_ENDPOINT_KEY: str = os.getenv("IMAGE_GEN_ENDPOINT_KEY", "placeholder")
IMAGE_EDIT_ENDPOINT_KEY: str = os.getenv("IMAGE_EDIT_ENDPOINT_KEY", "placeholder") IMAGE_EDIT_ENDPOINT_KEY: str = os.getenv("IMAGE_EDIT_ENDPOINT_KEY", "placeholder")
EMBEDDING_ENDPOINT_KEY: str = os.getenv("EMBEDDING_ENDPOINT_KEY", "placeholder") EMBEDDING_ENDPOINT_KEY: str = os.getenv("EMBEDDING_ENDPOINT_KEY", "placeholder")
# Models # Models
CHAT_MODEL: str = os.getenv("CHAT_MODEL", "") 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_MODEL: str = os.getenv("IMAGE_GEN_MODEL", "")
IMAGE_GEN_SIZE_SQUARE: str = os.getenv("IMAGE_GEN_SIZE_SQUARE", "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_PORTRAIT: str = os.getenv("IMAGE_GEN_SIZE_PORTRAIT", "1024x1536")
@@ -44,59 +36,54 @@ 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", "")
# Database and embeddings # Database and history
DB_PATH: str = os.getenv("DB_PATH", "chat_history.db") DB_PATH: str = os.getenv("DB_PATH", "chat_history.db")
EMBEDDING_DIMENSION: int = 2048
MAX_HISTORY_MESSAGES: int = int(os.getenv("MAX_HISTORY_MESSAGES", "1000")) MAX_HISTORY_MESSAGES: int = int(os.getenv("MAX_HISTORY_MESSAGES", "1000"))
SIMILARITY_THRESHOLD: float = float(os.getenv("SIMILARITY_THRESHOLD", "0.7")) SIMILARITY_THRESHOLD: float = float(os.getenv("SIMILARITY_THRESHOLD", "0.7"))
TOP_K_RESULTS: int = int(os.getenv("TOP_K_RESULTS", "5")) TOP_K_RESULTS: int = int(os.getenv("TOP_K_RESULTS", "5"))
# Check token
if not DISCORD_TOKEN:
msg = "DISCORD_TOKEN required."
raise RuntimeError(msg)
# Check endpoints def validate_config() -> None:
if not CHAT_ENDPOINT: """Raise ``RuntimeError`` if any required setting is missing."""
endpoint_msg = "CHAT_ENDPOINT required." # Check token
raise RuntimeError(endpoint_msg) if not DISCORD_TOKEN:
msg = "DISCORD_TOKEN required."
raise RuntimeError(msg)
if not COMPLETION_ENDPOINT: # Check endpoints
endpoint_msg = "COMPLETION_ENDPOINT required." if not CHAT_ENDPOINT:
raise RuntimeError(endpoint_msg) endpoint_msg = "CHAT_ENDPOINT required."
raise RuntimeError(endpoint_msg)
if not IMAGE_GEN_ENDPOINT: if not IMAGE_GEN_ENDPOINT:
endpoint_msg = "IMAGE_GEN_ENDPOINT required." endpoint_msg = "IMAGE_GEN_ENDPOINT required."
raise RuntimeError(endpoint_msg) raise RuntimeError(endpoint_msg)
if not IMAGE_EDIT_ENDPOINT: if not IMAGE_EDIT_ENDPOINT:
endpoint_msg = "IMAGE_EDIT_ENDPOINT required." endpoint_msg = "IMAGE_EDIT_ENDPOINT required."
raise RuntimeError(endpoint_msg) raise RuntimeError(endpoint_msg)
if not EMBEDDING_ENDPOINT: if not EMBEDDING_ENDPOINT:
endpoint_msg = "EMBEDDING_ENDPOINT required." endpoint_msg = "EMBEDDING_ENDPOINT required."
raise RuntimeError(endpoint_msg) raise RuntimeError(endpoint_msg)
# Check models # Check models
if not CHAT_MODEL: if not CHAT_MODEL:
model_msg = "CHAT_MODEL required." model_msg = "CHAT_MODEL required."
raise RuntimeError(model_msg) raise RuntimeError(model_msg)
if not COMPLETION_MODEL: if not IMAGE_GEN_MODEL:
model_msg = "COMPLETION_MODEL required." model_msg = "IMAGE_GEN_MODEL required."
raise RuntimeError(model_msg) raise RuntimeError(model_msg)
if not IMAGE_GEN_MODEL: if not IMAGE_EDIT_MODEL:
model_msg = "IMAGE_GEN_MODEL required." model_msg = "IMAGE_EDIT_MODEL required."
raise RuntimeError(model_msg) raise RuntimeError(model_msg)
if not IMAGE_EDIT_MODEL: if not EMBEDDING_MODEL:
model_msg = "IMAGE_EDIT_MODEL required." model_msg = "EMBEDDING_MODEL required."
raise RuntimeError(model_msg) raise RuntimeError(model_msg)
if not EMBEDDING_MODEL:
model_msg = "EMBEDDING_MODEL required."
raise RuntimeError(model_msg)
# TTS # TTS
TTS_MODEL_PATH: str = os.getenv("TTS_MODEL_PATH", "kokoro-v1.0.onnx") TTS_MODEL_PATH: str = os.getenv("TTS_MODEL_PATH", "kokoro-v1.0.onnx")
@@ -180,9 +167,3 @@ VOICES_LIST: dict[str, dict[str, str | list[str]]] = {
], ],
}, },
} }
logger.info("CHAT_ENDPOINT set to %s", CHAT_ENDPOINT)
logger.info("COMPLETION_ENDPOINT set to %s", COMPLETION_ENDPOINT)
logger.info("IMAGE_GEN_ENDPOINT set to %s", IMAGE_GEN_ENDPOINT)
logger.info("IMAGE_EDIT_ENDPOINT set to %s", IMAGE_EDIT_ENDPOINT)
logger.info("EMBEDDING_ENDPOINT set to %s", EMBEDDING_ENDPOINT)
vibe_bot/database.py
+16 -761
View File
@@ -1,599 +1,22 @@
"""SQLite database with RAG support for chat history and embeddings.""" """SQLite database with RAG support for chat history and embeddings.
Facade for the ``vibe_bot.db`` package; re-exports the public names so
imports of ``vibe_bot.database`` keep working unchanged.
"""
from __future__ import annotations from __future__ import annotations
import logging from vibe_bot.db.bots import CustomBotManager
import sqlite3 from vibe_bot.db.messages import ChatDatabase
from typing import TYPE_CHECKING from vibe_bot.db.timing import IMAGE_GEN_TIME_LIMIT, IMAGE_GEN_TIME_WINDOW
import numpy as np __all__ = [
from openai import OpenAI "IMAGE_GEN_TIME_LIMIT",
"IMAGE_GEN_TIME_WINDOW",
from vibe_bot import llama_wrapper "ChatDatabase",
from vibe_bot.config import ( "CustomBotManager",
DB_PATH, "get_database",
EMBEDDING_ENDPOINT, ]
EMBEDDING_ENDPOINT_KEY,
EMBEDDING_MODEL,
MAX_HISTORY_MESSAGES,
SIMILARITY_THRESHOLD,
TOP_K_RESULTS,
)
if TYPE_CHECKING:
from datetime import datetime
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
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:
"""SQLite database with RAG support for storing chat history
using OpenAI embeddings.
"""
def __init__(self, db_path: str = DB_PATH) -> None:
"""Initialize the database connection.
Args:
db_path: Path to the SQLite database file.
"""
logger.info("Initializing ChatDatabase with path: %s", db_path)
self.db_path = db_path
self.client = OpenAI(
base_url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
)
logger.info("Connecting to OpenAI API for embeddings")
self._initialize_database()
def _initialize_database(self) -> None:
"""Initialize the SQLite database with required tables."""
logger.info("Initializing SQLite database at %s", self.db_path)
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# Create messages table
logger.info("Creating chat_messages table if not exists")
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS chat_messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id TEXT UNIQUE,
user_id TEXT,
username TEXT,
content TEXT,
bot_name TEXT,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
channel_id TEXT,
guild_id TEXT
)
""",
)
logger.info("chat_messages table initialized successfully")
# Migrate: add bot_name column if it doesn't exist
logger.info("Checking for bot_name column migration")
cursor.execute("PRAGMA table_info(chat_messages)")
columns = [row[1] for row in cursor.fetchall()]
if "bot_name" not in columns:
logger.info("Adding bot_name column to chat_messages table")
cursor.execute(
"ALTER TABLE chat_messages ADD COLUMN bot_name TEXT",
)
logger.info("bot_name column added successfully")
# Create embeddings table for RAG
logger.info("Creating message_embeddings table if not exists")
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS message_embeddings (
message_id TEXT PRIMARY KEY,
embedding BLOB,
FOREIGN KEY (message_id) REFERENCES chat_messages(message_id)
)
""",
)
logger.info("message_embeddings table initialized successfully")
# Create index for faster lookups
logger.info("Creating idx_timestamp index if not exists")
cursor.execute(
"""
CREATE INDEX IF NOT EXISTS idx_timestamp ON chat_messages(timestamp)
""",
)
logger.info("idx_timestamp index created successfully")
logger.info("Creating idx_user_id index if not exists")
cursor.execute(
"""
CREATE INDEX IF NOT EXISTS idx_user_id ON chat_messages(user_id)
""",
)
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()
logger.info("Database initialization completed successfully")
conn.close()
def _vector_to_bytes(self, vector: list[float]) -> bytes:
"""Convert vector to bytes for SQLite storage."""
logger.debug("Converting vector (length: %d) to bytes", len(vector))
result = np.array(vector, dtype=np.float32).tobytes()
logger.debug("Vector converted to %d bytes", len(result))
return result
def _bytes_to_vector(self, blob: bytes) -> np.ndarray:
"""Convert bytes back to vector."""
logger.debug("Converting %d bytes back to vector", len(blob))
result = np.frombuffer(blob, dtype=np.float32)
logger.debug("Vector reconstructed with %d dimensions", len(result))
return result
def _calculate_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
"""Calculate cosine similarity between two vectors."""
vec1 = vec1.flatten()
vec2 = vec2.flatten()
logger.debug(
"Calculating cosine similarity between vectors of dimension %d",
len(vec1),
)
norm1 = np.linalg.norm(vec1)
norm2 = np.linalg.norm(vec2)
if norm1 == 0 or norm2 == 0:
return 0.0
result = float(np.dot(vec1, vec2) / (norm1 * norm2))
logger.debug("Similarity calculated: %.4f", result)
return result
def add_message(
self,
*,
message_id: str,
user_id: str,
username: str,
content: str,
bot_name: str | None = None,
channel_id: str | None = None,
guild_id: str | None = None,
) -> bool:
"""Add a message to the database and generate its embedding."""
logger.info("Adding message %s from user %s", message_id, username)
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
try:
# Insert message
logger.debug(
"Inserting message into chat_messages table: message_id=%s",
message_id,
)
cursor.execute(
"""
INSERT OR REPLACE INTO chat_messages
(message_id, user_id, username, content, bot_name, channel_id, guild_id)
VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(
message_id,
user_id,
username,
content,
bot_name,
channel_id,
guild_id,
),
)
logger.debug("Message %s inserted into chat_messages table", message_id)
# Generate and store embedding
logger.info("Generating embedding for message %s", message_id)
embedding = llama_wrapper.embedding(
content,
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
model=EMBEDDING_MODEL,
)
if embedding:
logger.debug(
"Embedding generated successfully for message %s, "
"storing in database",
message_id,
)
cursor.execute(
"""
INSERT OR REPLACE INTO message_embeddings
(message_id, embedding)
VALUES (?, ?)
""",
(message_id, self._vector_to_bytes(embedding)),
)
logger.debug(
"Embedding stored in message_embeddings table for message %s",
message_id,
)
else:
logger.warning(
"Failed to generate embedding for message %s, "
"skipping embedding storage",
message_id,
)
# Clean up old messages if exceeding limit
logger.info("Checking if cleanup of old messages is needed")
self._cleanup_old_messages(cursor)
conn.commit()
except Exception:
logger.exception("Error adding message %s", message_id)
conn.rollback()
return False
else:
logger.info("Successfully added message %s to database", message_id)
return True
finally:
conn.close()
def _cleanup_old_messages(self, cursor: sqlite3.Cursor) -> None:
"""Remove old messages to stay within the limit."""
cursor.execute(
"""
SELECT COUNT(*) FROM chat_messages
""",
)
count = cursor.fetchone()[0]
if count > MAX_HISTORY_MESSAGES:
cursor.execute(
"""
DELETE FROM chat_messages
WHERE id IN (
SELECT id FROM chat_messages
ORDER BY timestamp ASC
LIMIT ?
)
""",
(count - MAX_HISTORY_MESSAGES,),
)
# Also remove corresponding embeddings
cursor.execute(
"""
DELETE FROM message_embeddings
WHERE message_id IN (
SELECT message_id FROM chat_messages
ORDER BY timestamp ASC
LIMIT ?
)
""",
(count - MAX_HISTORY_MESSAGES,),
)
def get_recent_messages(
self,
limit: int = 10,
) -> list[tuple[str, str, str, datetime]]:
"""Get recent messages from the database."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute(
"""
SELECT message_id, username, content, timestamp
FROM chat_messages
ORDER BY timestamp DESC
LIMIT ?
""",
(limit,),
)
messages = cursor.fetchall()
conn.close()
return messages
def search_similar_messages(
self,
query: str,
top_k: int = TOP_K_RESULTS,
min_similarity: float = SIMILARITY_THRESHOLD,
) -> list[tuple[str, str, float]]:
"""Search for messages similar to the query using embeddings."""
query_embedding = llama_wrapper.embedding(
text=query,
model=EMBEDDING_MODEL,
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
)
if not query_embedding:
return []
query_vector = np.array(query_embedding, dtype=np.float32)
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# Join chat_messages and message_embeddings to get content and embeddings
cursor.execute(
"""
SELECT cm.message_id, cm.content, me.embedding
FROM chat_messages cm
JOIN message_embeddings me ON cm.message_id = me.message_id
WHERE cm.username != 'vibe-bot'
""",
)
rows = cursor.fetchall()
results: list[tuple[str, str, float]] = []
for message_id, content, embedding_blob in rows:
embedding_vector = self._bytes_to_vector(embedding_blob)
similarity = self._calculate_similarity(query_vector, embedding_vector)
if similarity >= min_similarity:
cursor.execute(
"""
SELECT content
FROM chat_messages
WHERE message_id = ?
ORDER BY timestamp DESC
""",
(f"{message_id}_response",),
)
response_row = cursor.fetchone()
if response_row:
results.append((content, response_row[0], similarity))
conn.close()
# Sort by similarity and return top results
results.sort(key=lambda x: x[2], reverse=True)
return results[:top_k]
def get_bot_history(self, bot_name: str, limit: int = 20) -> list[tuple[str, str]]:
"""Get message history for a specific custom bot.
Args:
bot_name: The name of the custom bot.
limit: Maximum number of messages to retrieve.
Returns:
List of (user_message, bot_response) tuples.
"""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
logger.info(
"Fetching last %d messages for bot %r",
limit,
bot_name,
)
cursor.execute(
"""
SELECT message_id, content, timestamp
FROM chat_messages
WHERE bot_name = ? AND message_id NOT LIKE '%%_response'
ORDER BY timestamp DESC
LIMIT ?
""",
(bot_name, limit),
)
messages = cursor.fetchall()
conversations: list[tuple[str, str]] = []
for message in messages:
msg_content = message[1]
logger.debug("Finding response for %s...", msg_content[:50])
cursor.execute(
"""
SELECT content
FROM chat_messages
WHERE message_id = ?
ORDER BY timestamp DESC
""",
(f"{message[0]}_response",),
)
response_row = cursor.fetchone()
if response_row:
logger.debug("Found response: %s...", response_row[0][:50])
conversations.append((msg_content, response_row[0]))
else:
logger.debug("No response found")
conn.close()
return conversations
def get_user_history(self, user_id: str, limit: int = 20) -> list[tuple[str, str]]:
"""Get message history for a specific user."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
logger.info("Fetching last %d user messages", limit)
cursor.execute(
"""
SELECT message_id, content, timestamp
FROM chat_messages
WHERE user_id = ? AND username != 'vibe-bot'
ORDER BY timestamp DESC
LIMIT ?
""",
(user_id, limit),
)
messages = cursor.fetchall()
# Format is [user message, bot response]
conversations: list[tuple[str, str]] = []
for message in messages:
msg_content = message[1]
logger.debug("Finding response for %s...", msg_content[:50])
cursor.execute(
"""
SELECT content
FROM chat_messages
WHERE message_id = ?
ORDER BY timestamp DESC
""",
(f"{message[0]}_response",),
)
response_row = cursor.fetchone()
if response_row:
logger.debug("Found response: %s...", response_row[0][:50])
conversations.append((msg_content, response_row[0]))
else:
logger.debug("No response found")
conn.close()
return conversations
def get_conversation_context(
self,
user_id: str,
current_message: str,
max_context: int = 5,
) -> list[dict[str, str]]:
"""Get relevant conversation context for RAG."""
# Get recent messages from the user
recent_messages = self.get_user_history(user_id, limit=max_context * 2)
# Search for similar messages
similar_messages = self.search_similar_messages(
current_message,
top_k=max_context,
)
# Combine contexts
context_parts: list[dict[str, str]] = []
# Add recent messages
for user_message, bot_message in recent_messages:
context_parts.append({"role": "assistant", "content": bot_message})
context_parts.append({"role": "user", "content": user_message})
# Add similar messages
for user_message, bot_message, _similarity in similar_messages:
context_parts.append({"role": "assistant", "content": bot_message})
context_parts.append({"role": "user", "content": user_message})
# Conversation history needs to be delivered in "newest context last" order
context_parts.reverse()
return context_parts
def clear_all_messages(self) -> None:
"""Clear all messages and embeddings from the database."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute("DELETE FROM message_embeddings")
cursor.execute("DELETE FROM chat_messages")
conn.commit()
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
@@ -605,171 +28,3 @@ def get_database() -> ChatDatabase:
if _chat_db is None: if _chat_db is None:
_chat_db = ChatDatabase() _chat_db = ChatDatabase()
return _chat_db return _chat_db
class CustomBotManager:
"""Manages custom bot configurations stored in SQLite database."""
def __init__(self, db_path: str = DB_PATH) -> None:
"""Initialize the custom bot manager.
Args:
db_path: Path to the SQLite database file.
"""
self.db_path = db_path
self._initialize_custom_bots_table()
def _initialize_custom_bots_table(self) -> None:
"""Initialize the custom bots table in SQLite."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
# Create table to hold custom bots
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS custom_bots (
bot_name TEXT PRIMARY KEY,
system_prompt TEXT NOT NULL,
created_by TEXT NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
is_active INTEGER DEFAULT 1
)
""",
)
conn.commit()
conn.close()
def create_custom_bot(
self,
bot_name: str,
system_prompt: str,
created_by: str,
) -> bool:
"""Create a new custom bot configuration."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
try:
cursor.execute(
"""
INSERT OR REPLACE INTO custom_bots
(bot_name, system_prompt, created_by, is_active)
VALUES (?, ?, ?, 1)
""",
(bot_name.lower(), system_prompt, created_by),
)
conn.commit()
except Exception:
logger.exception("Error creating custom bot")
conn.rollback()
return False
else:
return True
finally:
conn.close()
def get_custom_bot(self, bot_name: str) -> tuple[str, str, str, datetime] | None:
"""Get a custom bot configuration by name."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
cursor.execute(
"""
SELECT bot_name, system_prompt, created_by, created_at
FROM custom_bots
WHERE bot_name = ? AND is_active = 1
""",
(bot_name.lower(),),
)
result = cursor.fetchone()
conn.close()
if result is None:
return None
return (result[0], result[1], result[2], result[3])
def list_custom_bots(
self,
user_id: str | None = None,
) -> list[tuple[str, str, str]]:
"""List all custom bots, optionally filtered by creator."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
if user_id:
cursor.execute(
"""
SELECT bot_name, system_prompt, created_by
FROM custom_bots
WHERE is_active = 1 AND created_by = ?
ORDER BY created_at DESC
""",
(user_id,),
)
else:
cursor.execute(
"""
SELECT bot_name, system_prompt, created_by
FROM custom_bots
WHERE is_active = 1
ORDER BY created_at DESC
""",
)
bots = cursor.fetchall()
conn.close()
return bots
def delete_custom_bot(self, bot_name: str) -> bool:
"""Delete a custom bot configuration."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
try:
cursor.execute(
"""
DELETE FROM custom_bots
WHERE bot_name = ?
""",
(bot_name.lower(),),
)
conn.commit()
except Exception:
logger.exception("Error deleting custom bot")
conn.rollback()
return False
else:
return cursor.rowcount > 0
finally:
conn.close()
def deactivate_custom_bot(self, bot_name: str) -> bool:
"""Deactivate a custom bot (soft delete)."""
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
try:
cursor.execute(
"""
UPDATE custom_bots
SET is_active = 0
WHERE bot_name = ?
""",
(bot_name.lower(),),
)
conn.commit()
except Exception:
logger.exception("Error deactivating custom bot")
conn.rollback()
return False
else:
return cursor.rowcount > 0
finally:
conn.close()
vibe_bot/db/__init__.py
+1
View File
@@ -0,0 +1 @@
"""SQLite storage layer: connection, schema, message store, RAG, custom bots."""
vibe_bot/db/bots.py
+148
View File
@@ -0,0 +1,148 @@
"""Custom bot configuration store (create, read, list, delete)."""
from __future__ import annotations
import logging
from datetime import datetime
from vibe_bot.config import DB_PATH
from vibe_bot.db.connection import connect
from vibe_bot.db.schema import initialize_custom_bots_table
logger = logging.getLogger(__name__)
class CustomBotManager:
"""Manages custom bot configurations stored in SQLite database."""
def __init__(self, db_path: str = DB_PATH) -> None:
"""Initialize the custom bot manager.
Args:
db_path: Path to the SQLite database file.
"""
self.db_path = db_path
self._initialize_custom_bots_table()
def _initialize_custom_bots_table(self) -> None:
"""Initialize the custom bots table in SQLite."""
initialize_custom_bots_table(self.db_path)
def create_custom_bot(
self,
bot_name: str,
system_prompt: str,
created_by: str,
) -> str | bool:
"""Create a custom bot configuration.
Returns "created" if the name was new, "replaced" if a bot with that
name already existed (the namespace is shared), or False on error.
"""
conn = connect(self.db_path)
cursor = conn.cursor()
name = bot_name.lower()
try:
cursor.execute(
"SELECT 1 FROM custom_bots WHERE bot_name = ?",
(name,),
)
exists = cursor.fetchone() is not None
cursor.execute(
"""
INSERT OR REPLACE INTO custom_bots
(bot_name, system_prompt, created_by, is_active)
VALUES (?, ?, ?, 1)
""",
(name, system_prompt, created_by),
)
conn.commit()
except Exception:
logger.exception("Error creating custom bot")
conn.rollback()
return False
else:
return "replaced" if exists else "created"
finally:
conn.close()
def get_custom_bot(self, bot_name: str) -> tuple[str, str, str, datetime] | None:
"""Get a custom bot configuration by name."""
conn = connect(self.db_path)
cursor = conn.cursor()
cursor.execute(
"""
SELECT bot_name, system_prompt, created_by, created_at
FROM custom_bots
WHERE bot_name = ? AND is_active = 1
""",
(bot_name.lower(),),
)
result = cursor.fetchone()
conn.close()
if result is None:
return None
return (result[0], result[1], result[2], result[3])
def list_custom_bots(
self,
user_id: str | None = None,
) -> list[tuple[str, str, str]]:
"""List all custom bots, optionally filtered by creator."""
conn = connect(self.db_path)
cursor = conn.cursor()
if user_id:
cursor.execute(
"""
SELECT bot_name, system_prompt, created_by
FROM custom_bots
WHERE is_active = 1 AND created_by = ?
ORDER BY created_at DESC
""",
(user_id,),
)
else:
cursor.execute(
"""
SELECT bot_name, system_prompt, created_by
FROM custom_bots
WHERE is_active = 1
ORDER BY created_at DESC
""",
)
bots = cursor.fetchall()
conn.close()
return bots
def delete_custom_bot(self, bot_name: str) -> bool:
"""Delete a custom bot configuration."""
conn = connect(self.db_path)
cursor = conn.cursor()
try:
cursor.execute(
"""
DELETE FROM custom_bots
WHERE bot_name = ?
""",
(bot_name.lower(),),
)
conn.commit()
except Exception:
logger.exception("Error deleting custom bot")
conn.rollback()
return False
else:
return cursor.rowcount > 0
finally:
conn.close()
vibe_bot/db/connection.py
+36
View File
@@ -0,0 +1,36 @@
"""SQLite connection plumbing shared by the database layer."""
from __future__ import annotations
import sqlite3
from datetime import datetime
# Per-connection busy timeout (ms) so concurrent writers wait instead of
# failing with "database is locked".
SQLITE_BUSY_TIMEOUT_MS = 5000
def _parse_timestamp(value: str | bytes) -> datetime:
"""Decode a stored TIMESTAMP value into a naive datetime."""
if isinstance(value, bytes):
value = value.decode("utf-8")
return datetime.fromisoformat(value)
sqlite3.register_converter("TIMESTAMP", _parse_timestamp)
def connect(db_path: str) -> sqlite3.Connection:
"""Open a SQLite connection configured for concurrent access.
WAL journaling is persistent (set once per database file); the busy
timeout is per-connection, so it is applied on every connection here.
``PARSE_DECLTYPES`` plus the registered ``TIMESTAMP`` converter decode
declared ``TIMESTAMP`` columns into ``datetime`` objects instead of
raw strings.
"""
conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
conn.execute("PRAGMA journal_mode=WAL")
conn.execute(f"PRAGMA busy_timeout={SQLITE_BUSY_TIMEOUT_MS}")
return conn
vibe_bot/db/messages.py
+306
View File
@@ -0,0 +1,306 @@
"""Chat message store: persistence, cleanup, and recency queries."""
from __future__ import annotations
import logging
import sqlite3
import numpy as np
from vibe_bot import llm_client
from vibe_bot.config import (
DB_PATH,
EMBEDDING_ENDPOINT,
EMBEDDING_ENDPOINT_KEY,
EMBEDDING_MODEL,
MAX_HISTORY_MESSAGES,
SIMILARITY_THRESHOLD,
TOP_K_RESULTS,
)
from vibe_bot.db.connection import connect
from vibe_bot.db.schema import initialize_chat_tables
from vibe_bot.db.search import (
get_bot_history,
get_user_history,
search_similar_messages,
)
from vibe_bot.db.timing import (
get_image_generation_time_estimate,
record_image_generation_time,
)
from vibe_bot.db.vectors import (
bytes_to_vector,
cosine_similarity,
vector_to_bytes,
)
logger = logging.getLogger(__name__)
class ChatDatabase:
"""SQLite store for chat history, embedding-backed RAG, and image timing."""
def __init__(self, db_path: str = DB_PATH) -> None:
"""Initialize the database connection.
Args:
db_path: Path to the SQLite database file.
"""
logger.info("Initializing ChatDatabase with path: %s", db_path)
self.db_path = db_path
initialize_chat_tables(db_path)
def _vector_to_bytes(self, vector: list[float]) -> bytes:
"""Convert vector to bytes for SQLite storage."""
return vector_to_bytes(vector)
def _bytes_to_vector(self, blob: bytes) -> np.ndarray:
"""Convert bytes back to a vector."""
return bytes_to_vector(blob)
def _calculate_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
"""Calculate cosine similarity between two vectors."""
return cosine_similarity(vec1, vec2)
def add_message(
self,
*,
message_id: str,
user_id: str,
username: str,
content: str,
bot_name: str | None = None,
channel_id: str | None = None,
guild_id: str | None = None,
role: str = "user",
embed: bool = True,
) -> bool:
"""Add a message to the database, optionally storing its embedding.
Args:
role: Either "user" (a human message) or "assistant" (a bot
response). Used to scope RAG retrieval instead of matching a
hard-coded bot username.
embed: Whether to generate and store an embedding for the message.
Response rows pass False: only user rows feed RAG retrieval.
"""
logger.debug("Adding message %s from user %s", message_id, user_id)
conn = connect(self.db_path)
cursor = conn.cursor()
try:
logger.debug(
"Inserting message into chat_messages table: message_id=%s",
message_id,
)
cursor.execute(
"""
INSERT OR REPLACE INTO chat_messages
(message_id, user_id, username, content, bot_name, channel_id,
guild_id, role)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
(
message_id,
user_id,
username,
content,
bot_name,
channel_id,
guild_id,
role,
),
)
logger.debug("Message %s inserted into chat_messages table", message_id)
if embed:
logger.debug("Generating embedding for message %s", message_id)
embedding = llm_client.embedding(
content,
model=EMBEDDING_MODEL,
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
)
if embedding:
logger.debug(
"Embedding generated successfully for message %s, "
"storing in database",
message_id,
)
vector = np.array(embedding, dtype=np.float32)
cursor.execute(
"""
INSERT OR REPLACE INTO message_embeddings
(message_id, embedding, norm)
VALUES (?, ?, ?)
""",
(message_id, vector.tobytes(), float(np.linalg.norm(vector))),
)
logger.debug(
"Embedding stored in message_embeddings table for message %s",
message_id,
)
else:
logger.warning(
"Failed to generate embedding for message %s, "
"skipping embedding storage",
message_id,
)
logger.debug("Checking if cleanup of old messages is needed")
self._cleanup_old_messages(cursor)
conn.commit()
except Exception:
logger.exception("Error adding message %s", message_id)
conn.rollback()
return False
else:
logger.debug("Successfully added message %s to database", message_id)
return True
finally:
conn.close()
def _cleanup_old_messages(self, cursor: sqlite3.Cursor) -> None:
"""Remove old messages to stay within the limit.
The rows to delete are captured up front. Deriving the embedding
message_ids from a fresh subquery *after* the chat_messages delete
would select the next-oldest live rows instead of the ones just
removed, orphaning the deleted rows' embeddings and deleting the
embeddings of rows that should survive.
"""
cursor.execute(
"""
SELECT COUNT(*) FROM chat_messages
""",
)
count = cursor.fetchone()[0]
if count <= MAX_HISTORY_MESSAGES:
return
excess = count - MAX_HISTORY_MESSAGES
cursor.execute(
"""
SELECT id, message_id FROM chat_messages
ORDER BY timestamp ASC
LIMIT ?
""",
(excess,),
)
oldest = cursor.fetchall()
if not oldest:
return
row_ids = [row[0] for row in oldest]
# Include each row's `_response` companion so a deleted user message
# also sheds its response embedding (and vice versa).
message_ids: list[str] = []
for _id, message_id in oldest:
message_ids.append(message_id)
message_ids.append(f"{message_id}_response")
id_placeholders = ", ".join("?" for _ in row_ids)
cursor.execute(
f"DELETE FROM chat_messages WHERE id IN ({id_placeholders})",
row_ids,
)
mid_placeholders = ", ".join("?" for _ in message_ids)
cursor.execute(
f"DELETE FROM message_embeddings WHERE message_id IN ({mid_placeholders})",
message_ids,
)
def search_similar_messages(
self,
query: str,
top_k: int = TOP_K_RESULTS,
min_similarity: float = SIMILARITY_THRESHOLD,
) -> list[tuple[str, str, float]]:
"""Search for messages similar to the query using embeddings."""
return search_similar_messages(
self.db_path,
query,
top_k=top_k,
min_similarity=min_similarity,
)
def get_bot_history(self, bot_name: str, limit: int = 20) -> list[tuple[str, str]]:
"""Get message history for a specific custom bot.
Args:
bot_name: The name of the custom bot.
limit: Maximum number of messages to retrieve.
Returns:
List of (user_message, bot_response) tuples.
"""
return get_bot_history(self.db_path, bot_name, limit)
def get_user_history(self, user_id: str, limit: int = 20) -> list[tuple[str, str]]:
"""Get message history for a specific user."""
return get_user_history(self.db_path, user_id, limit)
def get_conversation_context(
self,
user_id: str,
current_message: str,
max_context: int = 5,
) -> list[dict[str, str]]:
"""Get relevant conversation context for RAG."""
recent_messages = get_user_history(self.db_path, user_id, limit=max_context * 2)
similar_messages = search_similar_messages(
self.db_path,
current_message,
top_k=max_context,
)
context_parts: list[dict[str, str]] = []
for user_message, bot_message in recent_messages:
context_parts.append({"role": "assistant", "content": bot_message})
context_parts.append({"role": "user", "content": user_message})
for user_message, bot_message, _similarity in similar_messages:
context_parts.append({"role": "assistant", "content": bot_message})
context_parts.append({"role": "user", "content": user_message})
# Conversation history needs to be delivered in "newest context last" order
context_parts.reverse()
return context_parts
def clear_all_messages(self) -> None:
"""Clear all messages and embeddings from the database."""
conn = connect(self.db_path)
cursor = conn.cursor()
cursor.execute("DELETE FROM message_embeddings")
cursor.execute("DELETE FROM chat_messages")
conn.commit()
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.
"""
return record_image_generation_time(self.db_path, duration_seconds)
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.
"""
return get_image_generation_time_estimate(self.db_path)
vibe_bot/db/schema.py
+167
View File
@@ -0,0 +1,167 @@
"""Schema creation and column migrations for the chat and custom-bot tables."""
from __future__ import annotations
import logging
import sqlite3
import numpy as np
from vibe_bot.db.connection import connect
logger = logging.getLogger(__name__)
def initialize_chat_tables(db_path: str) -> None:
"""Create (and migrate) the chat history and embedding tables."""
logger.info("Initializing SQLite database at %s", db_path)
conn = connect(db_path)
cursor = conn.cursor()
logger.info("Creating chat_messages table if not exists")
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS chat_messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id TEXT UNIQUE,
user_id TEXT,
username TEXT,
content TEXT,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
channel_id TEXT,
guild_id TEXT
)
""",
)
logger.info("chat_messages table initialized successfully")
_migrate_chat_messages(cursor)
logger.info("Creating message_embeddings table if not exists")
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS message_embeddings (
message_id TEXT PRIMARY KEY,
embedding BLOB,
norm REAL,
FOREIGN KEY (message_id) REFERENCES chat_messages(message_id)
)
""",
)
logger.info("message_embeddings table initialized successfully")
_migrate_message_embeddings(cursor)
logger.info("Creating idx_timestamp index if not exists")
cursor.execute(
"""
CREATE INDEX IF NOT EXISTS idx_timestamp ON chat_messages(timestamp)
""",
)
logger.info("idx_timestamp index created successfully")
logger.info("Creating idx_user_id index if not exists")
cursor.execute(
"""
CREATE INDEX IF NOT EXISTS idx_user_id ON chat_messages(user_id)
""",
)
logger.info("idx_user_id index created successfully")
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()
logger.info("Database initialization completed successfully")
conn.close()
def _migrate_chat_messages(cursor: sqlite3.Cursor) -> None:
"""Add the bot_name and role columns to pre-existing databases."""
logger.info("Checking for chat_messages column migrations")
cursor.execute("PRAGMA table_info(chat_messages)")
columns = {row[1] for row in cursor.fetchall()}
if "bot_name" not in columns:
logger.info("Adding bot_name column to chat_messages table")
cursor.execute("ALTER TABLE chat_messages ADD COLUMN bot_name TEXT")
logger.info("bot_name column added successfully")
# role replaces the old convention of identifying bot responses by a
# hard-coded bot username.
if "role" not in columns:
logger.info("Adding role column to chat_messages table")
cursor.execute("ALTER TABLE chat_messages ADD COLUMN role TEXT")
cursor.execute(
"UPDATE chat_messages SET role = 'assistant' "
"WHERE message_id LIKE '%_response' AND role IS NULL",
)
cursor.execute(
"UPDATE chat_messages SET role = 'user' WHERE role IS NULL",
)
logger.info("role column added and backfilled")
# Backfill in batches so a large legacy table does not build one huge
# executemany parameter list in memory.
NORM_BACKFILL_BATCH = 500
def _migrate_message_embeddings(cursor: sqlite3.Cursor) -> None:
"""Add the norm column to pre-existing databases and backfill it.
The norm is the L2 norm of the stored float32 blob, so search can score
candidates with one matrix multiply and no per-vector renormalization.
"""
logger.info("Checking for message_embeddings column migrations")
cursor.execute("PRAGMA table_info(message_embeddings)")
columns = {row[1] for row in cursor.fetchall()}
if "norm" in columns:
return
logger.info("Adding norm column to message_embeddings table")
cursor.execute("ALTER TABLE message_embeddings ADD COLUMN norm REAL")
cursor.execute(
"SELECT message_id, embedding FROM message_embeddings "
"WHERE embedding IS NOT NULL",
)
rows = cursor.fetchall()
for start in range(0, len(rows), NORM_BACKFILL_BATCH):
cursor.executemany(
"UPDATE message_embeddings SET norm = ? WHERE message_id = ?",
[
(
float(np.linalg.norm(np.frombuffer(blob, dtype=np.float32))),
message_id,
)
for message_id, blob in rows[start : start + NORM_BACKFILL_BATCH]
],
)
logger.info("norm column added and backfilled for %d rows", len(rows))
def initialize_custom_bots_table(db_path: str) -> None:
"""Create the custom bots table in SQLite."""
conn = connect(db_path)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS custom_bots (
bot_name TEXT PRIMARY KEY,
system_prompt TEXT NOT NULL,
created_by TEXT NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
is_active INTEGER DEFAULT 1
)
""",
)
conn.commit()
conn.close()
vibe_bot/db/search.py
+204
View File
@@ -0,0 +1,204 @@
"""RAG retrieval: similarity search over user messages and history lookups."""
from __future__ import annotations
import logging
import numpy as np
from vibe_bot import llm_client
from vibe_bot.config import (
EMBEDDING_ENDPOINT,
EMBEDDING_ENDPOINT_KEY,
EMBEDDING_MODEL,
SIMILARITY_THRESHOLD,
TOP_K_RESULTS,
)
from vibe_bot.db.connection import connect
logger = logging.getLogger(__name__)
def search_similar_messages(
db_path: str,
query: str,
top_k: int = TOP_K_RESULTS,
min_similarity: float = SIMILARITY_THRESHOLD,
) -> list[tuple[str, str, float]]:
"""Search for messages similar to the query using embeddings.
A single JOIN pulls every user row, its stored embedding, the stored L2
norm, and its ``_response`` companion. Similarities are one matrix
multiply over the stored norms — no per-vector renormalization. Rows
with a missing or zero norm score 0 instead of dividing by zero.
"""
query_embedding = llm_client.embedding(
text=query,
model=EMBEDDING_MODEL,
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
)
if not query_embedding:
return []
query_vector = np.array(query_embedding, dtype=np.float32)
query_norm = float(np.linalg.norm(query_vector))
if query_norm == 0:
return []
conn = connect(db_path)
try:
cursor = conn.cursor()
cursor.execute(
"""
SELECT cm.content, r.content, me.embedding, me.norm
FROM chat_messages cm
JOIN message_embeddings me ON me.message_id = cm.message_id
LEFT JOIN chat_messages r ON r.message_id = cm.message_id || '_response'
WHERE cm.role = 'user'
""",
)
rows = cursor.fetchall()
finally:
conn.close()
if not rows:
return []
n_rows = len(rows)
blobs = [embedding_blob for _c, _r, embedding_blob, _n in rows]
dim = len(blobs[0]) // 4
if sum(len(blob) for blob in blobs) == n_rows * dim * 4:
vectors = np.frombuffer(b"".join(blobs), dtype=np.float32).reshape(n_rows, dim)
else:
# Mixed blob lengths (e.g. EMBEDDING_MODEL changed mid-life) can't be
# batched into one reshape; reconstruct per row, zero-padded (or
# truncated) to the query dim so the single matrix multiply still works.
vectors = np.zeros((n_rows, query_vector.size), dtype=np.float32)
for i, blob in enumerate(blobs):
row = np.frombuffer(blob, dtype=np.float32)
k = min(row.size, query_vector.size)
vectors[i, :k] = row[:k]
norms = np.array(
[
stored_norm if stored_norm is not None else 0.0
for _c, _r, _b, stored_norm in rows
],
dtype=np.float64,
)
safe_norms = np.where(norms > 0, norms, 1.0)
similarities = vectors @ query_vector / (safe_norms * query_norm)
similarities = np.where(norms > 0, similarities, 0.0)
results: list[tuple[str, str, float]] = []
for (content, response, _blob, _norm), similarity in zip(
rows, similarities, strict=True
):
if response is None or similarity < min_similarity:
continue
results.append((str(content), str(response), float(similarity)))
results.sort(key=lambda item: item[2], reverse=True)
return results[:top_k]
def get_bot_history(
db_path: str, bot_name: str, limit: int = 20
) -> list[tuple[str, str]]:
"""Get message history for a specific custom bot.
Args:
bot_name: The name of the custom bot.
limit: Maximum number of messages to retrieve.
Returns:
List of (user_message, bot_response) tuples.
"""
conn = connect(db_path)
cursor = conn.cursor()
logger.debug(
"Fetching last %d messages for bot %r",
limit,
bot_name,
)
cursor.execute(
"""
SELECT message_id, content
FROM chat_messages
WHERE bot_name = ? AND message_id NOT LIKE '%%_response'
ORDER BY timestamp DESC
LIMIT ?
""",
(bot_name, limit),
)
conversations: list[tuple[str, str]] = []
try:
for message_id, msg_content in cursor.fetchall():
logger.debug("Finding response for message_id=%s", message_id)
cursor.execute(
"""
SELECT content
FROM chat_messages
WHERE message_id = ?
ORDER BY timestamp DESC
""",
(f"{message_id}_response",),
)
response_row = cursor.fetchone()
if response_row:
logger.debug("Found response for message_id=%s", message_id)
conversations.append((str(msg_content), str(response_row[0])))
else:
logger.debug("No response found")
finally:
conn.close()
return conversations
def get_user_history(
db_path: str, user_id: str, limit: int = 20
) -> list[tuple[str, str]]:
"""Get message history for a specific user."""
conn = connect(db_path)
cursor = conn.cursor()
logger.debug("Fetching last %d user messages", limit)
cursor.execute(
"""
SELECT message_id, content
FROM chat_messages
WHERE user_id = ? AND role = 'user'
ORDER BY timestamp DESC
LIMIT ?
""",
(user_id, limit),
)
# Format is [user message, bot response]
conversations: list[tuple[str, str]] = []
try:
for message_id, msg_content in cursor.fetchall():
logger.debug("Finding response for message_id=%s", message_id)
cursor.execute(
"""
SELECT content
FROM chat_messages
WHERE message_id = ?
ORDER BY timestamp DESC
""",
(f"{message_id}_response",),
)
response_row = cursor.fetchone()
if response_row:
logger.debug("Found response for message_id=%s", message_id)
conversations.append((str(msg_content), str(response_row[0])))
else:
logger.debug("No response found")
finally:
conn.close()
return conversations
vibe_bot/db/timing.py
+95
View File
@@ -0,0 +1,95 @@
"""Image-generation timing statistics (moving-average estimate)."""
from __future__ import annotations
import logging
from vibe_bot.db.connection import connect
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
def record_image_generation_time(db_path: str, duration_seconds: float) -> bool:
"""Record how long an image generation took.
Args:
duration_seconds: Wall-clock seconds the generation took.
"""
logger.debug("Recording image generation time: %.2fs", duration_seconds)
conn = connect(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(db_path: str) -> 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 = connect(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])
vibe_bot/db/vectors.py
+42
View File
@@ -0,0 +1,42 @@
"""float32 embedding (de)serialization and cosine similarity."""
from __future__ import annotations
import logging
import numpy as np
logger = logging.getLogger(__name__)
def vector_to_bytes(vector: list[float]) -> bytes:
"""Convert a vector to bytes for SQLite storage."""
logger.debug("Converting vector (length: %d) to bytes", len(vector))
result = np.array(vector, dtype=np.float32).tobytes()
logger.debug("Vector converted to %d bytes", len(result))
return result
def bytes_to_vector(blob: bytes) -> np.ndarray:
"""Convert bytes back to a vector."""
logger.debug("Converting %d bytes back to vector", len(blob))
result = np.frombuffer(blob, dtype=np.float32)
logger.debug("Vector reconstructed with %d dimensions", len(result))
return result
def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
"""Calculate cosine similarity between two vectors."""
vec1 = vec1.flatten()
vec2 = vec2.flatten()
logger.debug(
"Calculating cosine similarity between vectors of dimension %d",
len(vec1),
)
norm1 = np.linalg.norm(vec1)
norm2 = np.linalg.norm(vec2)
if norm1 == 0 or norm2 == 0:
return 0.0
result = float(np.dot(vec1, vec2) / (norm1 * norm2))
logger.debug("Similarity calculated: %.4f", result)
return result
vibe_bot/llama_wrapper.py
-431
View File
@@ -1,431 +0,0 @@
"""Wraps the openai calls in generic functions.
Supports chat, image, edit, and embeddings.
Allows custom endpoints for each of the above supported functions.
"""
from __future__ import annotations
import json
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, cast
import openai
import requests
if TYPE_CHECKING:
from io import BufferedReader, BytesIO
from openai.types.chat import ChatCompletionMessageParam
def chat_completion(
system_prompt: str,
user_prompt: str,
*,
openai_url: str,
openai_api_key: str,
model: str,
max_tokens: int = 1000,
) -> str:
"""Send a chat completion request and return the response.
Args:
system_prompt: The system prompt to use.
user_prompt: The user prompt to send.
openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication.
model: The model to use for completion.
max_tokens: Maximum number of tokens to generate.
Returns:
The model's response text, stripped of whitespace.
"""
client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
messages: list[ChatCompletionMessageParam] = [
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": user_prompt,
},
]
response = client.chat.completions.create(
model=model,
messages=messages,
max_tokens=max_tokens,
timeout=60.0,
)
if not response.choices:
return ""
content = response.choices[0].message.content
if content:
return content.strip()
return ""
def chat_completion_with_history(
system_prompt: str,
prompts: list[dict[str, str]],
*,
openai_url: str,
openai_api_key: str,
model: str,
max_tokens: int = 1000,
) -> str:
"""Send a chat completion request with conversation history.
Args:
system_prompt: The system prompt to use.
prompts: List of prompt dicts with role and content.
openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication.
model: The model to use for completion.
max_tokens: Maximum number of tokens to generate.
Returns:
The model's response text, stripped of whitespace.
"""
client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
messages: list[ChatCompletionMessageParam] = [
cast(
"ChatCompletionMessageParam",
{
"role": "system",
"content": system_prompt,
},
),
]
messages.extend(cast("list[ChatCompletionMessageParam]", prompts))
response = client.chat.completions.create(
model=model,
messages=messages,
max_tokens=max_tokens,
seed=-1,
timeout=60.0,
)
if not response.choices:
return ""
content = response.choices[0].message.content
if content:
return content.strip()
return ""
def chat_completion_instruct(
system_prompt: str,
user_prompt: str,
*,
openai_url: str,
openai_api_key: str,
model: str,
max_tokens: int = 1000,
) -> str:
"""Send an instruction-based chat completion request.
Args:
system_prompt: The system prompt to use.
user_prompt: The user prompt to send.
openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication.
model: The model to use for completion.
max_tokens: Maximum number of tokens to generate.
Returns:
The model's response text, stripped of whitespace.
"""
client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
messages: list[ChatCompletionMessageParam] = [
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": user_prompt,
},
]
response = client.chat.completions.create(
model=model,
messages=messages,
max_tokens=max_tokens,
seed=-1,
timeout=60.0,
)
if not response.choices:
return ""
content = response.choices[0].message.content
if content:
return content.strip()
return ""
async def chat_completion_with_tools(
system_prompt: str,
prompts: list[dict[str, str]],
tools: list[dict[str, object]],
tool_executor: Callable[[str, dict[str, str]], str],
*,
openai_url: str,
openai_api_key: str,
model: str,
max_tokens: int = 1000,
max_tool_rounds: int = 5,
tool_call_notifier: (
Callable[[str, dict[str, str]], None]
| Callable[[str, dict[str, str]], Awaitable[None]]
| None
) = None,
) -> str:
"""Send a chat completion request with tool support and iterative tool calling.
Args:
system_prompt: The system prompt to use.
prompts: List of prompt dicts with role and content.
tools: List of tool definitions in OpenAI format.
tool_executor: A callable that takes (tool_name: str, tool_args: dict) -> str.
openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication.
model: The model to use for completion.
max_tokens: Maximum number of tokens to generate.
max_tool_rounds: Maximum number of tool call rounds before giving up.
tool_call_notifier: Optional callback invoked before each tool call
with (tool_name, tool_args).
Returns:
The model's final response text, stripped of whitespace.
"""
client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
messages: list[ChatCompletionMessageParam] = [
cast(
"ChatCompletionMessageParam",
{
"role": "system",
"content": system_prompt,
},
),
]
messages.extend(cast("list[ChatCompletionMessageParam]", prompts))
for _round in range(max_tool_rounds):
response = client.chat.completions.create(
model=model,
messages=messages,
tools=tools, # type: ignore[arg-type]
max_tokens=max_tokens,
seed=-1,
timeout=60.0,
)
if not response.choices:
return ""
message = response.choices[0].message
# Check if the model wants to call a tool
tool_calls = message.tool_calls
if tool_calls:
assistant_msg: dict[str, object] = {
"role": "assistant",
"content": message.content or "",
}
tool_call_dicts: list[dict[str, object]] = []
for tool_call in tool_calls:
tool_call_dicts.append(
{
"id": tool_call.id,
"type": "function",
"function": {
"name": tool_call.function.name, # type: ignore[union-attr]
"arguments": tool_call.function.arguments, # type: ignore[union-attr]
},
},
)
assistant_msg["tool_calls"] = tool_call_dicts
messages.append(cast("ChatCompletionMessageParam", assistant_msg))
# Execute each tool call and add results to messages
for tool_call in tool_calls:
tool_name = tool_call.function.name # type: ignore[union-attr]
tool_args = json.loads(tool_call.function.arguments) # type: ignore[union-attr]
if tool_call_notifier:
result = tool_call_notifier(tool_name, tool_args)
if hasattr(result, "__await__"):
await result # type: ignore[misc]
tool_result = tool_executor(tool_name, tool_args)
messages.append(
cast(
"ChatCompletionMessageParam",
{
"role": "tool",
"tool_call_id": tool_call.id,
"content": tool_result,
},
),
)
continue
# No more tool calls, return the final response
content = message.content
if content:
return content.strip()
return ""
return ""
def image_generation(
prompt: str,
*,
openai_url: str,
openai_api_key: str,
model: str = "gen",
n: int = 1,
size: str = "1024x1024",
) -> str:
"""Generate an image using the given prompt.
Args:
prompt: The image generation prompt.
openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication.
model: The model to use for image generation.
n: Number of images to generate.
size: The size of the generated image, e.g. "1024x1024".
Returns:
The base64 encoded image data. Decode and write to a file.
"""
client = openai.OpenAI(
base_url=openai_url,
api_key=openai_api_key,
max_retries=0,
)
try:
response = client.images.generate(
prompt=prompt,
n=n,
size=size,
model=model,
timeout=300.0,
)
except openai.APIConnectionError:
return ""
if response.data:
return response.data[0].b64_json or ""
return ""
def image_edit(
image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO],
prompt: str,
*,
openai_url: str,
openai_api_key: str,
model: str = "edit",
n: int = 1,
) -> str:
"""Edit an existing image using a prompt.
Args:
image: The source image as a file-like object or list thereof.
prompt: The edit instruction.
openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication.
model: The model to use for image editing.
n: Number of edited images to generate.
Returns:
The base64 encoded edited image data.
"""
client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key)
response = client.images.edit(
image=image,
prompt=prompt,
n=n,
size="768x768",
model=model,
)
if response.data:
return response.data[0].b64_json or ""
return ""
def embedding(
text: str,
*,
openai_url: str,
openai_api_key: str,
model: str,
) -> list[float]:
"""Generate an embedding vector for the given text.
Uses a raw HTTP request to avoid the OpenAI SDK injecting
unsupported parameters like encoding_format.
Args:
text: The text to embed.
openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication.
model: The embedding model to use.
Returns:
The embedding vector as a list of floats, or an empty list on failure.
"""
url = f"{openai_url.rstrip('/')}/embeddings"
headers = {
"Authorization": f"Bearer {openai_api_key}",
"Content-Type": "application/json",
}
payload = {"model": model, "input": [text]}
try:
resp = requests.post(url, headers=headers, json=payload, timeout=30)
resp.raise_for_status()
except requests.RequestException:
return []
data = resp.json()
# Handle both OpenAI-style response ({"data": [...]}) and
# Ollama-style response ([{...}]) where the API returns a list directly
if isinstance(data, list):
first = data[0]
if not isinstance(first, dict):
return []
raw = first.get("embedding")
elif isinstance(data, dict):
if not data.get("data"):
return []
raw = data["data"][0].get("embedding")
else:
return []
if raw is None:
return []
if isinstance(raw, str):
raw = json.loads(raw)
if not isinstance(raw, list):
raw = list(raw)
if not raw:
return []
return list[float](raw)
vibe_bot/llm/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Submodules of the OpenAI-compatible client layer (chat, images, registry)."""
vibe_bot/llm/chat.py
+206
View File
@@ -0,0 +1,206 @@
"""Core async chat completion with iterative tool calling, plus adapters."""
from __future__ import annotations
import asyncio
import json
import logging
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, cast
import openai
if TYPE_CHECKING:
from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
logger = logging.getLogger(__name__)
ToolExecutor = Callable[[str, dict[str, str]], str]
ToolCallNotifier = Callable[[str, dict[str, str]], None | Awaitable[None]]
_chat_client: openai.AsyncOpenAI | None = None
def get_chat_client() -> openai.AsyncOpenAI:
"""Return the shared async chat client, building it once."""
global _chat_client
if _chat_client is None:
from vibe_bot.config import CHAT_ENDPOINT, CHAT_ENDPOINT_KEY
_chat_client = openai.AsyncOpenAI(
base_url=CHAT_ENDPOINT, api_key=CHAT_ENDPOINT_KEY
)
return _chat_client
async def chat_complete(
messages: list[ChatCompletionMessageParam],
*,
model: str,
max_tokens: int,
seed: int | None = None,
tools: list[dict[str, object]] | None = None,
tool_executor: ToolExecutor | None = None,
tool_call_notifier: ToolCallNotifier | None = None,
max_tool_rounds: int = 5,
timeout: float = 60.0,
) -> str:
"""Send a chat completion, optionally with iterative tool calling.
Args:
messages: The conversation messages (system/user/assistant/tool).
model: The model to use for completion.
max_tokens: Maximum number of tokens to generate.
seed: Optional sampling seed.
tools: Optional list of tool definitions in OpenAI format.
tool_executor: Sync callable (tool_name, tool_args) -> result string.
tool_call_notifier: Optional sync-or-async callback invoked before each
tool call with (tool_name, tool_args).
max_tool_rounds: Maximum tool call rounds before giving up.
timeout: Per-request timeout in seconds.
Returns:
The model's final response text, stripped of whitespace ("" on failure).
"""
client = get_chat_client()
messages = list(messages)
for _round in range(max_tool_rounds):
kwargs: dict[str, Any] = {
"model": model,
"messages": messages,
"max_tokens": max_tokens,
"timeout": timeout,
}
if seed is not None:
kwargs["seed"] = seed
if tools:
kwargs["tools"] = cast("list[Any]", tools)
response = cast(
"ChatCompletion", await client.chat.completions.create(**kwargs)
)
if not response.choices:
return ""
message = response.choices[0].message
tool_calls = message.tool_calls
if tool_calls and tool_executor is not None:
assistant_msg: dict[str, object] = {
"role": "assistant",
"content": message.content or "",
}
tool_call_dicts: list[dict[str, object]] = []
for tool_call in tool_calls:
if tool_call.type != "function":
continue
tool_call_dicts.append(
{
"id": tool_call.id,
"type": "function",
"function": {
"name": tool_call.function.name,
"arguments": tool_call.function.arguments,
},
},
)
assistant_msg["tool_calls"] = tool_call_dicts
messages.append(cast("ChatCompletionMessageParam", assistant_msg))
for tool_call in tool_calls:
if tool_call.type != "function":
continue
tool_name = tool_call.function.name
tool_args = json.loads(tool_call.function.arguments)
if tool_call_notifier is not None:
result = tool_call_notifier(tool_name, tool_args)
if result is not None:
await result
tool_result = await asyncio.to_thread(
tool_executor, tool_name, tool_args
)
messages.append(
cast(
"ChatCompletionMessageParam",
{
"role": "tool",
"tool_call_id": tool_call.id,
"content": tool_result,
},
),
)
continue
content = message.content
if content:
return content.strip()
return ""
return ""
async def chat_completion_instruct(
system_prompt: str,
user_prompt: str,
*,
model: str,
max_tokens: int = 1000,
) -> str:
"""Instruction-based completion over :func:`chat_complete`."""
messages: list[ChatCompletionMessageParam] = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
]
return await chat_complete(messages, model=model, max_tokens=max_tokens, seed=-1)
async def chat_completion_with_history(
system_prompt: str,
prompts: list[dict[str, str]],
*,
model: str,
max_tokens: int = 1000,
) -> str:
"""Completion with conversation history over :func:`chat_complete`."""
messages: list[ChatCompletionMessageParam] = [
cast(
"ChatCompletionMessageParam",
{"role": "system", "content": system_prompt},
),
]
messages.extend(cast("list[ChatCompletionMessageParam]", prompts))
return await chat_complete(messages, model=model, max_tokens=max_tokens, seed=-1)
async def chat_completion_with_tools(
system_prompt: str,
prompts: list[dict[str, str]],
tools: list[dict[str, object]],
tool_executor: ToolExecutor,
*,
model: str,
max_tokens: int = 1000,
max_tool_rounds: int = 5,
tool_call_notifier: ToolCallNotifier | None = None,
) -> str:
"""Tool-capable completion over :func:`chat_complete`."""
messages: list[ChatCompletionMessageParam] = [
cast(
"ChatCompletionMessageParam",
{"role": "system", "content": system_prompt},
),
]
messages.extend(cast("list[ChatCompletionMessageParam]", prompts))
return await chat_complete(
messages,
model=model,
max_tokens=max_tokens,
seed=-1,
tools=tools,
tool_executor=tool_executor,
tool_call_notifier=tool_call_notifier,
max_tool_rounds=max_tool_rounds,
)
vibe_bot/llm/images.py
+92
View File
@@ -0,0 +1,92 @@
"""Async image generation and editing clients."""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import openai
if TYPE_CHECKING:
from io import BufferedReader, BytesIO
logger = logging.getLogger(__name__)
_image_gen_client: openai.AsyncOpenAI | None = None
_image_edit_client: openai.AsyncOpenAI | None = None
def get_image_gen_client() -> openai.AsyncOpenAI:
"""Return the shared async image-generation client, building it once."""
global _image_gen_client
if _image_gen_client is None:
from vibe_bot.config import IMAGE_GEN_ENDPOINT, IMAGE_GEN_ENDPOINT_KEY
_image_gen_client = openai.AsyncOpenAI(
base_url=IMAGE_GEN_ENDPOINT,
api_key=IMAGE_GEN_ENDPOINT_KEY,
max_retries=0,
)
return _image_gen_client
def get_image_edit_client() -> openai.AsyncOpenAI:
"""Return the shared async image-edit client, building it once."""
global _image_edit_client
if _image_edit_client is None:
from vibe_bot.config import IMAGE_EDIT_ENDPOINT, IMAGE_EDIT_ENDPOINT_KEY
_image_edit_client = openai.AsyncOpenAI(
base_url=IMAGE_EDIT_ENDPOINT, api_key=IMAGE_EDIT_ENDPOINT_KEY
)
return _image_edit_client
async def image_generation(
prompt: str,
*,
model: str = "gen",
n: int = 1,
size: str = "1024x1024",
) -> str:
"""Generate an image; return base64 data ("" on failure)."""
client = get_image_gen_client()
try:
response = await client.images.generate(
prompt=prompt,
n=n,
size=size,
model=model,
timeout=300.0,
)
except openai.OpenAIError as e:
logger.warning("Image generation failed: %s", e)
return ""
if response.data:
return response.data[0].b64_json or ""
return ""
async def image_edit(
image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO],
prompt: str,
*,
model: str = "edit",
n: int = 1,
) -> str:
"""Edit an image; return base64 data ("" on failure)."""
client = get_image_edit_client()
try:
response = await client.images.edit(
image=image,
prompt=prompt,
n=n,
size="768x768",
model=model,
)
except openai.OpenAIError as e:
logger.warning("Image edit failed: %s", e)
return ""
if response.data:
return response.data[0].b64_json or ""
return ""
vibe_bot/llm/registry.py
+106
View File
@@ -0,0 +1,106 @@
"""Tool registry: OpenAI schemas plus dispatch to sync implementations."""
from __future__ import annotations
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, cast
if TYPE_CHECKING:
from pydantic import BaseModel
ToolImpl = Callable[..., str]
class _RegisteredTool:
"""A registered tool: its OpenAI schema plus its synchronous impl."""
__slots__ = ("args_schema", "description", "impl", "name")
def __init__(
self,
name: str,
description: str,
args_schema: dict[str, object],
impl: ToolImpl,
) -> None:
self.name = name
self.description = description
self.args_schema = args_schema
self.impl = impl
class ToolRegistry:
"""Holds tool schemas and dispatches tool calls to their implementations."""
def __init__(self) -> None:
self._tools: dict[str, _RegisteredTool] = {}
def register(
self,
name: str,
description: str,
args_schema: dict[str, object],
impl: ToolImpl,
) -> None:
"""Register a tool under ``name`` with its OpenAI args schema."""
self._tools[name] = _RegisteredTool(name, description, args_schema, impl)
def to_openai_tools(self) -> list[dict[str, object]]:
"""Return the registered tools in OpenAI function-calling format."""
return [
{
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": tool.args_schema,
},
}
for tool in self._tools.values()
]
def execute(self, name: str, args: dict[str, str], **impl_kwargs: Any) -> str:
"""Dispatch a tool call; unknown tools yield a friendly message.
``impl_kwargs`` (e.g. ``channel``) are forwarded to the impl so tools
can access per-invocation context.
"""
tool = self._tools.get(name)
if tool is None:
return f"Unknown tool: {name}"
return tool.impl(name, args, **impl_kwargs)
_default_registry: ToolRegistry | None = None
def get_tool_registry() -> ToolRegistry:
"""Return the shared tool registry, seeded with the channel-members tool."""
global _default_registry
if _default_registry is None:
from vibe_bot.tools import get_channel_members
raw_schema = get_channel_members.args_schema
if isinstance(raw_schema, dict):
args_schema: dict[str, object] = raw_schema
else:
# A LangChain @tool exposes args_schema as a pydantic model class.
args_schema = cast("type[BaseModel]", raw_schema).model_json_schema()
registry = ToolRegistry()
registry.register(
get_channel_members.name,
get_channel_members.description or "",
args_schema,
_channel_members_tool,
)
_default_registry = registry
return _default_registry
def _channel_members_tool(name: str, args: dict[str, str], **kwargs: Any) -> str:
"""Adapt the registry dispatch to ``get_channel_members_impl(channel)``."""
from vibe_bot.tools import get_channel_members_impl
channel = kwargs.get("channel")
return get_channel_members_impl(channel)
vibe_bot/llm_client.py
+126
View File
@@ -0,0 +1,126 @@
"""Async OpenAI-compatible LLM, image, and embedding clients.
Public API facade: chat completion and the tool registry live in the
``vibe_bot.llm`` subpackage; the embedding HTTP plumbing stays in this
module.
``image_edit`` (``!retcon``) requests a fixed 768x768 output rather than
matching the source image's aspect ratio. Matching it would require
decoding the downloaded image (Pillow is not a dependency) and most
OpenAI-compatible edit endpoints only accept a fixed set of sizes anyway;
the square output bounds request cost and is universally honored.
"""
from __future__ import annotations
import json
import logging
from typing import Any
import requests
from vibe_bot.llm.chat import (
ToolCallNotifier,
ToolExecutor,
chat_complete,
chat_completion_instruct,
chat_completion_with_history,
chat_completion_with_tools,
get_chat_client,
)
from vibe_bot.llm.images import (
get_image_edit_client,
get_image_gen_client,
image_edit,
image_generation,
)
from vibe_bot.llm.registry import ToolRegistry, get_tool_registry
__all__ = [
"ToolCallNotifier",
"ToolExecutor",
"ToolRegistry",
"chat_complete",
"chat_completion_instruct",
"chat_completion_with_history",
"chat_completion_with_tools",
"embedding",
"get_chat_client",
"get_embedding_session",
"get_image_edit_client",
"get_image_gen_client",
"get_tool_registry",
"image_edit",
"image_generation",
]
logger = logging.getLogger(__name__)
_embedding_session: requests.Session | None = None
def get_embedding_session() -> requests.Session:
"""Return the shared requests session for embedding HTTP calls."""
global _embedding_session
if _embedding_session is None:
_embedding_session = requests.Session()
return _embedding_session
def embedding(
text: str,
*,
url: str,
api_key: str,
model: str,
) -> list[float]:
"""Generate an embedding vector for the given text (synchronous).
Uses a raw HTTP request (shared session) to avoid the SDK injecting
unsupported parameters like encoding_format.
"""
endpoint = f"{url.rstrip('/')}/embeddings"
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
payload = {"model": model, "input": [text]}
try:
resp = get_embedding_session().post(
endpoint, headers=headers, json=payload, timeout=30
)
resp.raise_for_status()
# A 2xx body can still be non-JSON (e.g. an HTML error page);
# resp.json() would raise JSONDecodeError (a ValueError).
data = resp.json()
except (requests.RequestException, ValueError):
return []
# Handle both OpenAI-style response ({"data": [...]}) and
# Ollama-style response ([{...}]) where the API returns a list directly
if isinstance(data, list):
first = data[0]
if not isinstance(first, dict):
return []
raw: Any = first.get("embedding")
elif isinstance(data, dict):
if not data.get("data"):
return []
raw = data["data"][0].get("embedding")
else:
return []
if raw is None:
return []
if isinstance(raw, str):
try:
raw = json.loads(raw)
except ValueError:
return []
if not isinstance(raw, list):
raw = list(raw)
if not raw:
return []
return list[float](raw)
vibe_bot/main.py
+10 -1268
View File
File diff suppressed because it is too large Load Diff
vibe_bot/prompts.py
+209
View File
@@ -0,0 +1,209 @@
"""Prompt constants and system-prompt assembly helpers.
Holds every prompt string the bot sends to the LLM (image layout selection,
image-prompt engineering, prompt verification) plus the shared response-length
hint and the ``build_system_prompt`` assembler used by the chat and
speak-with-bot paths.
"""
from __future__ import annotations
import re
from typing import TYPE_CHECKING
from vibe_bot.config import (
IMAGE_GEN_SIZE_LANDSCAPE,
IMAGE_GEN_SIZE_PORTRAIT,
IMAGE_GEN_SIZE_SQUARE,
)
if TYPE_CHECKING:
import discord
# 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,
}
# Shared response-length hint appended to bot system prompts.
RESPONSE_LENGTH_HINT = "Keep your responses under 2-3 sentences."
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"
)
IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE = (
"You are an expert art director and image-generation prompt engineer. "
"Convert the user's message into one single, extremely detailed image "
"generation prompt that will be passed directly to a text-to-image model. "
"The image model is weak: it guesses at compositions, fumbles rendered "
"text, and invents details on its own. Your prompt must therefore leave "
"nothing to interpretation - explicitly describe every visible aspect of "
"the image so it can be created with extreme precision and detail.\n"
"The final image will use a {layout} canvas, so compose the scene to fit "
"that orientation.\n"
"Your prompt must cover all of the following as one flowing, descriptive "
"passage. Begin with the main subject, described completely in the very "
"first sentence:\n"
"- Subject(s): every subject with concrete specifics (species or "
"character, age, build, clothing, colors, materials, accessories), its "
"exact pose, expression, gaze direction, and its precise position in the "
'frame (for example "centered in the foreground" or "small in the '
'upper-left background"). State the relative scale of subjects to each '
"other and to the frame. If the subject is a fusion, hybrid, or anything "
"unusual, the first sentence must state in full what is joined to what "
"and exactly how it looks, and that description must be repeated near "
"the end of the passage.\n"
"- Composition and framing: the camera angle (eye-level, low, high, "
"bird's-eye), the shot type (extreme close-up, portrait, full body, wide "
"establishing shot), the focal point, the arrangement of elements across "
"the {layout} canvas, and the depth of field.\n"
"- Text: if the image must contain readable text (titles, signs, "
"posters, labels, banners, watermarks, logos, captions), quote the EXACT "
"text verbatim in double quotes with precise capitalization and "
"punctuation, and specify its font style, color, size, and exact "
"placement. If the image should contain no text, state that explicitly "
'("no text anywhere in the image").\n'
"- Setting and background: the complete environment with concrete "
"details - location, time of day, weather, and every notable background "
"and foreground element with its position.\n"
"- Style and rendering: the art style or medium (for example "
"photorealistic 35mm photograph, oil painting, watercolor, cel-shaded "
"anime, pixel art, vector illustration), the color palette with specific "
"colors, the lighting (source, direction, quality, mood), the overall "
"atmosphere, and the level of detail.\n"
'- Finish with concise quality terms such as "highly detailed, sharp '
'focus".\n'
"Rules:\n"
'- Be concrete and specific. Never use vague words like "nice", '
'"cool", "epic", or "various" - name exact colors, objects, '
"positions, and quantities.\n"
"- Be literal. Interpret the user's request exactly as written. Never "
'rationalize, normalize, or "improve" it: surreal, absurd, or '
'anthropomorphic requests are intentional, not mistakes. A "fountain '
'pen wearing pants" is an anthropomorphized fountain pen character '
"wearing pants, not a pen lying next to a pair of pants.\n"
"- Preserve everything the user specified. Fill in details the user did "
"not specify with coherent choices that fit the request, but never "
"alter, drop, or reinterpret what the user did specify.\n"
"- Decompose concepts. The image model lacks world knowledge, so never "
"rely on a name alone for anything it might misrender (mythical "
"creatures, fictional characters, cultural items, animal breeds, "
"instruments, vehicles). Spell out the visual anatomy: silhouette, body "
"parts, materials, and distinguishing features, with explicit "
"disambiguation. A centaur is a single creature with a human torso, "
"arms, and head seamlessly fused to a horse's front half, the horse's "
"four legs extending from the human's waist - one fused body, not a "
"person riding a horse.\n"
"- 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.\n"
"- Respond with ONLY the image generation prompt itself. Do not affirm "
"the user, do not answer the user's questions, and do not add headings, "
"labels, numbered lists, or any other text."
)
IMAGE_PROMPT_VERIFY_SYSTEM_PROMPT = (
"You are the final quality check for an image generation prompt. The "
"image model that will use it has no world knowledge: it renders only "
"what is described literally and silently drops anything it does not "
"understand - a prompt that merely names a centaur without describing "
"the fused human-animal body will produce a plain horse. "
"Given the user's original request and the drafted prompt, judge "
"strictly: would the drafted prompt, taken completely literally, "
"produce exactly what the user asked for, including every unusual, "
"mythical, surreal, or anthropomorphic element? "
"Respond with ONLY the single word PASS if it would. Otherwise respond "
"with ONLY a corrected version of the prompt that would produce exactly "
"what the user asked for: one flowing descriptive passage, the "
"subject's full anatomy and every unusual element described explicitly "
"in the first sentence and repeated near the end, no other text."
)
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 build_system_prompt(personality: str, user_info: str) -> str:
"""Assemble a custom-bot system prompt with the length hint and user info.
Args:
personality: The base bot personality / system prompt.
user_info: Preformatted user information to append.
Returns:
The assembled system prompt: personality, a response-length hint, and
the user information block.
"""
return f"{personality}\n{RESPONSE_LENGTH_HINT}\n\nUser Information:\n{user_info}"
def get_user_info(user: discord.User | discord.Member) -> str:
"""Format user information for inclusion in bot prompts.
Reads only presentation attributes off the (User or Member) object, so it
has no runtime dependency on the discord package.
"""
parts: list[str] = []
if user.global_name:
parts.append(f"Global Name: {user.global_name}")
nick = getattr(user, "nick", None)
if nick:
parts.append(f"Nickname: {nick}")
top_role = getattr(user, "top_role", None)
if top_role and top_role.name != "@everyone":
parts.append(f"Top Role: {top_role.name}")
activities = getattr(user, "activities", None)
if activities:
activity_names = [
getattr(a, "name", str(a))
for a in activities
if getattr(a, "name", "") != "custom_status"
]
if activity_names:
parts.append(f"Activities: {', '.join(activity_names)}")
joined_at = getattr(user, "joined_at", None)
if joined_at:
parts.append(f"Joined: {joined_at.strftime('%Y-%m-%d')}")
parts.append(f"Username: {user.name}")
parts.append(f"User ID: {user.id}")
parts.append(
f"Account Created: {user.created_at.strftime('%Y-%m-%d') if user.created_at else 'Unknown'}"
)
return "\n".join(parts)
vibe_bot/services/__init__.py
+8
View File
@@ -0,0 +1,8 @@
"""Stateful bot services, each wired with its dependencies via constructor.
The four services own the logic for the four LLM-backed flows (custom-bot chat,
image gen/edit, speech, and bot-vs-bot conversation). They take a discord
``ctx`` only as an argument (typed under ``TYPE_CHECKING``) and never import the
``discord`` package at runtime; a file factory is injected so they can hand
audio/image bytes to ``ctx.send`` without touching ``discord.File``.
"""
vibe_bot/services/chat_service.py
+129
View File
@@ -0,0 +1,129 @@
"""Chat service: RAG context + tool-capped LLM completion + persistence + reply."""
from __future__ import annotations
import asyncio
import logging
from typing import TYPE_CHECKING
from vibe_bot import llm_client
from vibe_bot.config import CHAT_MODEL, MAX_COMPLETION_TOKENS
from vibe_bot.prompts import build_system_prompt, get_user_info
from vibe_bot.textutil import split_message
if TYPE_CHECKING:
from discord.ext.commands import Bot, Context
from vibe_bot.database import ChatDatabase
from vibe_bot.llm_client import ToolRegistry
logger = logging.getLogger(__name__)
class ChatService:
"""Handles one custom-bot chat turn: context -> LLM (with tools) -> persist -> reply."""
def __init__(self, db: ChatDatabase, registry: ToolRegistry) -> None:
self._db = db
self._registry = registry
async def handle(
self,
ctx: Context[Bot],
*,
bot_name: str,
message: str,
system_prompt: str,
response_prefix: str,
) -> None:
"""Run a single chat turn for ``bot_name`` and send the reply.
Args:
ctx: The Discord command context.
bot_name: The name of the custom bot.
message: The user message to process.
system_prompt: The base system prompt (personality) for the bot.
response_prefix: The prefix message sent before the reply.
"""
await ctx.send(f"{bot_name} is searching its databanks for {message[:50]}...")
# Get conversation context using RAG (SQLite + embedding HTTP call).
context = await asyncio.to_thread(
self._db.get_conversation_context,
user_id=str(ctx.author.id),
current_message=message,
max_context=5,
)
prompts: list[dict[str, str]] = [{"role": "user", "content": message}]
if context:
prompts = context + prompts
logger.info(
"chat: bot=%s user=%s context_msgs=%d",
bot_name,
ctx.author.id,
len(context),
)
system_prompt_edit = build_system_prompt(
system_prompt, get_user_info(ctx.author)
)
tools = self._registry.to_openai_tools()
def tool_executor(tool_name: str, tool_args: dict[str, str]) -> str:
"""Dispatch a tool call through the registry."""
return self._registry.execute(tool_name, tool_args, channel=ctx.channel)
async def tool_call_notifier(tool_name: str, tool_args: dict[str, str]) -> None:
"""Send a notification message when a tool is called."""
if tool_name == "get_channel_members":
await ctx.send(f"{bot_name} is looking at the channel members...")
try:
bot_response = await llm_client.chat_completion_with_tools(
system_prompt=system_prompt_edit,
prompts=prompts,
tools=tools,
tool_executor=tool_executor,
tool_call_notifier=tool_call_notifier,
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
# Store both the user message and the bot response in the database.
await asyncio.to_thread(
self._db.add_message,
message_id=f"{ctx.message.id}",
user_id=str(ctx.author.id),
username=ctx.author.name,
content=f"User: {message}",
bot_name=bot_name,
channel_id=str(ctx.channel.id),
guild_id=str(ctx.guild.id) if ctx.guild else None,
)
if ctx.bot.user is not None:
await asyncio.to_thread(
self._db.add_message,
message_id=f"{ctx.message.id}_response",
user_id=str(ctx.bot.user.id),
username=ctx.bot.user.name,
content=bot_response,
bot_name=bot_name,
channel_id=str(ctx.channel.id),
guild_id=str(ctx.guild.id) if ctx.guild else None,
role="assistant",
embed=False,
)
# Send the response back to the chat.
await ctx.send(response_prefix)
for send_chunk in split_message(bot_response, 1000):
await ctx.send(send_chunk)
except Exception:
logger.exception("Error in handle_chat")
await ctx.send("An error occurred while processing your request.")
vibe_bot/services/conversation_service.py
+147
View File
@@ -0,0 +1,147 @@
"""Conversation service: run a capped bot-vs-bot conversation (talkforme)."""
from __future__ import annotations
import asyncio
import logging
from typing import TYPE_CHECKING
from vibe_bot import llm_client
from vibe_bot.config import CHAT_MODEL, MAX_COMPLETION_TOKENS
from vibe_bot.prompts import RESPONSE_LENGTH_HINT
from vibe_bot.textutil import split_message
if TYPE_CHECKING:
from discord.ext.commands import Bot, Context
from vibe_bot.database import CustomBotManager
logger = logging.getLogger(__name__)
# Input size bound: reject oversized topics before they reach the LLM.
MAX_TOPIC_LENGTH = 500
# Hard cap on the number of replies a single !talkforme invocation produces.
TALK_LIMIT = 20
def flip_counter(counter: int) -> int:
"""Flip between 0 and 1 (the two conversing bots)."""
return 1 if counter == 0 else 0
class ConversationService:
"""Runs a two-bot conversation about a topic, chunking each reply."""
def __init__(self, manager: CustomBotManager) -> None:
self._manager = manager
async def run(
self,
ctx: Context[Bot],
bot1: str,
bot2: str,
limit: str,
topic: str,
) -> None:
"""Have ``bot1`` and ``bot2`` talk about ``topic`` for up to ``limit`` replies.
Args:
ctx: The Discord command context.
bot1: Name of the first custom bot.
bot2: Name of the second custom bot.
limit: Requested number of replies (string; parsed to an int).
topic: The conversation topic.
"""
if len(topic) > MAX_TOPIC_LENGTH:
logger.warning(
"Talkforme topic too long from user %s: length=%d",
ctx.author.id,
len(topic),
)
await ctx.send(f"Topic too long. Max {MAX_TOPIC_LENGTH} characters.")
return
bot1_info = await asyncio.to_thread(self._manager.get_custom_bot, bot1)
if not bot1_info:
await ctx.send(f"{bot1} is not a real bot...")
return
bot1_prompt = bot1_info[1]
bot2_info = await asyncio.to_thread(self._manager.get_custom_bot, bot2)
if not bot2_info:
await ctx.send(f"{bot2} is not a real bot...")
return
bot2_prompt = bot2_info[1]
try:
message_limit = int(limit)
except ValueError:
await ctx.send("Message limit must be an integer.")
return
effective_limit = min(message_limit, TALK_LIMIT)
await ctx.send(
f"{bot1} is going to talk to {bot2} "
f'about "{topic[:50]}" for {effective_limit} replies.',
)
bot_list = [(bot1, bot1_prompt), (bot2, bot2_prompt)]
async def send_chunked(text: str) -> None:
"""Send text in 1000-char chunks to stay under Discord's limit."""
for chunk in split_message(text, 1000):
await ctx.send(chunk)
message_counter = 0
bot_counter = 0
current_bot = bot_list[bot_counter]
prompt_histories: list[list[dict[str, str]]] = [
[{"role": "user", "content": topic}],
[{"role": "assistant", "content": topic}],
]
first_bot_response = await llm_client.chat_completion_with_history(
system_prompt=(
current_bot[1] + f"\n{RESPONSE_LENGTH_HINT} "
f"You are talking to {current_bot[flip_counter(bot_counter)][0]}"
),
prompts=prompt_histories[bot_counter],
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
await ctx.send(f"## {current_bot[0]}")
await send_chunked(first_bot_response)
prompt_histories[0].append({"role": "assistant", "content": first_bot_response})
prompt_histories[1].append({"role": "user", "content": first_bot_response})
bot_counter = flip_counter(counter=bot_counter)
while message_counter < effective_limit:
current_bot = bot_list[bot_counter]
logger.debug("Current bot is %s", current_bot[0])
bot_response = await llm_client.chat_completion_with_history(
system_prompt=(
current_bot[1] + f"\n{RESPONSE_LENGTH_HINT} "
f"You are talking to {current_bot[flip_counter(bot_counter)][0]}"
),
prompts=prompt_histories[bot_counter],
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
message_counter += 1
prompt_histories[bot_counter].append(
{"role": "assistant", "content": bot_response},
)
prompt_histories[flip_counter(bot_counter)].append(
{"role": "user", "content": bot_response},
)
await ctx.send(f"## {current_bot[0]}")
await send_chunked(bot_response)
bot_counter = flip_counter(counter=bot_counter)
logger.debug(
"Message counter is %d/%d",
message_counter,
effective_limit,
)
vibe_bot/services/image_service.py
+285
View File
@@ -0,0 +1,285 @@
"""Image service: doodlebob (generate) and retcon (edit), plus their helpers."""
from __future__ import annotations
import asyncio
import base64
import logging
import re
import time
from collections.abc import Callable
from io import BytesIO
from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse
import requests
from vibe_bot import llm_client
from vibe_bot.config import (
CHAT_MODEL,
IMAGE_EDIT_MODEL,
IMAGE_GEN_MODEL,
MAX_COMPLETION_TOKENS,
)
from vibe_bot.prompts import (
IMAGE_LAYOUT_SYSTEM_PROMPT,
IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE,
IMAGE_PROMPT_VERIFY_SYSTEM_PROMPT,
LAYOUT_SIZES,
parse_image_layout,
)
if TYPE_CHECKING:
from discord.ext.commands import Bot, Context
from vibe_bot.database import ChatDatabase
logger = logging.getLogger(__name__)
# Input size bound: reject oversized prompts before they reach the LLM.
MAX_IMAGE_PROMPT_LENGTH = 2000
# !retcon download safety: only fetch images from Discord's own CDN hosts and
# cap the size of a single download.
ALLOWED_IMAGE_HOSTS = (
"discord.com",
"discordapp.com",
"discordapp.net",
"discordcdn.com",
"discord.media",
)
MAX_IMAGE_DOWNLOAD_BYTES = 8 * 1024 * 1024
# Injected factory that turns an in-memory image/audio into a discord.File.
# Kept out of the service so it never imports discord at runtime.
FileFactory = Callable[[BytesIO, str], Any]
# Matches a bare "pass" verdict from the prompt-verification LLM.
_PASS_RE = re.compile(r"\bpass\b", re.IGNORECASE)
async def select_image_layout(user_message: str) -> str:
"""Ask the LLM to pick an image layout (a single-word answer)."""
response = await llm_client.chat_completion_instruct(
system_prompt=IMAGE_LAYOUT_SYSTEM_PROMPT,
user_prompt=user_message,
model=CHAT_MODEL,
max_tokens=2,
)
return parse_image_layout(response)
async def verify_image_prompt(user_message: str, image_prompt: str) -> str:
"""Check the drafted prompt literally produces the user's request."""
check_prompt = f"User request: {user_message}\n\nDrafted prompt: {image_prompt}"
response = await llm_client.chat_completion_instruct(
system_prompt=IMAGE_PROMPT_VERIFY_SYSTEM_PROMPT,
user_prompt=check_prompt,
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
if not response:
return image_prompt
# A passing check is the single word PASS; a correction is a full
# rewritten passage, which is always much longer.
if len(response) <= 50 and _PASS_RE.search(response):
return image_prompt
return response
def _allowed_image_url(url: str) -> bool:
"""Return True when the URL points at an allowed Discord CDN host."""
try:
host = urlparse(url).hostname or ""
except ValueError:
return False
return any(
host == allowed or host.endswith(f".{allowed}")
for allowed in ALLOWED_IMAGE_HOSTS
)
def _download_image_bytes(url: str) -> bytes | None:
"""Download a source image for !retcon.
Returns the image bytes, or None when the URL is not on the Discord CDN
allowlist, the download fails, or the image exceeds the size cap.
"""
if not _allowed_image_url(url):
logger.warning("Refusing to download image from non-Discord host: %s", url)
return None
try:
response = requests.get(url, timeout=30, stream=True)
response.raise_for_status()
except requests.RequestException as e:
logger.warning("Failed to download image from %s: %s", url, e)
return None
content_length = response.headers.get("Content-Length")
if content_length is not None:
try:
if int(content_length) > MAX_IMAGE_DOWNLOAD_BYTES:
logger.warning("Image from %s exceeds the size cap", url)
return None
except ValueError:
pass
chunks: list[bytes] = []
total = 0
try:
for chunk in response.iter_content(chunk_size=65536):
if not chunk:
continue
total += len(chunk)
if total > MAX_IMAGE_DOWNLOAD_BYTES:
logger.warning(
"Image from %s exceeds the size cap while streaming", url
)
return None
chunks.append(chunk)
except requests.RequestException as e:
logger.warning("Failed while streaming image from %s: %s", url, e)
return None
return b"".join(chunks)
class ImageService:
"""Generates (doodlebob) and edits (retcon) images via the LLM image APIs."""
def __init__(self, db: ChatDatabase, make_file: FileFactory) -> None:
self._db = db
self._make_file = make_file
async def generate(self, ctx: Context[Bot], *, message: str) -> None:
"""Convert a message into an image using Doodlebob."""
logger.info(
"Doodlebob command triggered by user %s: prompt_chars=%d",
ctx.author.id,
len(message),
)
if len(message) > MAX_IMAGE_PROMPT_LENGTH:
logger.warning(
"Doodlebob prompt too long from user %s: length=%d",
ctx.author.id,
len(message),
)
await ctx.send(
f"Prompt too long. Max {MAX_IMAGE_PROMPT_LENGTH} characters."
)
return
await ctx.send("**Doodlebob shopping for a canvas...**")
# Let the LLM pick the canvas orientation based on the content.
layout = await select_image_layout(message)
logger.info("Doodlebob selected layout %r for user %s", layout, ctx.author.id)
await ctx.send(f"**Doodlebob selected {layout}**")
system_prompt = IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE.format(layout=layout)
# Wait for the generated image prompt.
image_prompt = await llm_client.chat_completion_instruct(
system_prompt=system_prompt,
user_prompt=message,
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
# If the string is empty we had an error.
if image_prompt == "":
logger.warning("No image prompt supplied. Check for errors.")
return
# Verify the prompt literally produces the user's request; the check
# may return a corrected prompt.
image_prompt = await verify_image_prompt(message, image_prompt)
logger.debug(
"Doodlebob final image prompt ready: prompt_chars=%d layout=%s",
len(image_prompt),
layout,
)
# Alert the user we're generating the image.
estimated_seconds = await asyncio.to_thread(
self._db.get_image_generation_time_estimate
)
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 = await llm_client.image_generation(
prompt=image_prompt,
model=IMAGE_GEN_MODEL,
size=LAYOUT_SIZES[layout],
)
elapsed_seconds = time.monotonic() - start_time
if not image_b64:
logger.warning("Image generation returned empty response.")
await ctx.send("Failed to generate image. The server may be busy.")
return
await asyncio.to_thread(self._db.record_image_generation_time, elapsed_seconds)
try:
edited_image_data = BytesIO(base64.b64decode(image_b64))
send_img = self._make_file(edited_image_data, "image.png")
await ctx.send(file=send_img)
await ctx.send(
f"**Strike complete. Image generated in {elapsed_seconds:.1f} seconds.**",
)
except Exception:
logger.exception("Failed to decode image data")
await ctx.send("Failed to process the generated image.")
async def edit(self, ctx: Context[Bot], *, message: str) -> None:
"""Edit an attached image based on a text prompt."""
if len(message) > MAX_IMAGE_PROMPT_LENGTH:
logger.warning(
"Retcon prompt too long from user %s: length=%d",
ctx.author.id,
len(message),
)
await ctx.send(
f"Prompt too long. Max {MAX_IMAGE_PROMPT_LENGTH} characters."
)
return
image_data_list: list[BytesIO] = []
for discord_image in ctx.message.attachments:
image_url = discord_image.url
image_bytes = await asyncio.to_thread(_download_image_bytes, image_url)
if image_bytes is None:
continue
image_data_list.append(BytesIO(image_bytes))
if not image_data_list:
await ctx.send("Please attach an image to edit.")
return
await ctx.send(f"**Rewriting history to match {message[:100]}...**")
image_b64 = await llm_client.image_edit(
image=image_data_list,
prompt=message,
model=IMAGE_EDIT_MODEL,
)
if not image_b64:
await ctx.send("Failed to edit the image.")
return
try:
edited_image_data = BytesIO(base64.b64decode(image_b64))
except ValueError as e:
logger.warning("Failed to decode edited image data: %s", e)
await ctx.send("Failed to process the edited image.")
return
send_img = self._make_file(edited_image_data, "image.png")
await ctx.send(file=send_img)
vibe_bot/services/speech_service.py
+279
View File
@@ -0,0 +1,279 @@
"""Speech service: parse the voice flag, dispatch bot-vs-plain, run TTS."""
from __future__ import annotations
import asyncio
import logging
import re
from collections.abc import Callable
from io import BytesIO
from typing import TYPE_CHECKING, Any
from vibe_bot import llm_client
from vibe_bot.config import (
CHAT_MODEL,
MAX_COMPLETION_TOKENS,
TTS_SPEED,
TTS_VOICE,
VOICES_LIST,
)
from vibe_bot.prompts import build_system_prompt, get_user_info
from vibe_bot.tts import DEFAULT_LANG
if TYPE_CHECKING:
from discord.ext.commands import Bot, Context
from vibe_bot.database import ChatDatabase, CustomBotManager
from vibe_bot.tts import TTSEngine
logger = logging.getLogger(__name__)
# Input size bound: reject oversized text before it reaches the TTS engine.
MAX_SPEAK_LENGTH = 5000
# Injected factory that turns an in-memory audio buffer into a discord.File.
FileFactory = Callable[[BytesIO, str], Any]
# Precomputed voice -> language lookup (replaces a per-call VOICES_LIST scan).
VOICE_LANGUAGES: dict[str, str] = {
voice: str(category["language"])
for category in VOICES_LIST.values()
for voice in category["voices"]
}
# Trailing-anchored voice flag: only a `--voice <name>` at the very end of the
# message is treated as a flag, so the flag mid-text is preserved as speech.
_VOICE_FLAG_RE = re.compile(r"^(?P<text>.*)\s+--voice\s+(?P<voice>\S+)$")
def parse_voice_flag(message: str) -> tuple[str, str | None]:
"""Split a trailing `--voice <name>` flag off a speak message.
Returns:
(text, voice) where voice is None when no trailing flag is present and
text is then returned unchanged.
"""
match = _VOICE_FLAG_RE.match(message)
if not match:
return message, None
return match["text"].rstrip(), match["voice"]
class SpeechService:
"""Speaks text (plain or via a custom bot) using the Kokoro TTS engine."""
def __init__(
self,
db: ChatDatabase,
manager: CustomBotManager,
tts: TTSEngine | None,
make_file: FileFactory,
) -> None:
self._db = db
self._manager = manager
self._tts = tts
self._make_file = make_file
async def speak(self, ctx: Context[Bot], *, message: str) -> None:
"""Have the bot speak the given text, or have a custom bot respond+speaks."""
if self._tts is None:
await ctx.send(
"TTS engine not initialized. "
"Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.",
)
return
text, voice = parse_voice_flag(message)
if not text or not text.strip():
await ctx.send("Please provide text to speak.")
return
if len(text) > MAX_SPEAK_LENGTH:
logger.warning(
"Speak text too long from user %s: length=%d",
ctx.author.id,
len(text),
)
await ctx.send(
f"Text too long to speak. Max {MAX_SPEAK_LENGTH} characters."
)
return
# Validate the voice if one was requested.
if voice is not None and voice not in VOICE_LANGUAGES:
await ctx.send(
f"Unknown voice '{voice}'. Use `!voices` to see available voices."
)
return
custom_bots = await asyncio.to_thread(self._manager.list_custom_bots)
bot_names = [b[0] for b in custom_bots]
first_word = text.split(maxsplit=1)[0] if text.split() else ""
if first_word in bot_names:
await self._speak_with_bot(ctx, first_word, text, voice)
else:
await self._speak_plain(ctx, text, voice)
async def _speak_with_bot(
self,
ctx: Context[Bot],
bot_name: str,
message: str,
voice: str | None,
) -> None:
"""Have a custom bot respond to the message and speak the response."""
text_to_speak = message[len(bot_name) :].lstrip()
if not text_to_speak:
await ctx.send("Please provide text for the bot to respond to.")
return
await ctx.send(f"**{bot_name}** is thinking...")
bot_info = await asyncio.to_thread(self._manager.get_custom_bot, bot_name)
if not bot_info:
await ctx.send(f"Custom bot '{bot_name}' not found.")
return
_, system_prompt, _, _ = bot_info
system_prompt_edit = build_system_prompt(
system_prompt, get_user_info(ctx.author)
)
engine = self._tts
if engine is None:
await ctx.send(
"TTS engine not initialized. "
"Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.",
)
return
# Determine language for the chosen voice.
chosen_voice = voice or TTS_VOICE
lang = VOICE_LANGUAGES.get(chosen_voice, DEFAULT_LANG)
try:
context = await asyncio.to_thread(
self._db.get_conversation_context,
user_id=str(ctx.author.id),
current_message=text_to_speak,
max_context=5,
)
prompts: list[dict[str, str]] = [{"role": "user", "content": text_to_speak}]
if context:
prompts = context + prompts
# Tools come from the shared registry (schema + dispatch).
registry = llm_client.get_tool_registry()
speak_tools = registry.to_openai_tools()
def speak_tool_executor(tool_name: str, tool_args: dict[str, str]) -> str:
"""Dispatch a tool call through the registry."""
return registry.execute(tool_name, tool_args, channel=ctx.channel)
async def speak_tool_call_notifier(
tool_name: str, tool_args: dict[str, str]
) -> None:
"""Send a notification message when a tool is called."""
if tool_name == "get_channel_members":
await ctx.send(
f"**{bot_name}** is looking at the channel members..."
)
bot_response = await llm_client.chat_completion_with_tools(
system_prompt=system_prompt_edit,
prompts=prompts,
tools=speak_tools,
tool_executor=speak_tool_executor,
tool_call_notifier=speak_tool_call_notifier,
model=CHAT_MODEL,
max_tokens=MAX_COMPLETION_TOKENS,
)
if not bot_response:
await ctx.send(f"**{bot_name}** failed to generate a response.")
return
await asyncio.to_thread(
self._db.add_message,
message_id=f"{ctx.message.id}",
user_id=str(ctx.author.id),
username=ctx.author.name,
content=f"User: {text_to_speak}",
bot_name=bot_name,
channel_id=str(ctx.channel.id),
guild_id=str(ctx.guild.id) if ctx.guild else None,
)
if ctx.bot.user is not None:
await asyncio.to_thread(
self._db.add_message,
message_id=f"{ctx.message.id}_response",
user_id=str(ctx.bot.user.id),
username=ctx.bot.user.name,
content=bot_response,
bot_name=bot_name,
channel_id=str(ctx.channel.id),
guild_id=str(ctx.guild.id) if ctx.guild else None,
role="assistant",
embed=False,
)
await ctx.send(f"**{bot_name}**: {bot_response}")
await ctx.send(f"Generating speech for **{bot_name}**...")
result = await asyncio.to_thread(
engine.generate_audio,
bot_response,
voice=chosen_voice,
speed=TTS_SPEED,
lang=lang,
)
if result.partial:
await ctx.send("Some audio chunks failed; audio may be incomplete.")
audio_file = self._make_file(result.audio, "speech.mp3")
await ctx.send(file=audio_file)
except Exception:
logger.exception(
"Error in speak command with bot %r",
bot_name,
)
await ctx.send("Error generating speech.")
async def _speak_plain(
self,
ctx: Context[Bot],
message: str,
voice: str | None,
) -> None:
"""Speak plain text (no custom bot involved)."""
engine = self._tts
if engine is None:
await ctx.send(
"TTS engine not initialized. "
"Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.",
)
return
chosen_voice = voice or TTS_VOICE
lang = VOICE_LANGUAGES.get(chosen_voice, DEFAULT_LANG)
try:
await ctx.send("Generating speech...")
result = await asyncio.to_thread(
engine.generate_audio,
message,
voice=chosen_voice,
speed=TTS_SPEED,
lang=lang,
)
if result.partial:
await ctx.send("Some audio chunks failed; audio may be incomplete.")
audio_file = self._make_file(result.audio, "speech.mp3")
await ctx.send(file=audio_file)
except Exception:
logger.exception("Error in speak command")
await ctx.send("Error generating speech.")
vibe_bot/tests/_helpers.py
+23
View File
@@ -0,0 +1,23 @@
"""Shared helpers for wiring tests (command invocation, sent-text asserts)."""
from __future__ import annotations
import asyncio
from collections.abc import Callable
from typing import Any, cast
from unittest.mock import MagicMock
from discord.ext import commands
def invoke(bot: commands.Bot, name: str, *args: Any, **kwargs: Any) -> None:
"""Invoke a registered command's callback directly with the given args."""
cmd = bot.get_command(name)
assert cmd is not None
callback = cast("Callable[..., Any]", cmd.callback)
asyncio.run(callback(*args, **kwargs))
def sent_texts(ctx: MagicMock) -> list[str]:
"""All positional text messages sent through ctx.send."""
return [c.args[0] for c in ctx.send.call_args_list if c.args]
vibe_bot/tests/conftest.py
+77 -148
View File
@@ -5,12 +5,16 @@ from __future__ import annotations
import tempfile import tempfile
import warnings import warnings
from collections.abc import Generator from collections.abc import Generator
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
import numpy as np import numpy as np
import pytest import pytest
from discord.ext import commands
from vibe_bot.app import App, build_bot
warnings.filterwarnings( warnings.filterwarnings(
"ignore", "ignore",
@@ -21,41 +25,79 @@ if TYPE_CHECKING:
from vibe_bot.database import ChatDatabase, CustomBotManager from vibe_bot.database import ChatDatabase, CustomBotManager
@dataclass(frozen=True)
class AppMocks:
"""The App under test plus its mock dependencies, for setup and asserts."""
app: App
db: MagicMock
manager: MagicMock
registry: MagicMock
chat: MagicMock
image: MagicMock
speech: MagicMock
conversation: MagicMock
@pytest.fixture @pytest.fixture
def mock_env_vars() -> Generator[None]: def mock_ctx() -> MagicMock:
"""Provide minimal env vars for config loading.""" """Create a mock Discord command context."""
with patch.dict( ctx = MagicMock()
"os.environ", ctx.author.name = "testuser"
{ ctx.author.id = "12345"
"DISCORD_TOKEN": "test-token", ctx.author.global_name = "Test User"
"CHAT_ENDPOINT": "https://chat.example.com/v1", ctx.author.nick = "tester"
"COMPLETION_ENDPOINT": "https://completion.example.com/v1", ctx.author.top_role.name = "@everyone"
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", ctx.author.activities = []
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", ctx.author.joined_at = None
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", ctx.author.created_at = None
"CHAT_MODEL": "test-chat-model", ctx.channel.id = "channel-1"
"COMPLETION_MODEL": "test-completion-model", ctx.guild.id = "guild-1"
"IMAGE_GEN_MODEL": "test-image-model", ctx.message.id = "msg-1"
"IMAGE_EDIT_MODEL": "test-image-edit-model", ctx.message.attachments = []
"EMBEDDING_MODEL": "test-embedding-model", ctx.bot.user = MagicMock()
"CHAT_ENDPOINT_KEY": "test-key", ctx.bot.user.name = "test-bot"
"COMPLETION_ENDPOINT_KEY": "test-completion-key", ctx.bot.user.id = "bot-123"
"IMAGE_GEN_ENDPOINT_KEY": "test-image-key", ctx.send = AsyncMock()
"IMAGE_EDIT_ENDPOINT_KEY": "test-image-edit-key", return ctx
"EMBEDDING_ENDPOINT_KEY": "test-embedding-key",
"MAX_COMPLETION_TOKENS": "1000",
"MAX_HISTORY_MESSAGES": "1000", @pytest.fixture
"SIMILARITY_THRESHOLD": "0.7", def app_mocks() -> AppMocks:
"TOP_K_RESULTS": "5", """An App built entirely from mocks, with alfred in the bot cache."""
"TTS_MODEL_PATH": "/tmp/test-model.onnx", db = MagicMock()
"TTS_VOICES_PATH": "/tmp/test-voices.bin", manager = MagicMock()
"TTS_VOICE": "af_sarah", manager.list_custom_bots.return_value = [
"TTS_SPEED": "1.0", ("alfred", "british butler", "user123"),
"DB_PATH": ":memory:", ]
}, registry = MagicMock()
clear=False, chat = MagicMock()
): chat.handle = AsyncMock()
yield image = MagicMock()
image.generate = AsyncMock()
image.edit = AsyncMock()
speech = MagicMock()
speech.speak = AsyncMock()
conversation = MagicMock()
conversation.run = AsyncMock()
app = App(
db=db,
manager=manager,
registry=registry,
tts=MagicMock(),
chat=chat,
image=image,
speech=speech,
conversation=conversation,
bot_cache={"alfred": ("british butler", "user123")},
)
return AppMocks(app, db, manager, registry, chat, image, speech, conversation)
@pytest.fixture
def bot(app_mocks: AppMocks) -> commands.Bot:
"""A real Bot built from the mock App (never connected)."""
return build_bot(app_mocks.app)
@pytest.fixture @pytest.fixture
@@ -71,22 +113,13 @@ def temp_db_path() -> Generator[str]:
def mock_embedding() -> Generator[MagicMock]: def mock_embedding() -> Generator[MagicMock]:
"""Provide a mock embedding function returning a fixed vector.""" """Provide a mock embedding function returning a fixed vector."""
vector: list[float] = [0.1] * 2048 vector: list[float] = [0.1] * 2048
with patch("vibe_bot.llama_wrapper.embedding", return_value=vector) as mock: with patch("vibe_bot.llm_client.embedding", return_value=vector) as mock:
yield mock
@pytest.fixture
def mock_openai_client() -> Generator[MagicMock]:
"""Provide a mock OpenAI client."""
mock_client = MagicMock()
with patch("vibe_bot.database.OpenAI", return_value=mock_client) as mock:
yield mock yield mock
@pytest.fixture @pytest.fixture
def chat_db( def chat_db(
temp_db_path: str, temp_db_path: str,
mock_openai_client: MagicMock,
mock_embedding: MagicMock, mock_embedding: MagicMock,
) -> Generator[ChatDatabase]: ) -> Generator[ChatDatabase]:
"""Provide a ChatDatabase instance with a temp database.""" """Provide a ChatDatabase instance with a temp database."""
@@ -94,7 +127,6 @@ def chat_db(
db = ChatDatabase(db_path=temp_db_path) db = ChatDatabase(db_path=temp_db_path)
yield db yield db
db.client.close()
@pytest.fixture @pytest.fixture
@@ -133,106 +165,3 @@ def mock_kokoro_tts() -> Generator[dict[str, Any]]:
"mock_samples": mock_samples, "mock_samples": mock_samples,
"mock_sr": 24000, "mock_sr": 24000,
} }
@pytest.fixture
def mock_discord() -> Generator[dict[str, MagicMock]]:
"""Mock discord module components."""
mock_intents = MagicMock()
mock_intents.default.return_value = MagicMock()
mock_intents.default.return_value.message_content = True
mock_bot_class = MagicMock()
mock_bot_instance = MagicMock()
mock_bot_instance.user = MagicMock()
mock_bot_instance.user.name = "test-bot"
mock_bot_instance.user.id = "123456789"
with (
patch("vibe_bot.main.discord") as mock_discord_module,
patch("vibe_bot.main.commands", MagicMock()),
patch("vibe_bot.main.commands.Bot", mock_bot_class),
):
mock_bot_class.return_value = mock_bot_instance
mock_discord_module.Intents = mock_intents
mock_discord_module.Message = MagicMock
mock_discord_module.File = MagicMock
yield {
"Intents": mock_intents,
"Bot": mock_bot_class,
"bot_instance": mock_bot_instance,
}
@pytest.fixture
def mock_tts_engine() -> Generator[MagicMock]:
"""Provide a mock TTSEngine."""
mock_engine = MagicMock()
mock_engine.generate_audio.return_value = MagicMock()
with (
patch("vibe_bot.main.tts_engine", mock_engine),
patch("vibe_bot.main.tts.TTSEngine", return_value=mock_engine),
):
yield mock_engine
@pytest.fixture
def mock_requests() -> Generator[MagicMock]:
"""Provide mock requests module."""
with patch("vibe_bot.main.requests") as mock_requests_module:
mock_response = MagicMock()
mock_response.content = b"fake image data"
mock_requests_module.get.return_value = mock_response
yield mock_requests_module
@pytest.fixture
def mock_base64() -> Generator[MagicMock]:
"""Provide mock base64 module."""
with patch("vibe_bot.main.base64") as mock_base64_module:
mock_base64_module.b64decode.return_value = b"fake image data"
yield mock_base64_module
@pytest.fixture
def mock_llama_wrapper() -> Generator[MagicMock]:
"""Provide mock llama_wrapper module."""
with patch("vibe_bot.main.llama_wrapper") as mock_wrapper:
mock_wrapper.chat_completion_with_history.return_value = "Bot response"
mock_wrapper.chat_completion_with_tools = AsyncMock(return_value="Bot response")
mock_wrapper.chat_completion_instruct.return_value = "image prompt"
mock_wrapper.image_generation.return_value = ""
mock_wrapper.image_edit.return_value = ""
mock_wrapper.embedding.return_value = [0.1] * 2048
yield mock_wrapper
@pytest.fixture
def mock_database() -> Generator[MagicMock]:
"""Provide mock database module."""
with patch("vibe_bot.main.get_database") as mock_get_db:
mock_db = MagicMock()
mock_db.get_conversation_context.return_value = []
mock_db.add_message.return_value = True
mock_get_db.return_value = mock_db
yield mock_db
@pytest.fixture
def mock_custom_bot_manager() -> Generator[MagicMock]:
"""Provide mock CustomBotManager."""
with patch("vibe_bot.main.CustomBotManager") as mock_manager_class:
mock_manager = MagicMock()
mock_manager.create_custom_bot.return_value = True
mock_manager.get_custom_bot.return_value = (
"alfred",
"british butler personality",
"user123",
"2024-01-01",
)
mock_manager.list_custom_bots.return_value = [
("alfred", "british butler personality", "user123"),
]
mock_manager.delete_custom_bot.return_value = True
mock_manager_class.return_value = mock_manager
yield mock_manager
vibe_bot/tests/test_app.py
+322
View File
@@ -0,0 +1,322 @@
"""Tests for the app composition root (singletons, services, bot handlers)."""
from __future__ import annotations
import asyncio
import logging
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from discord.ext import commands
from vibe_bot.app import (
configure_logging,
create_app,
invalidate_bot_cache,
)
from vibe_bot.commands import speech as speech_commands
from vibe_bot.config import TTS_MODEL_PATH, TTS_VOICES_PATH
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.conversation_service import ConversationService
from vibe_bot.services.image_service import ImageService
from vibe_bot.services.speech_service import SpeechService
from vibe_bot.tests._helpers import sent_texts
from vibe_bot.tests.conftest import AppMocks
def test_create_app_wires_singletons(
chat_db: Any,
custom_bot_manager: Any,
) -> None:
"""create_app shares one db/manager with the services and seeds the cache."""
engine = MagicMock()
with (
patch("vibe_bot.app.ChatDatabase", return_value=chat_db),
patch("vibe_bot.app.CustomBotManager", return_value=custom_bot_manager),
patch("vibe_bot.app.TTSEngine", return_value=engine) as mock_tts,
patch(
"vibe_bot.llm_client.get_tool_registry",
return_value=MagicMock(),
) as mock_registry,
):
app = create_app()
mock_tts.assert_called_once_with(TTS_MODEL_PATH, TTS_VOICES_PATH)
assert app.db is chat_db
assert app.manager is custom_bot_manager
assert app.tts is engine
assert app.registry is mock_registry.return_value
assert isinstance(app.chat, ChatService)
assert isinstance(app.image, ImageService)
assert isinstance(app.speech, SpeechService)
assert isinstance(app.conversation, ConversationService)
assert app.chat._db is chat_db
assert app.chat._registry is mock_registry.return_value
assert app.image._db is chat_db
assert app.speech._db is chat_db
assert app.speech._manager is custom_bot_manager
assert app.speech._tts is engine
assert app.conversation._manager is custom_bot_manager
assert app.bot_cache == {}
def test_create_app_tts_failure_tolerant(
chat_db: Any,
custom_bot_manager: Any,
) -> None:
"""A failing TTS engine degrades to None instead of crashing startup."""
with (
patch("vibe_bot.app.ChatDatabase", return_value=chat_db),
patch("vibe_bot.app.CustomBotManager", return_value=custom_bot_manager),
patch("vibe_bot.app.TTSEngine", side_effect=OSError("no model file")),
):
app = create_app()
assert app.tts is None
def test_invalidate_bot_cache_rebuilds(app_mocks: AppMocks) -> None:
"""invalidate_bot_cache rebuilds the cache from the manager."""
app_mocks.manager.list_custom_bots.return_value = [
("newbot", "a personality", "user999"),
("alfred", "british butler", "user123"),
]
invalidate_bot_cache(app_mocks.app)
assert app_mocks.app.bot_cache == {
"newbot": ("a personality", "user999"),
"alfred": ("british butler", "user123"),
}
def test_configure_logging_configures_root() -> None:
"""configure_logging is the sole basicConfig: it adds a root handler."""
root = logging.getLogger()
original = root.handlers
root.handlers.clear()
try:
configure_logging()
assert len(root.handlers) == 1
handler = root.handlers[0]
assert handler.formatter is not None
assert handler.formatter._fmt == (
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
finally:
root.handlers = original
# ---------------------------------------------------------------------------
# build_bot
# ---------------------------------------------------------------------------
def test_build_bot_registers_commands(bot: commands.Bot) -> None:
"""build_bot returns a commands.Bot with every command registered."""
assert isinstance(bot, commands.Bot)
for name in (
"custom-bot",
"list-custom-bots",
"delete-custom-bot",
"lobotomize",
"debug",
"voices",
"speak",
"doodlebob",
"retcon",
"history",
"talkforme",
):
assert bot.get_command(name) is not None
for event in ("on_ready", "on_message", "on_command_error"):
assert event in bot.__dict__
def test_build_bot_intents(bot: commands.Bot) -> None:
"""message_content, members, and presences intents are enabled."""
assert bot.intents.message_content is True
assert bot.intents.members is True
assert bot.intents.presences is True
def test_handlers_refuse_to_run_before_build_bot() -> None:
"""Event and command handlers raise if invoked before the wiring exists."""
import vibe_bot.app as app_module
from vibe_bot.commands import _state as commands_state
saved_app, saved_bot = app_module._app, app_module._bot
saved_commands_app = commands_state._app
app_module._app = None
app_module._bot = None
commands_state._app = None
try:
with pytest.raises(RuntimeError, match="App is not initialized"):
asyncio.run(app_module.on_message(MagicMock()))
with pytest.raises(RuntimeError, match="Bot is not initialized"):
asyncio.run(app_module.on_ready())
with pytest.raises(RuntimeError, match="App is not initialized"):
asyncio.run(speech_commands.speak(MagicMock(), message="hello"))
finally:
app_module._app = saved_app
app_module._bot = saved_bot
commands_state._app = saved_commands_app
# ---------------------------------------------------------------------------
# on_message guards
# ---------------------------------------------------------------------------
def test_on_message_ignores_non_bang_messages(
app_mocks: AppMocks,
bot: commands.Bot,
) -> None:
"""A non-! message returns before any service or DB touch."""
message = MagicMock()
message.author = MagicMock()
message.content = "hello there"
with patch.object(bot, "process_commands", new=AsyncMock()) as mock_process:
asyncio.run(bot.on_message(message))
mock_process.assert_not_called()
app_mocks.chat.handle.assert_not_awaited()
app_mocks.speech.speak.assert_not_awaited()
app_mocks.image.generate.assert_not_awaited()
app_mocks.image.edit.assert_not_awaited()
app_mocks.conversation.run.assert_not_awaited()
assert app_mocks.manager.mock_calls == []
assert app_mocks.db.mock_calls == []
def test_on_message_skips_bot_authors(app_mocks: AppMocks, bot: commands.Bot) -> None:
"""A bot-authored message returns without calling any service or DB."""
message = MagicMock()
message.author = bot.user
message.content = "!alfred hi"
with patch.object(bot, "process_commands", new=AsyncMock()) as mock_process:
asyncio.run(bot.on_message(message))
mock_process.assert_not_called()
app_mocks.chat.handle.assert_not_awaited()
assert app_mocks.manager.mock_calls == []
assert app_mocks.db.mock_calls == []
def test_on_message_routes_custom_bot(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!alfred hi dispatches to the chat service with the parsed args."""
message = MagicMock()
message.author = MagicMock()
message.author.name = "testuser"
message.content = "!alfred hi"
with (
patch.object(bot, "get_context", new=AsyncMock(return_value=mock_ctx)),
patch.object(bot, "process_commands", new=AsyncMock()) as mock_process,
):
asyncio.run(bot.on_message(message))
app_mocks.chat.handle.assert_awaited_once_with(
ctx=mock_ctx,
bot_name="alfred",
message="hi",
system_prompt="british butler",
response_prefix="alfred response",
)
mock_process.assert_not_called()
def test_on_message_falls_through_to_commands(
app_mocks: AppMocks,
bot: commands.Bot,
) -> None:
"""An unmatched ! message falls through to process_commands."""
message = MagicMock()
message.author = MagicMock()
message.author.name = "testuser"
message.content = "!unknown hi"
with patch.object(bot, "process_commands", new=AsyncMock()) as mock_process:
asyncio.run(bot.on_message(message))
mock_process.assert_awaited_once_with(message)
app_mocks.chat.handle.assert_not_awaited()
def test_on_message_uses_cache_not_manager(
app_mocks: AppMocks,
bot: commands.Bot,
) -> None:
"""Bot-name matching reads bot_cache, never list_custom_bots()."""
message = MagicMock()
message.author = MagicMock()
message.author.name = "testuser"
message.content = "!alfred hi"
with patch.object(bot, "get_context", new=AsyncMock(return_value=MagicMock())):
asyncio.run(bot.on_message(message))
app_mocks.manager.list_custom_bots.assert_not_called()
# ---------------------------------------------------------------------------
# on_ready / on_command_error
# ---------------------------------------------------------------------------
def test_on_ready(bot: commands.Bot) -> None:
"""on_ready logs startup without raising."""
handler = bot.__dict__["on_ready"]
asyncio.run(handler())
def test_on_command_error_cooldown(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A cooldown error becomes a friendly message."""
from discord.ext.commands.cooldowns import Cooldown
error = commands.CommandOnCooldown(
Cooldown(1.0, 60.0), 0.5, commands.BucketType.user
)
asyncio.run(bot.on_command_error(mock_ctx, error))
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
def test_on_command_error_generic(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Any other command error is logged and not re-raised."""
error = commands.CommandError("boom")
asyncio.run(bot.on_command_error(mock_ctx, error))
# ---------------------------------------------------------------------------
# main.py entrypoint
# ---------------------------------------------------------------------------
def test_main_builds_and_runs_bot() -> None:
"""main() validates config, configures logging, and runs the built bot."""
from vibe_bot import main as main_module
with (
patch.object(main_module, "validate_config") as mock_validate,
patch.object(main_module, "configure_logging") as mock_logging,
patch.object(
main_module, "create_app", return_value=MagicMock()
) as mock_create,
patch.object(main_module, "build_bot", return_value=MagicMock()) as mock_build,
patch.object(main_module, "DISCORD_TOKEN", "test-token"),
):
main_module.main()
mock_validate.assert_called_once_with()
mock_logging.assert_called_once_with()
mock_create.assert_called_once_with()
mock_build.assert_called_once()
mock_build.return_value.run.assert_called_once_with("test-token")
vibe_bot/tests/test_commands.py
+643
View File
@@ -0,0 +1,643 @@
"""Wiring tests for the commands package: every handler reaches its service."""
from __future__ import annotations
import asyncio
from datetime import UTC, datetime
from unittest.mock import MagicMock
import pytest
from discord.ext import commands
from vibe_bot.commands.custom_bots import MAX_PERSONALITY_LENGTH
from vibe_bot.tests._helpers import invoke, sent_texts
from vibe_bot.tests.conftest import AppMocks
# ---------------------------------------------------------------------------
# custom_bots: create / list / delete
# ---------------------------------------------------------------------------
def test_custom_bot_success_updates_cache(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A successful create invalidates the cache with the new bot."""
app_mocks.manager.create_custom_bot.return_value = "created"
app_mocks.manager.list_custom_bots.return_value = [
("newbot", "a personality", "user123"),
]
invoke(
bot, "custom-bot", mock_ctx, "newbot", personality="you are a british butler"
)
app_mocks.manager.create_custom_bot.assert_called_once_with(
bot_name="newbot",
system_prompt="you are a british butler",
created_by="12345",
)
assert app_mocks.app.bot_cache == {"newbot": ("a personality", "user123")}
assert "alfred" not in app_mocks.app.bot_cache
texts = sent_texts(mock_ctx)
assert any("has been created" in t for t in texts)
assert any("You can now use this bot" in t for t in texts)
def test_custom_bot_replaced_message(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""When the name already exists, the command reports a replace, not a create."""
app_mocks.manager.create_custom_bot.return_value = "replaced"
invoke(
bot, "custom-bot", mock_ctx, "alfred", personality="you are a british butler"
)
texts = sent_texts(mock_ctx)
assert any("already existed" in t and "replaced" in t for t in texts)
def test_custom_bot_invalid_name_too_short(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A one-character name is rejected before any manager call."""
invoke(bot, "custom-bot", mock_ctx, "a", personality="this is a valid personality")
texts = sent_texts(mock_ctx)
assert any("Invalid bot name" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_invalid_name_empty(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""An empty name is rejected before any manager call."""
invoke(bot, "custom-bot", mock_ctx, "", personality="this is a valid personality")
texts = sent_texts(mock_ctx)
assert any("Invalid bot name" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_invalid_personality(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A personality under 10 characters is rejected."""
invoke(bot, "custom-bot", mock_ctx, "testbot", personality="short")
texts = sent_texts(mock_ctx)
assert any("Invalid personality" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_personality_too_long(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A personality over MAX_PERSONALITY_LENGTH is rejected before any DB call."""
invoke(
bot,
"custom-bot",
mock_ctx,
"testbot",
personality="a" * (MAX_PERSONALITY_LENGTH + 1),
)
texts = sent_texts(mock_ctx)
assert any("Personality too long" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_create_fails(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A failed create reports an error and does not invalidate the cache."""
app_mocks.manager.create_custom_bot.return_value = False
invoke(
bot, "custom-bot", mock_ctx, "alfred", personality="you are a british butler"
)
texts = sent_texts(mock_ctx)
assert any("Failed to create custom bot" in t for t in texts)
app_mocks.manager.list_custom_bots.assert_not_called()
assert app_mocks.app.bot_cache == {"alfred": ("british butler", "user123")}
def test_list_custom_bots_empty(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Listing with no bots suggests creating one."""
app_mocks.manager.list_custom_bots.return_value = []
invoke(bot, "list-custom-bots", mock_ctx)
texts = sent_texts(mock_ctx)
assert any("No custom bots" in t for t in texts)
def test_list_custom_bots_with_bots(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Listing shows every bot name."""
app_mocks.manager.list_custom_bots.return_value = [
("alfred", "british butler", "user-1"),
("jarvis", "ai assistant", "user-2"),
]
invoke(bot, "list-custom-bots", mock_ctx)
texts = sent_texts(mock_ctx)
assert any("Available Custom Bots" in t for t in texts)
assert any("* alfred" in t and "* jarvis" in t for t in texts)
def test_delete_custom_bot_success(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""The creator can delete their bot; the cache is invalidated."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"12345",
"2024-01-01",
)
app_mocks.manager.delete_custom_bot.return_value = True
app_mocks.manager.list_custom_bots.return_value = []
invoke(bot, "delete-custom-bot", mock_ctx, "alfred")
app_mocks.manager.delete_custom_bot.assert_called_once_with("alfred")
texts = sent_texts(mock_ctx)
assert any("has been deleted" in t for t in texts)
assert app_mocks.app.bot_cache == {}
def test_delete_custom_bot_not_found(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Deleting a non-existent bot reports not found."""
app_mocks.manager.get_custom_bot.return_value = None
invoke(bot, "delete-custom-bot", mock_ctx, "nonexistent")
texts = sent_texts(mock_ctx)
assert any("not found" in t for t in texts)
app_mocks.manager.delete_custom_bot.assert_not_called()
def test_delete_custom_bot_not_owner(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A non-owner cannot delete the bot."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"other-user-id",
"2024-01-01",
)
invoke(bot, "delete-custom-bot", mock_ctx, "alfred")
texts = sent_texts(mock_ctx)
assert any("You can only delete your own" in t for t in texts)
app_mocks.manager.delete_custom_bot.assert_not_called()
def test_delete_custom_bot_delete_fails(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A failed delete reports an error and keeps the cache."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"12345",
"2024-01-01",
)
app_mocks.manager.delete_custom_bot.return_value = False
invoke(bot, "delete-custom-bot", mock_ctx, "alfred")
texts = sent_texts(mock_ctx)
assert any("Failed to delete" in t for t in texts)
app_mocks.manager.list_custom_bots.assert_not_called()
assert app_mocks.app.bot_cache == {"alfred": ("british butler", "user123")}
# ---------------------------------------------------------------------------
# speech: speak / voices
# ---------------------------------------------------------------------------
def test_speak_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!speak forwards to SpeechService.speak."""
invoke(bot, "speak", mock_ctx, message="hello world")
app_mocks.speech.speak.assert_awaited_once_with(mock_ctx, message="hello world")
def test_speak_cooldown_blocks_third_invocation(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""speak allows 3 per 30s; the 4th immediate call is rate-limited."""
cmd = bot.get_command("speak")
assert cmd is not None
assert cmd.cooldown is not None
mock_ctx.message.edited_at = None
mock_ctx.message.created_at = datetime(2026, 1, 1, tzinfo=UTC)
for _ in range(3):
cmd._prepare_cooldowns(mock_ctx)
with pytest.raises(commands.CommandOnCooldown) as exc_info:
cmd._prepare_cooldowns(mock_ctx)
asyncio.run(bot.on_command_error(mock_ctx, exc_info.value))
app_mocks.speech.speak.assert_not_awaited()
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
def test_voices(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!voices lists the voice catalog."""
invoke(bot, "voices", mock_ctx)
full = "\n".join(sent_texts(mock_ctx))
assert "Available Voices" in full
assert "af_sarah" in full
assert "Use `!speak" in full
# ---------------------------------------------------------------------------
# images: doodlebob / retcon
# ---------------------------------------------------------------------------
def test_doodlebob_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!doodlebob forwards to ImageService.generate."""
invoke(bot, "doodlebob", mock_ctx, message="a centaur")
app_mocks.image.generate.assert_awaited_once_with(mock_ctx, message="a centaur")
def test_retcon_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!retcon forwards to ImageService.edit."""
invoke(bot, "retcon", mock_ctx, message="make it blue")
app_mocks.image.edit.assert_awaited_once_with(mock_ctx, message="make it blue")
def test_doodlebob_cooldown_blocks_second_invocation(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A second immediate doodlebob invocation is rate-limited before the LLM."""
cmd = bot.get_command("doodlebob")
assert cmd is not None
assert cmd.cooldown is not None
mock_ctx.message.edited_at = None
mock_ctx.message.created_at = datetime(2026, 1, 1, tzinfo=UTC)
# First invocation consumes the single token (no error).
cmd._prepare_cooldowns(mock_ctx)
# Second immediate invocation is rejected before the service is called.
with pytest.raises(commands.CommandOnCooldown) as exc_info:
cmd._prepare_cooldowns(mock_ctx)
# The error handler turns it into a friendly message.
asyncio.run(bot.on_command_error(mock_ctx, exc_info.value))
app_mocks.image.generate.assert_not_awaited()
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
# ---------------------------------------------------------------------------
# conversation: talkforme
# ---------------------------------------------------------------------------
def test_talkforme_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!talkforme parses its args, then forwards to ConversationService.run."""
invoke(bot, "talkforme", mock_ctx, message="a b 3 talking cats")
app_mocks.conversation.run.assert_awaited_once_with(
mock_ctx, "a", "b", "3", "talking cats"
)
def test_talkforme_invalid_args(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!talkforme with too few parts shows usage and calls no service."""
invoke(bot, "talkforme", mock_ctx, message="bot1 bot2")
texts = sent_texts(mock_ctx)
assert any("Usage" in t for t in texts)
app_mocks.conversation.run.assert_not_awaited()
def test_talkforme_cooldown_blocks_second_invocation(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A second immediate talkforme invocation is rate-limited before the LLM."""
cmd = bot.get_command("talkforme")
assert cmd is not None
assert cmd.cooldown is not None
mock_ctx.message.edited_at = None
mock_ctx.message.created_at = datetime(2026, 1, 1, tzinfo=UTC)
cmd._prepare_cooldowns(mock_ctx)
with pytest.raises(commands.CommandOnCooldown) as exc_info:
cmd._prepare_cooldowns(mock_ctx)
asyncio.run(bot.on_command_error(mock_ctx, exc_info.value))
app_mocks.conversation.run.assert_not_awaited()
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
# ---------------------------------------------------------------------------
# admin: lobotomize / debug / history
# ---------------------------------------------------------------------------
def test_lobotomize(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!lobotomize clears all messages on the shared db."""
invoke(bot, "lobotomize", mock_ctx)
app_mocks.db.clear_all_messages.assert_called_once_with()
texts = sent_texts(mock_ctx)
assert any("cleared" in t for t in texts)
def test_history_bot_not_found(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!history on an unknown bot reports not found."""
app_mocks.manager.get_custom_bot.return_value = None
invoke(bot, "history", mock_ctx, "nonexistent")
texts = sent_texts(mock_ctx)
assert any("not found" in t for t in texts)
app_mocks.db.get_bot_history.assert_not_called()
def test_history_no_history(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!history on a bot with no messages says so."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
app_mocks.db.get_bot_history.return_value = []
invoke(bot, "history", mock_ctx, "alfred")
texts = sent_texts(mock_ctx)
assert any("No chat history" in t and "**alfred**" in t for t in texts)
def test_history_with_data(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!history formats the stored exchange, newest last."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
app_mocks.db.get_bot_history.return_value = [
("hello", "yes master?"),
("what time is it", "it is currently 3pm"),
]
invoke(bot, "history", mock_ctx, "alfred")
full = "\n".join(sent_texts(mock_ctx))
assert "Chat History for **alfred**" in full
assert "what time is it" in full
assert "alfred: it is currently 3pm" in full
def test_history_long_response_chunked(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Long history payloads are split into multiple sends."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
app_mocks.db.get_bot_history.return_value = [
("x" * 2000, "y" * 2000),
]
invoke(bot, "history", mock_ctx, "alfred")
assert mock_ctx.send.call_count >= 2
def test_debug_no_subcommand(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug without a subcommand shows the menu."""
invoke(bot, "debug", mock_ctx, subcommand=None)
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "Debug Menu" in call_args
assert "members" in call_args
def test_debug_members_no_guild(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug members on a channel without guild members says so."""
mock_ctx.channel.guild = None
invoke(bot, "debug", mock_ctx, subcommand="members")
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "No members found in this channel." in call_args
def test_debug_members_with_members(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug members lists the guild members."""
mock_member = MagicMock()
mock_member.display_name = "Alice"
mock_member.name = "alice"
mock_member.nick = None
mock_member.global_name = None
mock_member.status = MagicMock(value="online")
mock_ctx.channel.guild.members = [mock_member]
invoke(bot, "debug", mock_ctx, subcommand="members")
assert mock_ctx.send.called
call_args = mock_ctx.send.call_args[0][0]
assert "Alice" in call_args
assert "1 total" in call_args
def test_debug_unknown_subcommand(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug with an unknown subcommand explains the usage."""
invoke(bot, "debug", mock_ctx, subcommand="unknown")
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "Unknown debug sub-command" in call_args
def test_debug_members_many_chunks(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug members with many members exceeds the chunk limit."""
mock_members = []
for i in range(50):
mock_member = MagicMock()
mock_member.display_name = f"User{i}_with_a_very_long_display_name"
mock_member.name = f"user{i}"
mock_member.nick = None
mock_member.global_name = None
mock_member.status = MagicMock(value="online")
mock_members.append(mock_member)
mock_ctx.channel.guild.members = mock_members
invoke(bot, "debug", mock_ctx, subcommand="members")
assert mock_ctx.send.call_count >= 2
first_chunk = mock_ctx.send.call_args_list[0][0][0]
assert "Members in this channel (50 total)" in first_chunk
def test_debug_whoami(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug whoami shows user info."""
invoke(bot, "debug", mock_ctx, subcommand="whoami")
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "Username: testuser" in call_args
assert "User ID: 12345" in call_args
assert "Global Name: Test User" in call_args
assert "Nickname: tester" in call_args
def test_debug_whoami_minimal(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug whoami omits fields the user does not have."""
mock_ctx.author.global_name = None
mock_ctx.author.nick = None
mock_ctx.author.top_role.name = "@everyone"
mock_ctx.author.activities = []
mock_ctx.author.joined_at = None
mock_ctx.author.created_at = None
invoke(bot, "debug", mock_ctx, subcommand="whoami")
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "Username: testuser" in call_args
assert "User ID: 12345" in call_args
assert "Global Name" not in call_args
assert "Nickname" not in call_args
assert "Activities" not in call_args
assert "Joined" not in call_args
def test_debug_tools(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug tools shows the LLM tool schema."""
invoke(bot, "debug", mock_ctx, subcommand="tools")
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "LLM Tools" in call_args
assert "get_channel_members" in call_args
assert "members" in call_args.lower()
vibe_bot/tests/test_config.py
+64 -71
View File
@@ -4,6 +4,9 @@ from __future__ import annotations
import subprocess import subprocess
import sys import sys
from unittest.mock import patch
import pytest
def test_config_defaults() -> None: def test_config_defaults() -> None:
@@ -12,17 +15,14 @@ def test_config_defaults() -> None:
for k, v in { for k, v in {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
"CHAT_ENDPOINT_KEY": "test-key", "CHAT_ENDPOINT_KEY": "test-key",
"COMPLETION_ENDPOINT_KEY": "test-completion-key",
"IMAGE_GEN_ENDPOINT_KEY": "test-image-key", "IMAGE_GEN_ENDPOINT_KEY": "test-image-key",
"IMAGE_EDIT_ENDPOINT_KEY": "test-image-edit-key", "IMAGE_EDIT_ENDPOINT_KEY": "test-image-edit-key",
"EMBEDDING_ENDPOINT_KEY": "test-embedding-key", "EMBEDDING_ENDPOINT_KEY": "test-embedding-key",
@@ -39,8 +39,6 @@ def test_config_defaults() -> None:
env_str += f'os.environ["{k}"] = "{v}"\n' env_str += f'os.environ["{k}"] = "{v}"\n'
code = f""" code = f"""
import sys
sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots")
import os import os
os.environ.clear() os.environ.clear()
os.environ["PATH"] = "/usr/bin:/bin" os.environ["PATH"] = "/usr/bin:/bin"
@@ -48,12 +46,10 @@ os.environ["PATH"] = "/usr/bin:/bin"
import vibe_bot.config import vibe_bot.config
assert vibe_bot.config.DISCORD_TOKEN == "test-token" assert vibe_bot.config.DISCORD_TOKEN == "test-token"
assert vibe_bot.config.CHAT_ENDPOINT == "https://chat.example.com/v1" assert vibe_bot.config.CHAT_ENDPOINT == "https://chat.example.com/v1"
assert vibe_bot.config.COMPLETION_ENDPOINT == "https://completion.example.com/v1"
assert vibe_bot.config.IMAGE_GEN_ENDPOINT == "https://image.example.com/v1" assert vibe_bot.config.IMAGE_GEN_ENDPOINT == "https://image.example.com/v1"
assert vibe_bot.config.IMAGE_EDIT_ENDPOINT == "https://image-edit.example.com/v1" assert vibe_bot.config.IMAGE_EDIT_ENDPOINT == "https://image-edit.example.com/v1"
assert vibe_bot.config.EMBEDDING_ENDPOINT == "https://embedding.example.com/v1" assert vibe_bot.config.EMBEDDING_ENDPOINT == "https://embedding.example.com/v1"
assert vibe_bot.config.CHAT_MODEL == "test-chat-model" assert vibe_bot.config.CHAT_MODEL == "test-chat-model"
assert vibe_bot.config.COMPLETION_MODEL == "test-completion-model"
assert vibe_bot.config.IMAGE_GEN_MODEL == "test-image-model" assert vibe_bot.config.IMAGE_GEN_MODEL == "test-image-model"
assert vibe_bot.config.IMAGE_EDIT_MODEL == "test-image-edit-model" assert vibe_bot.config.IMAGE_EDIT_MODEL == "test-image-edit-model"
assert vibe_bot.config.EMBEDDING_MODEL == "test-embedding-model" assert vibe_bot.config.EMBEDDING_MODEL == "test-embedding-model"
@@ -78,20 +74,23 @@ print("OK")
def _run_config_check(env_vars: dict[str, str], expected_error: str) -> None: def _run_config_check(env_vars: dict[str, str], expected_error: str) -> None:
"""Run a subprocess that imports config and checks for expected RuntimeError.""" """Run a subprocess that imports config and calls validate_config().
The import itself must never raise; only validate_config() may raise the
expected RuntimeError for the missing required setting.
"""
env_str = "" env_str = ""
for k, v in env_vars.items(): for k, v in env_vars.items():
env_str += f'os.environ["{k}"] = "{v}"\n' env_str += f'os.environ["{k}"] = "{v}"\n'
code = f""" code = f"""
import sys
sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots")
import os import os
os.environ.clear() os.environ.clear()
os.environ["PATH"] = "/usr/bin:/bin" os.environ["PATH"] = "/usr/bin:/bin"
{env_str} {env_str}
try: try:
import vibe_bot.config import vibe_bot.config
vibe_bot.config.validate_config()
print("NO_ERROR") print("NO_ERROR")
except RuntimeError as e: except RuntimeError as e:
print(f"ERROR: {{e}}") print(f"ERROR: {{e}}")
@@ -116,12 +115,10 @@ def test_config_missing_discord_token() -> None:
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "", "DISCORD_TOKEN": "",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -134,12 +131,10 @@ def test_config_missing_chat_endpoint() -> None:
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "", "CHAT_ENDPOINT": "",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -147,35 +142,15 @@ def test_config_missing_chat_endpoint() -> None:
_run_config_check(env, "CHAT_ENDPOINT required") _run_config_check(env, "CHAT_ENDPOINT required")
def test_config_missing_completion_endpoint() -> None:
"""Test that RuntimeError is raised when COMPLETION_ENDPOINT is missing."""
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
}
_run_config_check(env, "COMPLETION_ENDPOINT required")
def test_config_missing_image_gen_endpoint() -> None: def test_config_missing_image_gen_endpoint() -> None:
"""Test that RuntimeError is raised when IMAGE_GEN_ENDPOINT is missing.""" """Test that RuntimeError is raised when IMAGE_GEN_ENDPOINT is missing."""
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "", "IMAGE_GEN_ENDPOINT": "",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -188,12 +163,10 @@ def test_config_missing_image_edit_endpoint() -> None:
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "", "IMAGE_EDIT_ENDPOINT": "",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -206,12 +179,10 @@ def test_config_missing_embedding_endpoint() -> None:
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "", "EMBEDDING_ENDPOINT": "",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -224,12 +195,10 @@ def test_config_missing_chat_model() -> None:
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "", "CHAT_MODEL": "",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -237,35 +206,15 @@ def test_config_missing_chat_model() -> None:
_run_config_check(env, "CHAT_MODEL required") _run_config_check(env, "CHAT_MODEL required")
def test_config_missing_completion_model() -> None:
"""Test that RuntimeError is raised when COMPLETION_MODEL is missing."""
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
}
_run_config_check(env, "COMPLETION_MODEL required")
def test_config_missing_image_gen_model() -> None: def test_config_missing_image_gen_model() -> None:
"""Test that RuntimeError is raised when IMAGE_GEN_MODEL is missing.""" """Test that RuntimeError is raised when IMAGE_GEN_MODEL is missing."""
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "", "IMAGE_GEN_MODEL": "",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -278,12 +227,10 @@ def test_config_missing_image_edit_model() -> None:
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "", "IMAGE_EDIT_MODEL": "",
"EMBEDDING_MODEL": "test-embedding-model", "EMBEDDING_MODEL": "test-embedding-model",
@@ -296,12 +243,10 @@ def test_config_missing_embedding_model() -> None:
env: dict[str, str] = { env: dict[str, str] = {
"DISCORD_TOKEN": "test-token", "DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1", "CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1", "IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1", "IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1", "EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model", "CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model", "IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model", "IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "", "EMBEDDING_MODEL": "",
@@ -309,16 +254,64 @@ def test_config_missing_embedding_model() -> None:
_run_config_check(env, "EMBEDDING_MODEL required") _run_config_check(env, "EMBEDDING_MODEL required")
REQUIRED_VARS = (
"DISCORD_TOKEN",
"CHAT_ENDPOINT",
"IMAGE_GEN_ENDPOINT",
"IMAGE_EDIT_ENDPOINT",
"EMBEDDING_ENDPOINT",
"CHAT_MODEL",
"IMAGE_GEN_MODEL",
"IMAGE_EDIT_MODEL",
"EMBEDDING_MODEL",
)
def test_validate_config_passes_with_full_env() -> None:
"""With all required settings present, validate_config() is a no-op."""
from vibe_bot import config
config.validate_config()
@pytest.mark.parametrize("var_name", REQUIRED_VARS)
def test_validate_config_missing_var(var_name: str) -> None:
"""Blanking any single required setting makes validate_config() raise."""
from vibe_bot import config
with (
patch.object(config, var_name, ""),
pytest.raises(RuntimeError, match=f"{var_name} required"),
):
config.validate_config()
def test_import_config_empty_env_no_raise_no_logging() -> None:
"""In an empty env the import succeeds and leaves the root logger unconfigured."""
code = """
import logging
import os
os.environ.clear()
os.environ["PATH"] = "/usr/bin:/bin"
import vibe_bot.config
handlers = logging.getLogger().handlers
assert handlers == [], f"config configured logging: {handlers}"
print("OK")
"""
result = subprocess.run( # noqa: PLW1510
[sys.executable, "-c", code],
capture_output=True,
text=True,
timeout=30,
)
assert result.returncode == 0, f"Subprocess failed: {result.stderr}"
assert "OK" in result.stdout
def test_config_logging_exists() -> None: def test_config_logging_exists() -> None:
"""Test that logging is configured in config module.""" """Test that logging is configured in config module."""
from vibe_bot.config import logger from vibe_bot.config import logger
assert logger is not None assert logger is not None
assert logger.name == "vibe_bot.config" assert logger.name == "vibe_bot.config"
def test_config_embedding_dimension() -> None:
"""Test that EMBEDDING_DIMENSION has expected default value."""
from vibe_bot.config import EMBEDDING_DIMENSION
assert EMBEDDING_DIMENSION == 2048
vibe_bot/tests/test_database.py
+774 -44
View File
@@ -2,17 +2,40 @@
from __future__ import annotations from __future__ import annotations
from typing import TYPE_CHECKING, Any from datetime import datetime
from typing import TYPE_CHECKING, Any, cast
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
import numpy as np import numpy as np
import pytest import pytest
if TYPE_CHECKING: if TYPE_CHECKING:
import sqlite3
from vibe_bot.database import ChatDatabase from vibe_bot.database import ChatDatabase
def _recent_messages(
db_path: str,
limit: int,
) -> list[tuple[str, str, str, datetime]]:
"""Read the newest rows straight from the database (test-side helper)."""
from vibe_bot.db.connection import connect
conn = connect(db_path)
try:
rows = conn.execute(
"SELECT message_id, username, content, timestamp "
"FROM chat_messages ORDER BY timestamp DESC LIMIT ?",
(limit,),
).fetchall()
finally:
conn.close()
return [
(str(row[0]), str(row[1]), str(row[2]), cast(datetime, row[3])) for row in rows
]
def test_vector_to_bytes(chat_db: ChatDatabase) -> None: def test_vector_to_bytes(chat_db: ChatDatabase) -> None:
"""Test converting a vector to bytes and back.""" """Test converting a vector to bytes and back."""
vector: list[float] = [0.1, 0.2, 0.3, 0.4] vector: list[float] = [0.1, 0.2, 0.3, 0.4]
@@ -55,6 +78,13 @@ def test_calculate_similarity_negative(chat_db: ChatDatabase) -> None:
assert similarity == pytest.approx(-1.0, abs=1e-6) assert similarity == pytest.approx(-1.0, abs=1e-6)
def test_calculate_similarity_zero_norm(chat_db: ChatDatabase) -> None:
"""A zero vector has no direction, so its similarity is 0."""
zero = np.zeros(3, dtype=np.float32)
other = np.array([1.0, 0.0, 0.0], dtype=np.float32)
assert chat_db._calculate_similarity(zero, other) == 0.0
def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None: def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
"""Test adding a message to the database.""" """Test adding a message to the database."""
result = chat_db.add_message( result = chat_db.add_message(
@@ -67,7 +97,7 @@ def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
) )
assert result is True assert result is True
messages = chat_db.get_recent_messages(limit=10) messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 1 assert len(messages) == 1
assert messages[0][0] == "msg-1" assert messages[0][0] == "msg-1"
assert messages[0][1] == "testuser" assert messages[0][1] == "testuser"
@@ -76,7 +106,7 @@ def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
def test_add_message_no_embedding(chat_db: ChatDatabase) -> None: def test_add_message_no_embedding(chat_db: ChatDatabase) -> None:
"""Test adding a message when embedding generation fails.""" """Test adding a message when embedding generation fails."""
with patch("vibe_bot.llama_wrapper.embedding", return_value=None): with patch("vibe_bot.llm_client.embedding", return_value=None):
result = chat_db.add_message( result = chat_db.add_message(
message_id="msg-no-embed", message_id="msg-no-embed",
user_id="user-1", user_id="user-1",
@@ -106,14 +136,14 @@ def test_add_message_duplicate(
content="Second content", content="Second content",
) )
messages = chat_db.get_recent_messages(limit=10) messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 1 assert len(messages) == 1
assert messages[0][2] == "Second content" assert messages[0][2] == "Second content"
def test_add_message_failure(chat_db: ChatDatabase) -> None: def test_add_message_failure(chat_db: ChatDatabase) -> None:
"""Test that add_message returns False on database error.""" """Test that add_message returns False on database error."""
with patch.object(chat_db, "_vector_to_bytes", side_effect=Exception("fail")): with patch("vibe_bot.db.messages.connect", return_value=_broken_connection()):
result = chat_db.add_message( result = chat_db.add_message(
message_id="msg-fail", message_id="msg-fail",
user_id="user-1", user_id="user-1",
@@ -123,11 +153,304 @@ def test_add_message_failure(chat_db: ChatDatabase) -> None:
assert result is False assert result is False
def test_get_recent_messages( def _embedding_row_count(db_path: str) -> int:
"""Count the rows stored in message_embeddings."""
import sqlite3
conn = sqlite3.connect(db_path)
count = conn.execute("SELECT COUNT(*) FROM message_embeddings").fetchone()[0]
conn.close()
return int(count)
def test_add_message_embed_false_skips_embedding(chat_db: ChatDatabase) -> None:
"""embed=False neither calls the embedding API nor stores a row."""
with patch("vibe_bot.llm_client.embedding") as mock_embedding:
result = chat_db.add_message(
message_id="msg-assist",
user_id="bot-1",
username="some-bot",
content="assistant reply",
role="assistant",
embed=False,
)
assert result is True
mock_embedding.assert_not_called()
assert _embedding_row_count(chat_db.db_path) == 0
def test_add_message_stores_embedding_by_default(
chat_db: ChatDatabase, chat_db: ChatDatabase,
mock_embedding: MagicMock, mock_embedding: MagicMock,
) -> None: ) -> None:
"""Test retrieving recent messages.""" """The default embed=True still stores exactly one embedding row."""
assert chat_db.add_message(
message_id="msg-embed",
user_id="user-1",
username="testuser",
content="default embed",
)
assert _embedding_row_count(chat_db.db_path) == 1
def test_add_message_stores_embedding_norm(chat_db: ChatDatabase) -> None:
"""add_message stores the L2 norm of the float32-encoded embedding."""
import sqlite3
vector: list[float] = [0.6, 0.8]
with patch("vibe_bot.llm_client.embedding", return_value=vector):
assert chat_db.add_message(
message_id="norm-1",
user_id="u1",
username="alice",
content="normed message",
)
conn = sqlite3.connect(chat_db.db_path)
blob, norm = conn.execute(
"SELECT embedding, norm FROM message_embeddings WHERE message_id = 'norm-1'"
).fetchone()
conn.close()
stored = np.frombuffer(blob, dtype=np.float32)
assert float(norm) == pytest.approx(float(np.linalg.norm(stored)), abs=1e-6)
def test_cleanup_old_messages_no_orphaned_embeddings(chat_db: ChatDatabase) -> None:
"""Deleting the oldest rows must delete their embeddings, not the next ones.
Regression test for the bug where the embedding cleanup re-queried
``chat_messages`` *after* the message delete, so it stripped embeddings
from the next-oldest live rows while leaving the deleted rows' embeddings
behind as orphans.
"""
import sqlite3
with patch("vibe_bot.db.messages.MAX_HISTORY_MESSAGES", 5):
# Seed 7 rows with distinct ascending timestamps and an embedding for
# each, so the "oldest" ordering is deterministic.
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
for i in range(1, 8):
ts = f"2024-01-0{i} 00:00:00"
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, bot_name, timestamp) "
"VALUES (?, ?, ?, ?, ?, ?)",
(f"m{i}", "u1", "alice", f"content {i}", "bot", ts),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding) "
"VALUES (?, ?)",
(f"m{i}", chat_db._vector_to_bytes([0.1, 0.2, 0.3])),
)
conn.commit()
conn.close()
# The 8th insert pushes the count to 8 (> 5), so add_message's
# cleanup must delete exactly the 3 oldest rows and their embeddings.
assert chat_db.add_message(
message_id="m8",
user_id="u1",
username="alice",
content="content 8",
)
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
cursor.execute("SELECT message_id FROM chat_messages")
live = {row[0] for row in cursor.fetchall()}
cursor.execute("SELECT message_id FROM message_embeddings")
embedded = {row[0] for row in cursor.fetchall()}
conn.close()
# The three oldest rows are gone; the rest (incl. the new one) remain.
assert live == {"m4", "m5", "m6", "m7", "m8"}
# No orphaned embeddings and no stripped survivors: the embedding set
# exactly matches the live message rows.
assert embedded == live
def test_role_scopes_history_and_search(chat_db: ChatDatabase) -> None:
"""get_user_history excludes responses; search matches only user rows."""
chat_db.add_message(
message_id="r-1",
user_id="u1",
username="alice",
content="User asks about the weather",
role="user",
)
chat_db.add_message(
message_id="r-1_response",
user_id="bot",
username="some-bot",
content="Bot answers the weather",
role="assistant",
)
# get_user_history returns only the user row (paired with its response).
conversations = chat_db.get_user_history("u1")
assert len(conversations) == 1
assert conversations[0][0] == "User asks about the weather"
assert conversations[0][1] == "Bot answers the weather"
# search_similar_messages only considers user rows, never responses.
results = chat_db.search_similar_messages(
"User asks about the weather", top_k=5, min_similarity=0.0
)
assert len(results) == 1
assert results[0][0] == "User asks about the weather"
assert results[0][1] == "Bot answers the weather"
def test_role_migration_backfills_legacy_rows(
temp_db_path: str,
) -> None:
"""Legacy rows (no role column) are backfilled when ChatDatabase inits."""
import sqlite3
from vibe_bot.database import ChatDatabase
# Create the legacy schema (no role column) with a user row and its
# response row inserted directly.
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT)"
)
cursor.execute(
"INSERT INTO chat_messages (message_id, user_id, username, content) "
"VALUES ('legacy-1', 'u1', 'alice', 'old question')"
)
cursor.execute(
"INSERT INTO chat_messages (message_id, user_id, username, content) "
"VALUES ('legacy-1_response', 'bot', 'old-bot', 'old answer')"
)
conn.commit()
conn.close()
# Initializing ChatDatabase should add and backfill the role column.
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute("SELECT message_id, role FROM chat_messages ORDER BY message_id")
roles = {row[0]: row[1] for row in cursor.fetchall()}
conn.close()
assert roles["legacy-1"] == "user"
assert roles["legacy-1_response"] == "assistant"
def test_bot_name_migration_adds_column(
temp_db_path: str,
) -> None:
"""A pre-bot_name schema gets the column added when ChatDatabase inits."""
import sqlite3
from vibe_bot.database import ChatDatabase
conn = sqlite3.connect(temp_db_path)
conn.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT, role TEXT)"
)
conn.commit()
conn.close()
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute("PRAGMA table_info(chat_messages)")
columns = {row[1] for row in cursor.fetchall()}
conn.close()
assert "bot_name" in columns
def _seed_pre_norm_db(db_path: str) -> list[tuple[str, str, list[float]]]:
"""Create a pre-norm-schema database (no norm column) with seeded rows."""
import sqlite3
rows: list[tuple[str, str, list[float]]] = [
("n-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-5", "ask about nothing", [0.0] * 8),
]
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT, role TEXT)"
)
cursor.execute(
"CREATE TABLE message_embeddings ("
"message_id TEXT PRIMARY KEY, embedding BLOB,"
"FOREIGN KEY (message_id) REFERENCES chat_messages(message_id))"
)
for i, (message_id, content, vector) in enumerate(rows):
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'u1', 'alice', ?, 'user', ?)",
(message_id, content, f"2024-01-0{i + 1} 00:00:00"),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', ?)",
(
f"{message_id}_response",
f"response {i + 1}",
f"2024-01-0{i + 1} 00:00:01",
),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding) VALUES (?, ?)",
(message_id, np.array(vector, dtype=np.float32).tobytes()),
)
conn.commit()
conn.close()
return rows
def test_norm_migration_backfills_stored_norms(temp_db_path: str) -> None:
"""A pre-norm database gets the norm column added and backfilled on init."""
import sqlite3
from vibe_bot.database import ChatDatabase
rows = _seed_pre_norm_db(temp_db_path)
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
stored = dict(conn.execute("SELECT message_id, norm FROM message_embeddings"))
conn.close()
for message_id, _content, vector in rows:
expected = float(np.linalg.norm(np.array(vector, dtype=np.float32)))
assert stored[message_id] == pytest.approx(expected, abs=1e-5)
def test_recent_messages_desc_order(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""Newest-first ordering of stored messages."""
chat_db.add_message( chat_db.add_message(
message_id="msg-1", message_id="msg-1",
user_id="u1", user_id="u1",
@@ -147,17 +470,17 @@ def test_get_recent_messages(
content="Third", content="Third",
) )
messages = chat_db.get_recent_messages(limit=2) messages = _recent_messages(chat_db.db_path, 2)
assert len(messages) == 2 assert len(messages) == 2
assert messages[0][2] == "Third" assert messages[0][2] == "Third"
assert messages[1][2] == "Second" assert messages[1][2] == "Second"
def test_get_recent_messages_limit( def test_recent_messages_limit(
chat_db: ChatDatabase, chat_db: ChatDatabase,
mock_embedding: MagicMock, mock_embedding: MagicMock,
) -> None: ) -> None:
"""Test that get_recent_messages respects the limit.""" """The newest-rows query respects the limit."""
for i in range(5): for i in range(5):
chat_db.add_message( chat_db.add_message(
message_id=f"msg-{i}", message_id=f"msg-{i}",
@@ -166,10 +489,25 @@ def test_get_recent_messages_limit(
content=f"Message {i}", content=f"Message {i}",
) )
messages = chat_db.get_recent_messages(limit=3) messages = _recent_messages(chat_db.db_path, 3)
assert len(messages) == 3 assert len(messages) == 3
def test_recent_messages_returns_datetime(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""The timestamp column comes back as a real datetime, not a string."""
chat_db.add_message(
message_id="dt-1",
user_id="u1",
username="alice",
content="fresh message",
)
messages = _recent_messages(chat_db.db_path, 1)
assert isinstance(messages[0][3], datetime)
def test_clear_all_messages( def test_clear_all_messages(
chat_db: ChatDatabase, chat_db: ChatDatabase,
mock_embedding: MagicMock, mock_embedding: MagicMock,
@@ -190,7 +528,7 @@ def test_clear_all_messages(
chat_db.clear_all_messages() chat_db.clear_all_messages()
messages = chat_db.get_recent_messages(limit=10) messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 0 assert len(messages) == 0
@@ -264,6 +602,25 @@ def test_image_generation_estimate_after_capping(
assert estimate == pytest.approx(99.5) assert estimate == pytest.approx(99.5)
def _broken_connection() -> MagicMock:
"""A mock connection whose first cursor.execute raises."""
fake_conn = MagicMock()
fake_conn.cursor.return_value.execute.side_effect = Exception("db error")
return fake_conn
def test_record_image_generation_time_failure(chat_db: ChatDatabase) -> None:
"""A database error while recording yields False, not an exception."""
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
assert chat_db.record_image_generation_time(1.0) is False
def test_image_generation_estimate_failure(chat_db: ChatDatabase) -> None:
"""A database error while reading yields None, not an exception."""
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
assert chat_db.get_image_generation_time_estimate() is None
def test_get_user_history( def test_get_user_history(
chat_db: ChatDatabase, chat_db: ChatDatabase,
mock_embedding: MagicMock, mock_embedding: MagicMock,
@@ -320,6 +677,39 @@ def test_get_user_history_excludes_bot(
assert len(conversations) == 0 assert len(conversations) == 0
def test_get_bot_history(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""get_bot_history pairs user messages with responses for one bot."""
chat_db.add_message(
message_id="bh-1",
user_id="u1",
username="alice",
content="bot question",
bot_name="alfred",
)
chat_db.add_message(
message_id="bh-1_response",
user_id="bot-1",
username="some-bot",
content="bot answer",
bot_name="alfred",
role="assistant",
embed=False,
)
chat_db.add_message(
message_id="bh-2",
user_id="u1",
username="alice",
content="unanswered question",
bot_name="alfred",
)
history = chat_db.get_bot_history("alfred")
assert history == [("bot question", "bot answer")]
def test_get_conversation_context( def test_get_conversation_context(
chat_db: ChatDatabase, chat_db: ChatDatabase,
mock_embedding: MagicMock, mock_embedding: MagicMock,
@@ -349,21 +739,333 @@ def test_get_conversation_context_empty(chat_db: ChatDatabase) -> None:
assert context == [] assert context == []
def _query_vector() -> list[float]:
"""A 8-dim query vector along the first axis."""
return [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
def _seed_search_rows(chat_db: ChatDatabase) -> list[tuple[str, str, list[float]]]:
"""Seed user/response rows with known embeddings (plus exclusion traps)."""
import sqlite3
rows: list[tuple[str, str, list[float]]] = [
("s-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
]
def stored_norm(vector: list[float]) -> float:
return float(np.linalg.norm(np.array(vector, dtype=np.float32)))
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
for i, (message_id, content, vector) in enumerate(rows):
timestamp = f"2024-01-0{i + 1} 00:00:00"
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, ?, ?, ?, 'user', ?)",
(message_id, "u1", "alice", content, timestamp),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, ?, ?, ?, 'assistant', ?)",
(
f"{message_id}_response",
"bot-1",
"some-bot",
f"response {i + 1}",
f"2024-01-0{i + 1} 00:00:01",
),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(message_id, chat_db._vector_to_bytes(vector), stored_norm(vector)),
)
# A user row with the top possible similarity but no response row: the
# JOIN must exclude it instead of returning a NULL response.
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES ('s-5', 'u1', 'alice', 'orphan question', 'user', "
"'2024-01-05 00:00:00')",
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(
"s-5",
chat_db._vector_to_bytes(_query_vector()),
stored_norm(_query_vector()),
),
)
# An assistant row carrying an embedding: the role filter must skip it.
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES ('s-6', 'bot-1', 'some-bot', 'assistant noise', 'assistant', "
"'2024-01-06 00:00:00')",
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(
"s-6",
chat_db._vector_to_bytes(_query_vector()),
stored_norm(_query_vector()),
),
)
conn.commit()
conn.close()
return rows
def test_search_matches_reference_topk_and_ordering(
chat_db: ChatDatabase,
) -> None:
"""The JOINed search matches the per-row reference: same top-k and order."""
rows = _seed_search_rows(chat_db)
query = _query_vector()
expected: list[tuple[str, str, float]] = []
query_arr = np.array(query, dtype=np.float32)
for i, (_message_id, content, vector) in enumerate(rows):
stored = np.frombuffer(chat_db._vector_to_bytes(vector), dtype=np.float32)
similarity = float(
np.dot(query_arr, stored)
/ (np.linalg.norm(query_arr) * np.linalg.norm(stored))
)
expected.append((content, f"response {i + 1}", similarity))
expected.sort(key=lambda item: item[2], reverse=True)
expected_top = expected[:3]
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = chat_db.search_similar_messages(
"query text", top_k=3, min_similarity=0.0
)
assert len(results) == 3
for (content, response, actual), (_ec, er, reference) in zip(
results, expected_top, strict=True
):
assert content == _ec
assert response == er
assert actual == pytest.approx(reference, abs=1e-5)
def test_search_with_stored_norms_matches_pre_norm_reference(
temp_db_path: str,
) -> None:
"""Stored-norm search is identical to per-vector renormalization.
Seeds a pre-norm database, migrates it, then asserts the stored-norm
search (top-k + ordering + similarities) matches a reference that
renormalizes every candidate vector inline — the pre-optimization
algorithm.
"""
from vibe_bot.database import ChatDatabase
rows = _seed_pre_norm_db(temp_db_path)
db = ChatDatabase(db_path=temp_db_path)
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=3, min_similarity=0.0)
query_arr = np.array(query, dtype=np.float32)
query_norm = float(np.linalg.norm(query_arr))
expected: list[tuple[str, str, float]] = []
for i, (_message_id, content, vector) in enumerate(rows):
stored = np.array(vector, dtype=np.float32)
stored_norm = float(np.linalg.norm(stored))
similarity = (
0.0
if stored_norm == 0
else float(np.dot(query_arr, stored) / (query_norm * stored_norm))
)
expected.append((content, f"response {i + 1}", similarity))
expected.sort(key=lambda item: item[2], reverse=True)
assert len(results) == 3
for (content, response, actual), (ec, er, reference) in zip(
results, expected[:3], strict=True
):
assert content == ec
assert response == er
assert actual == pytest.approx(reference, abs=1e-6)
def test_search_null_norm_row_scores_zero(temp_db_path: str) -> None:
"""A row whose norm was never written scores 0 instead of crashing."""
import sqlite3
from vibe_bot.database import ChatDatabase
_rows = _seed_pre_norm_db(temp_db_path)
db = ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
conn.execute("UPDATE message_embeddings SET norm = NULL WHERE message_id = 'n-1'")
conn.commit()
conn.close()
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
scores = {content: similarity for content, _response, similarity in results}
assert scores["ask about the sky"] == 0.0
# All four scored rows plus the zero-vector row (0.0 passes min_similarity=0.0).
assert len(results) == 5
def test_search_mixed_embedding_dims_falls_back_to_per_row(temp_db_path: str) -> None:
"""Rows with different blob lengths (mid-life model change) don't crash."""
import sqlite3
from vibe_bot.database import ChatDatabase
wide = [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
narrow = [0.5, 0.5, 0.0, 0.0]
db = ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
for message_id, content, vector in (
("m-1", "wide question", wide),
("m-2", "narrow question", narrow),
):
blob = np.array(vector, dtype=np.float32).tobytes()
stored_norm = float(np.linalg.norm(np.frombuffer(blob, dtype=np.float32)))
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'u1', 'alice', ?, 'user', '2024-01-01 00:00:00')",
(message_id, content),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', '2024-01-01 00:00:01')",
(f"{message_id}_response", f"{message_id} answer"),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(message_id, blob, stored_norm),
)
conn.commit()
conn.close()
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
scores = {content: similarity for content, _response, similarity in results}
# Both rows are returned: the 8-dim row scores exactly, the 4-dim row is
# zero-padded to the query dim and still scores sanely.
assert set(scores) == {"wide question", "narrow question"}
wide_norm = float(np.linalg.norm(np.array(wide, dtype=np.float32)))
narrow_norm = float(np.linalg.norm(np.array(narrow, dtype=np.float32)))
assert scores["wide question"] == pytest.approx(0.9 / wide_norm, abs=1e-5)
assert scores["narrow question"] == pytest.approx(0.5 / narrow_norm, abs=1e-5)
assert scores["wide question"] > scores["narrow question"]
def test_search_issues_single_select_per_call(chat_db: ChatDatabase) -> None:
"""search_similar_messages issues exactly one SELECT per call (no N+1)."""
import vibe_bot.db.search as db_search
from vibe_bot.db.connection import connect as realconnect
chat_db.add_message(
message_id="n-1",
user_id="u1",
username="alice",
content="one question",
)
chat_db.add_message(
message_id="n-1_response",
user_id="bot-1",
username="some-bot",
content="one answer",
role="assistant",
embed=False,
)
select_statements: list[str] = []
class _TracingCursor:
def __init__(self, cursor: sqlite3.Cursor) -> None:
self._cursor = cursor
def execute(self, sql: str, *args: Any) -> Any:
if sql.strip().upper().startswith("SELECT"):
select_statements.append(sql)
return self._cursor.execute(sql, *args)
def fetchall(self) -> Any:
return self._cursor.fetchall()
def fetchone(self) -> Any:
return self._cursor.fetchone()
class _TracingConnection:
def __init__(self, conn: sqlite3.Connection) -> None:
self._conn = conn
def cursor(self) -> _TracingCursor:
return _TracingCursor(self._conn.cursor())
def close(self) -> None:
self._conn.close()
def tracingconnect(db_path: str) -> _TracingConnection:
return _TracingConnection(realconnect(db_path))
with patch.object(db_search, "connect", side_effect=tracingconnect):
results = chat_db.search_similar_messages(
"one question", top_k=5, min_similarity=0.0
)
assert len(results) == 1
assert results[0][0] == "one question"
assert results[0][1] == "one answer"
assert len(select_statements) == 1
def test_search_empty_query_embedding_returns_empty(chat_db: ChatDatabase) -> None:
"""A failed query embedding yields no results."""
with patch("vibe_bot.llm_client.embedding", return_value=[]):
assert chat_db.search_similar_messages("anything") == []
def test_search_zero_query_vector_returns_empty(chat_db: ChatDatabase) -> None:
"""A zero-norm query vector yields no results (no division by zero)."""
with patch("vibe_bot.llm_client.embedding", return_value=[0.0] * 8):
assert chat_db.search_similar_messages("anything") == []
def test_custom_bot_create(custom_bot_manager: Any) -> None: def test_custom_bot_create(custom_bot_manager: Any) -> None:
"""Test creating a custom bot.""" """Test creating a custom bot returns "created" for a new name."""
result = custom_bot_manager.create_custom_bot( result = custom_bot_manager.create_custom_bot(
bot_name="alfred", bot_name="alfred",
system_prompt="You are a british butler", system_prompt="You are a british butler",
created_by="user-123", created_by="user-123",
) )
assert result is True assert result == "created"
def test_custom_bot_create_duplicate( def test_custom_bot_create_duplicate(
custom_bot_manager: Any, custom_bot_manager: Any,
) -> None: ) -> None:
"""Test creating a duplicate custom bot replaces the old one.""" """Test creating a duplicate custom bot replaces the old one."""
custom_bot_manager.create_custom_bot( first = custom_bot_manager.create_custom_bot(
bot_name="alfred", bot_name="alfred",
system_prompt="First personality", system_prompt="First personality",
created_by="user-1", created_by="user-1",
@@ -373,13 +1075,25 @@ def test_custom_bot_create_duplicate(
system_prompt="Second personality", system_prompt="Second personality",
created_by="user-1", created_by="user-1",
) )
assert result is True assert first == "created"
assert result == "replaced"
bot = custom_bot_manager.get_custom_bot("alfred") bot = custom_bot_manager.get_custom_bot("alfred")
assert bot is not None assert bot is not None
assert bot[1] == "Second personality" assert bot[1] == "Second personality"
def test_custom_bot_create_failure(custom_bot_manager: Any) -> None:
"""A database error while creating yields False, not an exception."""
with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
result = custom_bot_manager.create_custom_bot(
bot_name="failbot",
system_prompt="a long enough personality",
created_by="user-1",
)
assert result is False
def test_custom_bot_create_case_insensitive( def test_custom_bot_create_case_insensitive(
custom_bot_manager: Any, custom_bot_manager: Any,
) -> None: ) -> None:
@@ -399,6 +1113,18 @@ def test_custom_bot_get_not_found(custom_bot_manager: Any) -> None:
assert result is None assert result is None
def test_custom_bot_get_returns_datetime(custom_bot_manager: Any) -> None:
"""created_at comes back as a real datetime, not a string."""
custom_bot_manager.create_custom_bot(
bot_name="dtbot",
system_prompt="a long enough personality",
created_by="user-1",
)
result = custom_bot_manager.get_custom_bot("dtbot")
assert result is not None
assert isinstance(result[3], datetime)
def test_custom_bot_get_returns_correct_data( def test_custom_bot_get_returns_correct_data(
custom_bot_manager: Any, custom_bot_manager: Any,
) -> None: ) -> None:
@@ -413,8 +1139,7 @@ def test_custom_bot_get_returns_correct_data(
assert result[0] == "testbot" assert result[0] == "testbot"
assert result[1] == "test prompt" assert result[1] == "test prompt"
assert result[2] == "creator-1" assert result[2] == "creator-1"
assert result[3] is not None assert isinstance(result[3], datetime)
assert "20" in result[3]
def test_custom_bot_list_empty(custom_bot_manager: Any) -> None: def test_custom_bot_list_empty(custom_bot_manager: Any) -> None:
@@ -440,6 +1165,23 @@ def test_custom_bot_list(custom_bot_manager: Any) -> None:
assert len(bots) == 2 assert len(bots) == 2
def test_custom_bot_list_by_creator(custom_bot_manager: Any) -> None:
"""list_custom_bots filters by creator when user_id is given."""
custom_bot_manager.create_custom_bot(
bot_name="bot-x",
system_prompt="prompt x",
created_by="user-1",
)
custom_bot_manager.create_custom_bot(
bot_name="bot-y",
system_prompt="prompt y",
created_by="user-2",
)
bots = custom_bot_manager.list_custom_bots(user_id="user-1")
assert [bot[0] for bot in bots] == ["bot-x"]
def test_custom_bot_delete(custom_bot_manager: Any) -> None: def test_custom_bot_delete(custom_bot_manager: Any) -> None:
"""Test deleting a custom bot.""" """Test deleting a custom bot."""
custom_bot_manager.create_custom_bot( custom_bot_manager.create_custom_bot(
@@ -462,32 +1204,18 @@ def test_custom_bot_delete_nonexistent(
assert result is False assert result is False
def test_custom_bot_deactivate(custom_bot_manager: Any) -> None: def test_custom_bot_delete_failure(custom_bot_manager: Any) -> None:
"""Test deactivating a custom bot.""" """A database error while deleting yields False, not an exception."""
custom_bot_manager.create_custom_bot( with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
bot_name="inactive-bot", assert custom_bot_manager.delete_custom_bot("whatever") is False
system_prompt="will be deactivated",
created_by="user-1",
)
result = custom_bot_manager.deactivate_custom_bot("inactive-bot")
assert result is True
bot = custom_bot_manager.get_custom_bot("inactive-bot")
assert bot is None
def test_custom_bot_deactivate_nonexistent(
custom_bot_manager: Any,
) -> None:
"""Test deactivating a non-existent bot returns False."""
result = custom_bot_manager.deactivate_custom_bot("nonexistent")
assert result is False
def test_custom_bot_list_excludes_inactive( def test_custom_bot_list_excludes_inactive(
custom_bot_manager: Any, custom_bot_manager: Any,
) -> None: ) -> None:
"""Test that list_custom_bots excludes deactivated bots.""" """Test that list_custom_bots excludes bots with is_active = 0."""
import sqlite3
custom_bot_manager.create_custom_bot( custom_bot_manager.create_custom_bot(
bot_name="active-bot", bot_name="active-bot",
system_prompt="stays active", system_prompt="stays active",
@@ -498,7 +1226,12 @@ def test_custom_bot_list_excludes_inactive(
system_prompt="should not appear", system_prompt="should not appear",
created_by="user-1", created_by="user-1",
) )
custom_bot_manager.deactivate_custom_bot("deactivated-bot") conn = sqlite3.connect(custom_bot_manager.db_path)
conn.execute(
"UPDATE custom_bots SET is_active = 0 WHERE bot_name = 'deactivated-bot'"
)
conn.commit()
conn.close()
bots = custom_bot_manager.list_custom_bots() bots = custom_bot_manager.list_custom_bots()
assert len(bots) == 1 assert len(bots) == 1
@@ -532,16 +1265,13 @@ def test_database_get_database_singleton(temp_db_path: str) -> None:
db2 = get_database() db2 = get_database()
assert db1 is db2 assert db1 is db2
db1.client.close()
def test_database_init_creates_tables(temp_db_path: str) -> None: def test_database_init_creates_tables(temp_db_path: str) -> None:
"""Test that database initialization creates the expected tables.""" """Test that database initialization creates the expected tables."""
from vibe_bot.database import ChatDatabase, CustomBotManager from vibe_bot.database import ChatDatabase, CustomBotManager
db = ChatDatabase(db_path=temp_db_path) ChatDatabase(db_path=temp_db_path)
CustomBotManager(db_path=temp_db_path) CustomBotManager(db_path=temp_db_path)
db.client.close()
import sqlite3 import sqlite3
vibe_bot/tests/test_docs.py
+109
View File
@@ -0,0 +1,109 @@
"""Docs-as-tests: README.md stays in sync with the real commands and file tree."""
from __future__ import annotations
import re
import subprocess
from pathlib import Path
from discord.ext import commands
REPO_ROOT = Path(__file__).resolve().parents[2]
README = REPO_ROOT / "README.md"
BINARY_SUFFIXES = {
".bin",
".db",
".gif",
".ico",
".jpeg",
".jpg",
".mp3",
".onnx",
".png",
".wav",
}
def _git(args: list[str]) -> list[str]:
result = subprocess.run(
["git", *args],
cwd=REPO_ROOT,
capture_output=True,
text=True,
check=True,
)
return result.stdout.splitlines()
def _committed_files() -> set[str]:
"""Files that make up the current tree: index minus deletions, plus untracked."""
deleted: set[str] = set()
for line in _git(["status", "--porcelain"]):
if "D" in line[:2]:
deleted.add(line[3:])
tracked = set(_git(["ls-files"])) - deleted
untracked = set(_git(["ls-files", "--others", "--exclude-standard"]))
files: set[str] = set()
for path in tracked | untracked:
parts = Path(path).parts
if any(part.startswith(".") for part in parts):
continue
if Path(path).suffix.lower() in BINARY_SUFFIXES:
continue
files.add(path)
return files
def _readme_tree_paths() -> set[str]:
"""Reconstruct the relative paths (and directory entries) from the README tree."""
text = README.read_text(encoding="utf-8")
match = re.search(r"## File Structure\s*```text\n(.*?)```", text, re.DOTALL)
assert match is not None, "File Structure tree block not found in README"
lines = match.group(1).splitlines()
assert lines, "File Structure tree block is empty"
root = lines[0].split("#")[0].strip().rstrip("/")
assert root == REPO_ROOT.name, f"Tree root {root!r} != repo dir {REPO_ROOT.name!r}"
stack: list[str] = []
paths: set[str] = set()
entry_re = re.compile(
r"^(?P<prefix>(?:[│ ] )*)(?:├── |└── )(?P<name>\S.*?)(?:\s+#.*)?$"
)
for line in lines[1:]:
m = entry_re.match(line)
assert m is not None, f"Unparseable tree line: {line!r}"
depth = len(m.group("prefix")) // 4
name = m.group("name").strip()
if name.endswith("/"):
stack = stack[:depth] + [name.rstrip("/")]
paths.add("/".join(stack) + "/")
else:
paths.add("/".join(stack[:depth] + [name]))
return paths
def test_readme_documents_every_registered_command(bot: commands.Bot) -> None:
"""Every command registered on the bot is documented in README.md."""
readme = README.read_text(encoding="utf-8")
assert len(bot.commands) >= 11
missing = [cmd.name for cmd in bot.commands if cmd.name not in readme]
assert missing == [], f"Commands missing from README: {missing}"
def test_readme_file_tree_matches_actual_tree() -> None:
"""The README tree covers every committed non-dotfile, non-binary file."""
actual = _committed_files()
readme_paths = _readme_tree_paths()
missing_in_readme = actual - readme_paths
assert (
missing_in_readme == set()
), f"Files missing from the README tree: {sorted(missing_in_readme)}"
for path in readme_paths:
target = REPO_ROOT / path
if path.endswith("/"):
assert target.is_dir(), f"README tree lists missing directory: {path}"
else:
assert target.is_file(), f"README tree lists missing file: {path}"
vibe_bot/tests/test_llama_wrapper.py
-150
View File
@@ -1,150 +0,0 @@
"""Tests for the llama_wrapper module."""
from __future__ import annotations
import base64
import tempfile
from io import BytesIO
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch
import numpy as np
from vibe_bot.config import (
CHAT_ENDPOINT,
CHAT_ENDPOINT_KEY,
CHAT_MODEL,
EMBEDDING_ENDPOINT,
EMBEDDING_ENDPOINT_KEY,
IMAGE_EDIT_ENDPOINT,
IMAGE_EDIT_ENDPOINT_KEY,
IMAGE_GEN_ENDPOINT,
IMAGE_GEN_ENDPOINT_KEY,
)
from vibe_bot.llama_wrapper import (
chat_completion,
chat_completion_instruct,
embedding,
image_edit,
image_generation,
)
TEMPDIR = Path(tempfile.mkdtemp())
def test_chat_completion_think() -> None:
"""Test chat completion with think model."""
chat_completion(
system_prompt="You are a helpful assistant.",
user_prompt="Tell me about Everquest",
openai_url=CHAT_ENDPOINT,
openai_api_key=CHAT_ENDPOINT_KEY,
model=CHAT_MODEL,
max_tokens=100,
)
def test_chat_completion_instruct() -> None:
"""Test chat completion with instruct model."""
chat_completion_instruct(
system_prompt="You are a helpful assistant.",
user_prompt="Tell me about Everquest",
openai_url=CHAT_ENDPOINT,
openai_api_key=CHAT_ENDPOINT_KEY,
model=CHAT_MODEL,
max_tokens=100,
)
def test_image_generation() -> None:
"""Test image generation endpoint."""
with patch("vibe_bot.llama_wrapper.openai.OpenAI") as mock_openai:
mock_response = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake image data").decode()
mock_response.data = [mock_data]
mock_openai.return_value.images.generate.return_value = mock_response
result = image_generation(
prompt="Generate an image of a horse",
openai_url=IMAGE_GEN_ENDPOINT,
openai_api_key=IMAGE_GEN_ENDPOINT_KEY,
)
assert result == base64.b64encode(b"fake image data").decode()
def test_image_edit() -> None:
"""Test image edit endpoint."""
with patch("vibe_bot.llama_wrapper.openai.OpenAI") as mock_openai:
mock_response = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake edited image data").decode()
mock_response.data = [mock_data]
mock_openai.return_value.images.edit.return_value = mock_response
result = image_edit(
image=BytesIO(b"fake image"),
prompt="Paint the words 'horse' on the horse.",
openai_url=IMAGE_EDIT_ENDPOINT,
openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,
)
assert result == base64.b64encode(b"fake edited image data").decode()
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
"""Calculate cosine similarity between two arrays.
Returns a value close to 1 for similar vectors,
close to 0 for orthogonal vectors,
and close to -1 for opposite vectors.
"""
a_arr, b_arr = np.array(a), np.array(b)
return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))
EMBEDDING_SIMILARITY_HIGH = 0.9
EMBEDDING_SIMILARITY_LOW = 0.5
def test_embeddings() -> None:
"""Test embedding similarity for similar and different texts."""
mock_horse_vec = [0.8] * 1024 + [0.6] * 1024
mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024
mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024
def mock_post(*args: Any, **kwargs: Any) -> MagicMock:
json_data = kwargs.get("json", {})
text = json_data["input"][0]
if "horse" in text and "donkey" not in text and "also" not in text:
embedding_data = mock_horse_vec
elif "also" in text:
embedding_data = mock_horse_also_vec
else:
embedding_data = mock_donkey_vec
mock_resp = MagicMock()
mock_resp.json.return_value = {"data": [{"embedding": embedding_data}]}
return mock_resp
with patch("vibe_bot.llama_wrapper.requests.post", side_effect=mock_post):
result1 = embedding(
"this is a horse",
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result2 = embedding(
"this is a horse also",
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result3 = embedding(
"this is a donkey",
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
similarity_1 = _cosine_similarity(np.array(result1), np.array(result2))
assert similarity_1 > EMBEDDING_SIMILARITY_HIGH
similarity_2 = _cosine_similarity(np.array(result1), np.array(result3))
assert similarity_2 < EMBEDDING_SIMILARITY_LOW
vibe_bot/tests/test_llm_client.py
+419
View File
@@ -0,0 +1,419 @@
"""Tests for the llm_client module."""
from __future__ import annotations
import base64
from io import BytesIO
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock, patch
import numpy as np
import pytest
from vibe_bot.config import (
CHAT_MODEL,
EMBEDDING_ENDPOINT,
EMBEDDING_ENDPOINT_KEY,
)
from vibe_bot.llm_client import (
chat_complete,
chat_completion_instruct,
embedding,
image_edit,
image_generation,
)
@pytest.mark.live
def test_chat_complete_live() -> None:
"""Live call to the chat endpoint via the core async ``chat_complete``.
Unmocked: requires network access to the configured chat API. Ported from
the former ``test_chat_completion_think`` (its sync ``chat_completion``
wrapper was deleted).
"""
import asyncio
result = asyncio.run(
chat_complete(
[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Tell me about Everquest"},
],
model=CHAT_MODEL,
max_tokens=100,
)
)
assert isinstance(result, str)
@pytest.mark.live
def test_chat_completion_instruct_live() -> None:
"""Live call to the chat endpoint via the async instruct adapter.
Unmocked: requires network access to the configured chat API.
"""
import asyncio
result = asyncio.run(
chat_completion_instruct(
system_prompt="You are a helpful assistant.",
user_prompt="Tell me about Everquest",
model=CHAT_MODEL,
max_tokens=100,
)
)
assert isinstance(result, str)
def test_image_generation() -> None:
"""Image generation returns the first b64 payload from the API."""
import asyncio
mock_client = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake image data").decode()
mock_response = MagicMock()
mock_response.data = [mock_data]
mock_client.images.generate = AsyncMock(return_value=mock_response)
with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
result = asyncio.run(
image_generation(
prompt="Generate an image of a horse",
model="test-image-model",
)
)
assert result == base64.b64encode(b"fake image data").decode()
def test_image_generation_api_error_returns_empty() -> None:
"""A 4xx/5xx (APIStatusError) from the image API returns "" without raising."""
import asyncio
import openai
mock_client = MagicMock()
mock_client.images.generate = AsyncMock(
side_effect=openai.APIStatusError(
"boom",
response=MagicMock(),
body=None,
)
)
with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
result = asyncio.run(
image_generation(
prompt="Generate an image of a horse",
model="test-image-model",
)
)
assert result == ""
def test_image_edit() -> None:
"""Image edit returns the first b64 payload from the API."""
import asyncio
mock_client = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake edited image data").decode()
mock_response = MagicMock()
mock_response.data = [mock_data]
mock_client.images.edit = AsyncMock(return_value=mock_response)
with patch("vibe_bot.llm.images.get_image_edit_client", return_value=mock_client):
result = asyncio.run(
image_edit(
image=BytesIO(b"fake image"),
prompt="Paint the words 'horse' on the horse.",
model="test-image-edit-model",
)
)
assert result == base64.b64encode(b"fake edited image data").decode()
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
"""Calculate cosine similarity between two arrays.
Returns a value close to 1 for similar vectors,
close to 0 for orthogonal vectors,
and close to -1 for opposite vectors.
"""
a_arr, b_arr = np.array(a), np.array(b)
return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))
EMBEDDING_SIMILARITY_HIGH = 0.9
EMBEDDING_SIMILARITY_LOW = 0.5
def _mock_embedding_session(
post: MagicMock,
) -> MagicMock:
"""Build a mock requests.Session whose .post is ``post``."""
session = MagicMock()
session.post = post
return session
def test_embeddings() -> None:
"""Embedding similarity for similar and different texts."""
mock_horse_vec = [0.8] * 1024 + [0.6] * 1024
mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024
mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024
def mock_post(*args: Any, **kwargs: Any) -> MagicMock:
json_data = kwargs.get("json", {})
text = json_data["input"][0]
if "horse" in text and "donkey" not in text and "also" not in text:
embedding_data = mock_horse_vec
elif "also" in text:
embedding_data = mock_horse_also_vec
else:
embedding_data = mock_donkey_vec
mock_resp = MagicMock()
mock_resp.json.return_value = {"data": [{"embedding": embedding_data}]}
return mock_resp
session = _mock_embedding_session(MagicMock(side_effect=mock_post))
with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
result1 = embedding(
"this is a horse",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result2 = embedding(
"this is a horse also",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result3 = embedding(
"this is a donkey",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
similarity_1 = _cosine_similarity(np.array(result1), np.array(result2))
assert similarity_1 > EMBEDDING_SIMILARITY_HIGH
similarity_2 = _cosine_similarity(np.array(result1), np.array(result3))
assert similarity_2 < EMBEDDING_SIMILARITY_LOW
def test_embedding_non_json_2xx_returns_empty() -> None:
"""A 2xx response with a non-JSON body must return [] without raising.
Regression test for ``resp.json()`` sitting outside the try block, so an
HTML error page (or any non-JSON 2xx body) raised JSONDecodeError out of
``embedding`` and, through it, out of ``get_conversation_context``.
"""
mock_resp = MagicMock()
mock_resp.raise_for_status.return_value = None
mock_resp.json.side_effect = ValueError("<html>rate limited</html>")
session = _mock_embedding_session(MagicMock(return_value=mock_resp))
with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
result = embedding(
"this is a horse",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
assert result == []
def test_chat_client_singleton_identity() -> None:
"""The shared chat client is built once and reused across calls."""
from vibe_bot import llm_client
client1 = llm_client.get_chat_client()
client2 = llm_client.get_chat_client()
assert client1 is client2
def test_image_gen_client_singleton_identity() -> None:
"""The shared image-generation client is built once, from a cold start."""
import vibe_bot.llm.images as images_mod
saved = images_mod._image_gen_client
images_mod._image_gen_client = None
try:
client1 = images_mod.get_image_gen_client()
client2 = images_mod.get_image_gen_client()
assert client1 is client2
finally:
images_mod._image_gen_client = saved
def test_image_edit_client_singleton_identity() -> None:
"""The shared image-edit client is built once, from a cold start."""
import vibe_bot.llm.images as images_mod
saved = images_mod._image_edit_client
images_mod._image_edit_client = None
try:
client1 = images_mod.get_image_edit_client()
client2 = images_mod.get_image_edit_client()
assert client1 is client2
finally:
images_mod._image_edit_client = saved
def test_flows_build_no_new_clients_or_sessions(
mock_ctx: MagicMock,
temp_db_path: str,
) -> None:
"""A full !doodlebob + chat turn constructs no new clients or sessions.
Every shared client and the embedding session are built once at
"startup"; running the whole image flow and a chat turn through the real
singletons (with only HTTP mocked) must not construct another
AsyncOpenAI client or requests.Session.
"""
import asyncio
import openai
import requests
from vibe_bot import llm_client
from vibe_bot.database import ChatDatabase
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.image_service import ImageService
# Startup: build every shared client and the embedding session once.
chat_client = llm_client.get_chat_client()
gen_client = llm_client.get_image_gen_client()
edit_client = llm_client.get_image_edit_client()
llm_client.get_embedding_session()
counts = {"async_openai": 0, "session": 0}
real_session = requests.Session
class CountingAsyncOpenAI(openai.AsyncOpenAI):
def __init__(self, **kwargs: Any) -> None:
counts["async_openai"] += 1
super().__init__(**kwargs)
def counting_session() -> requests.Session:
counts["session"] += 1
return real_session()
# layout, image prompt, verify verdict, chat reply — in call order.
completions_create = AsyncMock(
side_effect=[
_make_response("square", None),
_make_response("a detailed prompt", None),
_make_response("PASS", None),
_make_response("a chat reply", None),
]
)
image_response = MagicMock()
image_response.data = [MagicMock(b64_json=base64.b64encode(b"fake image").decode())]
images_generate = AsyncMock(return_value=image_response)
registry = MagicMock()
registry.to_openai_tools.return_value = []
db = ChatDatabase(db_path=temp_db_path)
with (
patch.object(openai, "AsyncOpenAI", CountingAsyncOpenAI),
patch.object(requests, "Session", counting_session),
patch.object(chat_client.chat.completions, "create", completions_create),
patch.object(gen_client.images, "generate", images_generate),
patch("vibe_bot.llm_client.embedding", return_value=[0.25] * 32),
):
asyncio.run(
ImageService(db, MagicMock()).generate(mock_ctx, message="a centaur")
)
asyncio.run(
ChatService(db, registry).handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
assert counts["async_openai"] == 0
assert counts["session"] == 0
assert llm_client.get_chat_client() is chat_client
assert llm_client.get_image_gen_client() is gen_client
assert llm_client.get_image_edit_client() is edit_client
def _make_response(content: str | None, tool_calls: list[object] | None) -> MagicMock:
"""Build a mock chat completion response with the given message fields."""
message = MagicMock()
message.content = content
message.tool_calls = tool_calls
return MagicMock(choices=[MagicMock(message=message)])
def test_chat_complete_skips_non_function_tool_call() -> None:
"""A tool call that is not of type 'function' is skipped, not executed."""
import asyncio
from vibe_bot.llm_client import chat_complete
called = {"n": 0}
def tool_executor(name: str, args: dict[str, str]) -> str:
called["n"] += 1
return f"executed:{name}"
custom_tool_call = MagicMock()
custom_tool_call.type = "custom"
mock_client = MagicMock()
mock_client.chat.completions.create = AsyncMock(
side_effect=[
_make_response(content=None, tool_calls=[custom_tool_call]),
_make_response(content="final answer", tool_calls=None),
]
)
with patch("vibe_bot.llm.chat.get_chat_client", return_value=mock_client):
result = asyncio.run(
chat_complete(
[{"role": "user", "content": "hi"}],
model="m",
max_tokens=10,
tool_executor=tool_executor,
)
)
assert result == "final answer"
assert called["n"] == 0
def test_tool_registry_dispatch_and_unknown_tool() -> None:
"""The registry renders schemas, dispatches known tools, and names unknowns."""
from vibe_bot.llm_client import ToolRegistry
def echo_tool(name: str, args: dict[str, str], **kwargs: object) -> str:
return f"echo:{args.get('text', '')}"
registry = ToolRegistry()
registry.register(
"echo",
"Echoes the text argument back.",
{"type": "object", "properties": {"text": {"type": "string"}}},
echo_tool,
)
tools = registry.to_openai_tools()
first = tools[0]
assert first["type"] == "function"
function_def = cast("dict[str, object]", first["function"])
assert function_def["name"] == "echo"
assert function_def["description"] == "Echoes the text argument back."
assert registry.execute("echo", {"text": "hi"}) == "echo:hi"
assert registry.execute("does_not_exist", {}) == "Unknown tool: does_not_exist"
vibe_bot/tests/test_logging.py
+186
View File
@@ -0,0 +1,186 @@
"""No-content logging tests.
A recognizable secret string is seeded into message content, image prompts,
and bot personalities; the logs captured at DEBUG must never contain it,
while metadata (bot name, user id, message id, counts) must.
"""
from __future__ import annotations
import asyncio
import base64
import logging
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.image_service import ImageService
from vibe_bot.tests._helpers import invoke
SECRET = "SECRET-DO-NOT-LOG-12345"
def _registry() -> MagicMock:
"""A mock ToolRegistry."""
reg = MagicMock()
reg.to_openai_tools.return_value = []
reg.execute.return_value = "tool result"
return reg
def _assert_no_secret(
caplog: pytest.LogCaptureFixture,
*metadata: str,
) -> None:
"""The secret must be absent from every record; the metadata must be present."""
for record in caplog.records:
assert (
SECRET not in record.getMessage()
), f"Secret leaked in log record: {record.getMessage()!r}"
assert SECRET not in caplog.text
for token in metadata:
assert token in caplog.text, f"Expected metadata {token!r} missing from logs"
def test_chat_service_logs_metadata_not_content(
caplog: pytest.LogCaptureFixture,
chat_db: Any,
mock_ctx: MagicMock,
) -> None:
"""A full chat turn (RAG + persist + reply) never logs the message content."""
caplog.set_level(logging.DEBUG)
svc = ChatService(chat_db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="A perfectly fine response."),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message=f"please remember {SECRET} forever",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
_assert_no_secret(caplog, "alfred", "12345")
def test_add_message_logs_metadata_not_content(
caplog: pytest.LogCaptureFixture,
chat_db: Any,
) -> None:
"""add_message (with embedding) never logs the stored content."""
caplog.set_level(logging.DEBUG)
assert chat_db.add_message(
message_id="msg-1",
user_id="12345",
username="testuser",
content=f"User: tell me about {SECRET}",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
)
assert chat_db.add_message(
message_id="msg-1_response",
user_id="bot-123",
username="test-bot",
content=f"the bot knew {SECRET}",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
role="assistant",
embed=False,
)
_assert_no_secret(caplog, "msg-1", "12345")
def test_history_lookups_log_metadata_not_content(
caplog: pytest.LogCaptureFixture,
chat_db: Any,
) -> None:
"""get_user_history / get_bot_history never log message or response content."""
caplog.set_level(logging.DEBUG)
chat_db.add_message(
message_id="msg-2",
user_id="12345",
username="testuser",
content=f"User: what is {SECRET}",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
)
chat_db.add_message(
message_id="msg-2_response",
user_id="bot-123",
username="test-bot",
content=f"it is {SECRET}, clearly",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
role="assistant",
embed=False,
)
user_history = chat_db.get_user_history("12345", limit=5)
bot_history = chat_db.get_bot_history("alfred", limit=5)
assert len(user_history) == 1
assert len(bot_history) == 1
_assert_no_secret(caplog, "msg-2", "alfred")
def test_doodlebob_logs_metadata_not_content(
caplog: pytest.LogCaptureFixture,
mock_ctx: MagicMock,
) -> None:
"""A doodlebob generation never logs the prompt or the derived image prompt."""
caplog.set_level(logging.DEBUG)
db = MagicMock()
db.get_image_generation_time_estimate.return_value = 12.0
svc = ImageService(db, MagicMock())
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(
side_effect=[
"portrait",
f"a very detailed painting of {SECRET}",
"pass",
]
),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=base64.b64encode(b"img").decode()),
),
):
asyncio.run(svc.generate(mock_ctx, message=f"draw {SECRET} on the moon"))
_assert_no_secret(caplog, "Doodlebob", "12345")
def test_custom_bot_creation_logs_metadata_not_personality(
caplog: pytest.LogCaptureFixture,
bot: Any,
mock_ctx: MagicMock,
) -> None:
"""!custom-bot never logs the personality text."""
caplog.set_level(logging.DEBUG)
invoke(
bot,
"custom-bot",
mock_ctx,
"secretbot",
personality=f"a butler who knows {SECRET}",
)
_assert_no_secret(caplog, "secretbot", "12345")
vibe_bot/tests/test_main.py
-1321
View File
File diff suppressed because it is too large Load Diff
vibe_bot/tests/test_prompts.py
+176
View File
@@ -0,0 +1,176 @@
"""Tests for the prompts module (prompt constants, layout parsing, user info)."""
from __future__ import annotations
from datetime import UTC, datetime
from unittest.mock import MagicMock
import pytest
from vibe_bot.config import (
IMAGE_GEN_SIZE_LANDSCAPE,
IMAGE_GEN_SIZE_PORTRAIT,
IMAGE_GEN_SIZE_SQUARE,
)
from vibe_bot.prompts import (
IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE,
LAYOUT_SIZES,
RESPONSE_LENGTH_HINT,
build_system_prompt,
get_user_info,
parse_image_layout,
)
@pytest.fixture
def mock_ctx() -> MagicMock:
"""A minimal Discord user for get_user_info."""
author = MagicMock()
author.name = "testuser"
author.id = "12345"
author.global_name = "Test User"
author.nick = "tester"
author.top_role.name = "@everyone"
author.activities = []
author.joined_at = None
author.created_at = None
return author
@pytest.fixture
def mock_author_with_member_data() -> MagicMock:
"""A Discord user with full member data (role + activity + timestamps)."""
author = MagicMock()
author.name = "testuser"
author.id = "12345"
author.global_name = "Test User"
author.nick = "tester"
author.top_role.name = "Admin"
activity = MagicMock()
activity.name = "Chess"
author.activities = [activity]
author.joined_at = datetime(2024, 1, 15, tzinfo=UTC)
author.created_at = datetime(2023, 6, 1, tzinfo=UTC)
return author
def test_build_system_prompt_contains_personality_hint_and_user_info() -> None:
"""build_system_prompt assembles personality + length hint + user info block."""
result = build_system_prompt("you are a butler", "Username: alice")
assert result.startswith("you are a butler")
assert RESPONSE_LENGTH_HINT in result
assert "User Information:\nUsername: alice" in result
def test_layout_sizes_map_to_config() -> None:
"""LAYOUT_SIZES maps each layout to its configured canvas size."""
assert LAYOUT_SIZES["portrait"] == IMAGE_GEN_SIZE_PORTRAIT
assert LAYOUT_SIZES["landscape"] == IMAGE_GEN_SIZE_LANDSCAPE
assert LAYOUT_SIZES["square"] == IMAGE_GEN_SIZE_SQUARE
@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:
"""Valid LLM layout responses parse to the right layout."""
assert 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:
"""Empty or malformed responses fall back to square."""
assert parse_image_layout(response) == "square"
def test_image_prompt_system_prompt_covers_key_details() -> None:
"""The prompt-rewrite system prompt forces explicit detail on all aspects."""
prompt = IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE.format(layout="square")
lowered = prompt.lower()
assert "square" in prompt
assert "exact text" in lowered
assert "composition" in lowered
assert "style" in lowered
assert "canada goose" in lowered
assert "only the image generation prompt" in lowered
assert "exactly as written" in lowered
assert "fountain pen wearing pants" in lowered
assert "centaur" in lowered
assert "not a person riding a horse" in lowered
def test_get_user_info_minimal(mock_ctx: MagicMock) -> None:
"""get_user_info with minimal member data includes the core identity lines."""
result = get_user_info(mock_ctx)
assert "Username: testuser" in result
assert "User ID: 12345" in result
assert "Global Name: Test User" in result
assert "Nickname: tester" in result
def test_get_user_info_with_member_data(
mock_author_with_member_data: MagicMock,
) -> None:
"""get_user_info with full member data includes roles, activity, timestamps."""
result = get_user_info(mock_author_with_member_data)
assert "Global Name: Test User" in result
assert "Nickname: tester" in result
assert "Username: testuser" in result
assert "User ID: 12345" in result
assert "Top Role: Admin" in result
assert "Activities: Chess" in result
assert "Joined: 2024-01-15" in result
assert "Account Created: 2023-06-01" in result
def test_get_user_info_no_global_name(mock_ctx: MagicMock) -> None:
"""Optional fields are omitted when they are empty."""
mock_ctx.global_name = None
mock_ctx.nick = None
mock_ctx.top_role.name = "@everyone"
mock_ctx.activities = []
result = get_user_info(mock_ctx)
assert "Global Name:" not in result
assert "Nickname:" not in result
assert "Top Role:" not in result
assert "Activities:" not in result
assert "Username: testuser" in result
assert "User ID: 12345" in result
def test_get_user_info_with_top_role_not_everyone(
mock_author_with_member_data: MagicMock,
) -> None:
"""Top role is included when it is not @everyone."""
result = get_user_info(mock_author_with_member_data)
assert "Top Role: Admin" in result
def test_get_user_info_no_activities(mock_ctx: MagicMock) -> None:
"""The activities line is omitted when there are none."""
mock_ctx.activities = []
result = get_user_info(mock_ctx)
assert "Activities:" not in result
vibe_bot/tests/test_services.py
+953
View File
@@ -0,0 +1,953 @@
"""Service-layer tests: chat, image, speech, and conversation services.
These exercise the LLM-backed logic directly (constructing each service with
mock dependencies) rather than going through the thin Discord command wrappers.
"""
from __future__ import annotations
import asyncio
import base64
from io import BytesIO
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import requests
from vibe_bot.config import TTS_VOICE
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.conversation_service import (
MAX_TOPIC_LENGTH,
ConversationService,
flip_counter,
)
from vibe_bot.services.image_service import (
MAX_IMAGE_DOWNLOAD_BYTES,
MAX_IMAGE_PROMPT_LENGTH,
ImageService,
_allowed_image_url,
_download_image_bytes,
select_image_layout,
verify_image_prompt,
)
from vibe_bot.services.speech_service import (
MAX_SPEAK_LENGTH,
SpeechService,
parse_voice_flag,
)
@pytest.fixture
def mock_ctx() -> MagicMock:
"""A mock Discord command context."""
ctx = MagicMock()
ctx.author.name = "testuser"
ctx.author.id = "12345"
ctx.author.global_name = "Test User"
ctx.author.nick = "tester"
ctx.author.top_role.name = "@everyone"
ctx.author.activities = []
ctx.author.joined_at = None
ctx.author.created_at = None
ctx.channel.id = "channel-1"
ctx.guild.id = "guild-1"
ctx.message.id = "msg-1"
ctx.message.attachments = []
ctx.bot.user = MagicMock()
ctx.bot.user.name = "test-bot"
ctx.bot.user.id = "bot-123"
ctx.send = AsyncMock()
return ctx
def _file_factory() -> MagicMock:
"""A File factory that records (buffer, filename) as a tuple."""
factory = MagicMock()
def make_file(buf: BytesIO, name: str) -> tuple[str, BytesIO, str]:
return ("FILE", buf, name)
factory.side_effect = make_file
return factory
def _sent_texts(ctx: MagicMock) -> list[str]:
"""All positional text messages sent through ctx.send."""
return [c.args[0] for c in ctx.send.call_args_list if c.args]
def _registry() -> MagicMock:
"""A mock ToolRegistry."""
reg = MagicMock()
reg.to_openai_tools.return_value = []
reg.execute.return_value = "tool result"
return reg
def _fake_bot(name: str) -> tuple[str, str, str, str]:
"""A stand-in custom bot tuple for manager.get_custom_bot."""
return (name, "a personality", "user-123", "2024-01-01")
# ---------------------------------------------------------------------------
# ChatService
# ---------------------------------------------------------------------------
def test_chat_success(mock_ctx: MagicMock) -> None:
"""A normal turn persists the exchange and sends the reply."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="This is a bot response"),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
db.add_message.assert_called()
assert mock_ctx.send.call_count >= 2
def test_chat_turn_embedding_budget(
mock_ctx: MagicMock,
temp_db_path: str,
) -> None:
"""One chat turn embeds exactly twice: the RAG query and the user row.
The assistant row is persisted with embed=False, so it costs no
embedding call and stores no embedding row.
"""
import sqlite3
from vibe_bot.database import ChatDatabase
db = ChatDatabase(db_path=temp_db_path)
svc = ChatService(db, _registry())
with (
patch(
"vibe_bot.llm_client.embedding",
return_value=[0.25] * 32,
) as mock_embedding,
patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="This is a bot response"),
),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
assert mock_embedding.call_count == 2
conn = sqlite3.connect(temp_db_path)
embedding_rows = conn.execute("SELECT COUNT(*) FROM message_embeddings").fetchone()
conn.close()
assert embedding_rows[0] == 1
def test_chat_error(mock_ctx: MagicMock) -> None:
"""An LLM error surfaces a friendly message."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(side_effect=Exception("API error")),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
call_args = mock_ctx.send.call_args[0][0]
assert "error occurred" in call_args.lower()
db.add_message.assert_not_called()
def test_chat_long_response_chunked(mock_ctx: MagicMock) -> None:
"""Long responses are split into multiple sends."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="x" * 2500),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
assert mock_ctx.send.call_count >= 3
def test_chat_includes_user_info(mock_ctx: MagicMock) -> None:
"""The system prompt sent to the LLM includes the requester's info."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
mock_llm = AsyncMock(return_value="resp")
with patch("vibe_bot.llm_client.chat_completion_with_tools", new=mock_llm):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
system_prompt = mock_llm.call_args.kwargs["system_prompt"]
assert "testuser" in system_prompt
def test_chat_with_context_and_tools(mock_ctx: MagicMock) -> None:
"""Prior RAG context is prepended and tool calls reach the registry."""
db = MagicMock()
db.get_conversation_context.return_value = [
{"role": "user", "content": "old question"},
{"role": "assistant", "content": "old answer"},
]
registry = _registry()
captured: dict[str, Any] = {}
async def fake_llm(**kwargs: Any) -> str:
captured.update(kwargs)
kwargs["tool_executor"]("get_channel_members", {})
await kwargs["tool_call_notifier"]("get_channel_members", {})
return "resp"
with patch("vibe_bot.llm_client.chat_completion_with_tools", new=fake_llm):
asyncio.run(
ChatService(db, registry).handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
prompts = captured["prompts"]
assert prompts[0] == {"role": "user", "content": "old question"}
assert prompts[-1] == {"role": "user", "content": "hello"}
registry.execute.assert_called_once_with(
"get_channel_members", {}, channel=mock_ctx.channel
)
assert any("looking at the channel members" in t for t in _sent_texts(mock_ctx))
# ---------------------------------------------------------------------------
# SpeechService
# ---------------------------------------------------------------------------
def _speech_service(
tts: MagicMock | None,
manager: MagicMock | None = None,
make_file: MagicMock | None = None,
) -> SpeechService:
return SpeechService(
MagicMock(), manager or MagicMock(), tts, make_file or _file_factory()
)
def test_speak_tts_not_initialized(mock_ctx: MagicMock) -> None:
"""No TTS engine means a clear error, no LLM or TTS calls."""
svc = _speech_service(None)
asyncio.run(svc.speak(mock_ctx, message="hello world"))
call_args = mock_ctx.send.call_args[0][0]
assert "TTS engine not initialized" in call_args
def test_speak_empty_message(mock_ctx: MagicMock) -> None:
"""Empty text is rejected before any TTS work."""
svc = _speech_service(MagicMock())
asyncio.run(svc.speak(mock_ctx, message=""))
call_args = mock_ctx.send.call_args[0][0]
assert "Please provide text" in call_args
def test_speak_too_long(mock_ctx: MagicMock) -> None:
"""Oversized text is rejected without calling the TTS engine."""
tts = MagicMock()
svc = _speech_service(tts)
asyncio.run(svc.speak(mock_ctx, message="a" * (MAX_SPEAK_LENGTH + 1)))
tts.generate_audio.assert_not_called()
call_args = mock_ctx.send.call_args[0][0]
assert "Text too long to speak" in call_args
def test_speak_partial_audio_warns(mock_ctx: MagicMock) -> None:
"""Partial audio triggers a warning line."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(
audio=MagicMock(), partial=True, failed_chunks=1
)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello world"))
assert any("audio may be incomplete" in t for t in _sent_texts(mock_ctx))
def test_speak_plain_text(mock_ctx: MagicMock) -> None:
"""Plain text is spoken and the audio file is sent."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello world"))
tts.generate_audio.assert_called_once()
assert mock_ctx.send.call_count >= 2
def test_speak_with_custom_bot(mock_ctx: MagicMock) -> None:
"""A bot prefix routes through the LLM, then speaks the response."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = [
("alfred", "british butler", "user-123"),
]
manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
svc = _speech_service(tts, manager=manager)
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="The time is 3pm"),
):
asyncio.run(svc.speak(mock_ctx, message="alfred what time is it"))
tts.generate_audio.assert_called_once()
assert any("**alfred**:" in t for t in _sent_texts(mock_ctx))
def test_speak_uses_requested_voice(mock_ctx: MagicMock) -> None:
"""A trailing --voice flag selects that voice for the TTS call."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello world --voice af_bella"))
assert tts.generate_audio.call_args.kwargs["voice"] == "af_bella"
def test_speak_mid_text_voice_flag_spoken_verbatim(mock_ctx: MagicMock) -> None:
"""A --voice mid-message is preserved as speech and the default voice is used."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
message = "hello --voice af_bella world"
asyncio.run(svc.speak(mock_ctx, message=message))
assert tts.generate_audio.call_args.args[0] == message
assert tts.generate_audio.call_args.kwargs["voice"] == TTS_VOICE
def test_speak_unknown_voice(mock_ctx: MagicMock) -> None:
"""An unknown voice is rejected before any TTS call."""
tts = MagicMock()
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello --voice not_a_real_voice"))
tts.generate_audio.assert_not_called()
call_args = mock_ctx.send.call_args[0][0]
assert "Unknown voice" in call_args
def test_speak_language_lookup_uses_precomputed_dict(
mock_ctx: MagicMock,
) -> None:
"""The speak hot path resolves the language via VOICE_LANGUAGES.get().
The dict is built once at import from VOICES_LIST (covering every
catalog voice) and the per-speak lookup is a single dict get — no
per-call scan of the category list.
"""
from vibe_bot.config import VOICES_LIST
from vibe_bot.services import speech_service
class LanguageLookupSpy:
"""Counts .get() lookups on the voice->language mapping."""
def __init__(self, data: dict[str, str]) -> None:
self.data = data
self.lookups = 0
def get(self, key: str, default: str | None = None) -> str | None:
self.lookups += 1
return self.data.get(key, default)
def __contains__(self, key: object) -> bool:
return key in self.data
counting = LanguageLookupSpy(speech_service.VOICE_LANGUAGES)
assert set(counting.data) == {
voice for category in VOICES_LIST.values() for voice in category["voices"]
}
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
with patch.object(speech_service, "VOICE_LANGUAGES", counting):
asyncio.run(svc.speak(mock_ctx, message="hello --voice bf_alice"))
assert counting.lookups == 1
assert tts.generate_audio.call_args.kwargs["lang"] == "en-gb"
# ---------------------------------------------------------------------------
# ConversationService
# ---------------------------------------------------------------------------
def _conversation_service(
manager: MagicMock | None = None,
) -> ConversationService:
return ConversationService(manager or MagicMock())
def test_flip_counter() -> None:
"""flip_counter toggles between 0 and 1."""
assert flip_counter(0) == 1
assert flip_counter(1) == 0
def test_talkforme_topic_too_long(mock_ctx: MagicMock) -> None:
"""Oversized topics are rejected before any LLM call."""
svc = _conversation_service()
asyncio.run(svc.run(mock_ctx, "a", "b", "3", "x" * (MAX_TOPIC_LENGTH + 1)))
call_args = mock_ctx.send.call_args[0][0]
assert "Topic too long" in call_args
def test_talkforme_bot1_not_found(mock_ctx: MagicMock) -> None:
"""A missing first bot is reported and the run stops."""
manager = MagicMock()
manager.get_custom_bot.return_value = None
svc = _conversation_service(manager=manager)
asyncio.run(svc.run(mock_ctx, "ghost", "alfred", "3", "cats"))
call_args = mock_ctx.send.call_args[0][0]
assert "ghost is not a real bot" in call_args
def test_talkforme_invalid_limit(mock_ctx: MagicMock) -> None:
"""A non-integer limit is rejected after both bots are found."""
manager = MagicMock()
manager.get_custom_bot.side_effect = _fake_bot
svc = _conversation_service(manager=manager)
asyncio.run(svc.run(mock_ctx, "a", "b", "abc", "cats"))
call_args = mock_ctx.send.call_args[0][0]
assert "Message limit must be an integer" in call_args
def test_talkforme_first_reply_chunked(mock_ctx: MagicMock) -> None:
"""Long first replies are sent in multiple chunks."""
manager = MagicMock()
manager.get_custom_bot.side_effect = _fake_bot
svc = _conversation_service(manager=manager)
with patch(
"vibe_bot.llm_client.chat_completion_with_history",
new=AsyncMock(return_value="y" * 2500),
):
asyncio.run(svc.run(mock_ctx, "a", "b", "1", "cats"))
assert mock_ctx.send.call_count >= 3
# ---------------------------------------------------------------------------
# ImageService (doodlebob / retcon)
# ---------------------------------------------------------------------------
def _image_service(db: MagicMock | None = None) -> ImageService:
return ImageService(db or MagicMock(), _file_factory())
def test_doodlebob_prompt_too_long(mock_ctx: MagicMock) -> None:
"""Oversized prompts are rejected before any LLM call."""
svc = _image_service()
asyncio.run(svc.generate(mock_ctx, message="a" * (MAX_IMAGE_PROMPT_LENGTH + 1)))
call_args = mock_ctx.send.call_args[0][0]
assert "Prompt too long" in call_args
def test_doodlebob_generate_success(mock_ctx: MagicMock) -> None:
"""A full generate flow ends with an image file and a completion line."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur in a field"))
assert db.record_image_generation_time.called
assert any("Strike complete" in t for t in _sent_texts(mock_ctx))
def test_doodlebob_failed_generation(mock_ctx: MagicMock) -> None:
"""An empty image-generation response is reported as a failure."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=""),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert any("Failed to generate image" in t for t in _sent_texts(mock_ctx))
assert not db.record_image_generation_time.called
def test_doodlebob_reports_estimate(mock_ctx: MagicMock) -> None:
"""A prior-history estimate produces a Drone ETA line."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = 12.5
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert any("Drone ETA" in t for t in _sent_texts(mock_ctx))
def test_doodlebob_no_estimate_without_history(mock_ctx: MagicMock) -> None:
"""No estimate means no Drone ETA line."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert not any("Drone ETA" in t for t in _sent_texts(mock_ctx))
def test_doodlebob_empty_prompt_stops(mock_ctx: MagicMock) -> None:
"""An empty image-prompt response stops the flow without generating."""
db = MagicMock()
svc = _image_service(db)
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=""),
),
patch("vibe_bot.llm_client.image_generation") as mock_gen,
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
mock_gen.assert_not_called()
assert not db.record_image_generation_time.called
def test_doodlebob_decode_failure(mock_ctx: MagicMock) -> None:
"""Invalid base64 from the generation API is reported as a failure."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value="abcde!!!"),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert any(
"Failed to process the generated image" in t for t in _sent_texts(mock_ctx)
)
@pytest.mark.parametrize(
("response", "expected"),
[
("portrait", "portrait"),
("landscape", "landscape"),
("square", "square"),
("PORTRAIT", "portrait"),
("I think landscape", "landscape"),
],
)
def test_select_image_layout_returns_parsed(
response: str,
expected: str,
) -> None:
"""select_image_layout parses the LLM's layout choice."""
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=response),
):
result = asyncio.run(select_image_layout("a tall tree"))
assert result == expected
def test_select_image_layout_uses_minimal_token_budget() -> None:
"""Layout selection is a one-word answer, so max_tokens is 2."""
mock_llm = AsyncMock(return_value="square")
with patch("vibe_bot.llm_client.chat_completion_instruct", new=mock_llm):
assert asyncio.run(select_image_layout("a tall tree")) == "square"
assert mock_llm.call_args.kwargs["max_tokens"] == 2
def test_doodlebob_latency_within_budget(mock_ctx: MagicMock) -> None:
"""End-to-end doodlebob latency with 50ms simulated per LLM/image call.
Hermetic latency figure: four mocked calls (layout, prompt, verify,
generate) at 50ms each must dominate the wall time; the overhead on
top of the simulated 200ms stays far below the budget.
"""
import time
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
responses = ["square", "a detailed prompt", "PASS"]
async def slow_instruct(**_kwargs: Any) -> str:
await asyncio.sleep(0.05)
return responses.pop(0)
async def slow_generate(**_kwargs: Any) -> str:
await asyncio.sleep(0.05)
return b64
with (
patch("vibe_bot.llm_client.chat_completion_instruct", new=slow_instruct),
patch("vibe_bot.llm_client.image_generation", new=slow_generate),
):
start = time.monotonic()
asyncio.run(svc.generate(mock_ctx, message="a centaur in a field"))
elapsed = time.monotonic() - start
assert elapsed >= 0.2
assert elapsed < 5.0
assert any("Strike complete" in t for t in _sent_texts(mock_ctx))
def test_verify_image_prompt_pass_keeps_prompt() -> None:
"""A PASS verdict keeps the original prompt unchanged."""
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value="PASS"),
):
result = asyncio.run(verify_image_prompt("a centaur", "a detailed prompt"))
assert result == "a detailed prompt"
def test_verify_image_prompt_correction_replaces() -> None:
"""A non-passing verdict is used as the corrected prompt."""
correction = "a rewritten prompt that is definitely long enough to be a fix"
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=correction),
):
result = asyncio.run(verify_image_prompt("a centaur", "a detailed prompt"))
assert result == correction
def test_verify_image_prompt_empty_falls_back() -> None:
"""An empty verdict falls back to the original prompt."""
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=""),
):
result = asyncio.run(verify_image_prompt("a centaur", "a detailed prompt"))
assert result == "a detailed prompt"
def test_retcon_no_attachments(mock_ctx: MagicMock) -> None:
"""retcon with no attachments asks the user to attach an image."""
svc = _image_service()
mock_ctx.message.attachments = []
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
call_args = mock_ctx.send.call_args[0][0]
assert "Please attach an image" in call_args
def test_retcon_rejected_url_not_downloaded(mock_ctx: MagicMock) -> None:
"""A non-Discord attachment URL is refused before any download happens."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://evil.example.com/img.png"
mock_ctx.message.attachments = [attachment]
mock_edit = AsyncMock(return_value="")
with (
patch("vibe_bot.llm_client.image_edit", new=mock_edit),
patch("requests.get") as mock_get,
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
mock_get.assert_not_called()
mock_edit.assert_not_called()
assert any("Please attach an image" in t for t in _sent_texts(mock_ctx))
def test_retcon_prompt_too_long(mock_ctx: MagicMock) -> None:
"""Oversized retcon prompts are rejected before any download."""
svc = _image_service()
asyncio.run(svc.edit(mock_ctx, message="a" * (MAX_IMAGE_PROMPT_LENGTH + 1)))
call_args = mock_ctx.send.call_args[0][0]
assert "Prompt too long" in call_args
def test_retcon_image_edit_empty(mock_ctx: MagicMock) -> None:
"""An empty edit response is reported as a failure."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://cdn.discordapp.com/attachments/1/2/3/img.png"
mock_ctx.message.attachments = [attachment]
with (
patch(
"vibe_bot.services.image_service._download_image_bytes",
return_value=b"fake image bytes",
),
patch(
"vibe_bot.llm_client.image_edit",
new=AsyncMock(return_value=""),
),
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
call_args = mock_ctx.send.call_args[0][0]
assert "Failed to edit the image" in call_args
def test_retcon_success(mock_ctx: MagicMock) -> None:
"""A successful edit sends the edited image file."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://cdn.discordapp.com/attachments/1/2/3/img.png"
mock_ctx.message.attachments = [attachment]
b64 = base64.b64encode(b"edited").decode()
with (
patch(
"vibe_bot.services.image_service._download_image_bytes",
return_value=b"fake image bytes",
),
patch(
"vibe_bot.llm_client.image_edit",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
assert any("Rewriting history" in t for t in _sent_texts(mock_ctx))
def test_retcon_edit_decode_failure(mock_ctx: MagicMock) -> None:
"""Invalid base64 from the edit API is reported as a processing failure."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://cdn.discordapp.com/attachments/1/2/3/img.png"
mock_ctx.message.attachments = [attachment]
with (
patch(
"vibe_bot.services.image_service._download_image_bytes",
return_value=b"fake image bytes",
),
patch(
"vibe_bot.llm_client.image_edit",
new=AsyncMock(return_value="abcde!!!"),
),
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
call_args = mock_ctx.send.call_args[0][0]
assert "Failed to process the edited image" in call_args
def test_allowed_image_url() -> None:
"""Only Discord CDN hosts are allowed for retcon downloads."""
assert _allowed_image_url("https://cdn.discordapp.com/a/b/c.png")
assert _allowed_image_url("https://media.discordapp.net/a/b/c.png")
assert not _allowed_image_url("https://example.com/a/b/c.png")
assert not _allowed_image_url("https://evilcdn.com/a/b/c.png")
assert not _allowed_image_url("http://[::1")
def test_download_image_bytes_success() -> None:
"""An allowed Discord URL is downloaded and its chunks joined."""
response = MagicMock()
response.headers = {}
response.iter_content.return_value = [b"abc", b"", b"def"]
response.raise_for_status.return_value = None
with patch(
"vibe_bot.services.image_service.requests.get",
return_value=response,
) as mock_get:
data = _download_image_bytes("https://cdn.discordapp.com/a/b/c.png")
mock_get.assert_called_once()
assert data == b"abcdef"
def test_download_image_bytes_request_failure() -> None:
"""A failing download returns None instead of raising."""
with patch(
"vibe_bot.services.image_service.requests.get",
side_effect=requests.RequestException("boom"),
):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
def test_download_image_bytes_content_length_cap() -> None:
"""A declared Content-Length above the cap is refused before streaming."""
response = MagicMock()
response.headers = {"Content-Length": str(MAX_IMAGE_DOWNLOAD_BYTES + 1)}
response.iter_content = MagicMock()
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
response.iter_content.assert_not_called()
def test_download_image_bytes_streaming_cap() -> None:
"""Streaming past the size cap is refused even without a Content-Length."""
response = MagicMock()
response.headers = {}
response.iter_content.return_value = [b"a" * (MAX_IMAGE_DOWNLOAD_BYTES + 1)]
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
def test_download_image_bytes_bad_content_length() -> None:
"""A non-numeric Content-Length is ignored and streaming proceeds."""
response = MagicMock()
response.headers = {"Content-Length": "not-a-number"}
response.iter_content.return_value = [b"xyz"]
response.raise_for_status.return_value = None
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") == b"xyz"
def test_download_image_bytes_stream_error() -> None:
"""A mid-stream request error returns None instead of raising."""
response = MagicMock()
response.headers = {}
response.iter_content.side_effect = requests.RequestException("stream died")
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
# ---------------------------------------------------------------------------
# parse_voice_flag
# ---------------------------------------------------------------------------
def test_parse_voice_flag_no_flag() -> None:
"""Without a trailing flag the message is returned unchanged."""
assert parse_voice_flag("hello world") == ("hello world", None)
def test_parse_voice_flag_trailing() -> None:
"""A trailing --voice flag is split off."""
assert parse_voice_flag("hello world --voice af_bella") == (
"hello world",
"af_bella",
)
def test_parse_voice_flag_mid_text_preserved() -> None:
"""A --voice that is not at the end is treated as speech, not a flag."""
assert parse_voice_flag("hello --voice af_bella world") == (
"hello --voice af_bella world",
None,
)
def test_parse_voice_flag_missing_value() -> None:
"""A --voice with no value (or only trailing spaces) is not a flag."""
assert parse_voice_flag("hello --voice") == ("hello --voice", None)
assert parse_voice_flag("hello --voice ") == ("hello --voice ", None)
vibe_bot/tests/test_textutil.py
+109
View File
@@ -0,0 +1,109 @@
"""Tests for the textutil module."""
from __future__ import annotations
import pytest
from vibe_bot.textutil import split_message
def test_split_message_empty() -> None:
assert split_message("") == []
def test_split_message_short_text_single_chunk() -> None:
assert split_message("hello") == ["hello"]
def test_split_message_exact_limit() -> None:
text = "a" * 1900
chunks = split_message(text)
assert chunks == [text]
assert len(chunks) == 1
def test_split_message_just_over_limit_no_newline() -> None:
text = "a" * 1901
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
assert chunks == ["a" * 1900, "a"]
def test_split_message_no_newlines_long() -> None:
text = "x" * 5000
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
def test_split_message_respects_newlines() -> None:
# Many short lines spanning several chunks: boundaries fall on newlines.
text = "\n".join(f"line {i}" for i in range(1, 501))
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
assert len(chunks) > 1
# Every chunk except the last ends on a newline (split at a line boundary).
for chunk in chunks[:-1]:
assert chunk.endswith("\n")
def test_split_message_long_line_hard_split() -> None:
# A single line longer than the limit must be hard-split.
text = "a" * 5000 + "\n" + "short"
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
def test_split_message_emoji_round_trip() -> None:
# Multi-codepoint emoji at split boundaries must not corrupt on join.
text = "hello 🌍 world 👨‍👩‍👧‍👦 end " * 500
for limit in (1, 10, 100, 1900):
chunks = split_message(text, limit)
assert "".join(chunks) == text
assert all(len(c) <= limit for c in chunks)
def test_split_message_property_round_trip_corpus() -> None:
corpus = [
"",
"a",
"a" * 1900,
"a" * 1901,
"line1\nline2\nline3",
"para one\n\npara two\n\npara three",
"no newline " * 1000,
"emoji 🎉 " * 1000,
"👨‍👩‍👧‍👦" * 2000,
"mixed 🌍 text and 日本語 and emoji 🎊 here",
# NFD combining marks (e + U+0301, a + U+0308) straddle split points.
"cafe\u0301 " * 1000,
"a\u0308\u0301 b\u0327\u0301 c\u0308 " * 1000,
# Code spans with backticks and spaces.
"`inline code` and `x = 1` spans " * 500,
"``double backtick`` `single` " * 500,
]
for text in corpus:
for limit in (1, 5, 100, 1900):
chunks = split_message(text, limit)
assert "".join(chunks) == text, f"round-trip failed for limit={limit}"
assert all(len(c) <= limit for c in chunks)
def test_split_message_exact_multiple_of_limit() -> None:
"""Input that is an exact multiple of the limit splits into full chunks."""
for limit in (1, 5, 100, 1900):
text = "z" * (limit * 4)
chunks = split_message(text, limit)
assert "".join(chunks) == text
assert len(chunks) == 4
assert all(len(c) == limit for c in chunks)
def test_split_message_limit_must_be_positive() -> None:
with pytest.raises(ValueError):
split_message("hello", 0)
with pytest.raises(ValueError):
split_message("hello", -5)
vibe_bot/tests/test_tts.py
+21 -11
View File
@@ -22,7 +22,7 @@ def test_tts_engine_init(mock_kokoro_tts: MagicMock) -> None:
def test_generate_audio(mock_kokoro_tts: MagicMock) -> None: def test_generate_audio(mock_kokoro_tts: MagicMock) -> None:
"""Test audio generation returns a BytesIO object.""" """Test audio generation returns a full (non-partial) AudioResult."""
from io import BytesIO from io import BytesIO
from vibe_bot.tts import TTSEngine from vibe_bot.tts import TTSEngine
@@ -30,9 +30,11 @@ def test_generate_audio(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin") engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("hello world this is a test") result = engine.generate_audio("hello world this is a test")
assert isinstance(result, BytesIO) assert isinstance(result.audio, BytesIO)
result.seek(0) assert result.partial is False
data = result.read() assert result.failed_chunks == 0
result.audio.seek(0)
data = result.audio.read()
assert len(data) > 0 assert len(data) > 0
@@ -57,7 +59,8 @@ def test_generate_audio_single_chunk(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin") engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("single chunk text") result = engine.generate_audio("single chunk text")
assert isinstance(result, BytesIO) assert isinstance(result.audio, BytesIO)
assert result.partial is False
mock_kokoro_tts["process_chunk_sequential"].assert_called_once() mock_kokoro_tts["process_chunk_sequential"].assert_called_once()
@@ -77,7 +80,9 @@ def test_generate_audio_multiple_chunks(mock_kokoro_tts: MagicMock) -> None:
"this text is long enough to be split into multiple chunks", "this text is long enough to be split into multiple chunks",
) )
assert isinstance(result, BytesIO) assert isinstance(result.audio, BytesIO)
assert result.partial is False
assert result.failed_chunks == 0
assert mock_kokoro_tts["process_chunk_sequential"].call_count == 3 assert mock_kokoro_tts["process_chunk_sequential"].call_count == 3
@@ -108,7 +113,12 @@ def test_generate_audio_chunk_failure(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin") engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("good chunk bad chunk another good") result = engine.generate_audio("good chunk bad chunk another good")
assert isinstance(result, BytesIO) # Audio is still produced for the good chunks, but flagged as partial.
assert isinstance(result.audio, BytesIO)
assert result.partial is True
assert result.failed_chunks == 1
result.audio.seek(0)
assert len(result.audio.read()) > 0
def test_generate_audio_all_chunks_fail(mock_kokoro_tts: MagicMock) -> None: def test_generate_audio_all_chunks_fail(mock_kokoro_tts: MagicMock) -> None:
@@ -145,13 +155,13 @@ def test_generate_audio_returns_seekable(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin") engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("hello world") result = engine.generate_audio("hello world")
result.seek(0) result.audio.seek(0)
data = result.read() data = result.audio.read()
assert len(data) > 0 assert len(data) > 0
# Should be able to seek and read again # Should be able to seek and read again
result.seek(0) result.audio.seek(0)
data2 = result.read() data2 = result.audio.read()
assert data == data2 assert data == data2
vibe_bot/textutil.py
+41
View File
@@ -0,0 +1,41 @@
"""Text chunking utilities for sending long content within Discord's limit."""
from __future__ import annotations
def split_message(text: str, limit: int = 1900) -> list[str]:
"""Split ``text`` into chunks of at most ``limit`` characters.
Splits on newlines first so a line is never broken mid-text when
avoidable; a single line longer than ``limit`` is hard-split by plain
code-point slicing (Python ``str`` slicing is code-point safe, so no
lone surrogates are produced). Multi-codepoint grapheme clusters (e.g.
some emoji) are not protected — this matches the previous behavior and
keeps the dependency footprint at zero.
Guarantees ``"".join(split_message(text, limit)) == text``.
"""
if not text:
return []
if limit <= 0:
raise ValueError("limit must be a positive integer")
chunks: list[str] = []
current = ""
for raw_line in text.splitlines(keepends=True):
line = raw_line
while len(line) > limit:
if current:
chunks.append(current)
current = ""
chunks.append(line[:limit])
line = line[limit:]
if len(current) + len(line) <= limit:
current += line
else:
if current:
chunks.append(current)
current = line
if current:
chunks.append(current)
return chunks
vibe_bot/tools.py
+6 -4
View File
@@ -33,13 +33,15 @@ def _format_member(member: Any) -> str:
def get_channel_members_impl(channel: Any) -> str: def get_channel_members_impl(channel: Any) -> str:
"""Get a list of all members in the Discord channel the bot is part of. """Get a list of the members of the guild the bot is a member of.
Use this tool when asked about who is in the channel, who the members are, Discord has no per-channel membership, so the channel argument only locates
or to get a roster of people present in the current channel. the guild; the roster covers the whole guild, not just that channel. Use
this tool when asked who is around, who the members are, or to get a
roster of the people in the server.
Returns: Returns:
A formatted string listing all members in the channel with their usernames, A formatted string listing the guild's members with their usernames,
display names, and nicknames. display names, and nicknames.
""" """
vibe_bot/tts.py
+48 -13
View File
@@ -3,9 +3,12 @@
from __future__ import annotations from __future__ import annotations
import logging import logging
from dataclasses import dataclass
from io import BytesIO from io import BytesIO
import numpy as np import numpy as np
# kokoro-tts and soundfile ship no type stubs upstream.
import soundfile as sf # type: ignore[import-untyped] import soundfile as sf # type: ignore[import-untyped]
from kokoro_tts import ( # type: ignore[import-untyped] from kokoro_tts import ( # type: ignore[import-untyped]
Kokoro, Kokoro,
@@ -13,14 +16,31 @@ from kokoro_tts import ( # type: ignore[import-untyped]
process_chunk_sequential, process_chunk_sequential,
) )
from vibe_bot.config import TTS_SPEED, TTS_VOICE
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Default voice settings # Default voice settings (single source of truth: vibe_bot.config).
DEFAULT_VOICE = "af_sarah" DEFAULT_VOICE = TTS_VOICE
DEFAULT_SPEED = 1.0 DEFAULT_SPEED = TTS_SPEED
DEFAULT_LANG = "en-us" DEFAULT_LANG = "en-us"
@dataclass
class AudioResult:
"""Audio output from a TTS generation.
Attributes:
audio: The encoded audio (MP3) as a seekable BytesIO.
partial: True if one or more text chunks failed to produce audio.
failed_chunks: Number of text chunks that failed.
"""
audio: BytesIO
partial: bool
failed_chunks: int
class TTSEngine: class TTSEngine:
"""Text-to-speech engine wrapper around Kokoro TTS.""" """Text-to-speech engine wrapper around Kokoro TTS."""
@@ -43,13 +63,14 @@ class TTSEngine:
voice: str = DEFAULT_VOICE, voice: str = DEFAULT_VOICE,
speed: float = DEFAULT_SPEED, speed: float = DEFAULT_SPEED,
lang: str = DEFAULT_LANG, lang: str = DEFAULT_LANG,
) -> BytesIO: ) -> AudioResult:
"""Convert text to audio and return as BytesIO (MP3 format).""" """Convert text to audio and return an AudioResult (MP3 in .audio)."""
all_samples: list[np.ndarray] = [] all_samples: list[np.ndarray] = []
sample_rate: int | None = None sample_rate: int | None = None
failed_chunks = 0
chunks: list[str] = list(chunk_text(text)) chunks: list[str] = list(chunk_text(text))
logger.info("Split text into %d chunks", len(chunks)) logger.debug("Split text into %d chunks", len(chunks))
for i, chunk in enumerate(chunks): for i, chunk in enumerate(chunks):
try: try:
@@ -60,15 +81,21 @@ class TTSEngine:
speed, speed,
lang, lang,
) )
if samples is not None:
if sample_rate is None:
sample_rate = sr
all_samples.append(np.asarray(samples))
logger.info("Processed chunk %d/%d", i + 1, len(chunks))
except Exception: except Exception:
logger.exception("Error processing chunk %d", i + 1) logger.exception("Error processing chunk %d", i + 1)
failed_chunks += 1
continue continue
if samples is None:
logger.warning("Chunk %d/%d produced no audio", i + 1, len(chunks))
failed_chunks += 1
continue
if sample_rate is None:
sample_rate = sr
all_samples.append(np.asarray(samples))
logger.debug("Processed chunk %d/%d", i + 1, len(chunks))
if not all_samples: if not all_samples:
msg = "No audio samples generated - text may be invalid or too long" msg = "No audio samples generated - text may be invalid or too long"
raise ValueError(msg) raise ValueError(msg)
@@ -85,9 +112,17 @@ class TTSEngine:
) )
buffer.seek(0) buffer.seek(0)
logger.info( partial = failed_chunks > 0
if partial:
logger.warning(
"TTS produced partial audio: %d of %d chunks failed",
failed_chunks,
len(chunks),
)
logger.debug(
"Generated MP3 audio: %d samples at %dHz", "Generated MP3 audio: %d samples at %dHz",
len(combined), len(combined),
sample_rate or 0, sample_rate or 0,
) )
return buffer return AudioResult(audio=buffer, partial=partial, failed_chunks=failed_chunks)