1277 lines
44 KiB
Python
1277 lines
44 KiB
Python
"""Main Discord bot application."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import logging
|
|
import re
|
|
import time
|
|
from io import BytesIO
|
|
from typing import TYPE_CHECKING
|
|
|
|
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
|
|
|
|
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 <bot_name> <personality_description>
|
|
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} <your message>`",
|
|
)
|
|
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 <name> <personality>` 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 <bot_name>
|
|
"""
|
|
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 <subcommand>
|
|
Available sub-commands:
|
|
- members: List all members in the current channel
|
|
- whoami: Show all information the bot has about you
|
|
- tools: Show the LLM's available tools
|
|
"""
|
|
logger.info(
|
|
"Debug command triggered by %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 <text> --voice <voice_name>` 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 <text> --voice <voice_name> - plain text to speech
|
|
Usage: !speak <bot_name> <text> --voice <voice_name> - have a custom bot respond and speak
|
|
Example: !speak hello world
|
|
Example: !speak hello world --voice af_bella
|
|
Example: !speak alfred what time is it --voice am_puck
|
|
"""
|
|
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 <bot_name>
|
|
"""
|
|
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 <number> <topic>")
|
|
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)
|