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

This commit is contained in:
david committed 2026-10-14 21:49:41 -03:00
1 parent 9f82a0a223
commit 8f287c89e5
2 files changed
+197

No files matched your search

+22
View File
@@ -81,6 +81,28 @@ $ python3 admin.py unblock <user id>
```
`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 <blocked|prompt|settings|reactions|params> #opens $EDITOR, vim if you don't have one
$ python3 admin.py block <user id>
$ python3 admin.py unblock <user id>
```
`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 <name> 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
+175
View File
@@ -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()