faggot-chan/main.py

#!/usr/bin/env python3
import asyncio
import json
import re  # thank you claude
import ollama
import websockets
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"
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
SCRIPT_DIR = Path(__file__).resolve().parent
CONFIG_DIR = SCRIPT_DIR / "config"

try:
    AUTH_TOKEN = (CONFIG_DIR / "auth_token.txt").read_text(encoding="utf-8").strip()
except FileNotFoundError:
    raise SystemExit("no config/auth_token.txt found — run setup_wizard.py first")

# only one generation at a time, we've got one cpu running one model here
chat_semaphore = asyncio.Semaphore(1)


# THANK YOU GUY FROM HUGGING FACE THAT MADE ALL THIS SO THAT I DON'T HAVE TO THINK I LOVE YOU

# 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 = []

def add_chunk_to_database(chunk):
    embedding = ollama.embed(model=embedmodel, input=chunk)['embeddings'][0]
    return embedding

def cosine_similarity(a, b):
    dot_product = sum([x * y for x, y in zip(a, b)])
    norm_a = sum([x ** 2 for x in a]) ** 0.5
    norm_b = sum([x ** 2 for x in b]) ** 0.5
    return dot_product / (norm_a * norm_b)

def retrieve(query, top_n=5):
    query_embedding = ollama.embed(model=embedmodel, input=query)['embeddings'][0]
    similarities = []
    for chunk, embedding in VECTOR_DB:
        similarity = cosine_similarity(query_embedding, embedding)
        similarities.append((chunk, similarity))
    similarities.sort(key=lambda x: x[1], reverse=True)
    return similarities[:top_n]

def rebuild_vector_db(longterm_lines: list[str]):
    global VECTOR_DB
    existing = dict(VECTOR_DB)  # chunk -> embedding, so we don't re-embed stuff we already have
    new_db = []
    added = 0
    for i, chunk in enumerate(longterm_lines):
        embedding = existing.get(chunk)
        if embedding is None:
            embedding = add_chunk_to_database(chunk)
            added += 1
            print(f"[{now_str()}] embedded new chunk {added} ({i+1}/{len(longterm_lines)})")
        new_db.append((chunk, embedding))
    VECTOR_DB = new_db


async def ollamachat(websocket):
    try:
        # bot sends one json payload with everything main needs
        raw = await websocket.recv()
        payload = json.loads(raw)

        if payload.get("auth_token") != AUTH_TOKEN:
            await websocket.send(json.dumps({"error": "bad auth token"}))
            return

        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")

        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)

        # 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.insert(0, {"role": "system", "content": system_prompt + rag_block})

        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,
            )

        # 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")

        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


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):
        await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(main())