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
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
|
#!/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):
# loop over every message on this connection instead of handling one and returning —
# returning from this function is what makes the websockets lib close the connection,
# which was breaking bot.py's persistent connection after every single reply
async for raw in websocket:
try:
payload = json.loads(raw)
if payload.get("auth_token") != AUTH_TOKEN:
await websocket.send(json.dumps({"error": "bad auth token"}))
continue
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 line in shortterm:
if is_skip(line):
continue
speaker, text = parse_line(line)
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())
|
