fixed a couple bugs, moved a couple of this. also made main stateless so i can do funny things

1 files changed, 136 insertions, 0 deletions

diff --git a/main.py b/main.py
new file mode 100644
index 0000000..6091a7d
--- /dev/null
+++ b/main.py
@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+import asyncio
+import json
+import ollama
+import websockets
+from websockets.server import serve
+from datetime import datetime
+
+#chatmodel = "hf.co/Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2-GGUF:Llama-3.1-8B-Lexi-Uncensored_V2_Q4.gguf"
+embedmodel = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf'
+
+# these are test models just ignore
+chatmodel = "hf.co/tvall43/Qwen3.5-4B-heretic-gguf:Q4_K_M"
+#chatmodel = 'hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF'
+#chatmodel = 'hf.co/VibeStudio/Nidum-Llama-3.2-3B-Uncensored-GGUF'
+#chatmodel = "faggotchan"
+#chatmodel = "deepseek-r1:1.5b"
+
+NUM_THREADS = 6  # hardcoded since threads.txt is gone, change if needed
+
+def times() -> str:
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+def nots(line: str) -> str:
+    if line.startswith("[") and "]" in line:
+        return line.split("]", 1)[1].lstrip()
+    return line
+
+
+# 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]
+    VECTOR_DB.append((chunk, 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]):
+    VECTOR_DB.clear()
+    for i, chunk in enumerate(longterm_lines):
+        add_chunk_to_database(chunk)
+        print(f"[{times()}] Added chunk {i+1}/{len(longterm_lines)} to the database")
+
+interaction_count = 0
+
+
+async def ollamachat(websocket):
+    global interaction_count
+    try:
+        # bot sends one json payload with everything main needs
+        raw = await websocket.recv()
+        payload = json.loads(raw)
+
+        system_prompt = payload["system_prompt"]
+        user_input    = payload["user_input"]
+        shortterm     = payload.get("shortterm", [])   # list of strings
+        longterm      = payload.get("longterm", [])    # list of strings
+
+        interaction_count += 1
+
+        # rebuild vector db if bot signals it (every 150 interactions bot sends full longterm)
+        if payload.get("rebuild_vectors"):
+            rebuild_vector_db(longterm)
+            print(f"[{times()}] Vector DB rebuilt with {len(longterm)} chunks")
+
+        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)
+        )
+
+        messages = [
+            {"role": "system", "content": system_prompt + "\n\n" + context_message},
+            {"role": "user",   "content": user_input},
+        ]
+
+        response = ollama.chat(
+            model=chatmodel,
+            messages=messages,
+            options={
+                "temperature":    0.8,
+                "top_p":          0.9,
+                "top_k":          40,
+                "repeat_penalty": 1.1,
+                "num_predict":    4096,
+                "num_thread":     NUM_THREADS,
+            }
+        )
+
+        bot_reply = response["message"]["content"].strip()
+
+        print(messages)
+        print(f"[{times()}] reply: {bot_reply}\n")
+
+        await websocket.send(json.dumps({"ok": bot_reply}))
+
+    except Exception as e:
+        error_msg = f"[{times()}] 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"[{times()}] 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())