bunch of stuff

1 files changed, 28 insertions, 15 deletions

diff --git a/main.py b/main.py
index 88865a7..44cbaf4 100644
--- a/main.py
+++ b/main.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 import asyncio
+import base64
+import hmac
 import json
 import re  # thank you claude
 import ollama
@@ -8,7 +10,7 @@ from pathlib import Path
 from websockets.server import serve
 from utils import nots, now_str
 
-chatmodel = "hf.co/Stetsoncody/Llama-3.2-3B-Instruct-abliterated-GGUF:Q4_K_M"
+chatmodel = "hf.co/mradermacher/Huihui-gemma-4-E4B-it-abliterated-GGUF:IQ4_XS"
 embedmodel = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf'
 
 # resolve config relative to this script, not whatever dir it happens to get launched from
@@ -72,7 +74,7 @@ async def ollamachat(websocket):
         try:
             payload = json.loads(raw)
 
-            if payload.get("auth_token") != AUTH_TOKEN:
+            if not hmac.compare_digest(str(payload.get("auth_token", "")).encode(), AUTH_TOKEN.encode()):
                 await websocket.send(json.dumps({"error": "bad auth token"}))
                 continue
 
@@ -80,6 +82,7 @@ async def ollamachat(websocket):
             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
@@ -96,7 +99,7 @@ async def ollamachat(websocket):
                 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:")
+            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):
@@ -112,17 +115,19 @@ async def ollamachat(websocket):
             def parse_line(line):
                 c = nots(line).strip()
                 # "X said: Y", "X said in DM: Y", "X replied to ... with reply: Y"
-                if " said in DM: " in c:
-                    speaker, text = c.split(" said in DM: ", 1)
-                    return speaker.strip(), text.strip()
-                if " with reply: " in c:
+                # 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""
-                    speaker = c.split(" replied to ", 1)[0].strip()
-                    text = c.rsplit(' with reply: "', 1)[-1].rstrip('"').strip()
-                    return speaker.strip(), text.strip()
-                if " said: " in c:
-                    speaker, text = c.split(" said: ", 1)
+                    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/assistant turns from shortterm
@@ -158,13 +163,21 @@ async def ollamachat(websocket):
             # final isolated turn — the specific message she needs to reply to
             messages.append({"role": "user", "content": f"[responding to]: {nots(user_input)}"})
 
+            # 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/")]
+            if images:
+                messages[-1]["images"] = images
+
             # make sure we start with a user turn (model requirement)
             if messages and messages[0]["role"] != "user":
                 messages.insert(0, {"role": "user", "content": ""})
 
             messages.insert(0, {"role": "system", "content": system_prompt + rag_block})
 
-            params = json.loads((CONFIG_DIR / "params.json").read_text())
+            try:
+                params = json.loads((CONFIG_DIR / "params.json").read_text())
+            except FileNotFoundError:
+                params = {}  # no params.json, just use ollama's defaults
 
             async with chat_semaphore:
                 response = await asyncio.to_thread(
@@ -177,7 +190,7 @@ async def ollamachat(websocket):
             # strip <think> blocks in case the model ignores think=False, thank you claude
             bot_reply = re.sub(r'<think>.*?</think>', '', response["message"]["content"], flags=re.DOTALL).strip()
 
-            print(messages)
+            print([{k: v for k, v in m.items() if k != "images"} for m in messages])  # skip the raw image bytes
             print(f"[{now_str()}] reply: {bot_reply}\n")
 
             await websocket.send(json.dumps({"ok": bot_reply}))
@@ -193,7 +206,7 @@ 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):
+    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
         await asyncio.Future()