bunch of stuff

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 <think> blocks just in case, thank you claude
+            bot_reply = re.sub(r'<think>.*?</think>', '', 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())