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

everything working again after cleanup

Browse Source
This commit is contained in:
Reese Wells
2026-05-23 23:56:03 -04:00
parent 6ec9fbe85f
commit 87a578f1de
13 changed files with 380 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.vscode/launch.json
Vendored
+9 -4
View File
@@ -1,9 +1,14 @@
{ {
// Use IntelliSense to learn about possible attributes.
// Hover to view descriptions of existing attributes.
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
"version": "0.2.0", "version": "0.2.0",
"configurations": [ "configurations": [
{
"name": "Vibe Bot: Module",
"type": "debugpy",
"request": "launch",
"module": "vibe_bot.main",
"console": "integratedTerminal",
"envFile": "${workspaceFolder}/.env"
},
{ {
"name": "Python Debugger: Current File", "name": "Python Debugger: Current File",
"type": "debugpy", "type": "debugpy",
@@ -13,4 +18,4 @@
"envFile": "${workspaceFolder}/.env" "envFile": "${workspaceFolder}/.env"
} }
] ]
} }
README.md
+171 -117
View File
@@ -1,217 +1,271 @@
# Vibe Discord Bot with RAG Chat History # Vibe Discord Bot with RAG Chat History
A Discord bot that stores long-term chat history using SQLite database with RAG (Retrieval-Augmented Generation) capabilities powered by custom embedding models. A Discord bot that stores long-term chat history using SQLite with RAG (Retrieval-Augmented Generation) capabilities. It supports custom bots with personalities, text-to-speech via Kokoro, image generation, and image editing.
- [Vibe Discord Bot with RAG Chat History](#vibe-discord-bot-with-rag-chat-history) - [Vibe Discord Bot with RAG Chat History](#vibe-discord-bot-with-rag-chat-history)
- [Quick Start - Available Commands](#quick-start---available-commands) - [Available Commands](#available-commands)
- [Pre-built Bots](#pre-built-bots)
- [Custom Bot Management](#custom-bot-management) - [Custom Bot Management](#custom-bot-management)
- [Using Custom Bots](#using-custom-bots) - [Using Custom Bots](#using-custom-bots)
- [Text-to-Speech](#text-to-speech)
- [Image Commands](#image-commands)
- [Bot Conversations](#bot-conversations)
- [Features](#features) - [Features](#features)
- [Setup](#setup) - [Setup](#setup)
- [Prerequisites](#prerequisites) - [Prerequisites](#prerequisites)
- [Environment Variables](#environment-variables) - [Environment Variables](#environment-variables)
- [Installation](#installation) - [Installation](#installation)
- [Running the Bot](#running-the-bot)
- [How It Works](#how-it-works) - [How It Works](#how-it-works)
- [Database Structure](#database-structure) - [Database Structure](#database-structure)
- [RAG Process](#rag-process) - [RAG Process](#rag-process)
- [Configuration Options](#configuration-options)
- [Usage](#usage)
- [File Structure](#file-structure) - [File Structure](#file-structure)
- [Build](#build) - [Building](#building)
- [Using uv](#using-uv) - [Local](#local)
- [Container](#container) - [Container](#container)
- [Docs](#docs) - [Testing](#testing)
- [Open AI](#open-ai) - [Configuration](#configuration)
- [Models](#models)
- [Qwen3.5](#qwen35)
## Available Commands
## Quick Start - Available Commands
### Pre-built Bots
| Command | Description | Example Usage |
| ------------ | ----------------------------- | ------------------------------------------ |
| `!doodlebob` | Generate images from text | `!doodlebob a cat sitting on a moon` |
| `!retcon` | Edit images with text prompts | `!retcon <image attachment> Make it sunny` |
### Custom Bot Management ### Custom Bot Management
| Command | Description | Example Usage | | Command | Description | Example Usage |
| ------------------------------ | --------------------------------------------- | ------------------------------------------------ | | ---------------------------------- | -------------------------------------- | ---------------------------------------------------- |
| `!custom <name> <personality>` | Create a custom bot with specific personality | `!custom alfred you are a proper british butler` | | `!custom-bot <name> <personality>` | Create a custom bot with a personality | `!custom-bot alfred you are a proper british butler` |
| `!list-custom-bots` | List all available custom bots | `!list-custom-bots` | | `!list-custom-bots` | List all available custom bots | `!list-custom-bots` |
| `!delete-custom-bot <name>` | Delete your custom bot | `!delete-custom-bot alfred` | | `!delete-custom-bot <name>` | Delete your custom bot (owner only) | `!delete-custom-bot alfred` |
### Using Custom Bots ### Using Custom Bots
Once you create a custom bot, you can interact with it directly by prefixing your message with the bot name: Once you create a custom bot, interact with it by prefixing your message with the bot name:
```bash ```text
!<bot_name> <your message> !<bot_name> <your message>
``` ```
**Example:** **Example:**
1. Create a bot: `!custom alfred you are a proper british butler` 1. Create a bot: `!custom-bot alfred you are a proper british butler`
2. Use the bot: `alfred Could you fetch me some tea?` 2. Use the bot: `alfred Could you fetch me some tea?`
3. The bot will respond in character as a British butler 3. The bot will respond in character as a British butler
### Text-to-Speech
| Command | Description | Example Usage |
| -------------------------- | --------------------------------------- | ------------------------------- |
| `!speak <text>` | Convert text to speech (MP3 attachment) | `!speak hello world` |
| `!speak <bot_name> <text>` | Have a custom bot respond and speak | `!speak alfred what time is it` |
### Image Commands
| 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 <image attachment> Make it sunny` |
### Bot Conversations
| Command | Description | Example Usage |
| -------------------------------------- | ------------------------------------------- | ------------------------------------------------ |
| `!talkforme <bot1> <bot2> <n> <topic>` | Have two bots discuss a topic for n replies | `!talkforme alfred jarvis 4 the meaning of life` |
## Features ## Features
- **Long-term chat history storage**: Persistent storage of all bot interactions - **Long-term chat history storage**: Persistent storage of all bot interactions in SQLite
- **RAG-based context retrieval**: Smart retrieval of relevant conversation history using vector embeddings - **RAG-based context retrieval**: Smart retrieval of relevant conversation history using vector embeddings
- **Custom embedding model**: Uses qwen3-embed-4b for semantic search capabilities - **Custom bots**: Create unlimited bots with unique personalities
- **Efficient message management**: Automatic cleanup of old messages based on configurable limits - **Text-to-speech**: Kokoro TTS engine converts bot responses to MP3 audio
- **Image generation**: Generate images from text prompts via OpenAI-compatible API
- **Long-term chat history storage**: Persistent storage of all bot interactions - **Image editing**: Edit uploaded images with text instructions
- **RAG-based context retrieval**: Smart retrieval of relevant conversation history using vector embeddings - **Bot conversations**: Two custom bots can discuss a topic autonomously
- **Custom embedding model**: Uses qwen3-embed-4b for semantic search capabilities - **Automatic message cleanup**: Configurable limits on stored messages
- **Efficient message management**: Automatic cleanup of old messages based on configurable limits
## Setup ## Setup
### Prerequisites ### Prerequisites
- Python 3.10 or higher - Python 3.13 or higher
- [uv](https://docs.astral.sh/uv/) package manager - [uv](https://docs.astral.sh/uv/) package manager
- Embedding API key
- Discord bot token - Discord bot token
- OpenAI-compatible API endpoints (for chat, embeddings, and image generation)
### Environment Variables ### Environment Variables
Create a `.env` file or export the following variables: Create a `.env` file with the following variables:
```bash ```bash
# Discord Bot Token # Discord Bot Token (required)
export DISCORD_TOKEN=your_discord_bot_token DISCORD_TOKEN=your_discord_bot_token
# Embedding API Configuration # Chat/Completion API (required)
export OPENAI_API_KEY=your_embedding_api_key CHAT_ENDPOINT=https://your-api.com/v1
export OPENAI_API_ENDPOINT=https://llama-embed.reeselink.com/embedding 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 (optional) # Image Generation (required)
export IMAGE_GEN_ENDPOINT=http://toybox.reeselink.com:1234/v1 IMAGE_GEN_ENDPOINT=https://your-api.com/v1
export IMAGE_EDIT_ENDPOINT=http://toybox.reeselink.com:1235/v1 IMAGE_EDIT_ENDPOINT=https://your-api.com/v1
IMAGE_GEN_ENDPOINT_KEY=your_api_key
IMAGE_EDIT_ENDPOINT_KEY=your_api_key
IMAGE_GEN_MODEL=gen
IMAGE_EDIT_MODEL=edit
# Database Configuration (optional) # Embedding API (required)
export CHAT_DB_PATH=chat_history.db EMBEDDING_ENDPOINT=https://your-api.com/v1
export EMBEDDING_MODEL=qwen3-embed-4b EMBEDDING_ENDPOINT_KEY=your_api_key
export EMBEDDING_DIMENSION=2048 EMBEDDING_MODEL=your_embed_model
export MAX_HISTORY_MESSAGES=1000
export SIMILARITY_THRESHOLD=0.7 # Optional: TTS Configuration
export TOP_K_RESULTS=5 TTS_MODEL_PATH=kokoro-v1.0.onnx
TTS_VOICES_PATH=voices-v1.0.bin
TTS_VOICE=af_sarah
TTS_SPEED=1.0
# Optional: Database/Chat Settings
DB_PATH=chat_history.db
MAX_COMPLETION_TOKENS=1000
MAX_HISTORY_MESSAGES=1000
SIMILARITY_THRESHOLD=0.7
TOP_K_RESULTS=5
``` ```
### Installation ### Installation
1. Sync dependencies with uv: 1. Clone the repository and sync dependencies:
```bash
uv sync ```bash
``` uv sync
```
2. Ensure the TTS model files are present in the project root:
- `kokoro-v1.0.onnx`
- `voices-v1.0.bin`
### Running the Bot
2. Run the bot:
```bash ```bash
uv run main.py uv run python -m vibe_bot.main
``` ```
## How It Works ## How It Works
### Database Structure ### Database Structure
The system uses two SQLite tables: The system uses SQLite with three tables:
1. **chat_messages**: Stores message metadata 1. **chat_messages**: Stores message metadata
- message_id, user_id, username, content, timestamp, channel_id, guild_id - `message_id`, `user_id`, `username`, `content`, `timestamp`, `channel_id`, `guild_id`
2. **message_embeddings**: Stores vector embeddings for RAG 2. **message_embeddings**: Stores vector embeddings for RAG
- message_id, embedding (as binary blob) - `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`
### RAG Process ### RAG Process
1. When a message is received, it's stored in the database 1. When a message is sent to a custom bot, it's stored in `chat_messages`
2. An embedding is generated using OpenAI's embedding API 2. An embedding is generated via the configured embedding API and stored in `message_embeddings`
3. The embedding is stored alongside the message 3. When a new message is sent:
4. When a new message is sent to the bot: - The system retrieves recent messages from the same user
- The system searches for similar messages using vector similarity - It searches for semantically similar messages using cosine similarity on embeddings
- Relevant context is retrieved and added to the prompt - Relevant context (user + bot message pairs) is prepended to the prompt
- The LLM generates a response with awareness of past conversations - The LLM generates a response with awareness of past conversations
### Configuration Options
- **MAX_HISTORY_MESSAGES**: Maximum number of messages to keep (default: 1000)
- **SIMILARITY_THRESHOLD**: Minimum similarity score for context retrieval (default: 0.7)
- **TOP_K_RESULTS**: Number of similar messages to retrieve (default: 5)
- **EMBEDDING_MODEL**: OpenAI embedding model to use (default: text-embedding-3-small)
## Usage
The bot maintains conversation context automatically. When you ask a question, it will:
1. Search for similar past conversations
2. Include relevant context in the prompt
3. Generate responses that are aware of the conversation history
## File Structure ## File Structure
```text ```text
vibe_discord_bots/ vibe_discord_bots/
├── main.py # Main bot application ├── vibe_bot/
├── database.py # SQLite database with RAG support │ ├── __init__.py # Package marker
├── pyproject.toml # Project dependencies (uv) │ ├── main.py # Main bot application (commands, event handlers)
├── .env # Environment variables │ ├── config.py # Environment variable loading and validation
├── .venv/ # Virtual environment (created by uv) │ ├── database.py # SQLite database with RAG + CustomBotManager
└── README.md # This file │ ├── llama_wrapper.py # OpenAI-compatible API wrappers (chat, images, embeddings)
│ ├── 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
``` ```
## Build ## Building
### Using uv ### Local
```bash ```bash
# Set environment variables # Sync dependencies
export DISCORD_TOKEN=$(cat .token) uv sync
export OPENAI_API_KEY=your_api_key
export OPENAI_API_ENDPOINT="https://llama-cpp.reeselink.com"
export IMAGE_GEN_ENDPOINT="http://toybox.reeselink.com:1234/v1"
export IMAGE_EDIT_ENDPOINT="http://toybox.reeselink.com:1235/v1"
# Run with uv # Run the bot
uv run main.py uv run python -m vibe_bot.main
``` ```
### Container ### Container
```bash ```bash
# Build # Build the container image
podman build -t vibe-bot:latest . podman build -t vibe-bot:latest .
# Run # Run with environment file
podman run --env-file .env localhost/vibe-bot:latest podman run --env-file .env localhost/vibe-bot:latest
``` ```
## Docs ## Testing
### Open AI Run the full test suite:
Chat ```bash
uv run pytest vibe_bot/tests/ -v
```
<https://developers.openai.com/api/reference/resources/chat/subresources/completions/methods/create> Run linters:
Images ```bash
# Ruff (linter + formatter)
uv run ruff check vibe_bot/
<https://developers.openai.com/api/reference/python/resources/images/methods/edit> # Mypy (type checking)
uv run mypy vibe_bot/
## Models # Pyright (type checking)
uv run pyright vibe_bot/
### Qwen3.5 # Black (formatter check)
uv run black --check vibe_bot/
```
> We recommend using the following set of sampling parameters for generation ## Configuration
- Non-thinking mode for text tasks: temperature=1.0, top_p=1.00, top_k=20, min_p=0.0, presence_penalty=2.0, repetition_penalty=1.0 | Variable | Default | Description |
- Non-thinking mode for VL tasks: temperature=0.7, top_p=0.80, top_k=20, min_p=0.0, presence_penalty=1.5, repetition_penalty=1.0 | ----------------------- | ------------------ | ------------------------------------- |
- Thinking mode for text tasks: temperature=1.0, top_p=0.95, top_k=20, min_p=0.0, presence_penalty=1.5, repetition_penalty=1.0 | `DISCORD_TOKEN` | *(required)* | Discord bot authentication token |
- Thinking mode for VL or precise coding (e.g. WebDev) tasks : temperature=0.6, top_p=0.95, top_k=20, min_p=0.0, presence_penalty=0.0, repetition_penalty=1.0 | `CHAT_ENDPOINT` | *(required)* | OpenAI-compatible chat API URL |
| `CHAT_MODEL` | *(required)* | Model name for chat completions |
> Please note that the support for sampling parameters varies according to inference frameworks. | `IMAGE_GEN_ENDPOINT` | *(required)* | Image generation API URL |
| `IMAGE_EDIT_ENDPOINT` | *(required)* | Image editing API URL |
| `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 |
pyproject.toml
+3
View File
@@ -23,6 +23,7 @@ dev = [
"pyright>=1.1.398", "pyright>=1.1.398",
"mypy>=1.17.0", "mypy>=1.17.0",
"black>=25.1.0", "black>=25.1.0",
"debugpy>=1.8.0",
] ]
[tool.uv] [tool.uv]
@@ -56,6 +57,8 @@ ignore = [
"EM101", # string literals in exceptions "EM101", # string literals in exceptions
"TC003", # stdlib import in type-checking block "TC003", # stdlib import in type-checking block
"F401", # unused imports (bytesio used in isinstance) "F401", # unused imports (bytesio used in isinstance)
"SIM117", # nested with statements (needed for test fixtures)
"PT018", # assertion breakdown (test-specific patterns)
] ]
[tool.ruff.lint.isort] [tool.ruff.lint.isort]
uv.lock
Generated
+19
View File
@@ -434,6 +434,23 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/3a/88/9713d1ecac111742d60e1d9c2c15fec56fd121940f97a73d014dc9a7d521/csvw-4.0.0-py2.py3-none-any.whl", hash = "sha256:df875fcb1505afd15061b5f370268522bf162640de0662a724453dcb4db6a88b", size = 69424, upload-time = "2026-05-05T06:25:24.646Z" }, { url = "https://files.pythonhosted.org/packages/3a/88/9713d1ecac111742d60e1d9c2c15fec56fd121940f97a73d014dc9a7d521/csvw-4.0.0-py2.py3-none-any.whl", hash = "sha256:df875fcb1505afd15061b5f370268522bf162640de0662a724453dcb4db6a88b", size = 69424, upload-time = "2026-05-05T06:25:24.646Z" },
] ]
[[package]]
name = "debugpy"
version = "1.8.20"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e0/b7/cd8080344452e4874aae67c40d8940e2b4d47b01601a8fd9f44786c757c7/debugpy-1.8.20.tar.gz", hash = "sha256:55bc8701714969f1ab89a6d5f2f3d40c36f91b2cbe2f65d98bf8196f6a6a2c33", size = 1645207, upload-time = "2026-01-29T23:03:28.199Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/15/e2/fc500524cc6f104a9d049abc85a0a8b3f0d14c0a39b9c140511c61e5b40b/debugpy-1.8.20-cp313-cp313-macosx_15_0_universal2.whl", hash = "sha256:5dff4bb27027821fdfcc9e8f87309a28988231165147c31730128b1c983e282a", size = 2539560, upload-time = "2026-01-29T23:03:48.738Z" },
{ url = "https://files.pythonhosted.org/packages/90/83/fb33dcea789ed6018f8da20c5a9bc9d82adc65c0c990faed43f7c955da46/debugpy-1.8.20-cp313-cp313-manylinux_2_34_x86_64.whl", hash = "sha256:84562982dd7cf5ebebfdea667ca20a064e096099997b175fe204e86817f64eaf", size = 4293272, upload-time = "2026-01-29T23:03:50.169Z" },
{ url = "https://files.pythonhosted.org/packages/a6/25/b1e4a01bfb824d79a6af24b99ef291e24189080c93576dfd9b1a2815cd0f/debugpy-1.8.20-cp313-cp313-win32.whl", hash = "sha256:da11dea6447b2cadbf8ce2bec59ecea87cc18d2c574980f643f2d2dfe4862393", size = 5331208, upload-time = "2026-01-29T23:03:51.547Z" },
{ url = "https://files.pythonhosted.org/packages/13/f7/a0b368ce54ffff9e9028c098bd2d28cfc5b54f9f6c186929083d4c60ba58/debugpy-1.8.20-cp313-cp313-win_amd64.whl", hash = "sha256:eb506e45943cab2efb7c6eafdd65b842f3ae779f020c82221f55aca9de135ed7", size = 5372930, upload-time = "2026-01-29T23:03:53.585Z" },
{ url = "https://files.pythonhosted.org/packages/33/2e/f6cb9a8a13f5058f0a20fe09711a7b726232cd5a78c6a7c05b2ec726cff9/debugpy-1.8.20-cp314-cp314-macosx_15_0_universal2.whl", hash = "sha256:9c74df62fc064cd5e5eaca1353a3ef5a5d50da5eb8058fcef63106f7bebe6173", size = 2538066, upload-time = "2026-01-29T23:03:54.999Z" },
{ url = "https://files.pythonhosted.org/packages/c5/56/6ddca50b53624e1ca3ce1d1e49ff22db46c47ea5fb4c0cc5c9b90a616364/debugpy-1.8.20-cp314-cp314-manylinux_2_34_x86_64.whl", hash = "sha256:077a7447589ee9bc1ff0cdf443566d0ecf540ac8aa7333b775ebcb8ce9f4ecad", size = 4269425, upload-time = "2026-01-29T23:03:56.518Z" },
{ url = "https://files.pythonhosted.org/packages/c5/d9/d64199c14a0d4c476df46c82470a3ce45c8d183a6796cfb5e66533b3663c/debugpy-1.8.20-cp314-cp314-win32.whl", hash = "sha256:352036a99dd35053b37b7803f748efc456076f929c6a895556932eaf2d23b07f", size = 5331407, upload-time = "2026-01-29T23:03:58.481Z" },
{ url = "https://files.pythonhosted.org/packages/e0/d9/1f07395b54413432624d61524dfd98c1a7c7827d2abfdb8829ac92638205/debugpy-1.8.20-cp314-cp314-win_amd64.whl", hash = "sha256:a98eec61135465b062846112e5ecf2eebb855305acc1dfbae43b72903b8ab5be", size = 5372521, upload-time = "2026-01-29T23:03:59.864Z" },
{ url = "https://files.pythonhosted.org/packages/e0/c3/7f67dea8ccf8fdcb9c99033bbe3e90b9e7395415843accb81428c441be2d/debugpy-1.8.20-py2.py3-none-any.whl", hash = "sha256:5be9bed9ae3be00665a06acaa48f8329d2b9632f15fd09f6a9a8c8d9907e54d7", size = 5337658, upload-time = "2026-01-29T23:04:17.404Z" },
]
[[package]] [[package]]
name = "decorator" name = "decorator"
version = "5.3.1" version = "5.3.1"
@@ -2215,6 +2232,7 @@ dependencies = [
[package.optional-dependencies] [package.optional-dependencies]
dev = [ dev = [
{ name = "black" }, { name = "black" },
{ name = "debugpy" },
{ name = "mypy" }, { name = "mypy" },
{ name = "pyright" }, { name = "pyright" },
{ name = "ruff" }, { name = "ruff" },
@@ -2223,6 +2241,7 @@ dev = [
[package.metadata] [package.metadata]
requires-dist = [ requires-dist = [
{ name = "black", marker = "extra == 'dev'", specifier = ">=25.1.0" }, { name = "black", marker = "extra == 'dev'", specifier = ">=25.1.0" },
{ name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.8.0" },
{ name = "discord", specifier = ">=2.3.2" }, { name = "discord", specifier = ">=2.3.2" },
{ name = "kokoro-tts", specifier = ">=2.3.1" }, { name = "kokoro-tts", specifier = ">=2.3.1" },
{ name = "mypy", specifier = ">=2.1.0" }, { name = "mypy", specifier = ">=2.1.0" },
vibe_bot/database.py
+5 -4
View File
@@ -330,7 +330,7 @@ class ChatDatabase:
results.sort(key=lambda x: x[2], reverse=True) results.sort(key=lambda x: x[2], reverse=True)
return results[:top_k] return results[:top_k]
def get_user_history(self, _user_id: str, limit: int = 20) -> list[tuple[str, str]]: def get_user_history(self, user_id: str, limit: int = 20) -> list[tuple[str, str]]:
"""Get message history for a specific user.""" """Get message history for a specific user."""
conn = sqlite3.connect(self.db_path) conn = sqlite3.connect(self.db_path)
cursor = conn.cursor() cursor = conn.cursor()
@@ -340,11 +340,11 @@ class ChatDatabase:
""" """
SELECT message_id, content, timestamp SELECT message_id, content, timestamp
FROM chat_messages FROM chat_messages
WHERE username != 'vibe-bot' WHERE user_id = ? AND username != 'vibe-bot'
ORDER BY timestamp DESC ORDER BY timestamp DESC
LIMIT ? LIMIT ?
""", """,
(limit,), (user_id, limit),
) )
messages = cursor.fetchall() messages = cursor.fetchall()
@@ -528,9 +528,10 @@ class CustomBotManager:
""" """
SELECT bot_name, system_prompt, created_by SELECT bot_name, system_prompt, created_by
FROM custom_bots FROM custom_bots
WHERE is_active = 1 WHERE is_active = 1 AND created_by = ?
ORDER BY created_at DESC ORDER BY created_at DESC
""", """,
(user_id,),
) )
else: else:
cursor.execute( cursor.execute(
vibe_bot/llama_wrapper.py
+61 -21
View File
@@ -6,9 +6,11 @@ Allows custom endpoints for each of the above supported functions.
from __future__ import annotations from __future__ import annotations
import json
from typing import TYPE_CHECKING, cast from typing import TYPE_CHECKING, cast
import openai import openai
import requests
if TYPE_CHECKING: if TYPE_CHECKING:
from io import BufferedReader, BytesIO from io import BufferedReader, BytesIO
@@ -54,8 +56,12 @@ def chat_completion(
model=model, model=model,
messages=messages, messages=messages,
max_tokens=max_tokens, max_tokens=max_tokens,
timeout=60.0,
) )
if not response.choices:
return ""
content = response.choices[0].message.content content = response.choices[0].message.content
if content: if content:
return content.strip() return content.strip()
@@ -101,8 +107,12 @@ def chat_completion_with_history(
messages=messages, messages=messages,
max_tokens=max_tokens, max_tokens=max_tokens,
seed=-1, seed=-1,
timeout=60.0,
) )
if not response.choices:
return ""
content = response.choices[0].message.content content = response.choices[0].message.content
if content: if content:
return content.strip() return content.strip()
@@ -148,8 +158,12 @@ def chat_completion_instruct(
messages=messages, messages=messages,
max_tokens=max_tokens, max_tokens=max_tokens,
seed=-1, seed=-1,
timeout=60.0,
) )
if not response.choices:
return ""
content = response.choices[0].message.content content = response.choices[0].message.content
if content: if content:
return content.strip() return content.strip()
@@ -158,8 +172,10 @@ def chat_completion_instruct(
def image_generation( def image_generation(
prompt: str, prompt: str,
*,
openai_url: str, openai_url: str,
openai_api_key: str, openai_api_key: str,
model: str = "gen",
n: int = 1, n: int = 1,
) -> str: ) -> str:
"""Generate an image using the given prompt. """Generate an image using the given prompt.
@@ -168,19 +184,28 @@ def image_generation(
prompt: The image generation prompt. prompt: The image generation prompt.
openai_url: The OpenAI-compatible API URL. openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication. openai_api_key: The API key for authentication.
model: The model to use for image generation.
n: Number of images to generate. n: Number of images to generate.
Returns: Returns:
The base64 encoded image data. Decode and write to a file. The base64 encoded image data. Decode and write to a file.
""" """
client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key) client = openai.OpenAI(
response = client.images.generate( base_url=openai_url,
prompt=prompt, api_key=openai_api_key,
n=n, max_retries=0,
size="1024x1024",
model="gen",
) )
try:
response = client.images.generate(
prompt=prompt,
n=n,
size="1024x1024",
model=model,
timeout=120.0,
)
except openai.APIConnectionError:
return ""
if response.data: if response.data:
return response.data[0].b64_json or "" return response.data[0].b64_json or ""
return "" return ""
@@ -189,8 +214,10 @@ def image_generation(
def image_edit( def image_edit(
image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO], image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO],
prompt: str, prompt: str,
*,
openai_url: str, openai_url: str,
openai_api_key: str, openai_api_key: str,
model: str = "edit",
n: int = 1, n: int = 1,
) -> str: ) -> str:
"""Edit an existing image using a prompt. """Edit an existing image using a prompt.
@@ -200,6 +227,7 @@ def image_edit(
prompt: The edit instruction. prompt: The edit instruction.
openai_url: The OpenAI-compatible API URL. openai_url: The OpenAI-compatible API URL.
openai_api_key: The API key for authentication. openai_api_key: The API key for authentication.
model: The model to use for image editing.
n: Number of edited images to generate. n: Number of edited images to generate.
Returns: Returns:
@@ -212,7 +240,7 @@ def image_edit(
prompt=prompt, prompt=prompt,
n=n, n=n,
size="1024x1024", size="1024x1024",
model="edit", model=model,
) )
if response.data: if response.data:
return response.data[0].b64_json or "" return response.data[0].b64_json or ""
@@ -228,6 +256,9 @@ def embedding(
) -> list[float]: ) -> list[float]:
"""Generate an embedding vector for the given text. """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: Args:
text: The text to embed. text: The text to embed.
openai_url: The OpenAI-compatible API URL. openai_url: The OpenAI-compatible API URL.
@@ -238,17 +269,26 @@ def embedding(
The embedding vector as a list of floats, or an empty list on failure. The embedding vector as a list of floats, or an empty list on failure.
""" """
client = openai.OpenAI(base_url=openai_url, api_key=openai_api_key) url = f"{openai_url.rstrip('/')}/embeddings"
response = client.embeddings.create( headers = {
input=[text], "Authorization": f"Bearer {openai_api_key}",
model=model, "Content-Type": "application/json",
encoding_format="float", }
) payload = {"model": model, "input": [text]}
if response:
data = response.data try:
raw_data = data[0].embedding resp = requests.post(url, headers=headers, json=payload, timeout=30)
# The result could be an array of floats or a single float. resp.raise_for_status()
if not isinstance(raw_data, float): except requests.RequestException:
return list(raw_data) return []
return [raw_data]
return [] data = resp.json()
if not data.get("data"):
return []
raw = data["data"][0].get("embedding")
if isinstance(raw, str):
raw = json.loads(raw)
if not isinstance(raw, list):
raw = list(raw)
return raw
vibe_bot/main.py
+23 -9
View File
@@ -20,6 +20,10 @@ from vibe_bot.config import (
DISCORD_TOKEN, DISCORD_TOKEN,
IMAGE_EDIT_ENDPOINT, IMAGE_EDIT_ENDPOINT,
IMAGE_EDIT_ENDPOINT_KEY, IMAGE_EDIT_ENDPOINT_KEY,
IMAGE_EDIT_MODEL,
IMAGE_GEN_ENDPOINT,
IMAGE_GEN_ENDPOINT_KEY,
IMAGE_GEN_MODEL,
MAX_COMPLETION_TOKENS, MAX_COMPLETION_TOKENS,
TTS_MODEL_PATH, TTS_MODEL_PATH,
TTS_SPEED, TTS_SPEED,
@@ -415,7 +419,7 @@ async def _speak_with_bot(
message_id=f"{ctx.message.id}_response", message_id=f"{ctx.message.id}_response",
user_id=str(ctx.bot.user.id), user_id=str(ctx.bot.user.id),
username=ctx.bot.user.name, username=ctx.bot.user.name,
content=f"Bot: {bot_response}", content=bot_response,
channel_id=str(ctx.channel.id), channel_id=str(ctx.channel.id),
guild_id=str(ctx.guild.id) if ctx.guild else None, guild_id=str(ctx.guild.id) if ctx.guild else None,
) )
@@ -497,14 +501,23 @@ async def doodlebob(ctx: CommandsContext[Bot], *, message: str) -> None:
image_b64 = llama_wrapper.image_generation( image_b64 = llama_wrapper.image_generation(
prompt=image_prompt, prompt=image_prompt,
openai_url=IMAGE_EDIT_ENDPOINT, openai_url=IMAGE_GEN_ENDPOINT,
openai_api_key=IMAGE_EDIT_ENDPOINT_KEY, openai_api_key=IMAGE_GEN_ENDPOINT_KEY,
model=IMAGE_GEN_MODEL,
) )
# Save the image to a file if not image_b64:
edited_image_data = BytesIO(base64.b64decode(image_b64)) logger.warning("Image generation returned empty response.")
send_img = discord.File(edited_image_data, filename="image.png") await ctx.send("Failed to generate image. The server may be busy.")
await ctx.send(file=send_img) return
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)
except Exception:
logger.exception("Failed to decode image data")
await ctx.send("Failed to process the generated image.")
@bot.command(name="retcon") @bot.command(name="retcon")
@@ -529,6 +542,7 @@ async def retcon(ctx: CommandsContext[Bot], *, message: str) -> None:
prompt=message, prompt=message,
openai_url=IMAGE_EDIT_ENDPOINT, openai_url=IMAGE_EDIT_ENDPOINT,
openai_api_key=IMAGE_EDIT_ENDPOINT_KEY, openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,
model=IMAGE_EDIT_MODEL,
) )
# Save the image to a file # Save the image to a file
@@ -621,7 +635,7 @@ async def talkforme(ctx: CommandsContext[Bot], *, message: str) -> None:
bot_response = llama_wrapper.chat_completion_with_history( bot_response = llama_wrapper.chat_completion_with_history(
system_prompt=( system_prompt=(
current_bot[1] + f"\nKeep your responses under 2-3 sentences. " current_bot[1] + f"\nKeep your responses under 2-3 sentences. "
f"{current_bot[flip_counter(bot_counter)]}" f"You are talking to {current_bot[flip_counter(bot_counter)][0]}"
), ),
prompts=prompt_histories[bot_counter], prompts=prompt_histories[bot_counter],
openai_url=CHAT_ENDPOINT, openai_url=CHAT_ENDPOINT,
@@ -709,7 +723,7 @@ async def handle_chat(
message_id=f"{ctx.message.id}_response", message_id=f"{ctx.message.id}_response",
user_id=str(ctx.bot.user.id), user_id=str(ctx.bot.user.id),
username=ctx.bot.user.name, username=ctx.bot.user.name,
content=f"Bot: {bot_response}", content=bot_response,
channel_id=str(ctx.channel.id), channel_id=str(ctx.channel.id),
guild_id=str(ctx.guild.id) if ctx.guild else None, guild_id=str(ctx.guild.id) if ctx.guild else None,
) )
vibe_bot/tests/conftest.py
+18 -13
View File
@@ -117,17 +117,22 @@ def mock_kokoro_tts() -> Generator[dict[str, Any]]:
mock_samples = np.array([0.1, 0.2, 0.3], dtype=np.float32) mock_samples = np.array([0.1, 0.2, 0.3], dtype=np.float32)
mock_process = MagicMock(return_value=(mock_samples, 24000)) mock_process = MagicMock(return_value=(mock_samples, 24000))
with patch("vibe_bot.tts.Kokoro", return_value=mock_kokoro_instance): # noqa: SIM117 with (
with patch("vibe_bot.tts.chunk_text", mock_chunk): patch(
with patch("vibe_bot.tts.process_chunk_sequential", mock_process): "vibe_bot.tts.Kokoro",
yield { return_value=mock_kokoro_instance,
"Kokoro": mock_kokoro, ),
"chunk_text": mock_chunk, patch("vibe_bot.tts.chunk_text", mock_chunk),
"process_chunk_sequential": mock_process, ):
"kokoro_instance": mock_kokoro_instance, with patch("vibe_bot.tts.process_chunk_sequential", mock_process):
"mock_samples": mock_samples, yield {
"mock_sr": 24000, "Kokoro": mock_kokoro,
} "chunk_text": mock_chunk,
"process_chunk_sequential": mock_process,
"kokoro_instance": mock_kokoro_instance,
"mock_samples": mock_samples,
"mock_sr": 24000,
}
@pytest.fixture @pytest.fixture
@@ -143,7 +148,7 @@ def mock_discord() -> Generator[dict[str, MagicMock]]:
mock_bot_instance.user.name = "test-bot" mock_bot_instance.user.name = "test-bot"
mock_bot_instance.user.id = "123456789" mock_bot_instance.user.id = "123456789"
with patch("vibe_bot.main.discord") as mock_discord_module: # noqa: SIM117 with patch("vibe_bot.main.discord") as mock_discord_module:
with patch("vibe_bot.main.commands", MagicMock()): with patch("vibe_bot.main.commands", MagicMock()):
with patch("vibe_bot.main.commands.Bot", mock_bot_class): with patch("vibe_bot.main.commands.Bot", mock_bot_class):
mock_bot_class.return_value = mock_bot_instance mock_bot_class.return_value = mock_bot_instance
@@ -162,7 +167,7 @@ def mock_tts_engine() -> Generator[MagicMock]:
"""Provide a mock TTSEngine.""" """Provide a mock TTSEngine."""
mock_engine = MagicMock() mock_engine = MagicMock()
mock_engine.generate_audio.return_value = MagicMock() mock_engine.generate_audio.return_value = MagicMock()
with patch("vibe_bot.main.tts_engine", mock_engine): # noqa: SIM117 with patch("vibe_bot.main.tts_engine", mock_engine):
with patch("vibe_bot.main.tts.TTSEngine", return_value=mock_engine): with patch("vibe_bot.main.tts.TTSEngine", return_value=mock_engine):
yield mock_engine yield mock_engine
vibe_bot/tests/test_config.py
+3 -3
View File
@@ -106,9 +106,9 @@ except Exception as e:
timeout=30, timeout=30,
) )
output = result.stdout.strip() output = result.stdout.strip()
assert output.startswith("ERROR:") and expected_error in output, ( # noqa: PT018 assert (
f"Expected error '{expected_error}' but got: {output}" output.startswith("ERROR:") and expected_error in output
) ), f"Expected error '{expected_error}' but got: {output}"
def test_config_missing_discord_token() -> None: def test_config_missing_discord_token() -> None:
vibe_bot/tests/test_database.py
+29 -7
View File
@@ -129,13 +129,22 @@ def test_get_recent_messages(
) -> None: ) -> None:
"""Test retrieving recent messages.""" """Test retrieving recent messages."""
chat_db.add_message( chat_db.add_message(
message_id="msg-1", user_id="u1", username="alice", content="First", message_id="msg-1",
user_id="u1",
username="alice",
content="First",
) )
chat_db.add_message( chat_db.add_message(
message_id="msg-2", user_id="u2", username="bob", content="Second", message_id="msg-2",
user_id="u2",
username="bob",
content="Second",
) )
chat_db.add_message( chat_db.add_message(
message_id="msg-3", user_id="u1", username="alice", content="Third", message_id="msg-3",
user_id="u1",
username="alice",
content="Third",
) )
messages = chat_db.get_recent_messages(limit=2) messages = chat_db.get_recent_messages(limit=2)
@@ -167,10 +176,16 @@ def test_clear_all_messages(
) -> None: ) -> None:
"""Test clearing all messages.""" """Test clearing all messages."""
chat_db.add_message( chat_db.add_message(
message_id="msg-1", user_id="u1", username="alice", content="Hello", message_id="msg-1",
user_id="u1",
username="alice",
content="Hello",
) )
chat_db.add_message( chat_db.add_message(
message_id="msg-2", user_id="u2", username="bob", content="World", message_id="msg-2",
user_id="u2",
username="bob",
content="World",
) )
chat_db.clear_all_messages() chat_db.clear_all_messages()
@@ -185,7 +200,10 @@ def test_get_user_history(
) -> None: ) -> None:
"""Test retrieving user message history.""" """Test retrieving user message history."""
chat_db.add_message( chat_db.add_message(
message_id="msg-1", user_id="u1", username="alice", content="User question", message_id="msg-1",
user_id="u1",
username="alice",
content="User question",
) )
chat_db.add_message( chat_db.add_message(
message_id="msg-1_response", message_id="msg-1_response",
@@ -422,7 +440,9 @@ def test_custom_bot_delete_with_error(
) -> None: ) -> None:
"""Test that delete_custom_bot returns False on error.""" """Test that delete_custom_bot returns False on error."""
with patch.object( with patch.object(
custom_bot_manager, "_initialize_custom_bots_table", side_effect=Exception("db error"), # noqa: E501 custom_bot_manager,
"_initialize_custom_bots_table",
side_effect=Exception("db error"),
): ):
pass pass
result = custom_bot_manager.delete_custom_bot("nonexistent") result = custom_bot_manager.delete_custom_bot("nonexistent")
@@ -433,6 +453,7 @@ def test_database_get_database_singleton(temp_db_path: str) -> None:
"""Test that get_database returns the same instance.""" """Test that get_database returns the same instance."""
import vibe_bot.database as db_module import vibe_bot.database as db_module
from vibe_bot.database import ChatDatabase, get_database from vibe_bot.database import ChatDatabase, get_database
db_module._chat_db = None db_module._chat_db = None
db1 = get_database() db1 = get_database()
@@ -453,6 +474,7 @@ def test_database_init_creates_tables(temp_db_path: str) -> None:
db.client.close() db.client.close()
import sqlite3 import sqlite3
conn = sqlite3.connect(temp_db_path) conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute("SELECT name FROM sqlite_master WHERE type='table'") cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
vibe_bot/tests/test_llama_wrapper.py
+17 -16
View File
@@ -6,6 +6,7 @@ import base64
import tempfile import tempfile
from io import BytesIO from io import BytesIO
from pathlib import Path from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
import numpy as np import numpy as np
@@ -106,24 +107,24 @@ EMBEDDING_SIMILARITY_LOW = 0.5
def test_embeddings() -> None: def test_embeddings() -> None:
"""Test embedding similarity for similar and different texts.""" """Test embedding similarity for similar and different texts."""
with patch("vibe_bot.llama_wrapper.openai.OpenAI") as mock_openai: mock_horse_vec = [0.8] * 1024 + [0.6] * 1024
mock_horse_vec = [0.8] * 1024 + [0.6] * 1024 mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024
mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024 mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024
mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024
mock_response1 = MagicMock() def mock_post(*args: Any, **kwargs: Any) -> MagicMock:
mock_response1.data = [MagicMock(embedding=mock_horse_vec)] json_data = kwargs.get("json", {})
mock_response2 = MagicMock() text = json_data["input"][0]
mock_response2.data = [MagicMock(embedding=mock_horse_also_vec)] if "horse" in text and "donkey" not in text and "also" not in text:
mock_response3 = MagicMock() embedding_data = mock_horse_vec
mock_response3.data = [MagicMock(embedding=mock_donkey_vec)] elif "also" in text:
embedding_data = mock_horse_also_vec
mock_openai.return_value.embeddings.create.side_effect = [ else:
mock_response1, embedding_data = mock_donkey_vec
mock_response2, mock_resp = MagicMock()
mock_response3, 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( result1 = embedding(
"this is a horse", "this is a horse",
openai_url=EMBEDDING_ENDPOINT, openai_url=EMBEDDING_ENDPOINT,
vibe_bot/tests/test_main.py
+9 -3
View File
@@ -125,7 +125,9 @@ def test_custom_bot_command_success(
asyncio.run( asyncio.run(
main_module.custom_bot( main_module.custom_bot(
mock_ctx, bot_name="alfred", personality="you are a british butler", mock_ctx,
bot_name="alfred",
personality="you are a british butler",
), ),
) )
@@ -199,7 +201,9 @@ def test_custom_bot_command_create_fails(
asyncio.run( asyncio.run(
main_module.custom_bot( main_module.custom_bot(
mock_ctx, bot_name="alfred", personality="you are a british butler", mock_ctx,
bot_name="alfred",
personality="you are a british butler",
), ),
) )
call_args = mock_ctx.send.call_args[0][0] call_args = mock_ctx.send.call_args[0][0]
@@ -347,7 +351,9 @@ def test_handle_chat_success(
import vibe_bot.main as main_module import vibe_bot.main as main_module
mock_llama_wrapper.chat_completion_with_history.return_value = "This is a bot response" # noqa: E501 mock_llama_wrapper.chat_completion_with_history.return_value = (
"This is a bot response"
)
asyncio.run( asyncio.run(
main_module.handle_chat( main_module.handle_chat(
vibe_bot/tests/test_tts.py
+13 -3
View File
@@ -63,9 +63,15 @@ def test_generate_audio_multiple_chunks(mock_kokoro_tts: MagicMock) -> None:
from vibe_bot.tts import TTSEngine from vibe_bot.tts import TTSEngine
mock_kokoro_tts["chunk_text"].return_value = ["chunk one", "chunk two", "chunk three"] # noqa: E501 mock_kokoro_tts["chunk_text"].return_value = [
"chunk one",
"chunk two",
"chunk three",
]
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin") engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("this text is long enough to be split into multiple chunks") # noqa: E501 result = engine.generate_audio(
"this text is long enough to be split into multiple chunks",
)
assert isinstance(result, BytesIO) assert isinstance(result, BytesIO)
assert mock_kokoro_tts["process_chunk_sequential"].call_count == 3 assert mock_kokoro_tts["process_chunk_sequential"].call_count == 3
@@ -88,7 +94,11 @@ def test_generate_audio_chunk_failure(mock_kokoro_tts: MagicMock) -> None:
raise Exception("processing error") raise Exception("processing error")
return np.array([0.1, 0.2], dtype=np.float32), 24000 return np.array([0.1, 0.2], dtype=np.float32), 24000
mock_kokoro_tts["chunk_text"].return_value = ["good chunk", "bad chunk", "another good"] # noqa: E501 mock_kokoro_tts["chunk_text"].return_value = [
"good chunk",
"bad chunk",
"another good",
]
mock_kokoro_tts["process_chunk_sequential"].side_effect = process_with_failure mock_kokoro_tts["process_chunk_sequential"].side_effect = process_with_failure
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin") engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")