diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 43 |
1 files changed, 28 insertions, 15 deletions
@@ -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() |