hello

1 files changed, 54 insertions, 19 deletions

diff --git a/gemini.py b/gemini.py
index f7da375..e73f9d4 100644
--- a/gemini.py
+++ b/gemini.py
@@ -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:
-                response = await asyncio.to_thread(
-                    client.models.generate_content,
-                    model=chatmodel,
-                    contents=contents,
-                    config=config,
-                )
+                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()