Upgrade Knowledgebase to native pgvector (VECTOR type + ivfflat index)
- Migrated embedding column from double precision[] to VECTOR(1024) - Now uses native <=> cosine operator for SQL-level vector search - Added pgvector to requirements - Fixed collation mismatch on all DBs after pgvector/pgvector:pg18 image swap Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
mcp[cli]
|
||||
psycopg[binary]
|
||||
pgvector
|
||||
voyageai
|
||||
anthropic
|
||||
|
||||
+55
-24
@@ -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,6 +155,10 @@ def store_memory(
|
||||
conn = get_conn()
|
||||
try:
|
||||
with conn.transaction():
|
||||
if embedding:
|
||||
from pgvector.psycopg import register_vector
|
||||
import numpy as np
|
||||
register_vector(conn)
|
||||
row = conn.execute(
|
||||
"""
|
||||
INSERT INTO kb_memories
|
||||
@@ -174,7 +169,20 @@ def store_memory(
|
||||
""",
|
||||
(mem_type, title, content, summary,
|
||||
tags or [], project, source, session_id,
|
||||
importance, embedding, json.dumps(meta or {})),
|
||||
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:
|
||||
@@ -302,6 +310,11 @@ def store_conversation(
|
||||
|
||||
def _insert_memory_in_tx(conn, data: dict):
|
||||
"""Helper: insert memory within an existing transaction."""
|
||||
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
|
||||
@@ -314,7 +327,22 @@ def _insert_memory_in_tx(conn, data: dict):
|
||||
data.get("tags",[]), data.get("project"),
|
||||
data.get("source"), data.get("session_id"),
|
||||
data.get("importance",0.5),
|
||||
data.get("embedding"),
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user