This commit is contained in:
david committed 2026-10-14 20:45:52 -03:00
1 parent b8baac169d
commit 9f82a0a223
7 files changed
+409 -121

No files matched your search

+17
View File
@@ -64,6 +64,23 @@ $ 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
# 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
# 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
+107
View File
@@ -0,0 +1,107 @@
#!/usr/bin/env python3
"""
admin panel for faggot-chan. runs on the box that has the files, so there's nothing on discord to reach
(not even with admin perms). the bot re-reads all of these on every message, no restart needed.
admin.py list
admin.py edit <blocked|prompt|settings|reactions|params> opens it in $EDITOR (vim if unset)
admin.py block <user id> [note]
admin.py unblock <user id>
"""
import json
import os
import shlex
import subprocess
import sys
from pathlib import Path
SCRIPT_DIR = Path(__file__).resolve().parent
BOT_CONFIG = SCRIPT_DIR / "files" / "config"
TARGETS = {
"blocked": BOT_CONFIG / "blocked.txt",
"prompt": BOT_CONFIG / "system_prompt.txt",
"settings": BOT_CONFIG / "settings.json",
"reactions": BOT_CONFIG / "reactions.json",
"params": SCRIPT_DIR / "config" / "params.json", # main.py side, ollama options
}
def edit(name: str):
path = TARGETS[name]
path.parent.mkdir(parents=True, exist_ok=True)
original = path.read_bytes() if path.exists() else None # to put back if the json ends up broken and they give up
if original is None:
path.write_text("{}\n" if path.suffix == ".json" else "", encoding="utf-8")
editor = shlex.split(os.environ.get("EDITOR", "vim"))
while True:
subprocess.call([*editor, str(path)])
if path.suffix != ".json":
return
try:
json.loads(path.read_text(encoding="utf-8"))
return
except json.JSONDecodeError as e:
print(f"that's not valid json: {e}")
if input("open it again? (y/n): ").strip().lower() in ("y", "yes"):
continue
if original is None:
path.unlink()
else:
path.write_bytes(original)
print("reverted")
return
def blocked_ids(lines: list[str]) -> set[str]:
return {l.split("#")[0].strip() for l in lines if l.split("#")[0].strip()}
def block(user_id: str, note: str = ""):
path = TARGETS["blocked"]
lines = path.read_text(encoding="utf-8").splitlines() if path.exists() else []
if user_id in blocked_ids(lines):
print(f"{user_id} is already blocked")
return
lines.append(f"{user_id} # {note}" if note else user_id)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
print(f"blocked {user_id}")
def unblock(user_id: str):
path = TARGETS["blocked"]
lines = path.read_text(encoding="utf-8").splitlines() if path.exists() else []
kept = [l for l in lines if l.split("#")[0].strip() != user_id]
if len(kept) == len(lines):
print(f"{user_id} wasn't blocked")
return
path.write_text("\n".join(kept) + "\n", encoding="utf-8")
print(f"unblocked {user_id}")
def usage():
print(__doc__.strip())
sys.exit(1)
def main():
args = sys.argv[1:]
if not args:
usage()
cmd = args[0]
if cmd == "list":
for name, path in TARGETS.items():
print(f"{name:10} {path}{'' if path.exists() else ' (not made yet)'}")
elif cmd == "edit" and len(args) == 2 and args[1] in TARGETS:
edit(args[1])
elif cmd == "block" and len(args) >= 2 and args[1].isdigit():
block(args[1], " ".join(args[2:]))
elif cmd == "unblock" and len(args) == 2 and args[1].isdigit():
unblock(args[1])
else:
usage()
if __name__ == "__main__":
main()
+162 -94
View File
@@ -37,6 +37,11 @@ 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"
BLOCKED_FILE = CONFIG_DIR / "blocked.txt" # edited with admin.py, re-read on every message
SETTINGS_FILE = CONFIG_DIR / "settings.json"
REACTIONS_FILE = CONFIG_DIR / "reactions.json" # sticker/emoji/media descriptions she picks from
DEFAULT_SETTINGS = {"admin_id": "1000000000000001", "moralfag_id": "1000000000000002", "faglimit": 2} # faglimit: ignore every x reply
def init_files():
for d in (LOG_DIR, MEDIA_DIR, CONFIG_DIR, DM_DIR):
@@ -47,6 +52,10 @@ def init_files():
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")
if not BLOCKED_FILE.exists():
BLOCKED_FILE.write_text("# one discord user id per line, edit with admin.py edit blocked\n1000000000000003 # fuck these two in particular\n1000000000000004\n", encoding="utf-8")
if not SETTINGS_FILE.exists():
SETTINGS_FILE.write_text(json.dumps(DEFAULT_SETTINGS, indent=2) + "\n", encoding="utf-8")
init_files()
@@ -57,23 +66,27 @@ server = int(_config_lines[1])
_conn_lines = [l.strip() for l in (CONFIG_DIR / "connection.txt").read_text(encoding="utf-8").splitlines() if l.strip() and not l.strip().startswith("#")]
MAIN_URI = _conn_lines[0] # ws:// or wss:// address of main.py
MAIN_TOKEN = _conn_lines[1] # shared secret main.py checks on every payload
FALLBACK_URI = _conn_lines[2] if len(_conn_lines) > 2 else None # optional, local main.py (ollama) that takes over when gemini runs out of quota, same token
intents = discord.Intents.default()
intents.message_content = True
intents.members = True
client = discord.Client(intents=intents)
block_id = ["1000000000000003", "1000000000000004"] #fuck these two in particular
moralfag_id = "1000000000000002"
moralfag_counter = 0
moralfag_skip = False
faglimit = 2 #ignore every x reply
last_server_message = 0.0 # unix timestamp of last message anyone sent in the server
dm_sent = False
interaction_count = 0
QUOTA_MSG = "usage limit reached, see you in a month"
QUOTA_COOLDOWN = 3600 # after gemini says it's out of quota go straight to the fallback for this long (seconds), then try it again
quota_until = 0.0 # unix timestamp
MAX_REACTION_OPTIONS = 25 # per kind (stickers/emojis/media), random sample so the prompt stays small
REACTION_TAG = re.compile(r'\[(sticker|emoji|media):\s*([^\]]+?)\s*\]', re.I)
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",
@@ -85,10 +98,27 @@ def log_error(text: str):
with ERROR_FILE.open("a", encoding="utf-8") as f:
f.write(text + "\n")
def load_blocked() -> set[str]:
# re-read on every message so admin.py edits apply without restarting
try:
lines = [l.split("#")[0].strip() for l in BLOCKED_FILE.read_text(encoding="utf-8").splitlines()]
return {l for l in lines if l}
except Exception as e:
log_error(f"[{now_str()}] error reading blocked list: {e}")
return set()
def load_settings() -> dict:
try:
return {**DEFAULT_SETTINGS, **json.loads(SETTINGS_FILE.read_text(encoding="utf-8"))}
except Exception as e:
log_error(f"[{now_str()}] error reading settings: {e}")
return dict(DEFAULT_SETTINGS)
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")
return text.strip() # what shortterm will read back for this line, build_payload uses it to find the message she's answering
def load_shortterm(file: Path = MEMORY_FILE) -> list[str]:
try:
@@ -105,21 +135,23 @@ def load_longterm() -> list[str]:
log_error(f"[{now_str()}] error reading longterm: {e}")
return []
def build_payload(user_input: str, rebuild: bool = False, media: list = None, logged: bool = False, dm_file: Path = None) -> str:
def build_payload(user_input: str, rebuild: bool = False, media: list = None, logged: str = None, dm_file: Path = None, tools: str = "") -> str:
global vector_db_ready
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 logged and logged in shortterm:
# cut at the message she's answering (it goes in the final turn instead), found by content and not by position:
# people keep talking while she's busy so it's not always the last line, and whatever got logged after it isn't part of this reply
shortterm = shortterm[:len(shortterm) - 1 - shortterm[::-1].index(logged)]
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
longterm = load_longterm() if should_rebuild else [] # only sent when it's needed, the backend asks for it (need_longterm) if it lost it
if should_rebuild:
vector_db_ready = True
return json.dumps({
"system_prompt": load_system_prompt(),
"system_prompt": load_system_prompt() + tools,
"user_input": user_input,
"shortterm": shortterm,
"longterm": longterm,
@@ -128,34 +160,56 @@ def build_payload(user_input: str, rebuild: bool = False, media: list = None, lo
"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
_main_ws = {} # one connection per backend uri (gemini / fallback), kept open across messages instead of reconnecting every time
_main_lock = None # one send/recv at a time on the shared sockets, made on first use so it lives in discord's loop
async def call_backend(uri: str, payload: str) -> tuple[bool, str, bool]:
# last bool is True when the backend said it's out of quota
resent = False
for attempt in range(2): # retry once with a fresh connection if the cached one turned out dead
try:
ws = _main_ws.get(uri)
if ws is None:
ws = _main_ws[uri] = await websockets.connect(uri, open_timeout=None, ping_timeout=None)
await ws.send(payload)
data = json.loads(await ws.recv())
if data.get("need_longterm") and not resent: # backend has no vector db (restarted / first time as fallback), send it the longterm once
resent = True
p = json.loads(payload)
p["longterm"] = load_longterm()
p["rebuild_vectors"] = True
payload = json.dumps(p)
await ws.send(payload)
data = json.loads(await ws.recv())
if "error" in data:
return False, data["error"], bool(data.get("quota"))
return True, data["ok"], False
except Exception as e:
_main_ws.pop(uri, None)
if attempt == 0:
continue # one retry with a fresh connection before giving up
return False, f"[{now_str()}] websocket error: {e}", False
async def call_main(payload: str) -> tuple[bool, str]:
global _main_ws, _main_lock
global _main_lock, quota_until
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}"
if not (FALLBACK_URI and time.time() < quota_until): # skip gemini while it's known to be out of quota
ok, result, quota = await call_backend(MAIN_URI, payload)
if not quota:
return ok, result
quota_until = time.time() + QUOTA_COOLDOWN
if not FALLBACK_URI:
return False, QUOTA_MSG
print(f"[{now_str()}] gemini is out of quota, using the fallback for the next {QUOTA_COOLDOWN // 60} min")
ok, result, _ = await call_backend(FALLBACK_URI, payload)
return ok, result
async def handle_error(error_msg: str, notify_channel=None, notify_user=None):
print(error_msg)
log_error(error_msg)
notice = "someone tell <@1000000000000001> there is a problem with my ai"
notice = QUOTA_MSG if error_msg == QUOTA_MSG else f"someone tell <@{load_settings()['admin_id']}> there is a problem with my ai"
try:
if notify_channel:
await notify_channel.send(notice)
@@ -347,71 +401,83 @@ async def random_message():
@client.event
async def on_ready():
asyncio.get_running_loop().set_exception_handler(ignore_connection_closed_ok)
for guild in client.guilds:
sync_reactions(guild)
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.
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 media_files:
options.append(("media", 0.07)) # ~30% combined chance of firing when all three are available
if not options:
return False
roll = random.random()
cumulative = 0.0
for kind, weight in options:
cumulative += weight
if roll >= cumulative:
continue
if kind == "sticker":
chosen = random.choice(stickers)
def load_reactions() -> dict:
# re-read every time so admin.py edit reactions applies without a restart
try:
data = json.loads(REACTIONS_FILE.read_text(encoding="utf-8"))
except (FileNotFoundError, json.JSONDecodeError):
data = {}
for k in ("stickers", "emojis", "media"):
data.setdefault(k, {}) # name -> short description
return data
def sync_reactions(guild):
# anything new gets added with an empty description, fill them in with admin.py edit reactions
data = load_reactions()
changed = False
for key, names in (("stickers", [s.name for s in guild.stickers]), ("emojis", [e.name for e in guild.emojis]), ("media", sorted(saved_media))):
for n in names:
if n not in data[key]:
data[key][n] = ""
changed = True
if changed:
REACTIONS_FILE.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
def reaction_tool_prompt(guild) -> str:
# the "tool": instead of (or along with) words she can write a [sticker: name] / [emoji: name] / [media: name] tag and
# split_reactions picks it back out of the reply. same call as the reply, so no extra latency or gemini quota.
sync_reactions(guild)
data = load_reactions()
have = {
"stickers": {s.name for s in guild.stickers},
"emojis": {e.name for e in guild.emojis},
"media": {n for n in saved_media if (MEDIA_DIR / n).is_file()},
}
sections = []
for key in ("stickers", "emojis", "media"):
options = [(n, d) for n, d in data[key].items() if n in have[key] and (d or key != "media")] # media without a description tells her nothing
if not options:
continue
options = random.sample(options, min(len(options), MAX_REACTION_OPTIONS))
sections.append(f"{key}:\n" + "\n".join(f"- {n}: {d}" if d else f"- {n}" for n, d in options))
if not sections:
return ""
return ("\n\nREACTIONS: you can send a sticker, an emoji reaction or a media file instead of (or along with) words, but only when it really fits the moment, most of the time just reply normally. "
"to do it put ONE tag at the end of your message: [sticker: name], [emoji: name] or [media: name], using exactly one of the names below. "
"the tag alone is fine if you have nothing to say. never invent names.\n" + "\n".join(sections))
def split_reactions(text: str) -> tuple[str, list]:
return REACTION_TAG.sub("", text).strip(), [(k.lower(), n) for k, n in REACTION_TAG.findall(text)]
async def send_reaction(message, kind: str, name: str):
# looks the name up for real, a made up name just does nothing
name = name.strip().strip(":").lower()
try:
if kind == "sticker":
chosen = next((s for s in message.guild.stickers if s.name.lower() == name), None)
if chosen:
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 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"Sticker send failed: {e}")
return True
if kind == "emoji":
try:
await message.add_reaction(random.choice(emojis))
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"Emoji reaction failed: {e}")
return True
if kind == "media":
chosen = random.choice(media_files)
try:
elif kind == "emoji":
chosen = next((e for e in message.guild.emojis if e.name.lower() == name), None)
if chosen:
await message.add_reaction(chosen)
elif kind == "media":
chosen = next((n for n in saved_media if n.lower() == name and (MEDIA_DIR / n).is_file()), None) # only names we saved ourselves
if chosen:
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
await handle_error(f"[{now_str()}] failed to send {kind} {name}: {e}", notify_channel=message.channel)
@client.event
@@ -426,7 +492,7 @@ async def on_message(message):
# update last activity so random_message knows the server isn't quiet
last_server_message = time.time()
if str(message.author.id) in block_id:
if str(message.author.id) in load_blocked():
return
# drop guild messages outside allowed channels before doing anything else, thank you claude
@@ -452,15 +518,15 @@ async def on_message(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}"', mem_file)
logged_line = 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}', mem_file)
logged_line = 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}', mem_file)
logged_line = log_message(f'[{now_str()}] {message.author} said: {content}', mem_file)
else:
log_message(f'[{now_str()}] {message.author} said: {content}', mem_file)
logged_line = log_message(f'[{now_str()}] {message.author} said: {content}', mem_file)
interaction_count += 1
rebuild = (interaction_count % 150 == 0)
@@ -471,7 +537,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, await grab_media(message), True, mem_file) # a dm is always aimed at her
payload = build_payload(user_input, rebuild, await grab_media(message), logged_line, 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)
@@ -498,22 +564,20 @@ async def on_message(message):
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
settings = load_settings()
if str(message.author.id) == settings["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 await on_random(message): # ~30% chance she just sends a sticker/emoji/media instead
return
if replied_message:
user_input = f'{message.author} replied to {replied_message.author} saying: "{replied_content}" with reply: "{content}"'
else:
user_input = f"{message.author} said: {content}"
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)
payload = build_payload(user_input, rebuild, media, logged_line, tools=reaction_tool_prompt(message.guild))
stop_typing = asyncio.Event()
asyncio.create_task(keep_typing(message.channel, stop_typing))
ok, result = await call_main(payload)
@@ -521,18 +585,22 @@ async def on_message(message):
if not ok:
await handle_error(result, notify_channel=message.channel)
return
if not result.strip():
result, reactions = split_reactions(result) # she can answer with a [sticker: x] style tag instead of words
if not result.strip() and not reactions:
return
try:
if result.strip():
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:
if str(message.author.id) == settings["moralfag_id"]:
moralfag_counter += 1
if moralfag_counter >= faglimit:
if moralfag_counter >= settings["faglimit"]:
moralfag_skip = True
moralfag_counter = 0
for kind, name in reactions[:1]: # one is enough
await send_reaction(message, kind, name)
except Exception as e:
await handle_error(f"[{now_str()}] failed to send reply: {e}", notify_channel=message.channel)
+48 -13
View File
@@ -9,12 +9,16 @@ 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
from google.genai import types, errors
from utils import nots, now_str, make_chunks, pick_chunks
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
RAG_THRESHOLD = 0.6 # chunks scoring under this aren't injected at all, watch the "best rag score" print and tweak it
PORT = int(os.environ.get("PORT", 8765))
RETRIES = 5 # how many times we retry a 503 (5s apart) before giving up
QUOTA_MSG = "usage limit reached, see you in a month"
# resolve config relative to this script, not whatever dir it happens to get launched from
SCRIPT_DIR = Path(__file__).resolve().parent
@@ -51,6 +55,8 @@ SAFETY_SETTINGS = [
# Each element in the VECTOR_DB will be a tuple (chunk, embedding)
VECTOR_DB = []
RAG_CHUNKS = [] # chunks from the last rebuild, so a partial (quota limited) rebuild can be finished without the bot resending everything
rag_loaded = False # False until the bot sends longterm at least once, after a restart we ask for it
def embed(texts, task_type):
# batches of 100 since that's the api limit per call
@@ -77,7 +83,9 @@ def retrieve(query, top_n=5):
similarity = cosine_similarity(query_embedding, embedding)
similarities.append((chunk, similarity))
similarities.sort(key=lambda x: x[1], reverse=True)
return similarities[:top_n]
if similarities:
print(f"[{now_str()}] best rag score: {similarities[0][1]:.2f}")
return pick_chunks(similarities, top_n, RAG_THRESHOLD)
def rebuild_vector_db(longterm_lines: list[str]):
global VECTOR_DB
@@ -97,6 +105,7 @@ def rebuild_vector_db(longterm_lines: list[str]):
async def geminichat(websocket):
global RAG_CHUNKS, rag_loaded
# 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:
@@ -110,14 +119,21 @@ async def geminichat(websocket):
system_prompt = payload["system_prompt"]
user_input = payload["user_input"]
shortterm = payload.get("shortterm", []) # list of strings
longterm = payload.get("longterm", []) # list of strings
longterm = payload.get("longterm", []) # list of strings, only there on a rebuild
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
# bot only sends longterm on a rebuild, if we never got it (just restarted) ask for it and the bot resends
if payload.get("rebuild_vectors"):
RAG_CHUNKS = make_chunks(longterm)
rag_loaded = True
elif not rag_loaded:
await websocket.send(json.dumps({"need_longterm": True}))
continue
# rebuild vector db if bot signals it (every 150 interactions bot sends full longterm)
# 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)
if payload.get("rebuild_vectors") or len(VECTOR_DB) < len(RAG_CHUNKS):
await asyncio.to_thread(rebuild_vector_db, RAG_CHUNKS)
print(f"[{now_str()}] Vector DB reconciled, {len(VECTOR_DB)} chunks")
user_query = nots(user_input)
@@ -131,7 +147,7 @@ async def geminichat(websocket):
# 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)
rag_block = "\n\nOlder relevant messages for context:\n" + "\n".join(f"- {chunk.replace(chr(10), ' | ')}" for chunk, _ in retrieved_knowledge)
# lines to skip — not real dialogue
skip_fragments = ("sent sticker:", "sent media:", "sent a random dm to ", "saved media:")
@@ -160,7 +176,8 @@ async def geminichat(websocket):
text = rest.rsplit(' with reply: "', 1)[-1]
if text.endswith('"'):
text = text[:-1]
return speaker.strip(), text.strip()
target = rest[len("replied to "):].split(" saying: ", 1)[0] # keep who they were talking to, otherwise replies lose it
return speaker.strip(), f"(replying to {target}) {text.strip()}"
if rest.startswith("said: "):
return speaker.strip(), rest[len("said: "):].strip()
return None, None
@@ -196,7 +213,7 @@ async def geminichat(websocket):
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)}"})
messages.append({"role": "user", "content": f"[responding to]: {nots(user_input)}\n(that is the newest message and the only one you answer, the lines above are just context. only the name that starts it is speaking to you, everyone else is someone else)"})
# make sure we start with a user turn (gemini rejects empty parts so no empty string here)
if messages and messages[0]["role"] != "user":
@@ -220,13 +237,31 @@ async def geminichat(websocket):
safety_settings=SAFETY_SETTINGS,
)
quota_hit = False
async with chat_semaphore:
for attempt in range(RETRIES + 1):
try:
response = await asyncio.to_thread(
client.models.generate_content,
model=chatmodel,
contents=contents,
config=config,
)
break
except errors.APIError as e:
if e.code == 429 and "PerMinute" not in str(e): # out of quota for real, not just a rate limit blip
quota_hit = True
break
if (e.code == 503 or e.code == 429) and attempt < RETRIES: # overloaded / rate limited, usually clears up in a few secs
print(f"[{now_str()}] gemini {e.code}, retrying the same request in 5s ({attempt + 1}/{RETRIES})")
await asyncio.sleep(5) # sleeping with the semaphore held on purpose, the retry goes first
continue
raise
if quota_hit:
print(f"[{now_str()}] gemini is out of quota")
await websocket.send(json.dumps({"error": QUOTA_MSG, "quota": True})) # bot.py switches to its fallback on this
continue
# response.text is None when gemini blocks the reply anyway
if not response.text:
@@ -253,8 +288,8 @@ async def geminichat(websocket):
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
print(f"[{now_str()}] hosting on ws://0.0.0.0:{PORT}")
async with serve(geminichat, "0.0.0.0", PORT, ping_interval=None, ping_timeout=None, max_size=32 * 1024 * 1024): # default 1mb would kill any media payload
await asyncio.Future()
+23 -11
View File
@@ -3,15 +3,18 @@ import asyncio
import base64
import hmac
import json
import os
import re # thank you claude
import ollama
import websockets
from pathlib import Path
from websockets.server import serve
from utils import nots, now_str
from utils import nots, now_str, make_chunks, pick_chunks
chatmodel = "hf.co/mradermacher/Huihui-gemma-4-E4B-it-abliterated-GGUF:IQ4_XS"
embedmodel = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf'
RAG_THRESHOLD = 0.65 # chunks scoring under this aren't injected at all, watch the "best rag score" print and tweak it
PORT = int(os.environ.get("PORT", 8765)) # PORT=8766 python3 main.py when it runs as the fallback next to gemini.py
# resolve config relative to this script, not whatever dir it happens to get launched from
SCRIPT_DIR = Path(__file__).resolve().parent
@@ -31,6 +34,7 @@ chat_semaphore = asyncio.Semaphore(1)
# Each element in the VECTOR_DB will be a tuple (chunk, embedding)
# The embedding is a list of floats, for example: [0.1, 0.04, -0.34, 0.21, ...]
VECTOR_DB = []
rag_loaded = False # False until the bot sends longterm at least once, after a restart we ask for it
def add_chunk_to_database(chunk):
embedding = ollama.embed(model=embedmodel, input=chunk)['embeddings'][0]
@@ -49,7 +53,9 @@ def retrieve(query, top_n=5):
similarity = cosine_similarity(query_embedding, embedding)
similarities.append((chunk, similarity))
similarities.sort(key=lambda x: x[1], reverse=True)
return similarities[:top_n]
if similarities:
print(f"[{now_str()}] best rag score: {similarities[0][1]:.2f}")
return pick_chunks(similarities, top_n, RAG_THRESHOLD)
def rebuild_vector_db(longterm_lines: list[str]):
global VECTOR_DB
@@ -67,6 +73,7 @@ def rebuild_vector_db(longterm_lines: list[str]):
async def ollamachat(websocket):
global rag_loaded
# 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,
# which was breaking bot.py's persistent connection after every single reply
@@ -81,13 +88,17 @@ async def ollamachat(websocket):
system_prompt = payload["system_prompt"]
user_input = payload["user_input"]
shortterm = payload.get("shortterm", []) # list of strings
longterm = payload.get("longterm", []) # list of strings
longterm = payload.get("longterm", []) # list of strings, only there on a rebuild
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
if payload.get("rebuild_vectors") or (not VECTOR_DB and longterm):
await asyncio.to_thread(rebuild_vector_db, longterm)
# if we never got longterm (just restarted, or first time as a fallback) ask for it and the bot resends
if payload.get("rebuild_vectors"):
await asyncio.to_thread(rebuild_vector_db, make_chunks(longterm))
rag_loaded = True
elif not rag_loaded:
await websocket.send(json.dumps({"need_longterm": True}))
continue
print(f"[{now_str()}] Vector DB reconciled, {len(VECTOR_DB)} chunks")
user_query = nots(user_input)
@@ -96,7 +107,7 @@ async def ollamachat(websocket):
# 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)
rag_block = "\n\nOlder relevant messages for context:\n" + "\n".join(f"- {chunk.replace(chr(10), ' | ')}" for chunk, _ in retrieved_knowledge)
# lines to skip — not real dialogue
skip_fragments = ("sent sticker:", "sent media:", "sent a random dm to ", "saved media:")
@@ -125,7 +136,8 @@ async def ollamachat(websocket):
text = rest.rsplit(' with reply: "', 1)[-1]
if text.endswith('"'):
text = text[:-1]
return speaker.strip(), text.strip()
target = rest[len("replied to "):].split(" saying: ", 1)[0] # keep who they were talking to, otherwise replies lose it
return speaker.strip(), f"(replying to {target}) {text.strip()}"
if rest.startswith("said: "):
return speaker.strip(), rest[len("said: "):].strip()
return None, None
@@ -161,7 +173,7 @@ async def ollamachat(websocket):
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)}"})
messages.append({"role": "user", "content": f"[responding to]: {nots(user_input)}\n(that is the newest message and the only one you answer, the lines above are just context. only the name that starts it is speaking to you, everyone else is someone else)"})
# 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/")]
@@ -205,8 +217,8 @@ 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, max_size=32 * 1024 * 1024): # default 1mb would kill any media payload
print(f"[{now_str()}] hosting on ws://0.0.0.0:{PORT}")
async with serve(ollamachat, "0.0.0.0", PORT, ping_interval=None, ping_timeout=None, max_size=32 * 1024 * 1024): # default 1mb would kill any media payload
await asyncio.Future()
+11 -3
View File
@@ -63,11 +63,12 @@ def write_tokens_file(config_dir: Path):
print(f"wrote {config_dir / 'tokens.txt'}")
def write_connection_file(config_dir: Path, default_uri: str = "ws://127.0.0.1:8765"):
def write_connection_file(config_dir: Path, default_uri: str = "ws://127.0.0.1:8765", ask_fallback: bool = False):
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
fallback = old[2] if len(old) > 2 else None
print("\n-- connection to main.py --")
uri = ask_keep("main.py address (ws:// or wss://)", old_uri)
if old_token:
@@ -78,7 +79,14 @@ def write_connection_file(config_dir: Path, default_uri: str = "ws://127.0.0.1:8
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")
if ask_fallback:
print("\n-- fallback (optional) --")
if ask_yn("use a local ollama main.py as the fallback when gemini runs out of quota"):
fallback = ask_keep("fallback main.py address", fallback or "ws://127.0.0.1:8766")
print("start it with: PORT=8766 python3 main.py (it uses the same auth token, from config/auth_token.txt)")
else:
fallback = None
(config_dir / "connection.txt").write_text(f"{uri}\n{token}\n" + (f"{fallback}\n" if fallback else ""), encoding="utf-8")
print(f"wrote {config_dir / 'connection.txt'}")
return token
@@ -124,7 +132,7 @@ def main():
main_config = SCRIPT_DIR / "config"
backend = ask_backend()
write_tokens_file(bot_config)
token = write_connection_file(bot_config, default_uri="ws://127.0.0.1:8765")
token = write_connection_file(bot_config, default_uri="ws://127.0.0.1:8765", ask_fallback=(backend == "gemini"))
write_auth_token_file(main_config, token=token)
if backend == "gemini":
write_gemini_key_file(main_config)
+41
View File
@@ -36,3 +36,44 @@ def tail_lines(path: Path, n: int) -> list[str]:
break
block_size *= 2
return [l.decode("utf-8", errors="ignore").strip() for l in lines[-n:] if l.strip()]
CHUNK_LINES = 4 # consecutive messages per rag chunk
MIN_CHUNK_CHARS = 40 # windows shorter than this are basically "lol" "yeah", not worth embedding
SKIP_FRAGMENTS = ("sent sticker:", "sent media:", "sent a random dm to ", "saved media:")
def make_chunks(lines: list[str]) -> list[str]:
# rag chunks are windows of consecutive messages instead of single lines, one-liners embed badly.
# the step is 1 so when the tail of memory.txt slides, the old windows keep the exact same text and don't get re-embedded
kept = []
for l in lines:
c = l.strip()
if any(f in c for f in SKIP_FRAGMENTS):
continue
if "said:" in c and c.split("said:", 1)[1].strip() == "":
continue
kept.append(c)
if len(kept) <= CHUNK_LINES:
windows = ["\n".join(kept)] if kept else []
else:
windows = ["\n".join(kept[i:i + CHUNK_LINES]) for i in range(len(kept) - CHUNK_LINES + 1)]
return [w for w in windows if len(w) >= MIN_CHUNK_CHARS]
def pick_chunks(scored: list, top_n: int, threshold: float) -> list:
# scored is (chunk, similarity) sorted best first. drops anything under the threshold and skips windows
# that share a line with one already picked (they overlap, otherwise the top 5 would be the same messages)
picked = []
used = set()
for chunk, sim in scored:
if sim < threshold:
break
lines = set(chunk.split("\n"))
if lines & used:
continue
used |= lines
picked.append((chunk, sim))
if len(picked) >= top_n:
break
return picked