423 lines
15 KiB
Python
423 lines
15 KiB
Python
import asyncio
|
|
import os
|
|
import json
|
|
from collections import defaultdict, deque
|
|
from datetime import date
|
|
from pathlib import Path
|
|
|
|
import discord
|
|
from openai import AsyncOpenAI
|
|
from dotenv import load_dotenv
|
|
from mcp import ClientSession
|
|
from mcp.client.streamable_http import streamablehttp_client
|
|
|
|
load_dotenv()
|
|
|
|
DISCORD_TOKEN = os.environ["DISCORD_TOKEN"]
|
|
LITELLM_BASE_URL = os.environ["LITELLM_BASE_URL"]
|
|
LITELLM_API_KEY = os.getenv("LITELLM_API_KEY", "anything")
|
|
MODEL = os.getenv("MODEL", "gpt-4")
|
|
SYSTEM_PROMPT = os.getenv("SYSTEM_PROMPT", "You are a helpful assistant.")
|
|
MAX_CONTEXT = int(os.getenv("MAX_CONTEXT", "20"))
|
|
MCP_URL = os.getenv("MCP_URL", "http://mcp-searxng:4001/mcp")
|
|
DOSSIER_FILE = Path(os.getenv("DOSSIER_FILE", "dossiers.json"))
|
|
PERSONALITY_FILE = Path(os.getenv("PERSONALITY_FILE", "personality.txt"))
|
|
|
|
llm = AsyncOpenAI(base_url=LITELLM_BASE_URL, api_key=LITELLM_API_KEY)
|
|
|
|
# Per-channel rolling history
|
|
channel_history: dict[int, deque] = defaultdict(lambda: deque(maxlen=MAX_CONTEXT))
|
|
|
|
# MCP tools cached at startup in OpenAI function-calling format
|
|
openai_tools: list[dict] = []
|
|
|
|
# User dossiers: {str(user_id): {"name": str, "facts": [str, ...]}}
|
|
dossiers: dict[str, dict] = {}
|
|
|
|
# Current personality notes (persisted to disk)
|
|
personality_notes: str = ""
|
|
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
bot = discord.Client(intents=intents)
|
|
|
|
# Local tool definition for personality self-modification
|
|
PERSONALITY_TOOL = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "update_personality",
|
|
"description": (
|
|
"Update your own personality. Use this to add, change, or remove personality traits, "
|
|
"opinions, interests, or behavioural tendencies. Changes persist permanently across conversations."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"personality": {
|
|
"type": "string",
|
|
"description": "The full updated personality description, replacing the current one.",
|
|
}
|
|
},
|
|
"required": ["personality"],
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def load_personality():
|
|
global personality_notes
|
|
if PERSONALITY_FILE.exists():
|
|
personality_notes = PERSONALITY_FILE.read_text().strip()
|
|
print(f"Loaded personality ({len(personality_notes)} chars)")
|
|
else:
|
|
personality_notes = ""
|
|
|
|
|
|
def save_personality(text: str):
|
|
global personality_notes
|
|
personality_notes = text.strip()
|
|
PERSONALITY_FILE.write_text(personality_notes)
|
|
print(f"Personality updated ({len(personality_notes)} chars)")
|
|
|
|
|
|
def load_dossiers():
|
|
global dossiers
|
|
if DOSSIER_FILE.exists():
|
|
text = DOSSIER_FILE.read_text().strip()
|
|
dossiers = json.loads(text) if text else {}
|
|
print(f"Loaded dossiers for {len(dossiers)} users")
|
|
else:
|
|
dossiers = {}
|
|
|
|
|
|
def save_dossiers():
|
|
DOSSIER_FILE.write_text(json.dumps(dossiers, indent=2))
|
|
|
|
|
|
SELF_DOSSIER_KEY = "self"
|
|
COMPACTION_THRESHOLD = 15
|
|
|
|
|
|
def strip_thinking_tags(text: str) -> str:
|
|
"""Remove <think>...</think> blocks that some models (e.g. Qwen) emit before their response."""
|
|
import re
|
|
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
|
|
|
|
|
|
async def maybe_compact_dossier(key: str, name: str):
|
|
"""If a dossier exceeds the threshold, distill its facts into a compact summary paragraph."""
|
|
entry = dossiers.get(key)
|
|
if not entry or len(entry.get("facts", [])) < COMPACTION_THRESHOLD:
|
|
return
|
|
|
|
facts_str = "\n".join(f"- {f}" for f in entry["facts"])
|
|
try:
|
|
response = await llm.chat.completions.create(
|
|
model=MODEL,
|
|
messages=[
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"Distill the following facts about a person into a single compact summary paragraph "
|
|
"of 2-3 sentences. Be dense and preserve all meaningful details. No bullet points."
|
|
),
|
|
},
|
|
{"role": "user", "content": f"Person: {name}\n\nFacts:\n{facts_str}"},
|
|
],
|
|
)
|
|
summary = strip_thinking_tags(response.choices[0].message.content.strip())
|
|
entry["facts"] = [summary]
|
|
save_dossiers()
|
|
print(f"Compacted dossier for {name} into summary paragraph")
|
|
except Exception as exc:
|
|
print(f"Dossier compaction failed for {name}: {exc}")
|
|
|
|
|
|
def get_dossier_context(user_ids: list[int]) -> str:
|
|
"""Build a context string for known users and self to inject into the system prompt."""
|
|
parts = []
|
|
|
|
self_entry = dossiers.get(SELF_DOSSIER_KEY)
|
|
if self_entry and self_entry.get("facts"):
|
|
facts = "; ".join(self_entry["facts"])
|
|
parts.append(f"What you know about yourself:\n- {facts}")
|
|
|
|
lines = []
|
|
for uid in user_ids:
|
|
entry = dossiers.get(str(uid))
|
|
if entry and entry.get("facts"):
|
|
name = entry.get("name", str(uid))
|
|
facts = "; ".join(entry["facts"])
|
|
lines.append(f"- {name}: {facts}")
|
|
if lines:
|
|
parts.append("What you know about people in this conversation:\n" + "\n".join(lines))
|
|
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
async def extract_and_update_dossier(user: discord.Member | discord.User, conversation: list[dict]):
|
|
"""Fire-and-forget: ask the LLM to extract new facts about a user from the conversation."""
|
|
uid = str(user.id)
|
|
existing = dossiers.get(uid, {}).get("facts", [])
|
|
existing_str = "; ".join(existing) if existing else "nothing yet"
|
|
|
|
extraction_messages = [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"You maintain a fact sheet about a specific person based on conversations. "
|
|
"You will be given their existing known facts and a new conversation. "
|
|
"Return ONLY a JSON array of short fact strings representing the complete, up-to-date fact sheet. "
|
|
"Merge new facts in, update or remove any that are contradicted by newer information, and deduplicate. "
|
|
"Only include facts clearly stated or strongly implied. Return the existing facts unchanged if nothing new was learned."
|
|
),
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"Person: {user.display_name} (they appear as '{user.display_name}:' in messages)\n"
|
|
f"Existing facts: {existing_str}\n\n"
|
|
f"Conversation:\n"
|
|
+ "\n".join(f"{m['role']}: {m['content']}" for m in conversation)
|
|
),
|
|
},
|
|
]
|
|
|
|
try:
|
|
response = await llm.chat.completions.create(
|
|
model=MODEL,
|
|
messages=extraction_messages,
|
|
)
|
|
raw = strip_thinking_tags(response.choices[0].message.content.strip())
|
|
print(f"Dossier extraction raw response for {user.display_name}: {repr(raw[:200])}")
|
|
if raw.startswith("```"):
|
|
raw = raw.split("\n", 1)[-1]
|
|
raw = raw.rsplit("```", 1)[0].strip()
|
|
if not raw:
|
|
return
|
|
updated_facts = json.loads(raw)
|
|
if not isinstance(updated_facts, list) or not updated_facts:
|
|
return
|
|
|
|
entry = dossiers.setdefault(uid, {"name": user.display_name, "facts": []})
|
|
entry["name"] = user.display_name
|
|
entry["facts"] = updated_facts[:50]
|
|
save_dossiers()
|
|
print(f"Updated dossier for {user.display_name}: {len(updated_facts)} facts total")
|
|
await maybe_compact_dossier(uid, user.display_name)
|
|
except Exception as exc:
|
|
print(f"Dossier extraction failed for {user.display_name}: {exc}")
|
|
|
|
|
|
async def extract_and_update_self_dossier(conversation: list[dict]):
|
|
"""Fire-and-forget: extract facts DadBot expressed about itself in the conversation."""
|
|
existing = dossiers.get(SELF_DOSSIER_KEY, {}).get("facts", [])
|
|
existing_str = "; ".join(existing) if existing else "nothing yet"
|
|
|
|
extraction_messages = [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"You maintain a fact sheet about DadBot based on what it says about itself in conversations. "
|
|
"Extract opinions, preferences, stances, backstory, or personality traits DadBot expressed. "
|
|
"Return ONLY a JSON array of short fact strings representing the complete, up-to-date self fact sheet. "
|
|
"Merge new facts in, update or remove contradicted ones, and deduplicate. "
|
|
"Return the existing facts unchanged if DadBot revealed nothing new about itself."
|
|
),
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"Existing self-facts: {existing_str}\n\n"
|
|
f"Conversation:\n"
|
|
+ "\n".join(f"{m['role']}: {m['content']}" for m in conversation)
|
|
),
|
|
},
|
|
]
|
|
|
|
try:
|
|
response = await llm.chat.completions.create(
|
|
model=MODEL,
|
|
messages=extraction_messages,
|
|
)
|
|
raw = strip_thinking_tags(response.choices[0].message.content.strip())
|
|
if raw.startswith("```"):
|
|
raw = raw.split("\n", 1)[-1]
|
|
raw = raw.rsplit("```", 1)[0].strip()
|
|
if not raw:
|
|
return
|
|
updated_facts = json.loads(raw)
|
|
if not isinstance(updated_facts, list) or not updated_facts:
|
|
return
|
|
|
|
entry = dossiers.setdefault(SELF_DOSSIER_KEY, {"facts": []})
|
|
entry["facts"] = updated_facts[:50]
|
|
save_dossiers()
|
|
print(f"Updated self-dossier: {len(updated_facts)} facts total")
|
|
await maybe_compact_dossier(SELF_DOSSIER_KEY, "DadBot")
|
|
except Exception as exc:
|
|
print(f"Self-dossier extraction failed: {exc}")
|
|
|
|
|
|
async def fetch_mcp_tools() -> list[dict]:
|
|
try:
|
|
async with streamablehttp_client(MCP_URL) as (read, write, _):
|
|
async with ClientSession(read, write) as session:
|
|
await session.initialize()
|
|
result = await session.list_tools()
|
|
tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": t.name,
|
|
"description": t.description or "",
|
|
"parameters": t.inputSchema,
|
|
},
|
|
}
|
|
for t in result.tools
|
|
]
|
|
print(f"Loaded {len(tools)} MCP tools: {[t['function']['name'] for t in tools]}")
|
|
return tools
|
|
except Exception as exc:
|
|
print(f"Warning: could not load MCP tools: {exc}")
|
|
return []
|
|
|
|
|
|
async def call_mcp_tool(name: str, args: dict) -> str:
|
|
async with streamablehttp_client(MCP_URL) as (read, write, _):
|
|
async with ClientSession(read, write) as session:
|
|
await session.initialize()
|
|
result = await session.call_tool(name, args)
|
|
parts = []
|
|
for block in result.content:
|
|
if hasattr(block, "text"):
|
|
parts.append(block.text)
|
|
else:
|
|
parts.append(str(block))
|
|
return "\n".join(parts)
|
|
|
|
|
|
async def run_agent(messages: list) -> str:
|
|
"""Call the LLM in a loop, executing tool calls until a final reply is produced."""
|
|
all_tools = openai_tools + [PERSONALITY_TOOL]
|
|
kwargs = {"model": MODEL, "messages": list(messages), "tools": all_tools}
|
|
|
|
while True:
|
|
response = await llm.chat.completions.create(**kwargs)
|
|
msg = response.choices[0].message
|
|
|
|
if not msg.tool_calls:
|
|
return msg.content
|
|
|
|
kwargs["messages"].append(msg)
|
|
|
|
for tc in msg.tool_calls:
|
|
args = json.loads(tc.function.arguments)
|
|
print(f"Tool call: {tc.function.name}({args})")
|
|
|
|
if tc.function.name == "update_personality":
|
|
save_personality(args["personality"])
|
|
result = "Personality updated successfully."
|
|
else:
|
|
result = await call_mcp_tool(tc.function.name, args)
|
|
|
|
kwargs["messages"].append({
|
|
"role": "tool",
|
|
"tool_call_id": tc.id,
|
|
"content": result,
|
|
})
|
|
|
|
|
|
@bot.event
|
|
async def on_ready():
|
|
global openai_tools
|
|
print(f"Logged in as {bot.user} (id: {bot.user.id})")
|
|
load_dossiers()
|
|
load_personality()
|
|
openai_tools = await fetch_mcp_tools()
|
|
# Compact any dossiers that are already over the threshold (run in background)
|
|
for key, entry in list(dossiers.items()):
|
|
if len(entry.get("facts", [])) >= COMPACTION_THRESHOLD:
|
|
name = entry.get("name", key)
|
|
print(f"Scheduling startup compaction for '{key}'")
|
|
asyncio.create_task(maybe_compact_dossier(key, name))
|
|
|
|
|
|
@bot.event
|
|
async def on_message(message: discord.Message):
|
|
if message.author.bot:
|
|
return
|
|
|
|
history = channel_history[message.channel.id]
|
|
|
|
content = message.content
|
|
for mention in (f"<@{bot.user.id}>", f"<@!{bot.user.id}>"):
|
|
content = content.replace(mention, "")
|
|
content = content.strip()
|
|
|
|
history.append({"role": "user", "content": f"{message.author.display_name}: {content}"})
|
|
|
|
if bot.user not in message.mentions:
|
|
return
|
|
|
|
if not content:
|
|
await message.reply("You mentioned me but didn't say anything!")
|
|
return
|
|
|
|
recent_user_ids = [message.author.id]
|
|
dossier_context = get_dossier_context(recent_user_ids)
|
|
today = date.today().strftime("%B %d, %Y")
|
|
date_context = (
|
|
f"Today's date is {today}. Your training data has a cutoff and may be months or years out of date. "
|
|
f"For anything time-sensitive — news, prices, sports, current events, who holds a position, etc. — "
|
|
f"you must use the web search tool. Don't guess or rely on memory for current facts."
|
|
)
|
|
|
|
system_content = SYSTEM_PROMPT
|
|
if personality_notes:
|
|
system_content = f"{system_content}\n\nYour current personality notes (you wrote these yourself and can update them anytime using the update_personality tool):\n{personality_notes}"
|
|
system_content = f"{system_content}\n\n{date_context}"
|
|
if dossier_context:
|
|
system_content = f"{system_content}\n\n{dossier_context}"
|
|
|
|
messages = [{"role": "system", "content": system_content}] + list(history)
|
|
|
|
async with message.channel.typing():
|
|
try:
|
|
reply = await run_agent(messages)
|
|
except Exception as exc:
|
|
reply = f"Error: {exc}"
|
|
history.pop()
|
|
|
|
if reply:
|
|
history.append({"role": "assistant", "content": reply})
|
|
|
|
asyncio.create_task(extract_and_update_dossier(message.author, list(history)))
|
|
asyncio.create_task(extract_and_update_self_dossier(list(history)))
|
|
|
|
normalized = reply.replace("[SPIT]", "[SPLIT]").replace("[SPLT]", "[SPLIT]").replace("[SPLITT]", "[SPLIT]")
|
|
parts = [p.strip() for p in normalized.split("[SPLIT]") if p.strip()]
|
|
first = True
|
|
for part in parts:
|
|
for chunk in _split(part, 2000):
|
|
if not first:
|
|
await asyncio.sleep(min(len(chunk) / 50, 3))
|
|
if first:
|
|
await message.reply(chunk)
|
|
first = False
|
|
else:
|
|
await message.channel.send(chunk)
|
|
|
|
|
|
def _split(text: str, limit: int) -> list[str]:
|
|
if len(text) <= limit:
|
|
return [text]
|
|
chunks = []
|
|
while text:
|
|
chunks.append(text[:limit])
|
|
text = text[limit:]
|
|
return chunks
|
|
|
|
|
|
bot.run(DISCORD_TOKEN)
|