i've been vibecoding this too much but it passes all my tests so it's ok

1 files changed, 175 insertions, 0 deletions

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 <name>   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'<think>.*?</think>', '', 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()