diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 204 |
1 files changed, 103 insertions, 101 deletions
@@ -65,128 +65,130 @@ def rebuild_vector_db(longterm_lines: list[str]): async def ollamachat(websocket): - try: - # bot sends one json payload with everything main needs - raw = await websocket.recv() - payload = json.loads(raw) + # 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 + async for raw in websocket: + try: + payload = json.loads(raw) - if payload.get("auth_token") != AUTH_TOKEN: - await websocket.send(json.dumps({"error": "bad auth token"})) - return + if payload.get("auth_token") != AUTH_TOKEN: + 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 + system_prompt = payload["system_prompt"] + user_input = payload["user_input"] + shortterm = payload.get("shortterm", []) # list of strings + longterm = payload.get("longterm", []) # list of strings - # 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) - print(f"[{now_str()}] Vector DB reconciled, {len(VECTOR_DB)} chunks") + # 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) + print(f"[{now_str()}] Vector DB reconciled, {len(VECTOR_DB)} chunks") - user_query = nots(user_input) - retrieved_knowledge = await asyncio.to_thread(retrieve, user_query) if VECTOR_DB else [] + user_query = nots(user_input) + retrieved_knowledge = await asyncio.to_thread(retrieve, user_query) if VECTOR_DB else [] - # 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 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 + # 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" + 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" - if " said in DM: " in c: - speaker, text = c.split(" said in DM: ", 1) - return speaker.strip(), text.strip() - if " with reply: " in c: - # "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) - return speaker.strip(), text.strip() - return None, None + # 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" + if " said in DM: " in c: + speaker, text = c.split(" said in DM: ", 1) + return speaker.strip(), text.strip() + if " with reply: " in c: + # "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) + return speaker.strip(), text.strip() + return None, None - # build alternating user/assistant turns from shortterm - # group consecutive non-bot lines into one user blob, bot lines become assistant turns - messages = [] - current_user_lines = [] - seen_last = None # for dedup + # build alternating user/assistant turns from shortterm + # group consecutive non-bot lines into one user blob, bot lines become assistant turns + messages = [] + current_user_lines = [] + seen_last = None # for dedup - for raw in shortterm: - if is_skip(raw): - continue - speaker, text = parse_line(raw) - if not speaker or not text: - continue - entry = f"{speaker}: {text}" - if entry == seen_last: # skip duplicate consecutive lines - continue - seen_last = entry + 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": "assistant", "content": text}) - else: - current_user_lines.append(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": "assistant", "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)}) + # 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)}"}) + # 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 (model requirement) - if messages and messages[0]["role"] != "user": - messages.insert(0, {"role": "user", "content": ""}) + # 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}) + messages.insert(0, {"role": "system", "content": system_prompt + rag_block}) - params = json.loads((CONFIG_DIR / "params.json").read_text()) + params = json.loads((CONFIG_DIR / "params.json").read_text()) - async with chat_semaphore: - response = await asyncio.to_thread( - ollama.chat, - model=chatmodel, - messages=messages, - options=params, - ) + async with chat_semaphore: + response = await asyncio.to_thread( + ollama.chat, + model=chatmodel, + messages=messages, + options=params, + ) - # 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() + # 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(f"[{now_str()}] reply: {bot_reply}\n") + print(messages) + print(f"[{now_str()}] reply: {bot_reply}\n") - await websocket.send(json.dumps({"ok": bot_reply})) + await websocket.send(json.dumps({"ok": bot_reply})) - except Exception as e: - error_msg = f"[{now_str()}] error in ollamachat: {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 + except Exception as e: + error_msg = f"[{now_str()}] error in ollamachat: {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(): |