237 lines
8.9 KiB
Python
237 lines
8.9 KiB
Python
"""
|
||
Discord bot that connects users to the LangGraph agent via private messages.
|
||
|
||
Architecture
|
||
------------
|
||
- The bot runs in-process alongside FastAPI (on a background asyncio task).
|
||
- Private messages (DMs) are routed through the same LangGraph graphs that
|
||
power the REST API — no HTTP loopback needed.
|
||
- Per-user conversation history is maintained so the LLM has context.
|
||
|
||
Environment
|
||
-----------
|
||
DISCORD_BOT_TOKEN – the bot token from the Discord Developer Portal
|
||
DISCORD_MAX_HISTORY – how many past messages to keep per user (default 7)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
|
||
import discord
|
||
|
||
from agents import list_all as list_all_agents
|
||
from bot.conversation import ConversationStore
|
||
from core.config import DEEPSEEK_API_KEY, get_config
|
||
from core.graph import create_agent_graph
|
||
from core.llm import create_client
|
||
|
||
logger = logging.getLogger("bot.discord")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Config
|
||
# ---------------------------------------------------------------------------
|
||
DISCORD_BOT_TOKEN = get_config("DISCORD_BOT_TOKEN") or ""
|
||
DISCORD_MAX_HISTORY = int(get_config("DISCORD_MAX_HISTORY", "7"))
|
||
DISCORD_DEFAULT_AGENT = get_config("DISCORD_DEFAULT_AGENT", "media-agent")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# LLM client shared by all agents (same as the REST API uses)
|
||
# ---------------------------------------------------------------------------
|
||
_llm_client = create_client(DEEPSEEK_API_KEY)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Conversation store — one per process
|
||
# ---------------------------------------------------------------------------
|
||
_conversations = ConversationStore(max_history=DISCORD_MAX_HISTORY)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Graph cache — lazy-compiled per agent, same pattern as api/dependencies.py
|
||
# ---------------------------------------------------------------------------
|
||
_agent_graphs: dict[str, object] = {}
|
||
|
||
|
||
def _get_graph(agent_id: str):
|
||
"""Return a compiled LangGraph for *agent_id*, building it on first use."""
|
||
if agent_id not in _agent_graphs:
|
||
agents = list_all_agents()
|
||
agent = agents.get(agent_id, agents.get("naked"))
|
||
_agent_graphs[agent_id] = create_agent_graph(
|
||
client=_llm_client,
|
||
agent_skills=agent.skills if agent else [],
|
||
system_prompt=agent.build_system_prompt() if agent else (
|
||
"You are a helpful, general-purpose assistant."
|
||
),
|
||
)
|
||
return _agent_graphs[agent_id]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Discord client
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class AgentBot(discord.Client):
|
||
"""A discord.py Client that connects users to the LangGraph agent."""
|
||
|
||
def __init__(self) -> None:
|
||
# message_content lets us read DM text.
|
||
# guilds is required so that mutual_guilds is populated — without it
|
||
# every DM is silently ignored.
|
||
intents = discord.Intents.default()
|
||
intents.message_content = True
|
||
intents.guilds = True
|
||
super().__init__(intents=intents)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Lifecycle
|
||
# ------------------------------------------------------------------
|
||
|
||
async def on_ready(self) -> None:
|
||
logger.info("Bot logged in as %s (ID %s)", self.user, self.user.id)
|
||
# Print a ready banner so the dev knows it's alive
|
||
print(f"\n🤖 Discord bot ready — logged in as {self.user}\n")
|
||
|
||
# ------------------------------------------------------------------
|
||
# Shared-guild helper — uses the REST API, no privileged intents
|
||
# ------------------------------------------------------------------
|
||
|
||
async def _shares_guild(self, user: discord.User) -> bool:
|
||
"""Return True if *user* and the bot share at least one guild."""
|
||
for guild in self.guilds:
|
||
try:
|
||
member = await guild.fetch_member(user.id)
|
||
if member is not None:
|
||
return True
|
||
except (discord.NotFound, discord.Forbidden, discord.HTTPException):
|
||
continue
|
||
return False
|
||
|
||
# ------------------------------------------------------------------
|
||
# Message handler — DMs only
|
||
# ------------------------------------------------------------------
|
||
|
||
async def on_message(self, message: discord.Message) -> None:
|
||
# Never reply to ourselves
|
||
if message.author == self.user:
|
||
return
|
||
|
||
# |-- DM channel only for now ----------------------------------|
|
||
if not isinstance(message.channel, discord.DMChannel):
|
||
logger.debug("Ignoring message from #%s (not a DM)", message.channel)
|
||
return
|
||
# |--------------------------------------------------------------|
|
||
|
||
# |-- Shared-server gate — only users who share at least one --|
|
||
# | guild with the bot can interact via DM. --|
|
||
# | We use fetch_member (REST API) instead of --|
|
||
# | User.mutual_guilds because the latter requires the --|
|
||
# | privileged "members" intent. This way no privileged --|
|
||
# | intents are needed. --|
|
||
if not await self._shares_guild(message.author):
|
||
logger.warning(
|
||
"Blocking DM from %s — no mutual guilds.",
|
||
message.author.name,
|
||
)
|
||
return
|
||
# |--------------------------------------------------------------|
|
||
|
||
user_id = message.author.id
|
||
|
||
# Show typing indicator while the graph runs
|
||
async with message.channel.typing():
|
||
try:
|
||
reply = await self._run_agent(
|
||
user_id=user_id,
|
||
user_msg=message.content,
|
||
)
|
||
await message.channel.send(reply)
|
||
except Exception:
|
||
logger.exception("Agent run failed for user %s", user_id)
|
||
await message.channel.send(
|
||
"Sorry, something went wrong processing your request. "
|
||
"Please try again in a moment."
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Agent invocation
|
||
# ------------------------------------------------------------------
|
||
|
||
async def _run_agent(self, *, user_id: int, user_msg: str) -> str:
|
||
"""Build the message list from history, invoke the graph, store the
|
||
reply, and return the assistant's final text."""
|
||
|
||
# 1. Pick agent — defaults to DISCORD_DEFAULT_AGENT env var.
|
||
# Change DISCORD_DEFAULT_AGENT in .env to switch agents.
|
||
agent_id = DISCORD_DEFAULT_AGENT
|
||
|
||
# 2. Build message list from stored history + new user message
|
||
history = _conversations.get_history(user_id)
|
||
messages = [*history, {"role": "user", "content": user_msg}]
|
||
|
||
# 3. Run the LangGraph (tools execute inline if needed)
|
||
graph = _get_graph(agent_id)
|
||
state = {"messages": messages}
|
||
result = await graph.ainvoke(state)
|
||
|
||
last_msg = result["messages"][-1]
|
||
reply = last_msg.content or ""
|
||
|
||
# 4. Persist the conversation
|
||
_conversations.append(user_id, user_msg, reply)
|
||
|
||
return reply
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Bootstrap helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _start_bot_sync(token: str) -> None:
|
||
"""Synchronous entry-point that runs the bot in a new asyncio event loop.
|
||
|
||
Called from a background thread so the main thread can keep running the
|
||
FastAPI / uvicorn server.
|
||
"""
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
|
||
async def _run() -> None:
|
||
bot = AgentBot()
|
||
try:
|
||
await bot.start(token)
|
||
except discord.LoginFailure:
|
||
logger.error(
|
||
"Discord login failed — check DISCORD_BOT_TOKEN in your .env file."
|
||
)
|
||
except Exception:
|
||
logger.exception("Unhandled exception in bot event loop.")
|
||
|
||
loop.run_until_complete(_run())
|
||
|
||
|
||
def start_in_background(token: str | None = None) -> None:
|
||
"""Launch the Discord bot in a daemon thread.
|
||
|
||
Pass *token* explicitly if you already have it; otherwise it is read
|
||
from the DISCORD_BOT_TOKEN env variable.
|
||
"""
|
||
token = token or DISCORD_BOT_TOKEN
|
||
if not token:
|
||
logger.warning(
|
||
"DISCORD_BOT_TOKEN is not set — Discord bot will NOT start."
|
||
)
|
||
return
|
||
|
||
import threading
|
||
|
||
t = threading.Thread(
|
||
target=_start_bot_sync,
|
||
args=(token,),
|
||
daemon=True,
|
||
name="discord-bot",
|
||
)
|
||
t.start()
|
||
logger.info("Discord bot thread started.")
|