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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
|
#!/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 datetime import datetime
chatmodel = "hf.co/Stetsoncody/Llama-3.2-3B-Instruct-abliterated-GGUF:Q4_K_M"
embedmodel = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf'
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 []
# 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(Path("config/params.json").read_text())
response = 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"[{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())
|
