diff options
Diffstat (limited to 'bot.py')
| -rw-r--r-- | bot.py | 281 |
1 files changed, 193 insertions, 88 deletions
@@ -1,14 +1,17 @@ #!/usr/bin/env python3 -import os import re # thank you claude +import base64 +import hashlib import asyncio import json import time import random import traceback +import aiohttp # already installed, discord.py needs it import discord import websockets from pathlib import Path +from urllib.parse import urlparse from utils import nots, now_str, tail_lines def ignore_connection_closed_ok(loop, context): #this doesn't fucking work @@ -19,7 +22,7 @@ def ignore_connection_closed_ok(loop, context): #this doesn't fucking work def load_system_prompt() -> str: - return (CONFIG_DIR / "system_prompt.txt").read_text(encoding="utf-8") + return PROMPT_FILE.read_text(encoding="utf-8") # resolve files relative to this script, not whatever dir it happens to get launched from SCRIPT_DIR = Path(__file__).resolve().parent @@ -27,19 +30,23 @@ FILES_DIR = SCRIPT_DIR / "files" LOG_DIR = FILES_DIR / "log" MEDIA_DIR = FILES_DIR / "media" CONFIG_DIR = FILES_DIR / "config" +DM_DIR = LOG_DIR / "dm" # one memory file per person so dms never end up in the shared memory ERROR_FILE = LOG_DIR / "errorlog.txt" MEMORY_FILE = LOG_DIR / "memory.txt" MEDIA_FILE = MEDIA_DIR / "media.txt" THREADS_FILE = CONFIG_DIR / "threads.txt" +PROMPT_FILE = CONFIG_DIR / "system_prompt.txt" def init_files(): - for d in (LOG_DIR, MEDIA_DIR, CONFIG_DIR): + for d in (LOG_DIR, MEDIA_DIR, CONFIG_DIR, DM_DIR): d.mkdir(parents=True, exist_ok=True) for f in (ERROR_FILE, MEMORY_FILE, MEDIA_FILE): f.touch(exist_ok=True) if not THREADS_FILE.exists(): THREADS_FILE.write_text("6\n") + if not PROMPT_FILE.exists(): # otherwise every single message dies on a missing file + PROMPT_FILE.write_text("you are faggot-chan, a chaotic discord bot. edit files/config/system_prompt.txt\n", encoding="utf-8") init_files() @@ -69,17 +76,23 @@ interaction_count = 0 QUIET_COOLDOWN = 600 # seconds of silence before she's allowed to send a random message (10 mins) +MEDIA_MIMES = ("image/png", "image/jpeg", "image/webp", "image/heic", "image/heif", + "audio/", "video/", "application/pdf", "text/plain") # what gemini can actually read +MEDIA_MAX_BYTES = 10 * 1024 * 1024 # total per message, gemini caps inline request size (~20mb and base64 adds a third) +MEDIA_SAVE_MAX_BYTES = 8 * 1024 * 1024 # per file we keep, stays under discord's 10mb upload limit so she can send it back + def log_error(text: str): with ERROR_FILE.open("a", encoding="utf-8") as f: f.write(text + "\n") -def log_message(text: str): - with MEMORY_FILE.open("a", encoding="utf-8") as f: +def log_message(text: str, file: Path = MEMORY_FILE): + text = " ".join(text.splitlines()) # one message = one line, otherwise multiline messages break the history parsing + with file.open("a", encoding="utf-8") as f: f.write(text + "\n") -def load_shortterm() -> list[str]: +def load_shortterm(file: Path = MEMORY_FILE) -> list[str]: try: - return tail_lines(MEMORY_FILE, 15) # reads bottom 15 lines without loading the whole file + return tail_lines(file, 15) # reads bottom 15 lines without loading the whole file except Exception as e: log_error(f"[{now_str()}] error reading shortterm: {e}") return [] @@ -92,39 +105,52 @@ def load_longterm() -> list[str]: log_error(f"[{now_str()}] error reading longterm: {e}") return [] -def build_payload(user_input: str, rebuild: bool = False) -> str: +def build_payload(user_input: str, rebuild: bool = False, media: list = None, logged: bool = False, dm_file: Path = None) -> str: global vector_db_ready - should_rebuild = rebuild or not vector_db_ready # always rebuild on first ever message - if should_rebuild: - vector_db_ready = True + shortterm = load_shortterm(dm_file or MEMORY_FILE) + if logged and shortterm: + shortterm = shortterm[:-1] # the message she's answering is already the last line, it goes in the final turn instead + if dm_file: + longterm = [] # dms stay out of the shared rag so private stuff can't come up in a public channel + should_rebuild = False + else: + longterm = load_longterm() + should_rebuild = rebuild or not vector_db_ready # always rebuild on first ever message + if should_rebuild: + vector_db_ready = True return json.dumps({ "system_prompt": load_system_prompt(), "user_input": user_input, - "shortterm": load_shortterm(), - "longterm": load_longterm(), + "shortterm": shortterm, + "longterm": longterm, "rebuild_vectors": should_rebuild, + "media": media or [], # [{"mime_type", "data" (base64)}], gemini.py reads everything, main.py only the images "auth_token": MAIN_TOKEN, }) _main_ws = None # kept open across messages instead of reconnecting every time +_main_lock = None # one send/recv at a time on the shared socket, made on first use so it lives in discord's loop async def call_main(payload: str) -> tuple[bool, str]: - global _main_ws - for attempt in range(2): # retry once with a fresh connection if the cached one turned out dead - try: - if _main_ws is None: - _main_ws = await websockets.connect(MAIN_URI, open_timeout=None, ping_timeout=None) - await _main_ws.send(payload) - raw = await _main_ws.recv() - data = json.loads(raw) - if "error" in data: - return False, data["error"] - return True, data["ok"] - except Exception as e: - _main_ws = None - if attempt == 0: - continue # one retry with a fresh connection before giving up - return False, f"[{now_str()}] websocket error: {e}" + global _main_ws, _main_lock + if _main_lock is None: + _main_lock = asyncio.Lock() + async with _main_lock: + for attempt in range(2): # retry once with a fresh connection if the cached one turned out dead + try: + if _main_ws is None: + _main_ws = await websockets.connect(MAIN_URI, open_timeout=None, ping_timeout=None) + await _main_ws.send(payload) + raw = await _main_ws.recv() + data = json.loads(raw) + if "error" in data: + return False, data["error"] + return True, data["ok"] + except Exception as e: + _main_ws = None + if attempt == 0: + continue # one retry with a fresh connection before giving up + return False, f"[{now_str()}] websocket error: {e}" async def handle_error(error_msg: str, notify_channel=None, notify_user=None): print(error_msg) @@ -144,41 +170,103 @@ async def keep_typing(channel, stop_event): await channel.typing() await asyncio.sleep(8) # refresh every 8s before discord's 10s timeout +def dm_file_for(user) -> Path: + f = DM_DIR / f"{user.id}.txt" + f.touch(exist_ok=True) + return f + +def split_reply(text: str, limit: int = 2000) -> list[str]: + # discord rejects anything over 2000 chars (and empty messages), split on a newline/space when possible + chunks = [] + text = text.strip() + while len(text) > limit: + cut = max(text.rfind("\n", 0, limit), text.rfind(" ", 0, limit)) + if cut <= 0: + cut = limit + chunks.append(text[:cut].strip()) + text = text[cut:].strip() + if text: + chunks.append(text) + return chunks + +async def send_reply(message, text: str): + # first chunk is a real reply, the rest just follow in the channel + chunks = split_reply(text) + await message.reply(chunks[0]) + for chunk in chunks[1:]: + await message.channel.send(chunk) + vector_db_ready = False # gets set to True after first rebuild so we don't keep rebuilding -saved_media_urls = set() +saved_media = set() # file names inside MEDIA_DIR with MEDIA_FILE.open("r", encoding="utf-8") as f: for line in f: if "saved media:" in line: - url = line.split("saved media:", 1)[1].strip() - saved_media_urls.add(url) + name = line.split("saved media:", 1)[1].strip() + if (MEDIA_DIR / name).is_file(): # old entries were urls (those expire), they don't match a file so they get dropped + saved_media.add(name) -def steal_media(text: str): - if "saved media:" not in text: +async def steal_media(url: str, name: str): + # discord cdn links are signed and die after a day or so, so we keep our own copy + if name in saved_media: + return + saved_media.add(name) # claim it first so two messages can't download the same thing at once + try: + data = b"" + async with aiohttp.ClientSession() as session: + async with session.get(url) as r: + r.raise_for_status() + async for chunk in r.content.iter_chunked(65536): + data += chunk + if len(data) > MEDIA_SAVE_MAX_BYTES: + break + if len(data) > MEDIA_SAVE_MAX_BYTES: # too big for her to send back anyway + saved_media.discard(name) + return + (MEDIA_DIR / name).write_bytes(data) + except Exception as e: + saved_media.discard(name) + log_error(f"[{now_str()}] failed to save media {url}: {e}") return - url = text.split("saved media:", 1)[1].strip() - if url in saved_media_urls: - return - saved_media_urls.add(url) with MEDIA_FILE.open("a", encoding="utf-8") as f: - f.write(text + "\n") + f.write(f"[{now_str()}] saved media: {name}\n") -def steal_attachments(message): +async def steal_attachments(message): #steal images and gifs and stuff (always do this regardless of what else fires) - media_urls = [ - a.url for a in message.attachments - if a.content_type and a.content_type.startswith(("image/", "video/")) + media = [ + (a.url, f"{a.id}_{Path(a.filename).name}") for a in message.attachments + if a.content_type and a.content_type.startswith(("image/", "video/")) and a.size <= MEDIA_SAVE_MAX_BYTES ] for embed in message.embeds: - if embed.image and embed.image.url: - media_urls.append(embed.image.url) - if embed.video and embed.video.url: - media_urls.append(embed.video.url) + for part in (embed.image, embed.video): + if part and part.url: + ext = Path(urlparse(part.url).path).suffix # no extension, no preview on discord, not worth keeping + if ext: + media.append((part.url, hashlib.sha1(part.url.encode()).hexdigest()[:16] + ext)) - for url in media_urls: - steal_media(f"[{now_str()}] saved media: {url}") + for url, name in media: + await steal_media(url, name) + +async def grab_media(message) -> list[dict]: + # downloads supported attachments to send to gemini, only call this when the bot was directly mentioned + media = [] + total = 0 + for a in message.attachments: + mime = (a.content_type or "").split(";")[0].strip().lower() + if not mime.startswith(MEDIA_MIMES): + continue + if total + a.size > MEDIA_MAX_BYTES: + continue + try: + data = await a.read() + except Exception as e: + log_error(f"[{now_str()}] failed to download attachment {a.filename}: {e}") + continue + total += len(data) + media.append({"mime_type": mime, "data": base64.b64encode(data).decode("ascii")}) + return media def resolve_mentions(content: str, guild) -> str: # thank you claude if not guild: @@ -208,24 +296,28 @@ async def random_dm(): victim = random.choice(members) payload = build_payload( f"System: you are sending a dm to {victim}, write a long, playful story or an over-the-top pun monologue about anything. " - "It should sound like something faggot-chan would post when she's bored — chaotic, funny, and full of personality (at least 20 words)." - "keep in mind this is a dm so make it a bit personal" + "It should sound like something faggot-chan would post when she's bored — chaotic, funny, and full of personality (at least 20 words). " + "keep in mind this is a dm so make it a bit personal", + dm_file=dm_file_for(victim) ) stop_typing = asyncio.Event() asyncio.create_task(keep_typing(victim.dm_channel or await victim.create_dm(), stop_typing)) ok, result = await call_main(payload) stop_typing.set() if not ok: - await handle_error(result, notify_user=victim) + await handle_error(result) # just log it, the victim has nothing to do with this + continue + if not result.strip(): continue try: #sends the dm dm = await victim.create_dm() - await dm.send(result) + for chunk in split_reply(result): + await dm.send(chunk) dm_sent = True - log_message(f"[{now_str()}]Faggot-chan sent a random dm to {victim}: {result}") + log_message(f"[{now_str()}] Faggot-chan said in DM: {result}", dm_file_for(victim)) print(f"[{now_str()}] random dm to {victim}: {result}") except Exception as e: - await handle_error(f"[{now_str()}] failed to send dm to {victim}: {e}", notify_user=victim) + await handle_error(f"[{now_str()}] failed to send dm to {victim}: {e}") async def random_message(): await client.wait_until_ready() while not client.is_closed(): @@ -237,40 +329,43 @@ async def random_message(): channel = discord.utils.get(guild.text_channels, name="general") if channel: payload = build_payload( - "System: write a long, playful story or an over-the-top pun monologue about cheese. or anything random" - "it should sound like something faggot-chan would post when she's bored." + "System: write a short, playful story or an over-the-top pun monologue about cheese. or anything random " + "it should sound like something faggot-chan would post when she's bored. " "chaotic, funny, and full of personality (at least 15 words and no more than 20)" ) ok, result = await call_main(payload) if not ok: await handle_error(result, notify_channel=channel) continue - await channel.send(result) + if not result.strip(): + continue + for chunk in split_reply(result): + await channel.send(chunk) log_message(f"[{now_str()}] random message: {result}") print(f"[{now_str()}] random message: {result}") @client.event async def on_ready(): + asyncio.get_running_loop().set_exception_handler(ignore_connection_closed_ok) client.loop.create_task(random_message()) client.loop.create_task(random_dm()) print(f"faggot-chan is active as {client.user}") async def on_random(message) -> bool: # single weighted pick between sticker/emoji/media, used as a substitute for an actual # llm reply — only called once on_message already decided this message gets a response. # returns True if it fired (caller should skip the llm call), False if it should proceed normally. - # skip locked/role-restricted ones, discord rejects those (50081 / 10014) - stickers = [s for s in message.guild.stickers if s.available] - emojis = [str(e) for e in message.guild.emojis if e.available and not e.roles] + stickers = list(message.guild.stickers) + emojis = [str(e) for e in message.guild.emojis] + media_files = [n for n in saved_media if (MEDIA_DIR / n).is_file()] # in case something got deleted by hand options = [] if stickers: options.append(("sticker", 0.10)) if emojis: options.append(("emoji", 0.13)) - if saved_media_urls: + if media_files: options.append(("media", 0.07)) # ~30% combined chance of firing when all three are available if not options: @@ -289,8 +384,6 @@ async def on_random(message) -> bool: await message.channel.send(stickers=[chosen]) log_message(f"[{now_str()}] Faggot-chan sent sticker: {chosen.name}") print(f"[{now_str()}] Faggot-chan sent sticker: {chosen.name}") - except websockets.exceptions.ConnectionClosedOK: - pass except Exception as e: log_error(f"[{now_str()}] error: {e}") await message.reply("someone tell <@1000000000000001> there is a problem with my ai") @@ -300,8 +393,6 @@ async def on_random(message) -> bool: if kind == "emoji": try: await message.add_reaction(random.choice(emojis)) - except websockets.exceptions.ConnectionClosedOK: - pass except Exception as e: log_error(f"[{now_str()}] error: {e}") await message.reply("someone tell <@1000000000000001> there is a problem with my ai") @@ -309,10 +400,15 @@ async def on_random(message) -> bool: return True if kind == "media": - chosen = random.choice(list(saved_media_urls)) - await message.channel.send(chosen) - log_message(f"[{now_str()}] Faggot-chan sent media: {chosen}") - print(f"[{now_str()}] Faggot-chan sent media: {chosen}") + chosen = random.choice(media_files) + try: + await message.channel.send(file=discord.File(MEDIA_DIR / chosen)) + log_message(f"[{now_str()}] Faggot-chan sent media: {chosen}") + print(f"[{now_str()}] Faggot-chan sent media: {chosen}") + except Exception as e: + log_error(f"[{now_str()}] error: {e}") + await message.reply("someone tell <@1000000000000001> there is a problem with my ai") + print(f"Media send failed: {e}") return True return False @@ -321,7 +417,7 @@ async def on_random(message) -> bool: @client.event async def on_message(message): global dm_sent, interaction_count, last_server_message - if message.author == client.user: + if message.author == client.user or message.author.bot: # other bots too, otherwise two bots can loop forever return if not is_member_in_server(message.author): @@ -338,9 +434,10 @@ async def on_message(message): return if message.guild: - steal_attachments(message) + await steal_attachments(message) content = resolve_mentions(message.content, message.guild) + mem_file = MEMORY_FILE if message.guild else dm_file_for(message.author) # dms get their own memory # always detect replies replied_message = None @@ -348,19 +445,22 @@ async def on_message(message): replied_content = "" if message.reference: try: - replied_message = message.reference.resolved or await message.channel.fetch_message(message.reference.message_id) + replied_message = message.reference.resolved + if not isinstance(replied_message, discord.Message): # deleted or not cached, resolved is a placeholder with no author + replied_message = await message.channel.fetch_message(message.reference.message_id) if replied_message: if replied_message.author == client.user: replied_to_bot = True replied_content = resolve_mentions(replied_message.content, message.guild) - log_message(f'[{now_str()}] {message.author} replied to {replied_message.author} saying: "{replied_content}" with reply: "{content}"') + log_message(f'[{now_str()}] {message.author} replied to {replied_message.author} saying: "{replied_content}" with reply: "{content}"', mem_file) else: - log_message(f'[{now_str()}] {message.author} said: {content}') + log_message(f'[{now_str()}] {message.author} said: {content}', mem_file) except Exception as e: + replied_message = None # fetch failed (deleted), treat it like a normal message log_error(f"[{now_str()}] failed to resolve reply: {e}") - log_message(f'[{now_str()}] {message.author} said: {content}') + log_message(f'[{now_str()}] {message.author} said: {content}', mem_file) else: - log_message(f'[{now_str()}] {message.author} said: {content}') + log_message(f'[{now_str()}] {message.author} said: {content}', mem_file) interaction_count += 1 rebuild = (interaction_count % 150 == 0) @@ -371,7 +471,7 @@ async def on_message(message): user_input = f'{message.author} replied to {replied_message.author} saying: "{replied_content}" with reply: "{content}"' else: user_input = f"{message.author} said in DM: {content}" - payload = build_payload(user_input, rebuild) + payload = build_payload(user_input, rebuild, await grab_media(message), True, mem_file) # a dm is always aimed at her stop_typing = asyncio.Event() asyncio.create_task(keep_typing(message.channel, stop_typing)) ok, result = await call_main(payload) @@ -379,29 +479,31 @@ async def on_message(message): if not ok: await handle_error(result, notify_user=message.author) return + if not result.strip(): + return try: - await message.reply(result) + await send_reply(message, result) print(f"[{now_str()}] DM from {message.author}: {content} Bot replied in DM: {result}") - log_message(f"[{now_str()}] Faggot-chan said in DM: {result}") + log_message(f"[{now_str()}] Faggot-chan said in DM: {result}", mem_file) except Exception as e: await handle_error(f"[{now_str()}] failed to send DM reply: {e}", notify_user=message.author) return - mentioned = client.user.mentioned_in(message) + mentioned = client.user in message.mentions # mentioned_in() also fires on @everyone/@here #if it replied to bot, treat it as a mention automatically if replied_to_bot: mentioned = True - if str(message.author.id) == moralfag_id: + if not mentioned and random.random() > 0.03: #3% chance of replying if not mentiond + return + + if str(message.author.id) == moralfag_id: # only checked once she was actually going to reply, otherwise the skip gets eaten by chatter global moralfag_skip, moralfag_counter if moralfag_skip: moralfag_skip = False return - if not mentioned and random.random() > 0.03: #3% chance of replying if not mentiond - return - if await on_random(message): # ~30% chance she just sends a sticker/emoji/media instead return @@ -410,7 +512,8 @@ async def on_message(message): else: user_input = f"{message.author} said: {content}" - payload = build_payload(user_input, rebuild) + media = await grab_media(message) if client.user in message.mentions else [] # only direct @mentions, not replies or random 3% ones + payload = build_payload(user_input, rebuild, media, True) stop_typing = asyncio.Event() asyncio.create_task(keep_typing(message.channel, stop_typing)) ok, result = await call_main(payload) @@ -418,9 +521,11 @@ async def on_message(message): if not ok: await handle_error(result, notify_channel=message.channel) return + if not result.strip(): + return try: - await message.reply(result) + await send_reply(message, result) print(f"[{now_str()}] {message.author} said: {content}\n[{now_str()}] Faggot-chan said: {result}") log_message(f"[{now_str()}] Faggot-chan said: {result}") if str(message.author.id) == moralfag_id: |