diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 84 |
1 files changed, 72 insertions, 12 deletions
@@ -82,19 +82,79 @@ async def ollamachat(websocket): user_query = nots(user_input) retrieved_knowledge = retrieve(user_query) if VECTOR_DB else [] - context_message = ( - "These are memory notes to help you continue naturally. " - "Do NOT imitate tone, style, structure, or wording.\n\n" - "Short-term memory, these are the most 15 recent chats sent by various people:\n" - + "\n".join(f"- {line}" for line in shortterm) - + "\n\nLong-term memory, these are retrieved older messages using cosine similarity:\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 + + 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 + + # 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 + + 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)}) + + # 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": ""}) - messages = [ - {"role": "system", "content": system_prompt + "\n\n" + context_message}, - {"role": "user", "content": user_input}, - ] + messages.insert(0, {"role": "system", "content": system_prompt + rag_block}) response = ollama.chat( model=chatmodel, |