diff --git a/Knowledgebase/requirements.txt b/Knowledgebase/requirements.txt index 4c5d3f7..329d267 100644 --- a/Knowledgebase/requirements.txt +++ b/Knowledgebase/requirements.txt @@ -1,4 +1,5 @@ mcp[cli] psycopg[binary] +pgvector voyageai anthropic diff --git a/Knowledgebase/server.py b/Knowledgebase/server.py index ee28e5e..35f3b15 100644 --- a/Knowledgebase/server.py +++ b/Knowledgebase/server.py @@ -15,7 +15,6 @@ Env proměnné: """ import json -import math import os import sys import traceback @@ -117,14 +116,6 @@ def _indent(text: str, n: int) -> str: pad = " " * n return "\n".join(pad + line for line in text.splitlines()) -def _cosine(a: list[float], b: list[float]) -> float: - dot = sum(x * y for x, y in zip(a, b)) - na = math.sqrt(sum(x * x for x in a)) - nb = math.sqrt(sum(x * x for x in b)) - if na == 0 or nb == 0: - return 0.0 - return dot / (na * nb) - # ─── MCP server ────────────────────────────────────────────────────────────── mcp = FastMCP("knowledgebase") @@ -164,18 +155,35 @@ def store_memory( conn = get_conn() try: with conn.transaction(): - row = conn.execute( - """ - INSERT INTO kb_memories - (mem_type, title, content, summary, tags, project, - source, session_id, importance, embedding, meta) - VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) - RETURNING id, created_at - """, - (mem_type, title, content, summary, - tags or [], project, source, session_id, - importance, embedding, json.dumps(meta or {})), - ).fetchone() + if embedding: + from pgvector.psycopg import register_vector + import numpy as np + register_vector(conn) + row = conn.execute( + """ + INSERT INTO kb_memories + (mem_type, title, content, summary, tags, project, + source, session_id, importance, embedding, meta) + VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + RETURNING id, created_at + """, + (mem_type, title, content, summary, + tags or [], project, source, session_id, + importance, np.array(embedding), json.dumps(meta or {})), + ).fetchone() + else: + row = conn.execute( + """ + INSERT INTO kb_memories + (mem_type, title, content, summary, tags, project, + source, session_id, importance, meta) + VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + RETURNING id, created_at + """, + (mem_type, title, content, summary, + tags or [], project, source, session_id, + importance, json.dumps(meta or {})), + ).fetchone() return f"Stored memory id={row['id']} at {row['created_at']}" except Exception as e: conn.rollback() @@ -302,21 +310,41 @@ def store_conversation( def _insert_memory_in_tx(conn, data: dict): """Helper: insert memory within an existing transaction.""" - conn.execute( - """ - INSERT INTO kb_memories - (mem_type, title, content, summary, tags, project, - source, session_id, importance, embedding, meta) - VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) - """, - (data.get("mem_type","fact"), data.get("title"), - data["content"], data.get("summary"), - data.get("tags",[]), data.get("project"), - data.get("source"), data.get("session_id"), - data.get("importance",0.5), - data.get("embedding"), - json.dumps(data.get("meta",{}))), - ) + embedding = data.get("embedding") + if embedding: + from pgvector.psycopg import register_vector + import numpy as np + register_vector(conn) + conn.execute( + """ + INSERT INTO kb_memories + (mem_type, title, content, summary, tags, project, + source, session_id, importance, embedding, meta) + VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + """, + (data.get("mem_type","fact"), data.get("title"), + data["content"], data.get("summary"), + data.get("tags",[]), data.get("project"), + data.get("source"), data.get("session_id"), + data.get("importance",0.5), + np.array(embedding), + json.dumps(data.get("meta",{}))), + ) + else: + conn.execute( + """ + INSERT INTO kb_memories + (mem_type, title, content, summary, tags, project, + source, session_id, importance, meta) + VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) + """, + (data.get("mem_type","fact"), data.get("title"), + data["content"], data.get("summary"), + data.get("tags",[]), data.get("project"), + data.get("source"), data.get("session_id"), + data.get("importance",0.5), + json.dumps(data.get("meta",{}))), + ) # ────────────────────────────────────────────────────────────────────────────── @@ -388,11 +416,14 @@ def search( fts_ids = {r["id"] for r in rows} results = [_row_to_dict(r) for r in rows] - # ── Vector reranking (Python-side cosine similarity) ── - # Fetch candidates with embeddings, compute cosine similarity, merge + # ── Vector search (nativní pgvector, <=> cosine distance) ── query_emb = get_embedding(query) if query_emb: try: + from pgvector.psycopg import register_vector + import numpy as np + register_vector(conn) + vec_conditions = ["deleted = FALSE", "embedding IS NOT NULL"] vec_params2: list[Any] = [] @@ -414,24 +445,21 @@ def search( f""" SELECT id, mem_type, title, content, summary, tags, project, source, session_id, importance, created_at, - embedding + 1 - (embedding <=> %s::vector) AS score FROM kb_memories WHERE {vec_where} - LIMIT 200 + ORDER BY embedding <=> %s::vector + LIMIT %s """, - vec_params2, + [np.array(query_emb), np.array(query_emb)] + vec_params2 + [limit], ).fetchall() for r in vec_rows: - if r["id"] not in fts_ids and r["embedding"]: - sim = _cosine(query_emb, r["embedding"]) - if sim > 0.5: # threshold - d = _row_to_dict(r) - d["score"] = sim - results.append(d) + if r["id"] not in fts_ids: + results.append(_row_to_dict(r)) except Exception as e: - log(f"Vector reranking error: {e}") + log(f"Vector search error: {e}") # deduplicate & sort by score seen = set() @@ -672,8 +700,11 @@ def update_memory( params.append(content) new_emb = get_embedding(f"{title or ''} {content}") if new_emb: + from pgvector.psycopg import register_vector + import numpy as np + register_vector(conn) updates.append("embedding = %s") - params.append(new_emb) + params.append(np.array(new_emb)) if title is not None: updates.append("title = %s") params.append(title)