aboutsummaryrefslogtreecommitdiff
path: root/main.py
blob: 6091a7df341a9e621c16c04958c4d803480965ad (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
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())