diff --git a/README.md b/README.md index 89d0e37..c825bfd 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,28 @@ $ python3 admin.py unblock ``` `reactions.json` is where the stickers / emojis / media she can send live, the bot adds new ones with an empty description and you fill them in. media with no description never gets offered +# fallback +if you're using gemini you can have a local ollama `main.py` take over when gemini says you're out of quota. the wizard asks about it (it's the third line of `./files/config/connection.txt`), it needs to run on another port so run it as +```console +$ PORT=8766 python3 main.py +``` +the bot goes back to trying gemini after an hour. if there's no fallback (or it's down) she just says "usage limit reached, see you in a month" + +# admin +`admin.py` is the admin panel, it only edits files on the box so nobody on discord can touch it. the bot re-reads everything on every message, no restart needed +```console +$ python3 admin.py list +$ python3 admin.py edit #opens $EDITOR, vim if you don't have one +$ python3 admin.py block +$ python3 admin.py unblock +``` +`reactions.json` is where the stickers / emojis / media she can send live, the bot adds new ones with an empty description and you fill them in. media with no description never gets offered + +you can let the local gemma-4 write the descriptions for you, it needs `pip install pillow` (and ffmpeg for videos), ollama running, and the model has to be able to see images +```console +$ python3 describe.py #only the empty ones, --all redoes everything, --model to use another model, --gemini to use the gemini api instead +``` + # 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 diff --git a/describe.py b/describe.py new file mode 100644 index 0000000..0e56550 --- /dev/null +++ b/describe.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +""" +fills in the empty descriptions in files/config/reactions.json by showing every sticker / emoji / media to the +local gemma-4 on ollama. run it on the box with the files, ollama has to be up (main.py doesn't). +stickers and emojis are downloaded from discord with the bot token (login only, the bot doesn't need to be running). + + describe.py only entries with no description yet + describe.py --all redo everything (overwrites the ones you wrote by hand too) + describe.py --model ollama model, defaults to chatmodel from main.py (it has to be able to see images) + describe.py --gemini use the gemini api instead (chatmodel from gemini.py, key from GEMINI_API_KEY or config/gemini_key.txt) +""" +import argparse +import asyncio +import io +import json +import os +import re +import shutil +import subprocess +import tempfile +import time +from pathlib import Path + +from PIL import Image + +SCRIPT_DIR = Path(__file__).resolve().parent +CONFIG_DIR = SCRIPT_DIR / "files" / "config" +MEDIA_DIR = SCRIPT_DIR / "files" / "media" +REACTIONS_FILE = CONFIG_DIR / "reactions.json" +VIDEO_EXTS = (".mp4", ".webm", ".mov", ".mkv") + +PROMPT = ("this is a {kind} from a discord server, its name is \"{hint}\" (the name is a hint, it may or may not help). " + "describe what is literally in the image in one short line, under 15 words: the character or object (name the character if you recognize them, " + "the name above can help), what they're doing, and any text written in the image. " + "don't guess when to use it, no preamble, no quotes.") + + +def model_from(script: str) -> str: + # same trick as colab_setup.py, read the model name straight out of main.py / gemini.py so they can't drift apart + text = (SCRIPT_DIR / script).read_text(encoding="utf-8") + return re.search(r'^chatmodel\s*=\s*["\'](.+?)["\']', text, re.M).group(1) + + +def read_lines(path: Path) -> list[str]: + return [l.strip() for l in path.read_text(encoding="utf-8").splitlines() if l.strip() and not l.strip().startswith("#")] + + +async def fetch_discord(wanted: dict) -> dict: + # {(kind, name): bytes}. login only, no gateway, so this works while the bot is running too + import discord + token, guild_id = read_lines(CONFIG_DIR / "tokens.txt")[:2] + found = {} + client = discord.Client(intents=discord.Intents.none()) + await client.login(token) + try: + guild = await client.fetch_guild(int(guild_id)) + for kind, items in (("stickers", await guild.fetch_stickers()), ("emojis", await guild.fetch_emojis())): + for item in items: + if item.name not in wanted[kind] or not item.available: # locked ones (boost level dropped) can't be sent anyway + continue + if kind == "stickers" and item.format == discord.StickerFormatType.lottie: # json animation, nothing to look at + continue + found[(kind, item.name)] = await item.read() + finally: + await client.close() + return found + + +def to_png(data: bytes, dest: Path) -> bool: + # first frame on a white background (transparent stickers come out black otherwise), small so gemma doesn't chew on it + try: + img = Image.open(io.BytesIO(data)) + img.seek(getattr(img, "n_frames", 1) // 2) # middle frame, the first one of an animation is often blank or half drawn + img = img.convert("RGBA") + except Exception: + return False + bg = Image.new("RGBA", img.size, "white") + bg.alpha_composite(img) + bg = bg.convert("RGB") + bg.thumbnail((512, 512)) + bg.save(dest, "PNG") + return True + + +def video_frame(path: Path, dest: Path) -> bool: + if not shutil.which("ffmpeg"): + return False + r = subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(path), "-frames:v", "1", "-vf", "thumbnail,scale=512:-1", str(dest)]) + return r.returncode == 0 and dest.exists() + + +def describe(model: str, png: Path, prompt: str) -> str: + import ollama + r = ollama.chat(model=model, messages=[{"role": "user", "content": prompt, "images": [str(png)]}], options={"temperature": 0.2}) + return clean(r["message"]["content"]) + + +def describe_gemini(model: str, png: Path, prompt: str) -> str: + from google import genai + from google.genai import types + key = os.environ.get("GEMINI_API_KEY", "").strip() or (SCRIPT_DIR / "config" / "gemini_key.txt").read_text(encoding="utf-8").strip() + # a sticker with a slur in it would get blocked by the default filters, same reason gemini.py turns them off + safety = [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")] + client = genai.Client(api_key=key) # has to stay in a variable, as a temporary it gets garbage collected mid call and closes the connection + r = client.models.generate_content( + model=model, + contents=[types.Part.from_bytes(data=png.read_bytes(), mime_type="image/png"), prompt], + config=types.GenerateContentConfig(safety_settings=safety, temperature=0.2), + ) + time.sleep(4) # free tier rate limit, rerun the script if some fail, it only does the empty ones + return clean(r.text or "") + + +def clean(text: str) -> str: + text = re.sub(r'.*?', '', text, flags=re.DOTALL) + return " ".join(text.split()).strip(' "') + + +def save_description(kind: str, name: str, text: str, force: bool): + # read it fresh every time, the bot adds new entries to this file while we run + data = json.loads(REACTIONS_FILE.read_text(encoding="utf-8")) + if name in data.get(kind, {}) and (force or not data[kind][name]): + data[kind][name] = text + REACTIONS_FILE.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--all", action="store_true") + ap.add_argument("--model") + ap.add_argument("--gemini", action="store_true") + args = ap.parse_args() + model = args.model or model_from("gemini.py" if args.gemini else "main.py") + + if not REACTIONS_FILE.exists(): + raise SystemExit("no files/config/reactions.json yet, start bot.py once so it lists your stickers/emojis/media") + data = json.loads(REACTIONS_FILE.read_text(encoding="utf-8")) + todo = {k: [n for n, d in data.get(k, {}).items() if args.all or not d] for k in ("stickers", "emojis", "media")} + if not any(todo.values()): + raise SystemExit("nothing without a description") + + remote = asyncio.run(fetch_discord(todo)) if todo["stickers"] or todo["emojis"] else {} + + with tempfile.TemporaryDirectory() as tmp: + png = Path(tmp) / "thing.png" + for kind in ("stickers", "emojis", "media"): + for name in todo[kind]: + png.unlink(missing_ok=True) + if kind == "media": + path = MEDIA_DIR / name + if not path.is_file(): + print(f"skip media/{name}: file is gone") + continue + ok = video_frame(path, png) if path.suffix.lower() in VIDEO_EXTS else to_png(path.read_bytes(), png) + else: + ok = (kind, name) in remote and to_png(remote[(kind, name)], png) + if not ok: + print(f"skip {kind}/{name}: couldn't get an image out of it") + continue + try: + hint = re.sub(r"^\d+_", "", Path(name).stem) if kind == "media" else name # media is saved as id_filename + prompt = PROMPT.format(kind=kind.rstrip("s") if kind != "media" else "media file", hint=hint) + text = (describe_gemini if args.gemini else describe)(model, png, prompt) + except Exception as e: + print(f"skip {kind}/{name}: describe failed: {e}") + continue + if not text: + continue + save_description(kind, name, text, args.all) + print(f"{kind}/{name}: {text}") + + +if __name__ == "__main__": + main()