Dad-bot/bot.py

import asyncio
import logging
import os
import json
import time
from collections import defaultdict, deque
from datetime import date
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("dadbot")

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")
BACKGROUND_MODEL = os.getenv("BACKGROUND_MODEL", MODEL)
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()
        log.info(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)
    log.info(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 {}
        log.info(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

# Semaphore to prevent background LLM tasks from running concurrently
background_llm_sem = asyncio.Semaphore(1)

# Cooldown tracking: skip background LLM calls for a period after a model error
BACKGROUND_COOLDOWN = 120  # seconds
BACKGROUND_LLM_TIMEOUT = 90  # seconds
_last_background_error: float = 0.0


def background_on_cooldown() -> bool:
    return (time.monotonic() - _last_background_error) < BACKGROUND_COOLDOWN


def record_background_error():
    global _last_background_error
    _last_background_error = time.monotonic()


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"])
    if background_on_cooldown():
        log.info(f"[compaction] Skipping '{name}' — model on cooldown")
        return
    log.info(f"[compaction] Waiting for background slot for '{name}'")
    async with background_llm_sem:
        log.info(f"[compaction] Starting LLM call for '{name}' ({len(entry['facts'])} facts)")
        t0 = time.monotonic()
        try:
            response = await asyncio.wait_for(
                llm.chat.completions.create(
                    model=BACKGROUND_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}"},
                    ],
                ),
                timeout=BACKGROUND_LLM_TIMEOUT,
            )
            elapsed = time.monotonic() - t0
            log.info(f"[compaction] LLM call completed in {elapsed:.1f}s for '{name}'")
            summary = strip_thinking_tags(response.choices[0].message.content.strip())
            entry["facts"] = [summary]
            save_dossiers()
            log.info(f"[compaction] Saved compacted dossier for '{name}'")
        except asyncio.TimeoutError:
            elapsed = time.monotonic() - t0
            log.warning(f"[compaction] Timed out after {elapsed:.1f}s for '{name}', will retry later")
            record_background_error()
        except Exception as exc:
            elapsed = time.monotonic() - t0
            log.error(f"[compaction] Failed for '{name}' after {elapsed:.1f}s: {exc}")
            record_background_error()


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)
            ),
        },
    ]

    if background_on_cooldown():
        log.info(f"[dossier] Skipping {user.display_name} — model on cooldown")
        return
    log.info(f"[dossier] Waiting for background slot for {user.display_name}")
    async with background_llm_sem:
        log.info(f"[dossier] Starting extraction for {user.display_name}")
        t0 = time.monotonic()
        try:
            response = await asyncio.wait_for(
                llm.chat.completions.create(model=BACKGROUND_MODEL, messages=extraction_messages),
                timeout=BACKGROUND_LLM_TIMEOUT,
            )
            elapsed = time.monotonic() - t0
            raw = strip_thinking_tags(response.choices[0].message.content.strip())
            log.info(f"[dossier] Extraction for {user.display_name} completed in {elapsed:.1f}s: {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()
            log.info(f"[dossier] Saved {len(updated_facts)} facts for {user.display_name}")
        except asyncio.TimeoutError:
            elapsed = time.monotonic() - t0
            log.warning(f"[dossier] Timed out after {elapsed:.1f}s for {user.display_name}, will retry later")
            record_background_error()
        except Exception as exc:
            log.error(f"[dossier] Extraction failed for {user.display_name}: {exc}")
            record_background_error()
    await maybe_compact_dossier(uid, user.display_name)


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)
            ),
        },
    ]

    if background_on_cooldown():
        log.info("[self-dossier] Skipping — model on cooldown")
        return
    log.info("[self-dossier] Waiting for background slot")
    async with background_llm_sem:
        log.info("[self-dossier] Starting extraction")
        t0 = time.monotonic()
        try:
            response = await asyncio.wait_for(
                llm.chat.completions.create(model=BACKGROUND_MODEL, messages=extraction_messages),
                timeout=BACKGROUND_LLM_TIMEOUT,
            )
            elapsed = time.monotonic() - t0
            log.info(f"[self-dossier] Extraction completed in {elapsed:.1f}s")
            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()
            log.info(f"[self-dossier] Saved {len(updated_facts)} facts")
        except asyncio.TimeoutError:
            elapsed = time.monotonic() - t0
            log.warning(f"[self-dossier] Timed out after {elapsed:.1f}s, will retry later")
            record_background_error()
        except Exception as exc:
            log.error(f"[self-dossier] Extraction failed: {exc}")
            record_background_error()
    await maybe_compact_dossier(SELF_DOSSIER_KEY, "DadBot")


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
                ]
                log.info(f"Loaded {len(tools)} MCP tools: {[t['function']['name'] for t in tools]}")
                return tools
    except Exception as exc:
        log.info(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}

    iteration = 0
    while True:
        iteration += 1
        log.info(f"[agent] LLM call #{iteration} starting")
        t0 = time.monotonic()
        response = await llm.chat.completions.create(**kwargs)
        elapsed = time.monotonic() - t0
        msg = response.choices[0].message
        log.info(f"[agent] LLM call #{iteration} completed in {elapsed:.1f}s, tool_calls={bool(msg.tool_calls)}")

        if not msg.tool_calls:
            return msg.content

        kwargs["messages"].append(msg)

        for tc in msg.tool_calls:
            args = json.loads(tc.function.arguments)
            log.info(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
    log.info(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)
            log.info(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)