File size: 6,573 Bytes
2a735cc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import discord
import asyncio
import logging
import os
from typing import Dict, List, Any
import threading
from discord.ext import commands
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class DiscordBot:
"""
Discord bot integration for the AI second brain.
Handles message ingestion, responses, and synchronization with the main app.
"""
def __init__(self, agent, token=None, channel_whitelist=None):
"""
Initialize the Discord bot.
Args:
agent: The AssistantAgent instance to use for processing queries
token: Discord bot token (defaults to environment variable)
channel_whitelist: List of channel IDs to listen to (None for all)
"""
self.agent = agent
self.token = token or os.getenv("DISCORD_BOT_TOKEN")
self.channel_whitelist = channel_whitelist or []
self.message_history = []
# Set up Discord client with intents
intents = discord.Intents.default()
intents.message_content = True # Required to read message content
self.client = commands.Bot(command_prefix="!", intents=intents)
# Register event handlers
self.setup_event_handlers()
# Thread for bot
self.bot_thread = None
logger.info("Discord bot initialized")
def setup_event_handlers(self):
"""Register event handlers for the Discord client."""
@self.client.event
async def on_ready():
logger.info(f"Discord bot logged in as {self.client.user}")
@self.client.event
async def on_message(message):
# Don't respond to self
if message.author == self.client.user:
return
# Check if this is a command
await self.client.process_commands(message)
# Only process messages in whitelisted channels if whitelist exists
if self.channel_whitelist and message.channel.id not in self.channel_whitelist:
return
# Only respond to messages that mention the bot or are DMs
is_dm = isinstance(message.channel, discord.DMChannel)
is_mentioned = self.client.user in message.mentions
if is_dm or is_mentioned:
await self.process_message(message)
# Add a !help command
@self.client.command(name="help")
async def help_command(ctx):
help_text = """
**AI Assistant Commands**
- Mention me with a question to get an answer
- Send me a DM with your query
- Use `!search <query>` to search your knowledge base
- Use `!upload` with an attachment to add to your knowledge base
"""
await ctx.send(help_text)
# Add a search command
@self.client.command(name="search")
async def search_command(ctx, *, query):
async with ctx.typing():
response = await self.process_query(query)
await ctx.send(response["answer"])
# If there are sources, show them in a followup message
if response["sources"]:
sources_text = "**Sources:**\n" + "\n".join([
f"- {s['file_name']} ({s['source']})"
for s in response["sources"]
])
await ctx.send(sources_text)
async def process_message(self, message):
"""Process a Discord message and send a response."""
# Clean up mention and extract query
content = message.content
for mention in message.mentions:
content = content.replace(f'<@{mention.id}>', '').replace(f'<@!{mention.id}>', '')
query = content.strip()
if not query:
await message.channel.send("How can I help you?")
return
# Show typing indicator
async with message.channel.typing():
# Process the query and get a response
response = await self.process_query(query)
# Store in message history
self.message_history.append({
"user": str(message.author),
"query": query,
"response": response["answer"],
"timestamp": message.created_at.isoformat(),
"channel": str(message.channel)
})
# Send the response
await message.channel.send(response["answer"])
# If there are sources, send them in a followup message
if response["sources"]:
sources_text = "**Sources:**\n" + "\n".join([
f"- {s['file_name']} ({s['source']})"
for s in response["sources"]
])
await message.channel.send(sources_text)
async def process_query(self, query):
"""Process a query using the agent and return a response."""
# Run the query in a thread to avoid blocking the event loop
loop = asyncio.get_event_loop()
response = await loop.run_in_executor(None, self.agent.query, query)
# Add the conversation to the agent's memory
if "answer" in response:
await loop.run_in_executor(
None,
self.agent.add_conversation_to_memory,
query,
response["answer"]
)
return response
def start(self):
"""Start the Discord bot in a separate thread."""
if not self.token:
logger.error("Discord bot token not found")
return False
def run_bot():
asyncio.set_event_loop(asyncio.new_event_loop())
self.client.run(self.token)
self.bot_thread = threading.Thread(target=run_bot, daemon=True)
self.bot_thread.start()
logger.info("Discord bot started in background thread")
return True
def stop(self):
"""Stop the Discord bot."""
if self.client and self.client.is_ready():
asyncio.run_coroutine_threadsafe(self.client.close(), self.client.loop)
logger.info("Discord bot stopped")
def get_message_history(self):
"""Get the message history."""
return self.message_history |