diff --git a/README.md b/README.md index c057d5c..254d92d 100644 --- a/README.md +++ b/README.md @@ -1,31 +1,80 @@ # how to run ts bullshit -first you need ollama installed +two parts: `bot.py` talks to discord, `main.py` (ollama) or `gemini.py` (gemini api) is the brain. +`main.py` is stateless so it can live on another machine / colab and talk to `bot.py` through a websocket or tunnel, you never have to upload your files anywhere sketchy + +# installing +pick your distro. all of them use a venv because python doesn't let you `pip install` system wide anymore + +### ubuntu ```console -$ pacman -S ollama -#or whatever packet manager you have -#alternatively, -$curl -fsSL https://ollama.com/install.sh | sh +$ sudo apt install python3 python3-venv curl zstd +$ curl -fsSL https://ollama.com/install.sh | sh ``` +### arch +```console +$ sudo pacman -S python ollama +#ollama-cuda or ollama-rocm if you have a gpu +``` +### gentoo +```console +$ sudo emerge --ask dev-lang/python sci-ml/ollama +#sci-ml/ollama is ~amd64, so you need to keyword it first +$ echo "sci-ml/ollama ~amd64" | sudo tee /etc/portage/package.accept_keywords/ollama +#sci-ml/ollama-bin if you don't wanna compile it +``` +(skip the ollama part if you're using the gemini api) + then your python dependencies are -``` console -$ pip install ollama websockets discord +```console +$ python3 -m venv venv +$ source venv/bin/activate +$ pip install ollama websockets discord.py +#and this one only if you're using the gemini api +$ pip install google-genai ``` -# once you install your dependencies -you run this +# once you install your dependencies +you need ollama running (not needed for gemini) ```console -#i think ollama still needs to serve i'm not too sure -$ ollama serve & #or make it a service with systemd or whatever +#ollama needs to be serving, the packages above already ship a service +$ sudo systemctl enable --now ollama #ubuntu / arch / gentoo systemd +$ sudo rc-update add ollama default && sudo rc-service ollama start #gentoo openrc +#or just run it by hand +$ ollama serve & $ ollama pull $ ollama pull #those two are in main.py -#test it with -$ ollama run +#test it with +$ ollama run ``` -if that runs fine then you can -- add your discord bot token and discord server id to your .env which is not .env but ./files/config/tokens +if that runs fine (or you have a gemini key) then you can +- go to the discord developer portal and turn on **message content intent** and **server members intent** for your bot or it won't start - add the bot to your server -- you'll probably want to change the system prompt which is found in bot.py +- run the wizard, it writes all the config files for you +```console +$ python3 setup_wizard.py +``` +- edit your system prompt in `./files/config/system_prompt.txt` (it makes a placeholder the first time bot.py runs) +- optional: `./config/params.json` has the ollama options (temperature and stuff), without it ollama just uses its defaults +- run the brain first, then the bot +```console +$ python3 main.py #or gemini.py +$ python3 bot.py +``` +she only reads `general` and `announcements`, and every dm gets its own memory in `./files/log/dm/` so it never leaks into the server -# have fun \ No newline at end of file +# colab +if your pc can't run the model you can host `main.py` on colab, run `setup_wizard.py` on both sides and say yes to the colab question +```console +#on colab +$ python3 colab_setup.py +``` +it prints a wss:// address and a token, paste those into the wizard on the bot.py side. the address changes every time you restart the notebook + +# IF IT FAILS +- check `./files/log/errorlog.txt` +- `bad auth token` means the token in `./files/config/connection.txt` isn't the one main.py has in `./config/auth_token.txt` +- on colab main.py logs to `main.log` next to colab_setup.py + +# have fun diff --git a/bot.py b/bot.py index 7485ef2..3f5b2f5 100644 --- a/bot.py +++ b/bot.py @@ -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}") -@client.event 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: diff --git a/colab_setup.py b/colab_setup.py index 42d57d9..faf8f79 100644 --- a/colab_setup.py +++ b/colab_setup.py @@ -2,7 +2,7 @@ """ run this in a colab cell to stand up main.py + a public wss:// tunnel to it. session is ephemeral — restart the notebook, get a new url, re-paste it -into your bot.py-side config/connection.txt. +into your bot.py-side files/config/connection.txt (or rerun setup_wizard.py). """ import os import re @@ -14,8 +14,13 @@ from pathlib import Path SCRIPT_DIR = Path(__file__).resolve().parent CONFIG_DIR = SCRIPT_DIR / "config" -CHATMODEL = "hf.co/Stetsoncody/Llama-3.2-3B-Instruct-abliterated-GGUF:Q4_K_M" -EMBEDMODEL = "hf.co/CompendiumLabs/bge-base-en-v1.5-gguf" +def model_from_main(name): + # read the model names straight out of main.py so they can't drift apart + text = (SCRIPT_DIR / "main.py").read_text(encoding="utf-8") + return re.search(rf'^{name}\s*=\s*["\'](.+?)["\']', text, re.M).group(1) + +CHATMODEL = model_from_main("chatmodel") +EMBEDMODEL = model_from_main("embedmodel") def run(cmd, **kwargs): @@ -65,9 +70,13 @@ def ensure_params(): def start_main(): print("\n-- starting main.py --") + # same deal as the tunnel below, main.py prints a lot and raises BrokenPipe once this cell's stdout is gone + log_file = open(SCRIPT_DIR / "main.log", "w") return subprocess.Popen( f"python3 {SCRIPT_DIR / 'main.py'}", shell=True, + stdout=log_file, + stderr=subprocess.STDOUT, ) diff --git a/gemini.py b/gemini.py new file mode 100644 index 0000000..f7da375 --- /dev/null +++ b/gemini.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +import asyncio +import base64 +import hmac +import json +import os +import re # thank you claude +import websockets +from pathlib import Path +from websockets.server import serve +from google import genai +from google.genai import types +from utils import nots, now_str + +chatmodel = "gemini-3.5-flash-lite" +embedmodel = "gemini-embedding-001" +embeddim = 768 # default is 3072, the cosine loop below is pure python so smaller = way faster + +# resolve config relative to this script, not whatever dir it happens to get launched from +SCRIPT_DIR = Path(__file__).resolve().parent +CONFIG_DIR = SCRIPT_DIR / "config" + +try: + AUTH_TOKEN = (CONFIG_DIR / "auth_token.txt").read_text(encoding="utf-8").strip() +except FileNotFoundError: + raise SystemExit("no config/auth_token.txt found — run setup_wizard.py first") + +# api key comes from env var or config/gemini_key.txt, whichever exists +GEMINI_KEY = os.environ.get("GEMINI_API_KEY", "").strip() +if not GEMINI_KEY: + try: + GEMINI_KEY = (CONFIG_DIR / "gemini_key.txt").read_text(encoding="utf-8").strip() + except FileNotFoundError: + raise SystemExit("no gemini key — set GEMINI_API_KEY or put it in config/gemini_key.txt") + +client = genai.Client(api_key=GEMINI_KEY) + +# gemini free tier is rate limited, keeping this at one generation at a time +chat_semaphore = asyncio.Semaphore(1) + +# bot is edgy, don't let gemini's default filters eat the replies +SAFETY_SETTINGS = [ + types.SafetySetting(category=c, threshold="BLOCK_NONE") + for c in ( + "HARM_CATEGORY_HARASSMENT", + "HARM_CATEGORY_HATE_SPEECH", + "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "HARM_CATEGORY_DANGEROUS_CONTENT", + ) +] + +# Each element in the VECTOR_DB will be a tuple (chunk, embedding) +VECTOR_DB = [] + +def embed(texts, task_type): + # batches of 100 since that's the api limit per call + out = [] + for i in range(0, len(texts), 100): + res = client.models.embed_content( + model=embedmodel, + contents=texts[i:i + 100], + config=types.EmbedContentConfig(task_type=task_type, output_dimensionality=embeddim), + ) + out.extend(e.values for e in res.embeddings) + return out + +def cosine_similarity(a, b): + dot_product = sum([x * y for x, y in zip(a, b)]) + norm_a = sum([x ** 2 for x in a]) ** 0.5 + norm_b = sum([x ** 2 for x in b]) ** 0.5 + return dot_product / (norm_a * norm_b) + +def retrieve(query, top_n=5): + query_embedding = embed([query], "RETRIEVAL_QUERY")[0] + similarities = [] + for chunk, embedding in VECTOR_DB: + similarity = cosine_similarity(query_embedding, embedding) + similarities.append((chunk, similarity)) + similarities.sort(key=lambda x: x[1], reverse=True) + return similarities[:top_n] + +def rebuild_vector_db(longterm_lines: list[str]): + global VECTOR_DB + existing = dict(VECTOR_DB) # chunk -> embedding, so we don't re-embed stuff we already have + missing = [c for c in dict.fromkeys(longterm_lines) if c not in existing] + # free tier counts every chunk in a batch as a request (100/min), so only do 90 per rebuild + # and catch up on the next messages instead of failing the whole reply + batch = missing[:90] + if batch: + print(f"[{now_str()}] embedding {len(batch)} new chunks ({len(missing)} missing total)") + try: + for chunk, embedding in zip(batch, embed(batch, "RETRIEVAL_DOCUMENT")): + existing[chunk] = embedding + except Exception as e: + print(f"[{now_str()}] embed failed, will retry next message: {e}") + VECTOR_DB = [(chunk, existing[chunk]) for chunk in longterm_lines if chunk in existing] + + +async def geminichat(websocket): + # loop over every message on this connection instead of handling one and returning — + # returning from this function is what makes the websockets lib close the connection + async for raw in websocket: + try: + payload = json.loads(raw) + + if not hmac.compare_digest(str(payload.get("auth_token", "")).encode(), AUTH_TOKEN.encode()): + await websocket.send(json.dumps({"error": "bad auth token"})) + continue + + system_prompt = payload["system_prompt"] + user_input = payload["user_input"] + shortterm = payload.get("shortterm", []) # list of strings + longterm = payload.get("longterm", []) # list of strings + media = payload.get("media", []) # list of {"mime_type", "data" (base64)} + + # rebuild vector db if bot signals it (every 150 interactions bot sends full longterm), + # or if we come up empty but longterm isn't — e.g. we just restarted and lost VECTOR_DB + # also keeps going while the db is still catching up after a partial (quota limited) rebuild + if payload.get("rebuild_vectors") or len(VECTOR_DB) < len(longterm): + await asyncio.to_thread(rebuild_vector_db, longterm) + print(f"[{now_str()}] Vector DB reconciled, {len(VECTOR_DB)} chunks") + + user_query = nots(user_input) + try: + retrieved_knowledge = await asyncio.to_thread(retrieve, user_query) if VECTOR_DB else [] + except Exception as e: + # quota hit on the query embed, just reply without rag this time + print(f"[{now_str()}] retrieve failed, skipping rag: {e}") + retrieved_knowledge = [] + + # rag chunks go in system so they don't pollute turn history + rag_block = "" + if retrieved_knowledge: + rag_block = "\n\nOlder relevant messages for context:\n" + "\n".join(f"- {chunk}" for chunk, _ in retrieved_knowledge) + + # lines to skip — not real dialogue + skip_fragments = ("sent sticker:", "sent media:", "sent a random dm to ", "saved media:") + def is_skip(line): + c = nots(line) + if any(f in c for f in skip_fragments): + return True + # empty said: like "username said: " + if "said:" in c and c.split("said:", 1)[1].strip() == "": + return True + return False + + BOT_NAME = "Faggot-chan" + + # parse shortterm lines into (speaker, text) pairs + def parse_line(line): + c = nots(line).strip() + # "X said: Y", "X said in DM: Y", "X replied to ... with reply: Y" + # look at what comes right after the username (no spaces in those) instead of searching the whole line, + # otherwise a message that happens to contain " with reply: " gets read as a reply + speaker, _, rest = c.partition(" ") + if rest.startswith("said in DM: "): + return speaker.strip(), rest[len("said in DM: "):].strip() + if rest.startswith("replied to ") and ' with reply: "' in rest: + # "X replied to Y saying: "Z" with reply: "W"" + text = rest.rsplit(' with reply: "', 1)[-1] + if text.endswith('"'): + text = text[:-1] + return speaker.strip(), text.strip() + if rest.startswith("said: "): + return speaker.strip(), rest[len("said: "):].strip() + return None, None + + # build alternating user/model turns from shortterm + # group consecutive non-bot lines into one user blob, bot lines become model turns + messages = [] + current_user_lines = [] + seen_last = None # for dedup + + for line in shortterm: + if is_skip(line): + continue + speaker, text = parse_line(line) + if not speaker or not text: + continue + entry = f"{speaker}: {text}" + if entry == seen_last: # skip duplicate consecutive lines + continue + seen_last = entry + + if speaker == BOT_NAME: + # flush accumulated user lines first + if current_user_lines: + messages.append({"role": "user", "content": "\n".join(current_user_lines)}) + current_user_lines = [] + messages.append({"role": "model", "content": text}) + else: + current_user_lines.append(entry) + + # flush any trailing user lines before the final turn + if current_user_lines: + messages.append({"role": "user", "content": "\n".join(current_user_lines)}) + + # final isolated turn — the specific message she needs to reply to + messages.append({"role": "user", "content": f"[responding to]: {nots(user_input)}"}) + + # make sure we start with a user turn (gemini rejects empty parts so no empty string here) + if messages and messages[0]["role"] != "user": + messages.insert(0, {"role": "user", "content": "..."}) + + contents = [ + types.Content(role=m["role"], parts=[types.Part(text=m["content"])]) + for m in messages + ] + + # media only goes on the final turn (history stays text), put it before the text like gemini recommends + if media: + media_parts = [ + types.Part.from_bytes(data=base64.b64decode(m["data"]), mime_type=m["mime_type"]) + for m in media + ] + contents[-1].parts = media_parts + contents[-1].parts + + config = types.GenerateContentConfig( + system_instruction=system_prompt + rag_block, + safety_settings=SAFETY_SETTINGS, + ) + + async with chat_semaphore: + response = await asyncio.to_thread( + client.models.generate_content, + model=chatmodel, + contents=contents, + config=config, + ) + + # response.text is None when gemini blocks the reply anyway + if not response.text: + reason = response.prompt_feedback.block_reason if response.prompt_feedback else None + if not reason and response.candidates: + reason = response.candidates[0].finish_reason + raise RuntimeError(f"gemini returned no text ({reason})") + + # strip blocks just in case, thank you claude + bot_reply = re.sub(r'.*?', '', response.text, flags=re.DOTALL).strip() + + print(messages) + print(f"[{now_str()}] reply: {bot_reply}\n") + + await websocket.send(json.dumps({"ok": bot_reply})) + + except Exception as e: + error_msg = f"[{now_str()}] error in geminichat: {e}" + print(error_msg) + try: + await websocket.send(json.dumps({"error": error_msg})) + except Exception: + pass # if we can't even send the error, nothing we can do + + +async def main(): + print(f"[{now_str()}] hosting on ws://0.0.0.0:8765") + async with serve(geminichat, "0.0.0.0", 8765, ping_interval=None, ping_timeout=None, max_size=32 * 1024 * 1024): # default 1mb would kill any media payload + await asyncio.Future() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/main.py b/main.py index 88865a7..44cbaf4 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,7 @@ #!/usr/bin/env python3 import asyncio +import base64 +import hmac import json import re # thank you claude import ollama @@ -8,7 +10,7 @@ from pathlib import Path from websockets.server import serve from utils import nots, now_str -chatmodel = "hf.co/Stetsoncody/Llama-3.2-3B-Instruct-abliterated-GGUF:Q4_K_M" +chatmodel = "hf.co/mradermacher/Huihui-gemma-4-E4B-it-abliterated-GGUF:IQ4_XS" embedmodel = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf' # resolve config relative to this script, not whatever dir it happens to get launched from @@ -72,7 +74,7 @@ async def ollamachat(websocket): try: payload = json.loads(raw) - if payload.get("auth_token") != AUTH_TOKEN: + if not hmac.compare_digest(str(payload.get("auth_token", "")).encode(), AUTH_TOKEN.encode()): await websocket.send(json.dumps({"error": "bad auth token"})) continue @@ -80,6 +82,7 @@ async def ollamachat(websocket): user_input = payload["user_input"] shortterm = payload.get("shortterm", []) # list of strings longterm = payload.get("longterm", []) # list of strings + media = payload.get("media", []) # list of {"mime_type", "data" (base64)} # rebuild vector db if bot signals it (every 150 interactions bot sends full longterm), # or if we come up empty but longterm isn't — e.g. we just restarted and lost VECTOR_DB @@ -96,7 +99,7 @@ async def ollamachat(websocket): rag_block = "\n\nOlder relevant messages for context:\n" + "\n".join(f"- {chunk}" for chunk, _ in retrieved_knowledge) # lines to skip — not real dialogue - skip_fragments = ("sent sticker:", "sent media:", "sent a random dm to:", "saved media:") + skip_fragments = ("sent sticker:", "sent media:", "sent a random dm to ", "saved media:") def is_skip(line): c = nots(line) if any(f in c for f in skip_fragments): @@ -112,17 +115,19 @@ async def ollamachat(websocket): def parse_line(line): c = nots(line).strip() # "X said: Y", "X said in DM: Y", "X replied to ... with reply: Y" - if " said in DM: " in c: - speaker, text = c.split(" said in DM: ", 1) - return speaker.strip(), text.strip() - if " with reply: " in c: + # look at what comes right after the username (no spaces in those) instead of searching the whole line, + # otherwise a message that happens to contain " with reply: " gets read as a reply + speaker, _, rest = c.partition(" ") + if rest.startswith("said in DM: "): + return speaker.strip(), rest[len("said in DM: "):].strip() + if rest.startswith("replied to ") and ' with reply: "' in rest: # "X replied to Y saying: "Z" with reply: "W"" - speaker = c.split(" replied to ", 1)[0].strip() - text = c.rsplit(' with reply: "', 1)[-1].rstrip('"').strip() - return speaker.strip(), text.strip() - if " said: " in c: - speaker, text = c.split(" said: ", 1) + text = rest.rsplit(' with reply: "', 1)[-1] + if text.endswith('"'): + text = text[:-1] return speaker.strip(), text.strip() + if rest.startswith("said: "): + return speaker.strip(), rest[len("said: "):].strip() return None, None # build alternating user/assistant turns from shortterm @@ -158,13 +163,21 @@ async def ollamachat(websocket): # final isolated turn — the specific message she needs to reply to messages.append({"role": "user", "content": f"[responding to]: {nots(user_input)}"}) + # ollama only takes images (no audio/video/pdf), and only on the final turn so history stays text + images = [base64.b64decode(m["data"]) for m in media if m["mime_type"].startswith("image/")] + if images: + messages[-1]["images"] = images + # make sure we start with a user turn (model requirement) if messages and messages[0]["role"] != "user": messages.insert(0, {"role": "user", "content": ""}) messages.insert(0, {"role": "system", "content": system_prompt + rag_block}) - params = json.loads((CONFIG_DIR / "params.json").read_text()) + try: + params = json.loads((CONFIG_DIR / "params.json").read_text()) + except FileNotFoundError: + params = {} # no params.json, just use ollama's defaults async with chat_semaphore: response = await asyncio.to_thread( @@ -177,7 +190,7 @@ async def ollamachat(websocket): # strip blocks in case the model ignores think=False, thank you claude bot_reply = re.sub(r'.*?', '', response["message"]["content"], flags=re.DOTALL).strip() - print(messages) + print([{k: v for k, v in m.items() if k != "images"} for m in messages]) # skip the raw image bytes print(f"[{now_str()}] reply: {bot_reply}\n") await websocket.send(json.dumps({"ok": bot_reply})) @@ -193,7 +206,7 @@ async def ollamachat(websocket): async def main(): print(f"[{now_str()}] hosting on ws://0.0.0.0:8765") - async with serve(ollamachat, "0.0.0.0", 8765, ping_interval=None, ping_timeout=None): + async with serve(ollamachat, "0.0.0.0", 8765, ping_interval=None, ping_timeout=None, max_size=32 * 1024 * 1024): # default 1mb would kill any media payload await asyncio.Future() diff --git a/setup_wizard.py b/setup_wizard.py index 8575b27..743db51 100644 --- a/setup_wizard.py +++ b/setup_wizard.py @@ -25,27 +25,59 @@ def ask_yn(prompt: str) -> bool: print("just y or n please") +def read_lines(path: Path) -> list[str]: + # same parsing bot.py does, empty list if the file isn't there yet + if not path.exists(): + return [] + return [l.strip() for l in path.read_text(encoding="utf-8").splitlines() if l.strip() and not l.strip().startswith("#")] + + +def ask_keep(prompt: str, current: str = None, secret: bool = False) -> str: + # enter on an empty line keeps whatever the file already had (secrets don't get printed) + if current: + shown = "keep current" if secret else f"keep: {current}" + return input(f"{prompt} [{shown}]: ").strip() or current + return input(f"{prompt}: ").strip() + + +def ask_backend() -> str: + while True: + val = ask("is main.py using 'ollama' or the 'gemini' api", "ollama").lower() + if val in ("ollama", "gemini"): + return val + print("just ollama or gemini please") + + def write_tokens_file(config_dir: Path): config_dir.mkdir(parents=True, exist_ok=True) + old = read_lines(config_dir / "tokens.txt") + old_token = old[0] if len(old) > 0 else None + old_server = old[1] if len(old) > 1 else None print("\n-- discord bot setup --") - token = ask("discord bot token") - server_id = ask("discord server (guild) id") + token = ask_keep("discord bot token", old_token, secret=True) + server_id = ask_keep("discord server (guild) id", old_server) while not server_id.isdigit(): print("that's not a number") - server_id = ask("discord server (guild) id") + server_id = ask_keep("discord server (guild) id", old_server) (config_dir / "tokens.txt").write_text(f"{token}\n{server_id}\n", encoding="utf-8") print(f"wrote {config_dir / 'tokens.txt'}") -def write_connection_file(config_dir: Path, default_uri: str = "ws://0.0.0.0:8765"): +def write_connection_file(config_dir: Path, default_uri: str = "ws://127.0.0.1:8765"): config_dir.mkdir(parents=True, exist_ok=True) + old = read_lines(config_dir / "connection.txt") + old_uri = old[0] if len(old) > 0 else default_uri + old_token = old[1] if len(old) > 1 else None print("\n-- connection to main.py --") - uri = ask("main.py address (ws:// or wss://)", default_uri) - print("paste the auth token main.py is using (from colab_setup.py output),") - token = ask("or leave blank to generate a new one") - if not token: - token = secrets.token_urlsafe(32) - print(f"generated token: {token}") + uri = ask_keep("main.py address (ws:// or wss://)", old_uri) + if old_token: + token = ask_keep("auth token main.py is using (from colab_setup.py output)", old_token, secret=True) + else: + print("paste the auth token main.py is using (from colab_setup.py output),") + token = ask("or leave blank to generate a new one") + if not token: + token = secrets.token_urlsafe(32) + print(f"generated token: {token}") (config_dir / "connection.txt").write_text(f"{uri}\n{token}\n", encoding="utf-8") print(f"wrote {config_dir / 'connection.txt'}") return token @@ -54,17 +86,33 @@ def write_connection_file(config_dir: Path, default_uri: str = "ws://0.0.0.0:876 def write_auth_token_file(config_dir: Path, token: str = None): config_dir.mkdir(parents=True, exist_ok=True) if not token: + old = read_lines(config_dir / "auth_token.txt") print("\n-- websocket auth for main.py --") - print("this is the secret bot.py has to send. leave blank to generate one.") - token = ask("auth token") - if not token: - token = secrets.token_urlsafe(32) - print(f"generated token: {token}") + if old: + token = ask_keep("auth token", old[0], secret=True) + else: + print("this is the secret bot.py has to send. leave blank to generate one.") + token = ask("auth token") + if not token: + token = secrets.token_urlsafe(32) + print(f"generated token: {token}") (config_dir / "auth_token.txt").write_text(f"{token}\n", encoding="utf-8") print(f"wrote {config_dir / 'auth_token.txt'}") print(f"\ngive this token (and this machine's public wss:// url) to whoever's setting up bot.py:\n {token}") +def write_gemini_key_file(config_dir: Path): + config_dir.mkdir(parents=True, exist_ok=True) + old = read_lines(config_dir / "gemini_key.txt") + print("\n-- gemini api key for gemini.py --") + key = ask_keep("gemini api key", old[0] if old else None, secret=True) + while not key: + print("need a key") + key = ask_keep("gemini api key") + (config_dir / "gemini_key.txt").write_text(f"{key}\n", encoding="utf-8") + print(f"wrote {config_dir / 'gemini_key.txt'}") + + def main(): print("faggot-chan setup wizard\n") @@ -74,10 +122,13 @@ def main(): # everything lives together, plain local setup bot_config = SCRIPT_DIR / "files" / "config" main_config = SCRIPT_DIR / "config" + backend = ask_backend() write_tokens_file(bot_config) - token = write_connection_file(bot_config, default_uri="ws://0.0.0.0:8765") + token = write_connection_file(bot_config, default_uri="ws://127.0.0.1:8765") write_auth_token_file(main_config, token=token) - print("\nall set. run main.py, then bot.py.") + if backend == "gemini": + write_gemini_key_file(main_config) + print(f"\nall set. run {'gemini.py' if backend == 'gemini' else 'main.py'}, then bot.py.") return where = ask("running this on the 'vps' (bot.py host) or 'colab' box (main.py host)", "vps") @@ -86,8 +137,13 @@ def main(): # this box just needs the auth token main.py will check — connection details # (discord token, server id) live on whichever box runs bot.py main_config = SCRIPT_DIR / "config" + backend = ask_backend() write_auth_token_file(main_config) - print("\nnow run colab_setup.py to bring main.py + the tunnel up.") + if backend == "gemini": + write_gemini_key_file(main_config) + print("\nnow run gemini.py (colab_setup.py would start the ollama main.py and pull models you don't need).") + else: + print("\nnow run colab_setup.py to bring main.py + the tunnel up.") return # vps / bot.py host: skip discord-client setup entirely (handled separately),