191 lines
6.4 KiB
Python
191 lines
6.4 KiB
Python
"""AI Chat cog - handles chat commands and mention responses."""
|
|
|
|
import logging
|
|
import re
|
|
|
|
import discord
|
|
from discord import app_commands
|
|
from discord.ext import commands
|
|
|
|
from daemon_boyfriend.config import settings
|
|
from daemon_boyfriend.services import AIService, ConversationManager, Message
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Discord message character limit
|
|
MAX_MESSAGE_LENGTH = 2000
|
|
|
|
|
|
def split_message(content: str, max_length: int = MAX_MESSAGE_LENGTH) -> list[str]:
|
|
"""Split a long message into chunks that fit Discord's limit.
|
|
|
|
Tries to split on paragraph breaks, then sentence breaks, then word breaks.
|
|
"""
|
|
if len(content) <= max_length:
|
|
return [content]
|
|
|
|
chunks: list[str] = []
|
|
remaining = content
|
|
|
|
while remaining:
|
|
if len(remaining) <= max_length:
|
|
chunks.append(remaining)
|
|
break
|
|
|
|
# Find a good split point
|
|
split_point = max_length
|
|
|
|
# Try to split on paragraph break
|
|
para_break = remaining.rfind("\n\n", 0, max_length)
|
|
if para_break > max_length // 2:
|
|
split_point = para_break + 2
|
|
else:
|
|
# Try to split on line break
|
|
line_break = remaining.rfind("\n", 0, max_length)
|
|
if line_break > max_length // 2:
|
|
split_point = line_break + 1
|
|
else:
|
|
# Try to split on sentence
|
|
sentence_end = max(
|
|
remaining.rfind(". ", 0, max_length),
|
|
remaining.rfind("! ", 0, max_length),
|
|
remaining.rfind("? ", 0, max_length),
|
|
)
|
|
if sentence_end > max_length // 2:
|
|
split_point = sentence_end + 2
|
|
else:
|
|
# Fall back to word break
|
|
word_break = remaining.rfind(" ", 0, max_length)
|
|
if word_break > 0:
|
|
split_point = word_break + 1
|
|
|
|
chunks.append(remaining[:split_point].rstrip())
|
|
remaining = remaining[split_point:].lstrip()
|
|
|
|
return chunks
|
|
|
|
|
|
class AIChatCog(commands.Cog):
|
|
"""AI conversation commands and mention handling."""
|
|
|
|
def __init__(self, bot: commands.Bot) -> None:
|
|
self.bot = bot
|
|
self.ai_service = AIService()
|
|
self.conversations = ConversationManager()
|
|
|
|
@app_commands.command(name="chat", description="Chat with Daemon Boyfriend")
|
|
@app_commands.describe(message="Your message to the bot")
|
|
async def chat(self, interaction: discord.Interaction, message: str) -> None:
|
|
"""Slash command to chat with the AI."""
|
|
await interaction.response.defer(thinking=True)
|
|
|
|
try:
|
|
response_text = await self._generate_response(interaction.user.id, message)
|
|
|
|
# Split long responses
|
|
chunks = split_message(response_text)
|
|
await interaction.followup.send(chunks[0])
|
|
|
|
# Send additional chunks as follow-up messages
|
|
for chunk in chunks[1:]:
|
|
await interaction.followup.send(chunk)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Chat error: {e}", exc_info=True)
|
|
await interaction.followup.send("Sorry, I encountered an error. Please try again.")
|
|
|
|
@app_commands.command(name="clear", description="Clear your conversation history")
|
|
async def clear_history(self, interaction: discord.Interaction) -> None:
|
|
"""Clear the user's conversation history."""
|
|
self.conversations.clear_history(interaction.user.id)
|
|
await interaction.response.send_message(
|
|
"Your conversation history has been cleared!", ephemeral=True
|
|
)
|
|
|
|
@commands.Cog.listener()
|
|
async def on_message(self, message: discord.Message) -> None:
|
|
"""Respond when the bot is mentioned."""
|
|
# Ignore messages from bots
|
|
if message.author.bot:
|
|
return
|
|
|
|
# Check if bot is mentioned
|
|
if self.bot.user is None or self.bot.user not in message.mentions:
|
|
return
|
|
|
|
# Extract message content without the mention
|
|
content = self._extract_message_content(message)
|
|
|
|
if not content:
|
|
# Just a mention with no message
|
|
await message.reply(f"Hey {message.author.display_name}! How can I help you?")
|
|
return
|
|
|
|
# Show typing indicator while generating response
|
|
async with message.channel.typing():
|
|
try:
|
|
response_text = await self._generate_response(message.author.id, content)
|
|
|
|
# Split and send response
|
|
chunks = split_message(response_text)
|
|
await message.reply(chunks[0])
|
|
|
|
for chunk in chunks[1:]:
|
|
await message.channel.send(chunk)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Mention response error: {e}", exc_info=True)
|
|
await message.reply("Sorry, I encountered an error. Please try again.")
|
|
|
|
def _extract_message_content(self, message: discord.Message) -> str:
|
|
"""Extract the actual message content, removing bot mentions."""
|
|
content = message.content
|
|
|
|
# Remove all mentions of the bot
|
|
if self.bot.user:
|
|
# Remove <@BOT_ID> and <@!BOT_ID> patterns
|
|
content = re.sub(
|
|
rf"<@!?{self.bot.user.id}>",
|
|
"",
|
|
content,
|
|
)
|
|
|
|
return content.strip()
|
|
|
|
async def _generate_response(self, user_id: int, user_message: str) -> str:
|
|
"""Generate an AI response for a user message.
|
|
|
|
Args:
|
|
user_id: Discord user ID
|
|
user_message: The user's message
|
|
|
|
Returns:
|
|
The AI's response text
|
|
"""
|
|
# Get conversation history
|
|
history = self.conversations.get_history(user_id)
|
|
|
|
# Add current message to history for the API call
|
|
messages = history + [Message(role="user", content=user_message)]
|
|
|
|
# Generate response
|
|
response = await self.ai_service.chat(
|
|
messages=messages,
|
|
system_prompt=self.ai_service.get_system_prompt(),
|
|
)
|
|
|
|
# Save the exchange to history
|
|
self.conversations.add_exchange(user_id, user_message, response.content)
|
|
|
|
logger.debug(
|
|
f"Generated response for user {user_id}: "
|
|
f"{len(response.content)} chars, {response.usage}"
|
|
)
|
|
|
|
return response.content
|
|
|
|
|
|
async def setup(bot: commands.Bot) -> None:
|
|
"""Load the AI Chat cog."""
|
|
await bot.add_cog(AIChatCog(bot))
|