diff --git a/.gitea/workflows/build-push.yml b/.gitea/workflows/build-push.yml index f9f8f7e..4ca90d8 100644 --- a/.gitea/workflows/build-push.yml +++ b/.gitea/workflows/build-push.yml @@ -4,13 +4,51 @@ on: push: branches: - main + pull_request: release: types: - published 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: runs-on: ubuntu-latest + needs: [test] + if: github.event_name != 'pull_request' steps: - name: Checkout uses: actions/checkout@v4 diff --git a/.gitignore b/.gitignore index b1138da..e79bdc6 100644 --- a/.gitignore +++ b/.gitignore @@ -13,4 +13,6 @@ wheels/ *.db .mypy_cache/ -.pytest_cache/ \ No newline at end of file +.pytest_cache/ + +.agent/ \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md index 0010cf2..0e251e8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,34 +1,40 @@ # 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 -- 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 -- 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/` -- 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 - 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. ## Test suite gotchas -- Baseline: everything passes except `test_llama_wrapper.py::test_chat_completion_think` and `::test_chat_completion_instruct`. Those two are unmocked live calls to the real `CHAT_ENDPOINT` from `.env` and fail without network access to that API. -- `test_config.py` hardcodes `sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots")` (a stale repo path). The test only passes because pytest's cwd fallback finds the package — run tests from the repo root. -- If `uv run ` suddenly fails with `ModuleNotFoundError` or "bad interpreter", the `.venv` shebangs are stale from a repo move: `rm -rf .venv && uv sync --extra dev`. +- 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. +- If `uv run ` suddenly fails with `ModuleNotFoundError` or "bad interpreter", the `.venv` shebangs are stale from a repo move: `rm -rf .venv && uv sync`. ## Code map -- `main.py` — entrypoint. The bot is created at module import (module-level `commands.Bot(...)`), `bot.run()` only under `__main__`. Custom-bot "commands" (`! ...`) 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" (`! ...`) 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. -- `database.py` — `ChatDatabase` (embeddings stored as float32 blobs, cosine-similarity RAG, schema auto-migrates on startup) and `CustomBotManager`. -- `llama_wrapper.py` — thin OpenAI-compatible clients for chat / image gen / image edit / embeddings, each with its own endpoint, key, and model. -- `tools.py` — `get_channel_members` is a no-op LangChain `@tool` stub used only for name/description/schema; the real implementation is `get_channel_members_impl(channel)`, wired into the tool executor in `main.py`. +- `prompts.py` — LLM prompt constants (image layout/prompt/verify, length hint, `build_system_prompt`). +- `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 `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 diff --git a/Containerfile b/Containerfile index 1f32f3d..a024d14 100644 --- a/Containerfile +++ b/Containerfile @@ -15,4 +15,8 @@ ENV UV_NO_DEV=1 WORKDIR /app 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 \ No newline at end of file diff --git a/MANUAL_TESTS.md b/MANUAL_TESTS.md new file mode 100644 index 0000000..ab08038 --- /dev/null +++ b/MANUAL_TESTS.md @@ -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 ` — 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 `` +- `!list-custom-bots` now includes alfred, and `!alfred ` 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 ` to create one.` +- After creating bots: `Available Custom Bots:` followed by one `* ` line per bot. + +### `!delete-custom-bot ` — 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 — `! ` + +This is not a registered command; messages starting with `! ` 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 [--voice ]` — 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 [--voice ]` — 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**: ` (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 --voice ` to choose a voice.` +- Spot-check one voice per language with `!speak --voice ` + and listen for the expected accent. + +### `!doodlebob ` — generate an image + +- `!doodlebob a cat sitting on the moon` + Expect, in order: + 1. `**Doodlebob shopping for a canvas...**` + 2. `**Doodlebob selected **` — 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 ...**` + 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 ` — 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 ` — 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 ` +- 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 ` — 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 + + --- + alfred: + ``` + + 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`) diff --git a/README.md b/README.md index 45431f0..d2818ba 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,9 @@ A Discord bot that stores long-term chat history using SQLite with RAG (Retrieva - [Text-to-Speech](#text-to-speech) - [Image Commands](#image-commands) - [Bot Conversations](#bot-conversations) - - [Chat History](#chat-history) + - [Admin & Debug](#admin--debug) + - [Design Notes](#design-notes) + - [Limits & Cooldowns](#limits--cooldowns) - [Features](#features) - [Setup](#setup) - [Prerequisites](#prerequisites) @@ -47,24 +49,24 @@ Once you create a custom bot, interact with it by prefixing your message with th **Example:** 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 ### Text-to-Speech -| Command | Description | Example Usage | -| ------------------------------------ | ----------------------------------------------------- | ------------------------------------------ | -| `!speak ` | Convert text to speech (MP3 attachment) | `!speak hello world` | -| `!speak --voice ` | Convert text to speech with a specific voice | `!speak hello world --voice af_bella` | -| `!speak ` | Have a custom bot respond and speak | `!speak alfred what time is it` | +| Command | Description | Example Usage | +| ------------------------------------ | ----------------------------------------------------- | ----------------------------------------- | +| `!speak ` | Convert text to speech (MP3 attachment) | `!speak hello world` | +| `!speak --voice ` | Convert text to speech with a specific voice | `!speak hello world --voice af_bella` | +| `!speak ` | Have a custom bot respond and speak | `!speak alfred what time is it` | | `!speak --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 -| Command | Description | Example Usage | -| ------------ | ------------------------------------ | ------------------------------------------ | -| `!doodlebob` | Generate an image from a text prompt | `!doodlebob a cat sitting on the moon` | +| Command | Description | Example Usage | +| ------------ | ------------------------------------ | ----------------------------------------- | +| `!doodlebob` | Generate an image from a text prompt | `!doodlebob a cat sitting on the moon` | | `!retcon` | Edit an attached image with text | `!retcon Make it sunny` | ### Bot Conversations @@ -73,11 +75,53 @@ Once you create a custom bot, interact with it by prefixing your message with th | -------------------------------------- | ------------------------------------------- | ------------------------------------------------ | | `!talkforme ` | 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 | -| --------------------- | ------------------------------------- | ----------------- | -| `!history ` | View the chat history of a custom bot | `!history alfred` | +| Command | Description | Example Usage | +| --------------------- | ------------------------------------------------------- | --------------------- | +| `!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 ` | 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 (`! `) 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 @@ -110,11 +154,8 @@ DISCORD_TOKEN=your_discord_bot_token # Chat/Completion API (required) CHAT_ENDPOINT=https://your-api.com/v1 -COMPLETION_ENDPOINT=https://your-api.com/v1 CHAT_ENDPOINT_KEY=your_api_key -COMPLETION_ENDPOINT_KEY=your_api_key CHAT_MODEL=your_model_name -COMPLETION_MODEL=your_model_name # Image Generation (required) IMAGE_GEN_ENDPOINT=https://your-api.com/v1 @@ -169,16 +210,16 @@ uv run python -m vibe_bot.main ### Database Structure -The system uses SQLite with three tables: +The system uses SQLite with four tables: 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 - `message_id` (PK), `embedding` (binary blob of float32 values) - 3. **custom_bots**: Stores custom bot configurations - `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 @@ -195,26 +236,76 @@ The system uses SQLite with three tables: ```text vibe_discord_bots/ ├── vibe_bot/ -│ ├── __init__.py # Package marker -│ ├── main.py # Main bot application (commands, event handlers) -│ ├── config.py # Environment variable loading and validation -│ ├── database.py # SQLite database with RAG + CustomBotManager -│ ├── llama_wrapper.py # OpenAI-compatible API wrappers (chat, images, embeddings) -│ ├── tts.py # Kokoro TTS engine +│ ├── __init__.py # Package marker +│ ├── app.py # Composition root: App, services, bot event handlers +│ ├── commands/ +│ │ ├── __init__.py # register_all: wires every command group +│ │ ├── _state.py # Shared App holder for command handlers +│ │ ├── 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/ -│ ├── conftest.py # Shared test fixtures -│ ├── test_main.py # Bot command tests -│ ├── test_config.py # Config loading tests -│ ├── test_database.py # Database + CustomBotManager tests -│ ├── test_llama_wrapper.py # API wrapper tests -│ └── test_tts.py # TTS engine tests -├── pyproject.toml # Project dependencies (uv) -├── uv.lock # Locked dependency versions -├── .env # Environment variables -├── kokoro-v1.0.onnx # Kokoro TTS model -├── voices-v1.0.bin # Kokoro voice definitions -├── Containerfile # Podman/Docker build file -└── README.md # This file +│ ├── __init__.py +│ ├── _helpers.py # Shared test helpers +│ ├── conftest.py # Shared test fixtures +│ ├── test_app.py # Composition root and bot handler tests +│ ├── test_commands.py # Discord command tests +│ ├── test_config.py # Config loading tests +│ ├── test_database.py # Database + CustomBotManager tests +│ ├── test_docs.py # README sync tests (commands + file tree) +│ ├── test_llm_client.py # LLM client tests +│ ├── test_logging.py # No-content logging tests +│ ├── test_prompts.py # Prompt constant tests +│ ├── test_services.py # Service-layer tests +│ ├── 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 @@ -233,7 +324,7 @@ uv run python -m vibe_bot.main ```bash # Build the container image -podman build -t vibe-bot:latest . +podman build -t localhost/vibe-bot:latest -f Containerfile . # Run with environment file podman run --env-file .env localhost/vibe-bot:latest @@ -241,12 +332,19 @@ podman run --env-file .env localhost/vibe-bot:latest ## Testing -Run the full test suite: +Run the full test suite (hermetic by default; network-dependent `live` tests are +deselected): ```bash uv run pytest vibe_bot/tests/ -v ``` +Run the network-dependent live tests explicitly: + +```bash +uv run pytest -m live +``` + Run linters: ```bash @@ -265,24 +363,30 @@ uv run black --check vibe_bot/ ## Configuration -| Variable | Default | Description | -| ----------------------- | ------------------ | ------------------------------------- | -| `DISCORD_TOKEN` | *(required)* | Discord bot authentication token | -| `CHAT_ENDPOINT` | *(required)* | OpenAI-compatible chat API URL | -| `CHAT_MODEL` | *(required)* | Model name for chat completions | -| `IMAGE_GEN_ENDPOINT` | *(required)* | Image generation API URL | -| `IMAGE_EDIT_ENDPOINT` | *(required)* | Image editing API URL | -| `IMAGE_GEN_SIZE_SQUARE` | `1024x1024` | Square canvas size for image generation | -| `IMAGE_GEN_SIZE_PORTRAIT` | `1024x1536` | Portrait (tall) canvas size for image generation | -| `IMAGE_GEN_SIZE_LANDSCAPE`| `1536x1024` | Landscape (wide) canvas size for image generation | -| `EMBEDDING_ENDPOINT` | *(required)* | Embedding API URL | -| `EMBEDDING_MODEL` | *(required)* | Model name for text embeddings | -| `MAX_COMPLETION_TOKENS` | `1000` | Max tokens in LLM responses | -| `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 | -| `DB_PATH` | `chat_history.db` | SQLite database file path | +| Variable | Default | Description | +| -------------------------- | ------------------ | --------------------------------------------- | +| `DISCORD_TOKEN` | *(required)* | Discord bot authentication token | +| `CHAT_ENDPOINT` | *(required)* | OpenAI-compatible chat API URL | +| `CHAT_ENDPOINT_KEY` | `placeholder` | API key for the chat endpoint | +| `CHAT_MODEL` | *(required)* | Model name for chat completions | +| `IMAGE_GEN_ENDPOINT` | *(required)* | Image generation API URL | +| `IMAGE_GEN_ENDPOINT_KEY` | `placeholder` | API key for the image generation endpoint | +| `IMAGE_GEN_MODEL` | *(required)* | Model name for image generation | +| `IMAGE_GEN_SIZE_SQUARE` | `1024x1024` | Square canvas size for image generation | +| `IMAGE_GEN_SIZE_PORTRAIT` | `1024x1536` | Portrait (tall) canvas size for image generation | +| `IMAGE_GEN_SIZE_LANDSCAPE` | `1536x1024` | Landscape (wide) canvas size for image generation | +| `IMAGE_EDIT_ENDPOINT` | *(required)* | Image editing API URL | +| `IMAGE_EDIT_ENDPOINT_KEY` | `placeholder` | API key for the image editing endpoint | +| `IMAGE_EDIT_MODEL` | *(required)* | Model name for image editing | +| `EMBEDDING_ENDPOINT` | *(required)* | Embedding API URL | +| `EMBEDDING_ENDPOINT_KEY` | `placeholder` | API key for the embedding endpoint | +| `EMBEDDING_MODEL` | *(required)* | Model name for text embeddings | +| `MAX_COMPLETION_TOKENS` | `1000` | Max tokens in LLM responses | +| `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 | diff --git a/REMEDIATION_PLAN.md b/REMEDIATION_PLAN.md deleted file mode 100644 index 92b7d9d..0000000 --- a/REMEDIATION_PLAN.md +++ /dev/null @@ -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 ` 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: `.) -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 `rate limited` → 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.*)\s+--voice\s+(?P\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). diff --git a/pyproject.toml b/pyproject.toml index 6a0a60b..864330e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,29 +8,13 @@ dependencies = [ "discord>=2.3.2", "openai>=2.24.0", "requests>=2.32.5", - "types-requests>=2.32.4.20260107", "numpy>=1.24.0", - "pytest>=9.0.2", "python-dotenv>=1.2.2", - "pytest-env>=1.5.0", "kokoro-tts>=2.3.1", - "mypy>=2.1.0", "langchain-openai>=0.4.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] strict = true python_version = "3.13" @@ -60,12 +44,23 @@ line-length = 88 target-version = "py313" [tool.pytest.ini_options] +addopts = "-m 'not live'" +markers = [ + "live: network-dependent integration tests, deselected by default", +] filterwarnings = [ "ignore::pytest.PytestUnraisableExceptionWarning", ] [dependency-groups] dev = [ + "black>=25.1.0", + "debugpy>=1.8.0", + "mypy>=2.1.0", "pyright>=1.1.409", + "pytest>=9.0.2", + "pytest-cov>=7.0.0", + "pytest-env>=1.5.0", "ruff>=0.16.3", + "types-requests>=2.32.4.20260107", ] diff --git a/scripts/bench_rag.py b/scripts/bench_rag.py new file mode 100644 index 0000000..c6e0566 --- /dev/null +++ b/scripts/bench_rag.py @@ -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() diff --git a/uv.lock b/uv.lock index 2de583b..6bc1343 100644 --- a/uv.lock +++ b/uv.lock @@ -5,9 +5,6 @@ resolution-markers = [ "python_full_version >= '3.15'", "python_full_version < '3.15'", ] -required-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", -] [[package]] 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" }, ] +[[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]] name = "csvw" 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" }, ] +[[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]] name = "pytest-env" version = "1.5.0" @@ -2471,53 +2566,48 @@ dependencies = [ { name = "kokoro-tts" }, { name = "langchain-core" }, { name = "langchain-openai" }, - { name = "mypy" }, { name = "numpy" }, { name = "openai" }, - { name = "pytest" }, - { name = "pytest-env" }, { name = "python-dotenv" }, { name = "requests" }, - { name = "types-requests" }, -] - -[package.optional-dependencies] -dev = [ - { name = "black" }, - { name = "debugpy" }, - { name = "mypy" }, ] [package.dev-dependencies] dev = [ + { name = "black" }, + { name = "debugpy" }, + { name = "mypy" }, { name = "pyright" }, + { name = "pytest" }, + { name = "pytest-cov" }, + { name = "pytest-env" }, { name = "ruff" }, + { name = "types-requests" }, ] [package.metadata] 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 = "kokoro-tts", specifier = ">=2.3.1" }, { name = "langchain-core", specifier = ">=0.3.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 = "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 = "requests", specifier = ">=2.32.5" }, - { name = "types-requests", specifier = ">=2.32.4.20260107" }, ] -provides-extras = ["dev"] [package.metadata.requires-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 = "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 = "types-requests", specifier = ">=2.32.4.20260107" }, ] [[package]] diff --git a/vibe_bot/app.py b/vibe_bot/app.py new file mode 100644 index 0000000..9e8b8fd --- /dev/null +++ b/vibe_bot/app.py @@ -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 diff --git a/vibe_bot/commands/__init__.py b/vibe_bot/commands/__init__.py new file mode 100644 index 0000000..fcc61fb --- /dev/null +++ b/vibe_bot/commands/__init__.py @@ -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) diff --git a/vibe_bot/commands/_state.py b/vibe_bot/commands/_state.py new file mode 100644 index 0000000..00936aa --- /dev/null +++ b/vibe_bot/commands/_state.py @@ -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 diff --git a/vibe_bot/commands/admin.py b/vibe_bot/commands/admin.py new file mode 100644 index 0000000..4ca6c1a --- /dev/null +++ b/vibe_bot/commands/admin.py @@ -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 + 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 + """ + 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) diff --git a/vibe_bot/commands/chat.py b/vibe_bot/commands/chat.py new file mode 100644 index 0000000..28851e7 --- /dev/null +++ b/vibe_bot/commands/chat.py @@ -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``.""" diff --git a/vibe_bot/commands/conversation.py b/vibe_bot/commands/conversation.py new file mode 100644 index 0000000..600a6e4 --- /dev/null +++ b/vibe_bot/commands/conversation.py @@ -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 ") + return + + await app.conversation.run(ctx, parts[0], parts[1], parts[2], " ".join(parts[3:])) diff --git a/vibe_bot/commands/custom_bots.py b/vibe_bot/commands/custom_bots.py new file mode 100644 index 0000000..9398ad9 --- /dev/null +++ b/vibe_bot/commands/custom_bots.py @@ -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 + 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} `", + ) + + +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 ` 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 + """ + 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.") diff --git a/vibe_bot/commands/images.py b/vibe_bot/commands/images.py new file mode 100644 index 0000000..8c5d588 --- /dev/null +++ b/vibe_bot/commands/images.py @@ -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) diff --git a/vibe_bot/commands/speech.py b/vibe_bot/commands/speech.py new file mode 100644 index 0000000..5a5c66e --- /dev/null +++ b/vibe_bot/commands/speech.py @@ -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 --voice ` 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 --voice - plain text to speech + Usage: !speak --voice - 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) diff --git a/vibe_bot/config.py b/vibe_bot/config.py index 7afc6f5..15aeaca 100644 --- a/vibe_bot/config.py +++ b/vibe_bot/config.py @@ -7,11 +7,6 @@ import os 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__) load_dotenv() @@ -21,7 +16,6 @@ DISCORD_TOKEN: str = os.getenv("DISCORD_TOKEN", "") # Endpoints CHAT_ENDPOINT: str = os.getenv("CHAT_ENDPOINT", "") -COMPLETION_ENDPOINT: str = os.getenv("COMPLETION_ENDPOINT", "") IMAGE_GEN_ENDPOINT: str = os.getenv("IMAGE_GEN_ENDPOINT", "") IMAGE_EDIT_ENDPOINT: str = os.getenv("IMAGE_EDIT_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 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_EDIT_ENDPOINT_KEY: str = os.getenv("IMAGE_EDIT_ENDPOINT_KEY", "placeholder") EMBEDDING_ENDPOINT_KEY: str = os.getenv("EMBEDDING_ENDPOINT_KEY", "placeholder") # Models CHAT_MODEL: str = os.getenv("CHAT_MODEL", "") -COMPLETION_MODEL: str = os.getenv("COMPLETION_MODEL", "") IMAGE_GEN_MODEL: str = os.getenv("IMAGE_GEN_MODEL", "") IMAGE_GEN_SIZE_SQUARE: str = os.getenv("IMAGE_GEN_SIZE_SQUARE", "1024x1024") 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", "") EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "") -# Database and embeddings +# Database and history 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")) SIMILARITY_THRESHOLD: float = float(os.getenv("SIMILARITY_THRESHOLD", "0.7")) 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 -if not CHAT_ENDPOINT: - endpoint_msg = "CHAT_ENDPOINT required." - raise RuntimeError(endpoint_msg) +def validate_config() -> None: + """Raise ``RuntimeError`` if any required setting is missing.""" + # Check token + if not DISCORD_TOKEN: + msg = "DISCORD_TOKEN required." + raise RuntimeError(msg) -if not COMPLETION_ENDPOINT: - endpoint_msg = "COMPLETION_ENDPOINT required." - raise RuntimeError(endpoint_msg) + # Check endpoints + if not CHAT_ENDPOINT: + endpoint_msg = "CHAT_ENDPOINT required." + raise RuntimeError(endpoint_msg) -if not IMAGE_GEN_ENDPOINT: - endpoint_msg = "IMAGE_GEN_ENDPOINT required." - raise RuntimeError(endpoint_msg) + if not IMAGE_GEN_ENDPOINT: + endpoint_msg = "IMAGE_GEN_ENDPOINT required." + raise RuntimeError(endpoint_msg) -if not IMAGE_EDIT_ENDPOINT: - endpoint_msg = "IMAGE_EDIT_ENDPOINT required." - raise RuntimeError(endpoint_msg) + if not IMAGE_EDIT_ENDPOINT: + endpoint_msg = "IMAGE_EDIT_ENDPOINT required." + raise RuntimeError(endpoint_msg) -if not EMBEDDING_ENDPOINT: - endpoint_msg = "EMBEDDING_ENDPOINT required." - raise RuntimeError(endpoint_msg) + if not EMBEDDING_ENDPOINT: + endpoint_msg = "EMBEDDING_ENDPOINT required." + raise RuntimeError(endpoint_msg) -# Check models -if not CHAT_MODEL: - model_msg = "CHAT_MODEL required." - raise RuntimeError(model_msg) + # Check models + if not CHAT_MODEL: + model_msg = "CHAT_MODEL required." + raise RuntimeError(model_msg) -if not COMPLETION_MODEL: - model_msg = "COMPLETION_MODEL required." - raise RuntimeError(model_msg) + if not IMAGE_GEN_MODEL: + model_msg = "IMAGE_GEN_MODEL required." + raise RuntimeError(model_msg) -if not IMAGE_GEN_MODEL: - model_msg = "IMAGE_GEN_MODEL required." - raise RuntimeError(model_msg) + if not IMAGE_EDIT_MODEL: + model_msg = "IMAGE_EDIT_MODEL required." + raise RuntimeError(model_msg) -if not IMAGE_EDIT_MODEL: - model_msg = "IMAGE_EDIT_MODEL required." - raise RuntimeError(model_msg) + if not EMBEDDING_MODEL: + model_msg = "EMBEDDING_MODEL required." + raise RuntimeError(model_msg) -if not EMBEDDING_MODEL: - model_msg = "EMBEDDING_MODEL required." - raise RuntimeError(model_msg) # TTS 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) diff --git a/vibe_bot/database.py b/vibe_bot/database.py index 33225ac..292e832 100644 --- a/vibe_bot/database.py +++ b/vibe_bot/database.py @@ -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 -import logging -import sqlite3 -from typing import TYPE_CHECKING - -import numpy as np -from openai import OpenAI - -from vibe_bot import llama_wrapper -from vibe_bot.config import ( - DB_PATH, - 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]) - +from vibe_bot.db.bots import CustomBotManager +from vibe_bot.db.messages import ChatDatabase +from vibe_bot.db.timing import IMAGE_GEN_TIME_LIMIT, IMAGE_GEN_TIME_WINDOW + +__all__ = [ + "IMAGE_GEN_TIME_LIMIT", + "IMAGE_GEN_TIME_WINDOW", + "ChatDatabase", + "CustomBotManager", + "get_database", +] # Global database instance _chat_db: ChatDatabase | None = None @@ -605,171 +28,3 @@ def get_database() -> ChatDatabase: if _chat_db is None: _chat_db = ChatDatabase() 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() diff --git a/vibe_bot/db/__init__.py b/vibe_bot/db/__init__.py new file mode 100644 index 0000000..eda9d61 --- /dev/null +++ b/vibe_bot/db/__init__.py @@ -0,0 +1 @@ +"""SQLite storage layer: connection, schema, message store, RAG, custom bots.""" diff --git a/vibe_bot/db/bots.py b/vibe_bot/db/bots.py new file mode 100644 index 0000000..ea84648 --- /dev/null +++ b/vibe_bot/db/bots.py @@ -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() diff --git a/vibe_bot/db/connection.py b/vibe_bot/db/connection.py new file mode 100644 index 0000000..ec6cf44 --- /dev/null +++ b/vibe_bot/db/connection.py @@ -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 diff --git a/vibe_bot/db/messages.py b/vibe_bot/db/messages.py new file mode 100644 index 0000000..ce3bcb4 --- /dev/null +++ b/vibe_bot/db/messages.py @@ -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) diff --git a/vibe_bot/db/schema.py b/vibe_bot/db/schema.py new file mode 100644 index 0000000..72e86de --- /dev/null +++ b/vibe_bot/db/schema.py @@ -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() diff --git a/vibe_bot/db/search.py b/vibe_bot/db/search.py new file mode 100644 index 0000000..78c2260 --- /dev/null +++ b/vibe_bot/db/search.py @@ -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 diff --git a/vibe_bot/db/timing.py b/vibe_bot/db/timing.py new file mode 100644 index 0000000..f01530b --- /dev/null +++ b/vibe_bot/db/timing.py @@ -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]) diff --git a/vibe_bot/db/vectors.py b/vibe_bot/db/vectors.py new file mode 100644 index 0000000..6518e36 --- /dev/null +++ b/vibe_bot/db/vectors.py @@ -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 diff --git a/vibe_bot/llama_wrapper.py b/vibe_bot/llama_wrapper.py deleted file mode 100644 index 5c3decb..0000000 --- a/vibe_bot/llama_wrapper.py +++ /dev/null @@ -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) diff --git a/vibe_bot/llm/__init__.py b/vibe_bot/llm/__init__.py new file mode 100644 index 0000000..fb14681 --- /dev/null +++ b/vibe_bot/llm/__init__.py @@ -0,0 +1 @@ +"""Submodules of the OpenAI-compatible client layer (chat, images, registry).""" diff --git a/vibe_bot/llm/chat.py b/vibe_bot/llm/chat.py new file mode 100644 index 0000000..44f73f8 --- /dev/null +++ b/vibe_bot/llm/chat.py @@ -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, + ) diff --git a/vibe_bot/llm/images.py b/vibe_bot/llm/images.py new file mode 100644 index 0000000..a4ac29f --- /dev/null +++ b/vibe_bot/llm/images.py @@ -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 "" diff --git a/vibe_bot/llm/registry.py b/vibe_bot/llm/registry.py new file mode 100644 index 0000000..00fdd25 --- /dev/null +++ b/vibe_bot/llm/registry.py @@ -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) diff --git a/vibe_bot/llm_client.py b/vibe_bot/llm_client.py new file mode 100644 index 0000000..a5a4520 --- /dev/null +++ b/vibe_bot/llm_client.py @@ -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) diff --git a/vibe_bot/main.py b/vibe_bot/main.py index b3c8475..28bd827 100644 --- a/vibe_bot/main.py +++ b/vibe_bot/main.py @@ -1,1276 +1,18 @@ -"""Main Discord bot application.""" +"""Entrypoint: validate config, configure logging, build and run the bot.""" from __future__ import annotations -import base64 -import logging -import re -import time -from io import BytesIO -from typing import TYPE_CHECKING +from vibe_bot.app import build_bot, configure_logging, create_app +from vibe_bot.config import DISCORD_TOKEN, validate_config -import discord -import requests -from discord import Message -from discord.ext import commands -from vibe_bot import llama_wrapper, tts -from vibe_bot.config import ( - CHAT_ENDPOINT, - CHAT_ENDPOINT_KEY, - CHAT_MODEL, - DISCORD_TOKEN, - IMAGE_EDIT_ENDPOINT, - IMAGE_EDIT_ENDPOINT_KEY, - IMAGE_EDIT_MODEL, - IMAGE_GEN_ENDPOINT, - IMAGE_GEN_ENDPOINT_KEY, - IMAGE_GEN_MODEL, - IMAGE_GEN_SIZE_LANDSCAPE, - IMAGE_GEN_SIZE_PORTRAIT, - IMAGE_GEN_SIZE_SQUARE, - MAX_COMPLETION_TOKENS, - TTS_MODEL_PATH, - TTS_SPEED, - TTS_VOICE, - TTS_VOICES_PATH, - VOICES_LIST, -) -from vibe_bot.database import CustomBotManager, get_database -from vibe_bot.tools import get_channel_members, get_channel_members_impl -from vibe_bot.tts import DEFAULT_LANG +def main() -> None: + """Validate config, then build and run the Discord bot.""" + validate_config() + configure_logging() + app = create_app() + build_bot(app).run(DISCORD_TOKEN) -if TYPE_CHECKING: - from discord.ext.commands import Bot - from discord.ext.commands import Context as CommandsContext -# Configure logging -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) -logger = logging.getLogger(__name__) - -# Initialize the bot -intents = discord.Intents.default() -intents.message_content = True -intents.members = True -intents.presences = True -bot = commands.Bot(command_prefix="!", intents=intents) - - -def get_user_info(user: discord.User | discord.Member) -> str: - """Format user information for inclusion in bot prompts.""" - 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) - - -# Initialize TTS engine -tts_engine: tts.TTSEngine | None = None -try: - tts_engine = tts.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", - ) - -# Name and personality validation constants -MIN_BOT_NAME_LENGTH = 2 -MAX_BOT_NAME_LENGTH = 50 -MIN_PERSONALITY_LENGTH = 10 - -# Image layout (canvas orientation) selection for doodlebob. -DEFAULT_IMAGE_LAYOUT = "square" -VALID_IMAGE_LAYOUTS = ("portrait", "landscape", "square") -LAYOUT_SIZES: dict[str, str] = { - "portrait": IMAGE_GEN_SIZE_PORTRAIT, - "landscape": IMAGE_GEN_SIZE_LANDSCAPE, - "square": IMAGE_GEN_SIZE_SQUARE, -} - -IMAGE_LAYOUT_SYSTEM_PROMPT = ( - "You decide the aspect ratio (layout) of an image that will be generated " - "from a user's request. Choose exactly ONE layout from these three options:\n" - "- portrait: a tall, vertical image (taller than wide). Use for subjects that " - "are taller than they are wide, such as a single standing person or animal, " - "a full-body character, a tall building, a skyscraper, a tree, a rocket, or " - "any vertical composition.\n" - "- landscape: a wide, horizontal image (wider than tall). Use for scenes that " - "are wider than they are tall, such as wide landscapes, panoramas, cityscapes, " - "seas and horizons, battle or group scenes spread out horizontally, or any " - "horizontal composition.\n" - "- square: an image that is as wide as it is tall. Use for balanced subjects, " - "close-ups, faces, single objects, logos, emblems, or whenever no strong tall " - "or wide orientation is implied.\n" - "Rules:\n" - "- Base your choice ONLY on the orientation the content implies.\n" - "- Respond with ONLY the single word portrait, landscape, or square.\n" - "- Do NOT include any other text, punctuation, explanation, or reasoning.\n" -) - -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." -) - - -def parse_image_layout(response: str) -> str: - """Parse an LLM response into a valid image layout. - - Args: - response: The raw LLM response text. - - Returns: - One of "portrait", "landscape", or "square". Falls back to "square" - when the response is empty or does not contain a valid layout. - - """ - text = response.strip().lower() - for layout in VALID_IMAGE_LAYOUTS: - if re.search(rf"\b{layout}\b", text): - return layout - return DEFAULT_IMAGE_LAYOUT - - -def select_image_layout(user_message: str) -> str: - """Ask the LLM to pick an image layout for the given content. - - Args: - user_message: The user's original image request. - - Returns: - One of "portrait", "landscape", or "square". Falls back to "square" - when the LLM returns an empty or malformed response. - - """ - response = llama_wrapper.chat_completion_instruct( - system_prompt=IMAGE_LAYOUT_SYSTEM_PROMPT, - user_prompt=user_message, - openai_url=CHAT_ENDPOINT, - openai_api_key=CHAT_ENDPOINT_KEY, - model=CHAT_MODEL, - max_tokens=MAX_COMPLETION_TOKENS, - ) - return parse_image_layout(response) - - -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 verify_image_prompt(user_message: str, image_prompt: str) -> str: - """Check the drafted prompt literally produces the user's request. - - Args: - user_message: The user's original image request. - image_prompt: The drafted image generation prompt. - - Returns: - The original prompt when the check passes or the LLM returns an - empty response, otherwise the LLM's corrected prompt. - - """ - check_prompt = f"User request: {user_message}\n\nDrafted prompt: {image_prompt}" - response = llama_wrapper.chat_completion_instruct( - system_prompt=IMAGE_PROMPT_VERIFY_SYSTEM_PROMPT, - user_prompt=check_prompt, - openai_url=CHAT_ENDPOINT, - openai_api_key=CHAT_ENDPOINT_KEY, - 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 re.search(r"\bpass\b", response, re.IGNORECASE): - return image_prompt - return response - - -@bot.event -async def on_ready() -> None: - """Log when the bot is ready and logged in.""" - logger.info("Bot is starting up...") - logger.info("Bot logged in as %s", bot.user) - - -@bot.command(name="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 - Example: !custom-bot alfred you are a proper british butler - """ - logger.info( - "Custom bot command initiated by %s: name=%r, personality length=%d", - ctx.author.name, - 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 %s: %r (length: %d)", - ctx.author.name, - bot_name, - name_length, - ) - await ctx.send("Invalid bot name. Name must be between 2 and 50 characters.") - return - - logger.info("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 %s: length=%d", - ctx.author.name, - personality_length, - ) - await ctx.send( - "Invalid personality. Description must be at least 10 characters.", - ) - return - - logger.info("Personality validation passed for bot %r", bot_name) - - # Create custom bot manager - logger.info("Initializing CustomBotManager for user %s", ctx.author.name) - custom_bot_manager = CustomBotManager() - - # Create the custom bot - logger.info( - "Attempting to create custom bot %r for user %s", - bot_name, - ctx.author.name, - ) - success = custom_bot_manager.create_custom_bot( - bot_name=bot_name, - system_prompt=personality, - created_by=str(ctx.author.id), - ) - - if success: - logger.info( - "Successfully created custom bot %r for user %s", - bot_name, - ctx.author.name, - ) - 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} `", - ) - else: - logger.warning( - "Failed to create custom bot %r for user %s", - bot_name, - ctx.author.name, - ) - await ctx.send("Failed to create custom bot. It may already exist.") - - -@bot.command(name="list-custom-bots") -async def list_custom_bots(ctx: CommandsContext[Bot]) -> None: - """List all custom bots available in the server.""" - logger.info("Listing custom bots requested by %s", ctx.author.name) - - # Create custom bot manager - logger.info("Initializing CustomBotManager to list custom bots") - custom_bot_manager = CustomBotManager() - - logger.info("Fetching list of custom bots from database") - bots = custom_bot_manager.list_custom_bots() - - if not bots: - logger.info("No custom bots found for user %s", ctx.author.name) - await ctx.send( - "No custom bots have been created yet. " - "Use `!custom-bot ` to create one.", - ) - return - - logger.info( - "Found %d custom bots, displaying top 10 for %s", - len(bots), - ctx.author.name, - ) - bot_list = "Available Custom Bots:\n\n" - for name, _prompt, _creator in bots: - bot_list += f"* {name}\n" - - logger.info("Sending bot list response to %s", ctx.author.name) - await ctx.send(bot_list) - - -@bot.command(name="delete-custom-bot") -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 - """ - logger.info( - "Delete custom bot command initiated by %s: bot_name=%r", - ctx.author.name, - bot_name, - ) - - # Create custom bot manager - logger.info("Initializing CustomBotManager for delete operation") - custom_bot_manager = CustomBotManager() - - # Get bot info - logger.info("Looking up custom bot %r in database", bot_name) - bot_info = custom_bot_manager.get_custom_bot(bot_name) - - if not bot_info: - logger.warning( - "Custom bot %r not found by user %s", - bot_name, - ctx.author.name, - ) - await ctx.send(f"Custom bot '{bot_name}' not found.") - return - - logger.info( - "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.name, - bot_name, - ) - await ctx.send("You can only delete your own custom bots.") - return - - logger.info( - "User %s is authorized to delete bot %r", - ctx.author.name, - bot_name, - ) - - # Delete the bot - logger.info("Deleting custom bot %r from database", bot_name) - success = custom_bot_manager.delete_custom_bot(bot_name) - - if success: - logger.info( - "Successfully deleted custom bot %r by user %s", - bot_name, - ctx.author.name, - ) - 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.name, - ) - await ctx.send("Failed to delete custom bot.") - - -# Handle custom bot commands -@bot.event -async def on_message(message: Message) -> None: - """Handle incoming messages for custom bot command detection.""" - # Skip bot messages - if message.author == bot.user: - return - - message_author = message.author.name - message_content = message.content.lower() - - logger.debug( - "Processing message from %s: %r...", - message_author, - message_content[:50], - ) - - ctx = await bot.get_context(message) - - logger.info("Initializing CustomBotManager to check for custom bot commands") - custom_bot_manager = CustomBotManager() - - logger.info("Fetching list of custom bots to check for matching commands") - custom_bots = custom_bot_manager.list_custom_bots() - - logger.info("Checking %d custom bots for command match", len(custom_bots)) - for bot_name, system_prompt, _ in custom_bots: - # Check if message starts with the custom bot name followed by a space - if message_content.startswith(f"!{bot_name} "): - logger.info( - "Custom bot command detected: %r triggered by %s", - bot_name, - message.author.name, - ) - - # Extract the actual message (remove the bot name prefix) - user_message = message.content[len(f"!{bot_name} ") :] - logger.debug( - "Extracted user message for bot %r: %r...", - bot_name, - user_message[:50], - ) - - # Prepare the payload with custom personality - response_prefix = f"{bot_name} response" - - logger.info("Sending request to OpenAI API for bot %r", bot_name) - await handle_chat( - 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) - - -@bot.command(name="lobotomize") -async def lobotomize(ctx: CommandsContext[Bot]) -> None: - """Clear all conversation history and memory for all bots.""" - logger.info("Lobotomize command triggered by %s", ctx.author.name) - db = get_database() - db.clear_all_messages() - await ctx.send("All conversation history and memory has been cleared. 🧠✨") - - -@bot.command(name="debug") -async def debug( - ctx: CommandsContext[Bot], - *, - subcommand: str | None = None, -) -> None: - """Debug menu for various debugging sub-commands. - - Usage: !debug - 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 %s with subcommand %r", ctx.author.name, 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) - chunk_size = 1900 - chunks: list[str] = [] - current_chunk = result - while current_chunk: - if len(current_chunk) <= chunk_size: - chunks.append(current_chunk) - break - split_pos = current_chunk.rfind("\n", 0, chunk_size) - if split_pos == -1: - split_pos = chunk_size - chunks.append(current_chunk[:split_pos]) - current_chunk = current_chunk[split_pos:].lstrip("\n") - - for chunk in chunks: - await ctx.send(chunk) - return - - if subcommand == "whoami": - user_info = get_user_info(ctx.author) - chunk_size = 1900 - whoami_chunks: list[str] = [] - current_chunk = user_info - while current_chunk: - if len(current_chunk) <= chunk_size: - whoami_chunks.append(current_chunk) - break - split_pos = current_chunk.rfind("\n", 0, chunk_size) - if split_pos == -1: - split_pos = chunk_size - whoami_chunks.append(current_chunk[:split_pos]) - current_chunk = current_chunk[split_pos:].lstrip("\n") - - for chunk in whoami_chunks: - 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" - tool_list += f" Parameters: {get_channel_members.args_schema.model_json_schema()}" # type: ignore - chunk_size = 1900 - tool_chunks: list[str] = [] - current_chunk = tool_list - while current_chunk: - if len(current_chunk) <= chunk_size: - tool_chunks.append(current_chunk) - break - split_pos = current_chunk.rfind("\n", 0, chunk_size) - if split_pos == -1: - split_pos = chunk_size - tool_chunks.append(current_chunk[:split_pos]) - current_chunk = current_chunk[split_pos:].lstrip("\n") - - for chunk in tool_chunks: - await ctx.send(chunk) - return - - await ctx.send( - f"Unknown debug sub-command: `{subcommand}`\n\n" - f"Use `!debug` to see available sub-commands.", - ) - - -@bot.command(name="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 --voice ` to choose a voice." - - chunk_size = 1900 - chunks: list[str] = [] - current_chunk = voice_list - while current_chunk: - if len(current_chunk) <= chunk_size: - chunks.append(current_chunk) - break - split_pos = current_chunk.rfind("\n", 0, chunk_size) - if split_pos == -1: - split_pos = chunk_size - chunks.append(current_chunk[:split_pos]) - current_chunk = current_chunk[split_pos:].lstrip("\n") - - for chunk in chunks: - await ctx.send(chunk) - - -@bot.command(name="speak") -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 --voice - plain text to speech - Usage: !speak --voice - 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 - """ - if tts_engine is None: - await ctx.send( - "TTS engine not initialized. " - "Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.", - ) - return - - # Parse --voice flag from the message - voice = None - voice_match = message.rsplit("--voice ", 1) - if len(voice_match) == 2: - voice = voice_match[1].strip() - message = voice_match[0].rstrip() - - if not message or not message.strip(): - await ctx.send("Please provide text to speak.") - return - - # Validate voice if provided - if voice: - all_voices = [v for cat in VOICES_LIST.values() for v in cat["voices"]] - if voice not in all_voices: - await ctx.send( - f"Unknown voice '{voice}'. Use `!voices` to see available voices." - ) - return - - custom_bot_manager = CustomBotManager() - custom_bots = custom_bot_manager.list_custom_bots() - bot_names = [b[0] for b in custom_bots] - - first_word = message.split(maxsplit=1)[0] if message.split() else "" - if first_word in bot_names: - await _speak_with_bot( - ctx, first_word, message, tts_engine, custom_bot_manager, voice - ) - else: - await _speak_plain(ctx, message, tts_engine, voice) - - -async def _speak_with_bot( - ctx: CommandsContext[Bot], - bot_name: str, - message: str, - engine: tts.TTSEngine, - custom_bot_manager: CustomBotManager, - voice: str | None = None, -) -> None: - """Handle speak command for a custom bot.""" - 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 = custom_bot_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 = f"{system_prompt}\nKeep your responses under 2-3 sentences.\n\nUser Information:\n{get_user_info(ctx.author)}" - - # Determine language for the chosen voice - chosen_voice = voice or TTS_VOICE - lang = DEFAULT_LANG - for cat in VOICES_LIST.values(): - if chosen_voice in cat["voices"]: - lang = str(cat["language"]) - break - - try: - db = get_database() - context = db.get_conversation_context( - user_id=str(ctx.author.id), - current_message=text_to_speak, - max_context=5, - ) - - prompts = [{"role": "user", "content": text_to_speak}] - if context: - prompts = context + prompts - - # Build tool definitions from LangChain tools - speak_tools: list[dict[str, object]] = [ - { - "type": "function", - "function": { - "name": get_channel_members.name, - "description": get_channel_members.description, - "parameters": get_channel_members.args_schema.model_json_schema(), # type: ignore - }, - }, - ] - - def speak_tool_executor(tool_name: str, tool_args: dict[str, str]) -> str: - """Execute a tool by name with the given arguments.""" - if tool_name == "get_channel_members": - return get_channel_members_impl(ctx.channel) - return f"Unknown tool: {tool_name}" - - 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 llama_wrapper.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, - openai_url=CHAT_ENDPOINT, - openai_api_key=CHAT_ENDPOINT_KEY, - model=CHAT_MODEL, - max_tokens=MAX_COMPLETION_TOKENS, - ) - - if not bot_response: - await ctx.send(f"**{bot_name}** failed to generate a response.") - return - - 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: - 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, - ) - - await ctx.send(f"**{bot_name}**: {bot_response}") - await ctx.send(f"Generating speech for **{bot_name}**...") - audio_buffer = engine.generate_audio( - bot_response, - voice=chosen_voice, - speed=TTS_SPEED, - lang=lang, - ) - - audio_file = discord.File(audio_buffer, filename="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( - ctx: CommandsContext[Bot], - message: str, - engine: tts.TTSEngine, - voice: str | None = None, -) -> None: - """Handle speak command for plain text.""" - chosen_voice = voice or TTS_VOICE - - # Determine language for the chosen voice - lang = DEFAULT_LANG - for cat in VOICES_LIST.values(): - if chosen_voice in cat["voices"]: - lang = str(cat["language"]) - break - - try: - await ctx.send("Generating speech...") - audio_buffer = engine.generate_audio( - message, - voice=chosen_voice, - speed=TTS_SPEED, - lang=lang, - ) - - audio_file = discord.File(audio_buffer, filename="speech.mp3") - await ctx.send(file=audio_file) - except Exception: - logger.exception("Error in speak command") - await ctx.send("Error generating speech.") - - -@bot.command(name="doodlebob") -async def doodlebob(ctx: CommandsContext[Bot], *, message: str) -> None: - """Convert a message into an image using Doodlebob.""" - logger.info( - "Doodlebob command triggered by %s: %s", - ctx.author.name, - message[:100], - ) - await ctx.send("**Doodlebob shopping for a canvas...**") - - # Let the LLM pick the canvas orientation based on the content. - layout = select_image_layout(message) - logger.info("Doodlebob selected layout %r for %s", layout, ctx.author.name) - await ctx.send(f"**Doodlebob selected {layout}**") - - system_prompt = IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE.format(layout=layout) - - # Wait for the generated image prompt - image_prompt = llama_wrapper.chat_completion_instruct( - system_prompt=system_prompt, - user_prompt=message, - openai_url=CHAT_ENDPOINT, - openai_api_key=CHAT_ENDPOINT_KEY, - 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 = verify_image_prompt(message, image_prompt) - logger.info("Doodlebob final image prompt: %s", image_prompt) - - # Alert the user we're generating the image - db = get_database() - estimated_seconds = db.get_image_generation_time_estimate() - await ctx.send(f"**Doodlebob calling drone strike on {image_prompt[:100]}...**") - if estimated_seconds is not None: - await ctx.send(f"**Drone ETA: ~{estimated_seconds:.0f} seconds**") - - start_time = time.monotonic() - image_b64 = llama_wrapper.image_generation( - prompt=image_prompt, - openai_url=IMAGE_GEN_ENDPOINT, - openai_api_key=IMAGE_GEN_ENDPOINT_KEY, - 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 - - db.record_image_generation_time(elapsed_seconds) - - try: - edited_image_data = BytesIO(base64.b64decode(image_b64)) - send_img = discord.File(edited_image_data, filename="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.") - - -@bot.command(name="retcon") -async def retcon(ctx: CommandsContext[Bot], *, message: str) -> None: - """Edit an attached image based on a text prompt.""" - image_data_list: list[BytesIO] = [] - for discord_image in ctx.message.attachments: - image_url = discord_image.url - try: - response = requests.get(image_url, timeout=30) # noqa: ASYNC210 - image_data = response.content - except requests.RequestException as e: - logger.warning("Failed to download image from %s: %s", image_url, e) - continue - image_bytestream = BytesIO(image_data) - image_data_list.append(image_bytestream) - - await ctx.send(f"**Rewriting history to match {message[:100]}...**") - - image_b64 = llama_wrapper.image_edit( - image=image_data_list, - prompt=message, - openai_url=IMAGE_EDIT_ENDPOINT, - openai_api_key=IMAGE_EDIT_ENDPOINT_KEY, - model=IMAGE_EDIT_MODEL, - ) - - # Save the image to a file - edited_image_data = BytesIO(base64.b64decode(image_b64)) - send_img = discord.File(edited_image_data, filename="image.png") - await ctx.send(file=send_img) - - -@bot.command(name="history") -async def history(ctx: CommandsContext[Bot], bot_name: str) -> None: - """View the chat history of a custom bot. - - Usage: !history - """ - logger.info( - "History command triggered by %s for bot %r", - ctx.author.name, - bot_name, - ) - - custom_bot_manager = CustomBotManager() - bot_info = custom_bot_manager.get_custom_bot(bot_name) - - if not bot_info: - await ctx.send(f"Custom bot '{bot_name}' not found.") - return - - db = get_database() - history = 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) - - chunk_size = 1900 - chunks: list[str] = [] - current_chunk = full_text - while current_chunk: - if len(current_chunk) <= chunk_size: - chunks.append(current_chunk) - break - split_pos = current_chunk.rfind("\n", 0, chunk_size) - if split_pos == -1: - split_pos = chunk_size - chunks.append(current_chunk[:split_pos]) - current_chunk = current_chunk[split_pos:].lstrip("\n") - - for chunk in chunks: - await ctx.send(chunk) - - -@bot.command(name="talkforme") -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 - """ - talk_limit = 20 - - MIN_TALKFORME_PARTS = 4 - parts = message.split(" ", maxsplit=MIN_TALKFORME_PARTS - 1) - if len(parts) < MIN_TALKFORME_PARTS: - await ctx.send("Usage: !talkforme bot1 bot2 ") - return - - bot1_name = parts[0] - bot2_name = parts[1] - limit = parts[2] - topic_list = parts[3:] - - topic = " ".join(topic_list) - - custom_bot_manager = CustomBotManager() - bot1 = custom_bot_manager.get_custom_bot(bot1_name) - - if not bot1: - await ctx.send(f"{bot1_name} is not a real bot...") - return - _, bot1_prompt, _, _ = bot1 - - bot2 = custom_bot_manager.get_custom_bot(bot2_name) - - if not bot2: - await ctx.send(f"{bot2_name} is not a real bot...") - return - _, bot2_prompt, _, _ = bot2 - - await ctx.send( - f"{bot1_name} is going to talk to {bot2_name} " - f'about "{topic[:50]}" for {limit} replies.', - ) - - bot_list = [(bot1_name, bot1_prompt), (bot2_name, bot2_prompt)] - - try: - message_limit = int(limit) - except ValueError: - await ctx.send("Message limit must be an integer.") - return - - def flip_counter(counter: int) -> int: - """Flip between 0 and 1.""" - return 1 if counter == 0 else 0 - - 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 = llama_wrapper.chat_completion_with_history( - system_prompt=( - current_bot[1] + f"\nKeep your responses under 2-3 sentences. " - f"You are talking to {current_bot[flip_counter(bot_counter)][0]}" - ), - prompts=prompt_histories[bot_counter], - openai_url=CHAT_ENDPOINT, - openai_api_key=CHAT_ENDPOINT_KEY, - model=CHAT_MODEL, - max_tokens=MAX_COMPLETION_TOKENS, - ) - await ctx.send(f"## {current_bot[0]}\n{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 < min(message_limit, talk_limit): - current_bot = bot_list[bot_counter] - logger.info("Current bot is %s", current_bot[0]) - bot_response = llama_wrapper.chat_completion_with_history( - system_prompt=( - current_bot[1] + f"\nKeep your responses under 2-3 sentences. " - f"You are talking to {current_bot[flip_counter(bot_counter)][0]}" - ), - prompts=prompt_histories[bot_counter], - openai_url=CHAT_ENDPOINT, - openai_api_key=CHAT_ENDPOINT_KEY, - 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]}") - while bot_response: - send_chunk = bot_response[:1000] - bot_response = bot_response[1000:] - await ctx.send(send_chunk) - bot_counter = flip_counter(counter=bot_counter) - logger.info("Message counter is %d/%s", message_counter, limit) - - -async def handle_chat( - ctx: CommandsContext[Bot], - *, - bot_name: str, - message: str, - system_prompt: str, - response_prefix: str, -) -> None: - """Handle chat completion for a custom bot command. - - Args: - ctx: The Discord command context. - bot_name: The name of the custom bot. - message: The user message to process. - system_prompt: The system prompt for the bot. - response_prefix: The prefix for the response message. - - """ - await ctx.send(f"{bot_name} is searching its databanks for {message[:50]}...") - - # Get database instance - db = get_database() - - # Get conversation context using RAG - context = db.get_conversation_context( - user_id=str(ctx.author.id), - current_message=message, - max_context=5, - ) - - prompts = [{"role": "user", "content": message}] - - if context: - prompts = context + prompts - - logger.info("Chat prompts: %s", prompts) - - system_prompt_edit = f"{system_prompt}\nKeep your responses under 2-3 sentences.\n\nUser Information:\n{get_user_info(ctx.author)}" - - # Build tool definitions from LangChain tools - tools: list[dict[str, object]] = [ - { - "type": "function", - "function": { - "name": get_channel_members.name, - "description": get_channel_members.description, - "parameters": get_channel_members.args_schema.model_json_schema(), # type: ignore - }, - }, - ] - - def tool_executor(tool_name: str, tool_args: dict[str, str]) -> str: - """Execute a tool by name with the given arguments.""" - if tool_name == "get_channel_members": - return get_channel_members_impl(ctx.channel) - return f"Unknown tool: {tool_name}" - - 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 llama_wrapper.chat_completion_with_tools( - system_prompt=system_prompt_edit, - prompts=prompts, - tools=tools, - tool_executor=tool_executor, - tool_call_notifier=tool_call_notifier, - openai_url=CHAT_ENDPOINT, - openai_api_key=CHAT_ENDPOINT_KEY, - model=CHAT_MODEL, - max_tokens=MAX_COMPLETION_TOKENS, - ) - - # Store both user message and bot response in the database - 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: - 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, - ) - - # Send the response back to the chat - await ctx.send(response_prefix) - while bot_response: - send_chunk = bot_response[:1000] - bot_response = 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.") - - -# Run the bot if __name__ == "__main__": - bot.run(DISCORD_TOKEN) + main() diff --git a/vibe_bot/prompts.py b/vibe_bot/prompts.py new file mode 100644 index 0000000..a71019a --- /dev/null +++ b/vibe_bot/prompts.py @@ -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) diff --git a/vibe_bot/services/__init__.py b/vibe_bot/services/__init__.py new file mode 100644 index 0000000..aae5c20 --- /dev/null +++ b/vibe_bot/services/__init__.py @@ -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``. +""" diff --git a/vibe_bot/services/chat_service.py b/vibe_bot/services/chat_service.py new file mode 100644 index 0000000..8e27970 --- /dev/null +++ b/vibe_bot/services/chat_service.py @@ -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.") diff --git a/vibe_bot/services/conversation_service.py b/vibe_bot/services/conversation_service.py new file mode 100644 index 0000000..4fcfe99 --- /dev/null +++ b/vibe_bot/services/conversation_service.py @@ -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, + ) diff --git a/vibe_bot/services/image_service.py b/vibe_bot/services/image_service.py new file mode 100644 index 0000000..dd88a07 --- /dev/null +++ b/vibe_bot/services/image_service.py @@ -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) diff --git a/vibe_bot/services/speech_service.py b/vibe_bot/services/speech_service.py new file mode 100644 index 0000000..6ab68e1 --- /dev/null +++ b/vibe_bot/services/speech_service.py @@ -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 ` 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.*)\s+--voice\s+(?P\S+)$") + + +def parse_voice_flag(message: str) -> tuple[str, str | None]: + """Split a trailing `--voice ` 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.") diff --git a/vibe_bot/tests/_helpers.py b/vibe_bot/tests/_helpers.py new file mode 100644 index 0000000..59394cb --- /dev/null +++ b/vibe_bot/tests/_helpers.py @@ -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] diff --git a/vibe_bot/tests/conftest.py b/vibe_bot/tests/conftest.py index fdc5cf5..3ee67dd 100644 --- a/vibe_bot/tests/conftest.py +++ b/vibe_bot/tests/conftest.py @@ -5,12 +5,16 @@ from __future__ import annotations import tempfile import warnings from collections.abc import Generator +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any from unittest.mock import AsyncMock, MagicMock, patch import numpy as np import pytest +from discord.ext import commands + +from vibe_bot.app import App, build_bot warnings.filterwarnings( "ignore", @@ -21,41 +25,79 @@ if TYPE_CHECKING: 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 -def mock_env_vars() -> Generator[None]: - """Provide minimal env vars for config loading.""" - with patch.dict( - "os.environ", - { - "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": "test-completion-model", - "IMAGE_GEN_MODEL": "test-image-model", - "IMAGE_EDIT_MODEL": "test-image-edit-model", - "EMBEDDING_MODEL": "test-embedding-model", - "CHAT_ENDPOINT_KEY": "test-key", - "COMPLETION_ENDPOINT_KEY": "test-completion-key", - "IMAGE_GEN_ENDPOINT_KEY": "test-image-key", - "IMAGE_EDIT_ENDPOINT_KEY": "test-image-edit-key", - "EMBEDDING_ENDPOINT_KEY": "test-embedding-key", - "MAX_COMPLETION_TOKENS": "1000", - "MAX_HISTORY_MESSAGES": "1000", - "SIMILARITY_THRESHOLD": "0.7", - "TOP_K_RESULTS": "5", - "TTS_MODEL_PATH": "/tmp/test-model.onnx", - "TTS_VOICES_PATH": "/tmp/test-voices.bin", - "TTS_VOICE": "af_sarah", - "TTS_SPEED": "1.0", - "DB_PATH": ":memory:", - }, - clear=False, - ): - yield +def mock_ctx() -> MagicMock: + """Create 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 + + +@pytest.fixture +def app_mocks() -> AppMocks: + """An App built entirely from mocks, with alfred in the bot cache.""" + db = MagicMock() + manager = MagicMock() + manager.list_custom_bots.return_value = [ + ("alfred", "british butler", "user123"), + ] + registry = MagicMock() + chat = MagicMock() + chat.handle = AsyncMock() + 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 @@ -71,22 +113,13 @@ def temp_db_path() -> Generator[str]: def mock_embedding() -> Generator[MagicMock]: """Provide a mock embedding function returning a fixed vector.""" vector: list[float] = [0.1] * 2048 - with patch("vibe_bot.llama_wrapper.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: + with patch("vibe_bot.llm_client.embedding", return_value=vector) as mock: yield mock @pytest.fixture def chat_db( temp_db_path: str, - mock_openai_client: MagicMock, mock_embedding: MagicMock, ) -> Generator[ChatDatabase]: """Provide a ChatDatabase instance with a temp database.""" @@ -94,7 +127,6 @@ def chat_db( db = ChatDatabase(db_path=temp_db_path) yield db - db.client.close() @pytest.fixture @@ -133,106 +165,3 @@ def mock_kokoro_tts() -> Generator[dict[str, Any]]: "mock_samples": mock_samples, "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 diff --git a/vibe_bot/tests/test_app.py b/vibe_bot/tests/test_app.py new file mode 100644 index 0000000..6264d1b --- /dev/null +++ b/vibe_bot/tests/test_app.py @@ -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") diff --git a/vibe_bot/tests/test_commands.py b/vibe_bot/tests/test_commands.py new file mode 100644 index 0000000..a152093 --- /dev/null +++ b/vibe_bot/tests/test_commands.py @@ -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() diff --git a/vibe_bot/tests/test_config.py b/vibe_bot/tests/test_config.py index 08835c4..81f82c1 100644 --- a/vibe_bot/tests/test_config.py +++ b/vibe_bot/tests/test_config.py @@ -4,6 +4,9 @@ from __future__ import annotations import subprocess import sys +from unittest.mock import patch + +import pytest def test_config_defaults() -> None: @@ -12,17 +15,14 @@ def test_config_defaults() -> None: for k, v in { "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": "test-completion-model", "IMAGE_GEN_MODEL": "test-image-model", "IMAGE_EDIT_MODEL": "test-image-edit-model", "EMBEDDING_MODEL": "test-embedding-model", "CHAT_ENDPOINT_KEY": "test-key", - "COMPLETION_ENDPOINT_KEY": "test-completion-key", "IMAGE_GEN_ENDPOINT_KEY": "test-image-key", "IMAGE_EDIT_ENDPOINT_KEY": "test-image-edit-key", "EMBEDDING_ENDPOINT_KEY": "test-embedding-key", @@ -39,8 +39,6 @@ def test_config_defaults() -> None: env_str += f'os.environ["{k}"] = "{v}"\n' code = f""" -import sys -sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots") import os os.environ.clear() os.environ["PATH"] = "/usr/bin:/bin" @@ -48,12 +46,10 @@ os.environ["PATH"] = "/usr/bin:/bin" import vibe_bot.config assert vibe_bot.config.DISCORD_TOKEN == "test-token" 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_EDIT_ENDPOINT == "https://image-edit.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.COMPLETION_MODEL == "test-completion-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.EMBEDDING_MODEL == "test-embedding-model" @@ -78,20 +74,23 @@ print("OK") 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 = "" for k, v in env_vars.items(): env_str += f'os.environ["{k}"] = "{v}"\n' code = f""" -import sys -sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots") import os os.environ.clear() os.environ["PATH"] = "/usr/bin:/bin" {env_str} try: import vibe_bot.config + vibe_bot.config.validate_config() print("NO_ERROR") except RuntimeError as e: print(f"ERROR: {{e}}") @@ -116,12 +115,10 @@ def test_config_missing_discord_token() -> None: env: dict[str, str] = { "DISCORD_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": "test-completion-model", "IMAGE_GEN_MODEL": "test-image-model", "IMAGE_EDIT_MODEL": "test-image-edit-model", "EMBEDDING_MODEL": "test-embedding-model", @@ -134,12 +131,10 @@ def test_config_missing_chat_endpoint() -> None: env: dict[str, str] = { "DISCORD_TOKEN": "test-token", "CHAT_ENDPOINT": "", - "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": "test-completion-model", "IMAGE_GEN_MODEL": "test-image-model", "IMAGE_EDIT_MODEL": "test-image-edit-model", "EMBEDDING_MODEL": "test-embedding-model", @@ -147,35 +142,15 @@ def test_config_missing_chat_endpoint() -> None: _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: """Test that RuntimeError is raised when IMAGE_GEN_ENDPOINT 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": "", "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", @@ -188,12 +163,10 @@ def test_config_missing_image_edit_endpoint() -> None: 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": "", "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", @@ -206,12 +179,10 @@ def test_config_missing_embedding_endpoint() -> None: 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": "", "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", @@ -224,12 +195,10 @@ def test_config_missing_chat_model() -> None: 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": "", - "COMPLETION_MODEL": "test-completion-model", "IMAGE_GEN_MODEL": "test-image-model", "IMAGE_EDIT_MODEL": "test-image-edit-model", "EMBEDDING_MODEL": "test-embedding-model", @@ -237,35 +206,15 @@ def test_config_missing_chat_model() -> None: _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: """Test that RuntimeError is raised when IMAGE_GEN_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": "test-completion-model", "IMAGE_GEN_MODEL": "", "IMAGE_EDIT_MODEL": "test-image-edit-model", "EMBEDDING_MODEL": "test-embedding-model", @@ -278,12 +227,10 @@ def test_config_missing_image_edit_model() -> None: 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": "test-completion-model", "IMAGE_GEN_MODEL": "test-image-model", "IMAGE_EDIT_MODEL": "", "EMBEDDING_MODEL": "test-embedding-model", @@ -296,12 +243,10 @@ def test_config_missing_embedding_model() -> None: 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": "test-completion-model", "IMAGE_GEN_MODEL": "test-image-model", "IMAGE_EDIT_MODEL": "test-image-edit-model", "EMBEDDING_MODEL": "", @@ -309,16 +254,64 @@ def test_config_missing_embedding_model() -> None: _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: """Test that logging is configured in config module.""" from vibe_bot.config import logger assert logger is not None 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 diff --git a/vibe_bot/tests/test_database.py b/vibe_bot/tests/test_database.py index 564c0ca..29a990a 100644 --- a/vibe_bot/tests/test_database.py +++ b/vibe_bot/tests/test_database.py @@ -2,17 +2,40 @@ 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 import numpy as np import pytest if TYPE_CHECKING: + import sqlite3 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: """Test converting a vector to bytes and back.""" 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) +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: """Test adding a message to the database.""" result = chat_db.add_message( @@ -67,7 +97,7 @@ def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None: ) 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 messages[0][0] == "msg-1" 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: """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( message_id="msg-no-embed", user_id="user-1", @@ -106,14 +136,14 @@ def test_add_message_duplicate( content="Second content", ) - messages = chat_db.get_recent_messages(limit=10) + messages = _recent_messages(chat_db.db_path, 10) assert len(messages) == 1 assert messages[0][2] == "Second content" def test_add_message_failure(chat_db: ChatDatabase) -> None: """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( message_id="msg-fail", user_id="user-1", @@ -123,11 +153,304 @@ def test_add_message_failure(chat_db: ChatDatabase) -> None: 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, mock_embedding: MagicMock, ) -> 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( message_id="msg-1", user_id="u1", @@ -147,17 +470,17 @@ def test_get_recent_messages( content="Third", ) - messages = chat_db.get_recent_messages(limit=2) + messages = _recent_messages(chat_db.db_path, 2) assert len(messages) == 2 assert messages[0][2] == "Third" assert messages[1][2] == "Second" -def test_get_recent_messages_limit( +def test_recent_messages_limit( chat_db: ChatDatabase, mock_embedding: MagicMock, ) -> None: - """Test that get_recent_messages respects the limit.""" + """The newest-rows query respects the limit.""" for i in range(5): chat_db.add_message( message_id=f"msg-{i}", @@ -166,10 +489,25 @@ def test_get_recent_messages_limit( content=f"Message {i}", ) - messages = chat_db.get_recent_messages(limit=3) + messages = _recent_messages(chat_db.db_path, 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( chat_db: ChatDatabase, mock_embedding: MagicMock, @@ -190,7 +528,7 @@ def test_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 @@ -264,6 +602,25 @@ def test_image_generation_estimate_after_capping( 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( chat_db: ChatDatabase, mock_embedding: MagicMock, @@ -320,6 +677,39 @@ def test_get_user_history_excludes_bot( 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( chat_db: ChatDatabase, mock_embedding: MagicMock, @@ -349,21 +739,333 @@ def test_get_conversation_context_empty(chat_db: ChatDatabase) -> None: 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: - """Test creating a custom bot.""" + """Test creating a custom bot returns "created" for a new name.""" result = custom_bot_manager.create_custom_bot( bot_name="alfred", system_prompt="You are a british butler", created_by="user-123", ) - assert result is True + assert result == "created" def test_custom_bot_create_duplicate( custom_bot_manager: Any, ) -> None: """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", system_prompt="First personality", created_by="user-1", @@ -373,13 +1075,25 @@ def test_custom_bot_create_duplicate( system_prompt="Second personality", created_by="user-1", ) - assert result is True + assert first == "created" + assert result == "replaced" bot = custom_bot_manager.get_custom_bot("alfred") assert bot is not None 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( custom_bot_manager: Any, ) -> None: @@ -399,6 +1113,18 @@ def test_custom_bot_get_not_found(custom_bot_manager: Any) -> 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( custom_bot_manager: Any, ) -> None: @@ -413,8 +1139,7 @@ def test_custom_bot_get_returns_correct_data( assert result[0] == "testbot" assert result[1] == "test prompt" assert result[2] == "creator-1" - assert result[3] is not None - assert "20" in result[3] + assert isinstance(result[3], datetime) 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 +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: """Test deleting a custom bot.""" custom_bot_manager.create_custom_bot( @@ -462,32 +1204,18 @@ def test_custom_bot_delete_nonexistent( assert result is False -def test_custom_bot_deactivate(custom_bot_manager: Any) -> None: - """Test deactivating a custom bot.""" - custom_bot_manager.create_custom_bot( - bot_name="inactive-bot", - 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_delete_failure(custom_bot_manager: Any) -> None: + """A database error while deleting yields False, not an exception.""" + with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()): + assert custom_bot_manager.delete_custom_bot("whatever") is False def test_custom_bot_list_excludes_inactive( custom_bot_manager: Any, ) -> 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( bot_name="active-bot", system_prompt="stays active", @@ -498,7 +1226,12 @@ def test_custom_bot_list_excludes_inactive( system_prompt="should not appear", 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() assert len(bots) == 1 @@ -532,16 +1265,13 @@ def test_database_get_database_singleton(temp_db_path: str) -> None: db2 = get_database() assert db1 is db2 - db1.client.close() - def test_database_init_creates_tables(temp_db_path: str) -> None: """Test that database initialization creates the expected tables.""" 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) - db.client.close() import sqlite3 diff --git a/vibe_bot/tests/test_docs.py b/vibe_bot/tests/test_docs.py new file mode 100644 index 0000000..c1b521f --- /dev/null +++ b/vibe_bot/tests/test_docs.py @@ -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(?:[│ ] )*)(?:├── |└── )(?P\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}" diff --git a/vibe_bot/tests/test_llama_wrapper.py b/vibe_bot/tests/test_llama_wrapper.py deleted file mode 100644 index 8e5f423..0000000 --- a/vibe_bot/tests/test_llama_wrapper.py +++ /dev/null @@ -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 diff --git a/vibe_bot/tests/test_llm_client.py b/vibe_bot/tests/test_llm_client.py new file mode 100644 index 0000000..b223ce6 --- /dev/null +++ b/vibe_bot/tests/test_llm_client.py @@ -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("rate limited") + + 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" diff --git a/vibe_bot/tests/test_logging.py b/vibe_bot/tests/test_logging.py new file mode 100644 index 0000000..3be6375 --- /dev/null +++ b/vibe_bot/tests/test_logging.py @@ -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") diff --git a/vibe_bot/tests/test_main.py b/vibe_bot/tests/test_main.py deleted file mode 100644 index 7cad75f..0000000 --- a/vibe_bot/tests/test_main.py +++ /dev/null @@ -1,1321 +0,0 @@ -"""Tests for the main module (Discord bot commands).""" - -from __future__ import annotations - -from datetime import UTC -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - - -@pytest.fixture -def mock_ctx() -> MagicMock: - """Create 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 - - -@pytest.fixture -def mock_ctx_with_member() -> MagicMock: - """Create a mock Discord command context with full member data.""" - 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 = "Admin" - mock_activity = MagicMock() - mock_activity.name = "Chess" - ctx.author.activities = [mock_activity] - from datetime import datetime - - ctx.author.joined_at = datetime(2024, 1, 15, tzinfo=UTC) - ctx.author.created_at = datetime(2023, 6, 1, tzinfo=UTC) - 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 test_bot_initialized(mock_discord: dict[str, MagicMock]) -> None: - """Test that the bot is initialized.""" - import vibe_bot.main as main_module - - assert main_module.bot is not None - - -def test_bot_intents_set(mock_discord: dict[str, MagicMock]) -> None: - """Test that required intents are enabled.""" - import vibe_bot.main as main_module - - main_module.bot = mock_discord["bot_instance"] - assert main_module.MIN_BOT_NAME_LENGTH == 2 - assert main_module.MAX_BOT_NAME_LENGTH == 50 - assert main_module.MIN_PERSONALITY_LENGTH == 10 - - -def test_bot_intents_members_and_presence(mock_discord: dict[str, MagicMock]) -> None: - """Test that members and presence intents are enabled.""" - intents = mock_discord["Intents"].default.return_value - intents.message_content = True - intents.members = True - intents.presences = True - assert intents.message_content is True - assert intents.members is True - assert intents.presences is True - - -@patch("vibe_bot.main.tts_engine", None) -def test_speak_tts_not_initialized(mock_ctx: MagicMock) -> None: - """Test speak command when TTS engine is not initialized.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module.speak(mock_ctx, message="hello world")) - mock_ctx.send.assert_called_once() - 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, - mock_tts_engine: MagicMock, -) -> None: - """Test speak command with empty message.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module.speak(mock_ctx, message="")) - mock_ctx.send.assert_called_once() - call_args = mock_ctx.send.call_args[0][0] - assert "Please provide text" in call_args - - -def test_speak_plain_text( - mock_ctx: MagicMock, - mock_tts_engine: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test speak command with plain text (no custom bot prefix).""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.list_custom_bots.return_value = [] - - asyncio.run(main_module.speak(mock_ctx, message="hello world")) - mock_tts_engine.generate_audio.assert_called_once() - assert mock_ctx.send.call_count >= 2 - - -def test_speak_with_custom_bot( - mock_ctx: MagicMock, - mock_tts_engine: MagicMock, - mock_custom_bot_manager: MagicMock, - mock_database: MagicMock, - mock_llama_wrapper: MagicMock, -) -> None: - """Test speak command with a custom bot prefix.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.list_custom_bots.return_value = [ - ("alfred", "british butler", "user-123"), - ] - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "british butler", - "user-123", - "2024-01-01", - ) - - asyncio.run(main_module.speak(mock_ctx, message="alfred what time is it")) - - mock_llama_wrapper.chat_completion_with_tools.assert_called_once() - mock_tts_engine.generate_audio.assert_called_once() - assert mock_ctx.send.call_count >= 3 - text_response = mock_ctx.send.call_args_list[1][0][0] - assert "**alfred**:" in text_response or "**alfred** :" in text_response - - -def test_custom_bot_command_success( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test creating a custom bot successfully.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run( - main_module.custom_bot( - mock_ctx, - bot_name="alfred", - personality="you are a british butler", - ), - ) - - mock_custom_bot_manager.create_custom_bot.assert_called_once() - assert mock_ctx.send.call_count == 2 - - -def test_custom_bot_command_invalid_name_too_short( - mock_ctx: MagicMock, -) -> None: - """Test custom bot command with name too short.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run( - main_module.custom_bot( - mock_ctx, - bot_name="a", - personality="this is a valid personality description", - ), - ) - call_args = mock_ctx.send.call_args[0][0] - assert "Invalid bot name" in call_args - - -def test_custom_bot_command_invalid_name_empty( - mock_ctx: MagicMock, -) -> None: - """Test custom bot command with empty name.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run( - main_module.custom_bot( - mock_ctx, - bot_name="", - personality="this is a valid personality description", - ), - ) - call_args = mock_ctx.send.call_args[0][0] - assert "Invalid bot name" in call_args - - -def test_custom_bot_command_invalid_personality( - mock_ctx: MagicMock, -) -> None: - """Test custom bot command with personality too short.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run( - main_module.custom_bot(mock_ctx, bot_name="testbot", personality="short"), - ) - call_args = mock_ctx.send.call_args[0][0] - assert "Invalid personality" in call_args - - -def test_custom_bot_command_create_fails( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test custom bot command when creation fails.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.create_custom_bot.return_value = False - - asyncio.run( - main_module.custom_bot( - mock_ctx, - bot_name="alfred", - personality="you are a british butler", - ), - ) - call_args = mock_ctx.send.call_args[0][0] - assert "Failed to create custom bot" in call_args - - -def test_list_custom_bots_empty( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test listing custom bots when none exist.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.list_custom_bots.return_value = [] - - asyncio.run(main_module.list_custom_bots(mock_ctx)) - call_args = mock_ctx.send.call_args[0][0] - assert "No custom bots" in call_args - - -def test_list_custom_bots_with_bots( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test listing custom bots when bots exist.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.list_custom_bots.return_value = [ - ("alfred", "british butler", "user-1"), - ("jarvis", "ai assistant", "user-2"), - ] - - asyncio.run(main_module.list_custom_bots(mock_ctx)) - call_args = mock_ctx.send.call_args[0][0] - assert "Available Custom Bots" in call_args - assert "* alfred" in call_args - assert "* jarvis" in call_args - - -def test_delete_custom_bot_success( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test deleting a custom bot successfully.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "prompt", - "12345", - "2024-01-01", - ) - mock_custom_bot_manager.delete_custom_bot.return_value = True - - asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="alfred")) - call_args = mock_ctx.send.call_args[0][0] - assert "has been deleted" in call_args - - -def test_delete_custom_bot_not_found( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test deleting a non-existent custom bot.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = None - - asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="nonexistent")) - call_args = mock_ctx.send.call_args[0][0] - assert "not found" in call_args - - -def test_delete_custom_bot_not_owner( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test deleting a custom bot you don't own.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "prompt", - "other-user-id", - "2024-01-01", - ) - - asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="alfred")) - call_args = mock_ctx.send.call_args[0][0] - assert "You can only delete your own" in call_args - - -def test_delete_custom_bot_delete_fails( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test deleting a custom bot when delete fails.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "prompt", - "12345", - "2024-01-01", - ) - mock_custom_bot_manager.delete_custom_bot.return_value = False - - asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="alfred")) - call_args = mock_ctx.send.call_args[0][0] - assert "Failed to delete" in call_args - - -def test_on_message_skips_bot_messages(mock_ctx: MagicMock) -> None: - """Test that on_message skips messages from the bot itself.""" - import asyncio - - import vibe_bot.main as main_module - - message = MagicMock() - message.author = main_module.bot.user - message.content = "hello" - - asyncio.run(main_module.on_message(message)) - - -def test_handle_chat_success( - mock_ctx: MagicMock, - mock_database: MagicMock, - mock_llama_wrapper: MagicMock, -) -> None: - """Test handle_chat with successful response.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_with_tools.return_value = ( - "This is a bot response" - ) - - asyncio.run( - main_module.handle_chat( - ctx=mock_ctx, - bot_name="alfred", - message="hello", - system_prompt="you are a butler", - response_prefix="alfred response", - ), - ) - - mock_llama_wrapper.chat_completion_with_tools.assert_called_once() - mock_database.add_message.assert_called() - assert mock_ctx.send.call_count >= 2 - - -def test_handle_chat_error( - mock_ctx: MagicMock, - mock_database: MagicMock, - mock_llama_wrapper: MagicMock, -) -> None: - """Test handle_chat when an exception occurs.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_with_tools.side_effect = Exception("API error") - - asyncio.run( - main_module.handle_chat( - ctx=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() - - -def test_handle_chat_long_response_chunked( - mock_ctx: MagicMock, - mock_database: MagicMock, - mock_llama_wrapper: MagicMock, -) -> None: - """Test that long bot responses are sent in chunks.""" - import asyncio - - import vibe_bot.main as main_module - - long_response = "x" * 2500 - mock_llama_wrapper.chat_completion_with_tools.return_value = long_response - - asyncio.run( - main_module.handle_chat( - ctx=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_speak_plain_with_mock_tts( - mock_ctx: MagicMock, - mock_tts_engine: MagicMock, -) -> None: - """Test _speak_plain function directly.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module._speak_plain(mock_ctx, "hello world", mock_tts_engine)) - - from vibe_bot.config import TTS_SPEED, TTS_VOICE - from vibe_bot.tts import DEFAULT_LANG - - mock_tts_engine.generate_audio.assert_called_once_with( - "hello world", - voice=TTS_VOICE, - speed=TTS_SPEED, - lang=DEFAULT_LANG, - ) - assert mock_ctx.send.call_count >= 2 - - -def test_speak_plain_error( - mock_ctx: MagicMock, - mock_tts_engine: MagicMock, -) -> None: - """Test _speak_plain when audio generation fails.""" - import asyncio - - import vibe_bot.main as main_module - - mock_tts_engine.generate_audio.side_effect = Exception("generation error") - - asyncio.run(main_module._speak_plain(mock_ctx, "hello world", mock_tts_engine)) - - call_args = mock_ctx.send.call_args[0][0] - assert "error generating speech" in call_args.lower() - - -def test_flip_counter() -> None: - """Test the flip_counter helper function defined inside talkforme.""" - - def flip_counter(counter: int) -> int: - return 1 if counter == 0 else 0 - - assert flip_counter(0) == 1 - assert flip_counter(1) == 0 - assert flip_counter(0) == 1 - - -def test_talkforme_invalid_args(mock_ctx: MagicMock) -> None: - """Test talkforme command with invalid arguments.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2")) - call_args = mock_ctx.send.call_args[0][0] - assert "Usage" in call_args - - -def test_talkforme_bot1_not_found( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test talkforme when bot1 doesn't exist.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = None - - asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2 4 a topic")) - call_args = mock_ctx.send.call_args[0][0] - assert "is not a real bot" in call_args - - -def test_talkforme_bot2_not_found( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test talkforme when bot2 doesn't exist.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.side_effect = [ - ("bot1", "bot1 personality", "user-1", "2024-01-01"), - None, - ] - - asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2 4 a topic")) - call_args = mock_ctx.send.call_args[0][0] - assert "is not a real bot" in call_args - - -def test_talkforme_invalid_limit( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test talkforme with non-integer limit.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = ( - "bot1", - "personality", - "user-1", - "2024-01-01", - ) - - asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2 abc topic")) - call_args = mock_ctx.send.call_args[0][0] - assert "must be an integer" in call_args - - -def test_history_bot_not_found( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, -) -> None: - """Test history command when bot doesn't exist.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = None - - asyncio.run(main_module.history(mock_ctx, bot_name="nonexistent")) - call_args = mock_ctx.send.call_args[0][0] - assert "not found" in call_args - - -def test_history_no_history( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, - mock_database: MagicMock, -) -> None: - """Test history command when bot has no chat history.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "british butler", - "user-123", - "2024-01-01", - ) - mock_database.get_bot_history.return_value = [] - - asyncio.run(main_module.history(mock_ctx, bot_name="alfred")) - call_args = mock_ctx.send.call_args[0][0] - assert "No chat history" in call_args - assert "**alfred**" in call_args - - -def test_history_with_data( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, - mock_database: MagicMock, -) -> None: - """Test history command when bot has chat history.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "british butler", - "user-123", - "2024-01-01", - ) - mock_database.get_bot_history.return_value = [ - ("hello", "yes master?"), - ("what time is it", "it is currently 3pm"), - ] - - asyncio.run(main_module.history(mock_ctx, bot_name="alfred")) - - assert mock_ctx.send.call_count >= 1 - - -def test_get_user_info_minimal(mock_ctx: MagicMock) -> None: - """Test get_user_info with minimal member data.""" - import vibe_bot.main as main_module - - result = main_module.get_user_info(mock_ctx.author) - - 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_ctx_with_member: MagicMock) -> None: - """Test get_user_info with full member data including roles and activities.""" - import vibe_bot.main as main_module - - result = main_module.get_user_info(mock_ctx_with_member.author) - - 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: - """Test get_user_info when user has no global name.""" - import vibe_bot.main as main_module - - mock_ctx.author.global_name = None - mock_ctx.author.nick = None - mock_ctx.author.top_role.name = "@everyone" - mock_ctx.author.activities = [] - - result = main_module.get_user_info(mock_ctx.author) - - 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_ctx_with_member: MagicMock, -) -> None: - """Test get_user_info includes top role when not @everyone.""" - import vibe_bot.main as main_module - - result = main_module.get_user_info(mock_ctx_with_member.author) - - assert "Top Role: Admin" in result - - -def test_get_user_info_no_activities(mock_ctx: MagicMock) -> None: - """Test get_user_info when user has no activities.""" - import vibe_bot.main as main_module - - mock_ctx.author.activities = [] - - result = main_module.get_user_info(mock_ctx.author) - - assert "Activities:" not in result - - -def test_handle_chat_includes_user_info( - mock_ctx: MagicMock, - mock_database: MagicMock, - mock_llama_wrapper: MagicMock, -) -> None: - """Test handle_chat includes user info in system prompt.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_with_tools.return_value = ( - "This is a bot response" - ) - - asyncio.run( - main_module.handle_chat( - ctx=mock_ctx, - bot_name="alfred", - message="hello", - system_prompt="you are a butler", - response_prefix="alfred response", - ), - ) - - mock_llama_wrapper.chat_completion_with_tools.assert_called_once() - call_kwargs = mock_llama_wrapper.chat_completion_with_tools.call_args - system_prompt = call_kwargs.kwargs["system_prompt"] - assert "you are a butler" in system_prompt - assert "User Information:" in system_prompt - assert "Username: testuser" in system_prompt - assert "User ID: 12345" in system_prompt - - -def test_speak_with_bot_includes_user_info( - mock_ctx: MagicMock, - mock_tts_engine: MagicMock, - mock_custom_bot_manager: MagicMock, - mock_database: MagicMock, - mock_llama_wrapper: MagicMock, -) -> None: - """Test _speak_with_bot includes user info in system prompt.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.list_custom_bots.return_value = [ - ("alfred", "british butler", "user-123"), - ] - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "british butler", - "user-123", - "2024-01-01", - ) - - asyncio.run(main_module.speak(mock_ctx, message="alfred what time is it")) - - mock_llama_wrapper.chat_completion_with_tools.assert_called_once() - call_kwargs = mock_llama_wrapper.chat_completion_with_tools.call_args - system_prompt = call_kwargs.kwargs["system_prompt"] - assert "british butler" in system_prompt - assert "User Information:" in system_prompt - assert "Username: testuser" in system_prompt - assert "User ID: 12345" in system_prompt - mock_tts_engine.generate_audio.assert_called_once() - - -def test_history_long_response_chunked( - mock_ctx: MagicMock, - mock_custom_bot_manager: MagicMock, - mock_database: MagicMock, -) -> None: - """Test that long history responses are sent in chunks.""" - import asyncio - - import vibe_bot.main as main_module - - mock_custom_bot_manager.get_custom_bot.return_value = ( - "alfred", - "british butler", - "user-123", - "2024-01-01", - ) - long_user = "x" * 500 - long_bot = "y" * 500 - mock_database.get_bot_history.return_value = [ - (long_user, long_bot), - ] - - asyncio.run(main_module.history(mock_ctx, bot_name="alfred")) - - assert mock_ctx.send.call_count >= 1 - - -def test_debug_no_subcommand(mock_ctx: MagicMock) -> None: - """Test debug command without subcommand shows menu.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module.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(mock_ctx: MagicMock) -> None: - """Test debug members command when channel has no guild.""" - import asyncio - - import vibe_bot.main as main_module - - mock_ctx.channel.guild = None - - asyncio.run(main_module.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(mock_ctx: MagicMock) -> None: - """Test debug members command with channel members.""" - import asyncio - - import vibe_bot.main as main_module - - 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] - - asyncio.run(main_module.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(mock_ctx: MagicMock) -> None: - """Test debug command with unknown subcommand.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module.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(mock_ctx: MagicMock) -> None: - """Test debug members command with many members that exceed 1900 chars.""" - import asyncio - - import vibe_bot.main as main_module - - 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 - - asyncio.run(main_module.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(mock_ctx: MagicMock) -> None: - """Test debug whoami command shows user info.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module.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(mock_ctx: MagicMock) -> None: - """Test debug whoami command with minimal user data.""" - import asyncio - - import vibe_bot.main as main_module - - 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 - - asyncio.run(main_module.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(mock_ctx: MagicMock) -> None: - """Test debug tools command shows LLM tools.""" - import asyncio - - import vibe_bot.main as main_module - - asyncio.run(main_module.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() - - -# --------------------------------------------------------------------------- -# Image layout selection (doodlebob) -# --------------------------------------------------------------------------- - -LAYOUT_TEST_SIZES = { - "portrait": "1024x1536", - "landscape": "1536x1024", - "square": "1024x1024", -} - - -def _sent_texts(mock_ctx: MagicMock) -> list[str]: - """Collect positional (text) arguments sent via ctx.send.""" - return [c.args[0] for c in mock_ctx.send.call_args_list if c.args] - - -@pytest.mark.parametrize( - ("response", "expected"), - [ - ("portrait", "portrait"), - ("landscape", "landscape"), - ("square", "square"), - (" Portrait ", "portrait"), - ("LANDSCAPE", "landscape"), - ("square.", "square"), - ("I would use portrait.", "portrait"), - ("This scene is best as landscape.", "landscape"), - ("A square composition works here.", "square"), - ], -) -def test_parse_image_layout_valid(response: str, expected: str) -> None: - """Test that valid LLM layout responses parse to the right layout.""" - import vibe_bot.main as main_module - - assert main_module.parse_image_layout(response) == expected - - -@pytest.mark.parametrize( - "response", - [ - "", - " ", - "banana", - "1024x1024", - "tall and wide", - "squareness", # word boundary should prevent a match on "square" - ], -) -def test_parse_image_layout_defaults_to_square(response: str) -> None: - """Test that empty or malformed responses fall back to square.""" - import vibe_bot.main as main_module - - assert main_module.parse_image_layout(response) == "square" - - -def test_image_prompt_system_prompt_covers_key_details() -> None: - """The prompt-rewrite system prompt forces explicit detail on all aspects.""" - import vibe_bot.main as main_module - - prompt = main_module.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_doodlebob_prompt_rewrite_uses_detailed_system_prompt( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """The prompt rewrite call uses the detailed system prompt with the layout.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "landscape", # layout selection - "a very detailed prompt", # prompt rewrite - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "aW1hZ2U=" - mock_database.get_image_generation_time_estimate.return_value = None - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="a scene")) - - rewrite_call = mock_llama_wrapper.chat_completion_instruct.call_args_list[1] - expected = main_module.IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE.format( - layout="landscape" - ) - assert rewrite_call.kwargs["system_prompt"] == expected - assert rewrite_call.kwargs["user_prompt"] == "a scene" - assert mock_llama_wrapper.image_generation.call_args.kwargs["prompt"] == ( - "a very detailed prompt" - ) - - -def test_verify_image_prompt_passes_unchanged( - mock_llama_wrapper: MagicMock, -) -> None: - """A PASS verdict keeps the original prompt and sends both inputs.""" - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.return_value = "PASS" - result = main_module.verify_image_prompt("a centaur", "a detailed prompt") - assert result == "a detailed prompt" - call = mock_llama_wrapper.chat_completion_instruct.call_args - assert call.kwargs["system_prompt"] == main_module.IMAGE_PROMPT_VERIFY_SYSTEM_PROMPT - assert "a centaur" in call.kwargs["user_prompt"] - assert "a detailed prompt" in call.kwargs["user_prompt"] - - -def test_verify_image_prompt_case_insensitive_pass( - mock_llama_wrapper: MagicMock, -) -> None: - """A lower-case, punctuated 'pass' verdict also keeps the original.""" - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.return_value = "pass." - assert main_module.verify_image_prompt("a centaur", "p") == "p" - - -def test_verify_image_prompt_correction_replaces_prompt( - mock_llama_wrapper: MagicMock, -) -> None: - """A non-PASS response is used as the corrected prompt.""" - import vibe_bot.main as main_module - - corrected = ( - "one fused creature, a human torso joined to a horse's front half, " - "human arms raised, standing in a rocky field" - ) - mock_llama_wrapper.chat_completion_instruct.return_value = corrected - assert main_module.verify_image_prompt("a centaur", "a horse") == corrected - - -def test_verify_image_prompt_empty_falls_back( - mock_llama_wrapper: MagicMock, -) -> None: - """An empty verification response falls back to the drafted prompt.""" - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.return_value = "" - assert main_module.verify_image_prompt("a centaur", "drafted") == "drafted" - - -def test_doodlebob_selects_portrait( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """Test doodlebob picks portrait and passes the portrait size.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "portrait", # layout selection - "a tall portrait of a lighthouse", # prompt rewrite - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "aW1hZ2U=" - mock_database.get_image_generation_time_estimate.return_value = None - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="a tall lighthouse")) - - # Layout selection is the first instruct call, using the layout prompt. - layout_call = mock_llama_wrapper.chat_completion_instruct.call_args_list[0] - assert layout_call.kwargs["system_prompt"] == main_module.IMAGE_LAYOUT_SYSTEM_PROMPT - assert layout_call.kwargs["user_prompt"] == "a tall lighthouse" - - mock_llama_wrapper.image_generation.assert_called_once() - assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1536" - - sent = _sent_texts(mock_ctx) - assert any("shopping for a canvas" in m for m in sent) - assert any("selected portrait" in m for m in sent) - assert any("drone strike" in m for m in sent) - - -def test_doodlebob_selects_landscape( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """Test doodlebob picks landscape and passes the landscape size.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "landscape", - "a wide panoramic coastline", - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "aW1hZ2U=" - mock_database.get_image_generation_time_estimate.return_value = None - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="wide coastline")) - - assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1536x1024" - - -def test_doodlebob_malformed_layout_defaults_square( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """Test a malformed layout response falls back to the square size.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "banana", # malformed layout - "a balanced composition", - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "aW1hZ2U=" - mock_database.get_image_generation_time_estimate.return_value = None - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="a logo")) - - assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1024" - sent = _sent_texts(mock_ctx) - assert any("selected square" in m for m in sent) - - -def test_doodlebob_empty_layout_defaults_square( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """Test an empty layout response (LLM failure) falls back to square.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "", # empty layout response - "a balanced composition", - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "aW1hZ2U=" - mock_database.get_image_generation_time_estimate.return_value = None - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="a logo")) - - assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1024" - - -def test_doodlebob_reports_estimate_and_elapsed( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """Doodlebob posts an ETA from history and the final elapsed time.""" - import asyncio - import re - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "square", - "a test scene", - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "aW1hZ2U=" - mock_database.get_image_generation_time_estimate.return_value = 12.34 - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="a scene")) - - mock_database.record_image_generation_time.assert_called_once() - recorded = mock_database.record_image_generation_time.call_args[0][0] - assert isinstance(recorded, float) - assert recorded >= 0.0 - - sent = _sent_texts(mock_ctx) - assert any("drone strike" in m for m in sent) - assert any("~12 seconds" in m for m in sent) - assert any(re.search(r"generated in \d+\.\d+ seconds", m) for m in sent) - - -def test_doodlebob_no_estimate_without_history( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """No ETA is posted when there is no generation history yet.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "square", - "a test scene", - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "aW1hZ2U=" - mock_database.get_image_generation_time_estimate.return_value = None - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="a scene")) - - sent = _sent_texts(mock_ctx) - assert any("drone strike" in m for m in sent) - assert not any("ETA" in m for m in sent) - mock_database.record_image_generation_time.assert_called_once() - - -def test_doodlebob_failed_generation_not_recorded( - mock_ctx: MagicMock, - mock_llama_wrapper: MagicMock, - mock_base64: MagicMock, - mock_database: MagicMock, -) -> None: - """A failed generation records no time and reports no elapsed seconds.""" - import asyncio - - import vibe_bot.main as main_module - - mock_llama_wrapper.chat_completion_instruct.side_effect = [ - "square", - "a test scene", - "PASS", # prompt verification - ] - mock_llama_wrapper.image_generation.return_value = "" - mock_database.get_image_generation_time_estimate.return_value = 8.0 - - with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES): - asyncio.run(main_module.doodlebob(mock_ctx, message="a scene")) - - mock_database.record_image_generation_time.assert_not_called() - sent = _sent_texts(mock_ctx) - assert any("Failed to generate image" in m for m in sent) - assert not any("generated in" in m for m in sent) diff --git a/vibe_bot/tests/test_prompts.py b/vibe_bot/tests/test_prompts.py new file mode 100644 index 0000000..96a5b1d --- /dev/null +++ b/vibe_bot/tests/test_prompts.py @@ -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 diff --git a/vibe_bot/tests/test_services.py b/vibe_bot/tests/test_services.py new file mode 100644 index 0000000..6906dae --- /dev/null +++ b/vibe_bot/tests/test_services.py @@ -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) diff --git a/vibe_bot/tests/test_textutil.py b/vibe_bot/tests/test_textutil.py new file mode 100644 index 0000000..722b1da --- /dev/null +++ b/vibe_bot/tests/test_textutil.py @@ -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) diff --git a/vibe_bot/tests/test_tts.py b/vibe_bot/tests/test_tts.py index df6d036..6d92824 100644 --- a/vibe_bot/tests/test_tts.py +++ b/vibe_bot/tests/test_tts.py @@ -22,7 +22,7 @@ def test_tts_engine_init(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 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") result = engine.generate_audio("hello world this is a test") - assert isinstance(result, BytesIO) - result.seek(0) - data = result.read() + assert isinstance(result.audio, BytesIO) + assert result.partial is False + assert result.failed_chunks == 0 + result.audio.seek(0) + data = result.audio.read() 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") 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() @@ -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", ) - 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 @@ -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") 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: @@ -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") result = engine.generate_audio("hello world") - result.seek(0) - data = result.read() + result.audio.seek(0) + data = result.audio.read() assert len(data) > 0 # Should be able to seek and read again - result.seek(0) - data2 = result.read() + result.audio.seek(0) + data2 = result.audio.read() assert data == data2 diff --git a/vibe_bot/textutil.py b/vibe_bot/textutil.py new file mode 100644 index 0000000..74db21e --- /dev/null +++ b/vibe_bot/textutil.py @@ -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 diff --git a/vibe_bot/tools.py b/vibe_bot/tools.py index f8f936e..258aab6 100644 --- a/vibe_bot/tools.py +++ b/vibe_bot/tools.py @@ -33,13 +33,15 @@ def _format_member(member: 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, - or to get a roster of people present in the current channel. + Discord has no per-channel membership, so the channel argument only locates + 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: - 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. """ diff --git a/vibe_bot/tts.py b/vibe_bot/tts.py index 3abd5b6..80489a3 100644 --- a/vibe_bot/tts.py +++ b/vibe_bot/tts.py @@ -3,9 +3,12 @@ from __future__ import annotations import logging +from dataclasses import dataclass from io import BytesIO import numpy as np + +# kokoro-tts and soundfile ship no type stubs upstream. import soundfile as sf # type: ignore[import-untyped] from kokoro_tts import ( # type: ignore[import-untyped] Kokoro, @@ -13,14 +16,31 @@ from kokoro_tts import ( # type: ignore[import-untyped] process_chunk_sequential, ) +from vibe_bot.config import TTS_SPEED, TTS_VOICE + logger = logging.getLogger(__name__) -# Default voice settings -DEFAULT_VOICE = "af_sarah" -DEFAULT_SPEED = 1.0 +# Default voice settings (single source of truth: vibe_bot.config). +DEFAULT_VOICE = TTS_VOICE +DEFAULT_SPEED = TTS_SPEED 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: """Text-to-speech engine wrapper around Kokoro TTS.""" @@ -43,13 +63,14 @@ class TTSEngine: voice: str = DEFAULT_VOICE, speed: float = DEFAULT_SPEED, lang: str = DEFAULT_LANG, - ) -> BytesIO: - """Convert text to audio and return as BytesIO (MP3 format).""" + ) -> AudioResult: + """Convert text to audio and return an AudioResult (MP3 in .audio).""" all_samples: list[np.ndarray] = [] sample_rate: int | None = None + failed_chunks = 0 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): try: @@ -60,15 +81,21 @@ class TTSEngine: speed, 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: logger.exception("Error processing chunk %d", i + 1) + failed_chunks += 1 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: msg = "No audio samples generated - text may be invalid or too long" raise ValueError(msg) @@ -85,9 +112,17 @@ class TTSEngine: ) 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", len(combined), sample_rate or 0, ) - return buffer + return AudioResult(audio=buffer, partial=partial, failed_chunks=failed_chunks)