diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 73 |
1 files changed, 43 insertions, 30 deletions
@@ -6,19 +6,22 @@ import ollama import websockets from pathlib import Path from websockets.server import serve -from datetime import datetime +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" -def times() -> str: - return datetime.now().strftime("%Y-%m-%d %H:%M:%S") +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") -def nots(line: str) -> str: - if line.startswith("[") and "]" in line: - return line.split("]", 1)[1].lstrip() - return line +# 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 @@ -29,7 +32,7 @@ VECTOR_DB = [] def add_chunk_to_database(chunk): embedding = ollama.embed(model=embedmodel, input=chunk)['embeddings'][0] - VECTOR_DB.append((chunk, embedding)) + return embedding def cosine_similarity(a, b): dot_product = sum([x * y for x, y in zip(a, b)]) @@ -47,35 +50,43 @@ def retrieve(query, top_n=5): return similarities[:top_n] def rebuild_vector_db(longterm_lines: list[str]): - VECTOR_DB.clear() + 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): - add_chunk_to_database(chunk) - print(f"[{times()}] Added chunk {i+1}/{len(longterm_lines)} to the database") - -interaction_count = 0 + 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): - global interaction_count 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 - 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") + # 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 = retrieve(user_query) if VECTOR_DB else [] + 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 = "" @@ -151,24 +162,26 @@ async def ollamachat(websocket): messages.insert(0, {"role": "system", "content": system_prompt + rag_block}) - params = json.loads(Path("config/params.json").read_text()) + params = json.loads((CONFIG_DIR / "params.json").read_text()) - response = ollama.chat( - model=chatmodel, - messages=messages, - options=params, - ) + 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"[{times()}] reply: {bot_reply}\n") + print(f"[{now_str()}] reply: {bot_reply}\n") await websocket.send(json.dumps({"ok": bot_reply})) except Exception as e: - error_msg = f"[{times()}] error in ollamachat: {e}" + error_msg = f"[{now_str()}] error in ollamachat: {e}" print(error_msg) try: await websocket.send(json.dumps({"error": error_msg})) @@ -177,7 +190,7 @@ async def ollamachat(websocket): async def main(): - print(f"[{times()}] hosting on ws://0.0.0.0:8765") + 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() |