hello

1 files changed, 23 insertions, 11 deletions

diff --git a/main.py b/main.py
index 44cbaf4..2c188a6 100644
--- a/main.py
+++ b/main.py
@@ -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()