308 lines
10 KiB
Python
308 lines
10 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Vector store — Qdrant remoto (DS920) con fallback SQLite locale su Pi ARM."""
|
|
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from typing import Optional
|
|
|
|
LOGGER = logging.getLogger("loogle_mcp.vector_store")
|
|
VECTOR_SIZE = 768
|
|
_local = sqlite3.connect(":memory:", check_same_thread=False) # placeholder
|
|
_qdrant_client = None
|
|
_next_qdrant_check = 0.0 # monotonic: prossimo tentativo di connessione mentre si usa il fallback
|
|
_qdrant_lock = threading.Lock()
|
|
|
|
|
|
def _fallback_path() -> str:
|
|
return os.environ.get("MCP_VECTOR_FALLBACK", "/data/vector_fallback.db")
|
|
|
|
|
|
def _fallback_conn() -> sqlite3.Connection:
|
|
path = _fallback_path()
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
conn = sqlite3.connect(path, timeout=30)
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS vectors (
|
|
id TEXT PRIMARY KEY,
|
|
collection TEXT NOT NULL,
|
|
vector TEXT NOT NULL,
|
|
payload TEXT NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_vectors_collection ON vectors(collection)")
|
|
conn.commit()
|
|
return conn
|
|
|
|
|
|
def _get_qdrant():
|
|
"""Client Qdrant, oppure None (fallback SQLite).
|
|
|
|
Se Qdrant non risponde si usa il fallback, ma si riprova ogni QDRANT_RETRY_S secondi
|
|
invece di restare sul fallback fino al riavvio.
|
|
"""
|
|
global _next_qdrant_check
|
|
if _qdrant_client is not None or time.monotonic() < _next_qdrant_check:
|
|
return _qdrant_client
|
|
with _qdrant_lock:
|
|
if _qdrant_client is not None or time.monotonic() < _next_qdrant_check:
|
|
return _qdrant_client
|
|
client = _connect_qdrant()
|
|
if client is None:
|
|
_next_qdrant_check = time.monotonic() + float(os.environ.get("QDRANT_RETRY_S", "60"))
|
|
return client
|
|
|
|
|
|
def _connect_qdrant():
|
|
global _qdrant_client
|
|
url = os.environ.get("QDRANT_URL", "").strip()
|
|
if not url:
|
|
LOGGER.warning("QDRANT_URL non impostato — fallback SQLite")
|
|
return None
|
|
try:
|
|
import httpx
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.http import models as qm
|
|
|
|
# Sonda breve: un tentativo con Qdrant giù non deve bloccare la ricerca per 20s
|
|
httpx.get(f"{url.rstrip('/')}/readyz", timeout=3.0).raise_for_status()
|
|
client = QdrantClient(url=url, timeout=20)
|
|
client.get_collections()
|
|
_qdrant_client = client
|
|
globals()["qm"] = qm
|
|
LOGGER.info("Qdrant connesso: %s", url)
|
|
return client
|
|
except Exception as exc:
|
|
LOGGER.warning("Qdrant non disponibile (%s) — fallback SQLite, nuovo tentativo tra %ss",
|
|
exc, os.environ.get("QDRANT_RETRY_S", "60"))
|
|
return None
|
|
|
|
|
|
def ensure_collection(name: str, vector_size: int = VECTOR_SIZE) -> None:
|
|
client = _get_qdrant()
|
|
if client is None:
|
|
return
|
|
from qdrant_client.http import models as qm
|
|
|
|
names = {c.name for c in client.get_collections().collections}
|
|
if name in names:
|
|
return
|
|
client.create_collection(
|
|
collection_name=name,
|
|
vectors_config=qm.VectorParams(size=vector_size, distance=qm.Distance.COSINE),
|
|
)
|
|
|
|
|
|
def kb_collection(username: str) -> str:
|
|
return f"kb_personal_{username}"
|
|
|
|
|
|
def ctx_collection(username: str) -> str:
|
|
return f"ctx_{username}"
|
|
|
|
|
|
SHARED_COLLECTION = "kb_shared_family"
|
|
GITEA_SHARED_COLLECTION = "gitea_shared_family"
|
|
APPS_SHARED_COLLECTION = "apps_shared_family"
|
|
CTX_SHARED_COLLECTION = "ctx_shared_family"
|
|
|
|
|
|
def gitea_collection(username: str) -> str:
|
|
return f"gitea_personal_{username}"
|
|
|
|
|
|
def _point_id(name: str) -> str:
|
|
return str(uuid.uuid5(uuid.NAMESPACE_URL, name))
|
|
|
|
|
|
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)) or 1.0
|
|
nb = math.sqrt(sum(x * x for x in b)) or 1.0
|
|
return dot / (na * nb)
|
|
|
|
|
|
def upsert_chunks(
|
|
collection: str,
|
|
ids: list[str],
|
|
vectors: list[list[float]],
|
|
payloads: list[dict],
|
|
) -> None:
|
|
if not vectors:
|
|
return
|
|
ensure_collection(collection, len(vectors[0]))
|
|
client = _get_qdrant()
|
|
if client is not None:
|
|
from qdrant_client.http import models as qm
|
|
|
|
points = [
|
|
qm.PointStruct(id=_point_id(pid), vector=vec, payload=payload)
|
|
for pid, vec, payload in zip(ids, vectors, payloads)
|
|
]
|
|
client.upsert(collection_name=collection, points=points)
|
|
return
|
|
conn = _fallback_conn()
|
|
for pid, vec, payload in zip(ids, vectors, payloads):
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO vectors(id,collection,vector,payload) VALUES (?,?,?,?)",
|
|
(_point_id(pid), collection, json.dumps(vec), json.dumps(payload, ensure_ascii=False)),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def delete_by_doc(collection: str, doc_id: int) -> None:
|
|
delete_where(collection, "doc_id", doc_id)
|
|
|
|
|
|
def delete_where(collection: str, key: str, value) -> None:
|
|
"""Elimina i punti il cui payload[key] == value."""
|
|
client = _get_qdrant()
|
|
if client is not None:
|
|
from qdrant_client.http import models as qm
|
|
|
|
ensure_collection(collection)
|
|
client.delete(
|
|
collection_name=collection,
|
|
points_selector=qm.FilterSelector(
|
|
filter=qm.Filter(
|
|
must=[qm.FieldCondition(key=key, match=qm.MatchValue(value=value))]
|
|
)
|
|
),
|
|
)
|
|
return
|
|
conn = _fallback_conn()
|
|
rows = conn.execute("SELECT id,payload FROM vectors WHERE collection=?", (collection,)).fetchall()
|
|
for row_id, payload_raw in rows:
|
|
payload = json.loads(payload_raw)
|
|
if payload.get(key) == value:
|
|
conn.execute("DELETE FROM vectors WHERE id=?", (row_id,))
|
|
conn.commit()
|
|
|
|
|
|
def delete_matching(collection: str, key: str, value, predicate) -> int:
|
|
"""Elimina i punti con payload[key] == value per cui predicate(payload) è vero."""
|
|
client = _get_qdrant()
|
|
if client is not None:
|
|
from qdrant_client.http import models as qm
|
|
|
|
ensure_collection(collection)
|
|
flt = qm.Filter(must=[qm.FieldCondition(key=key, match=qm.MatchValue(value=value))])
|
|
ids, offset = [], None
|
|
while True:
|
|
points, offset = client.scroll(
|
|
collection_name=collection, scroll_filter=flt, limit=256, offset=offset, with_payload=True
|
|
)
|
|
ids.extend(p.id for p in points if predicate(p.payload or {}))
|
|
if offset is None:
|
|
break
|
|
if ids:
|
|
client.delete(collection_name=collection, points_selector=qm.PointIdsList(points=ids))
|
|
return len(ids)
|
|
conn = _fallback_conn()
|
|
rows = conn.execute("SELECT id,payload FROM vectors WHERE collection=?", (collection,)).fetchall()
|
|
removed = 0
|
|
for row_id, payload_raw in rows:
|
|
payload = json.loads(payload_raw)
|
|
if payload.get(key) == value and predicate(payload):
|
|
conn.execute("DELETE FROM vectors WHERE id=?", (row_id,))
|
|
removed += 1
|
|
conn.commit()
|
|
return removed
|
|
|
|
|
|
def count_by_doc(collection: str, doc_id: int) -> int:
|
|
client = _get_qdrant()
|
|
if client is not None:
|
|
from qdrant_client.http import models as qm
|
|
|
|
ensure_collection(collection)
|
|
result = client.count(
|
|
collection_name=collection,
|
|
count_filter=qm.Filter(
|
|
must=[qm.FieldCondition(key="doc_id", match=qm.MatchValue(value=doc_id))]
|
|
),
|
|
exact=True,
|
|
)
|
|
return int(result.count)
|
|
conn = _fallback_conn()
|
|
rows = conn.execute("SELECT payload FROM vectors WHERE collection=?", (collection,)).fetchall()
|
|
count = 0
|
|
for (payload_raw,) in rows:
|
|
payload = json.loads(payload_raw)
|
|
if payload.get("doc_id") == doc_id:
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def collection_point_count(collection: str) -> int:
|
|
client = _get_qdrant()
|
|
if client is not None:
|
|
try:
|
|
info = client.get_collection(collection)
|
|
return int(info.points_count or 0)
|
|
except Exception:
|
|
return 0
|
|
conn = _fallback_conn()
|
|
row = conn.execute(
|
|
"SELECT COUNT(*) FROM vectors WHERE collection=?", (collection,)
|
|
).fetchone()
|
|
return int(row[0] if row else 0)
|
|
|
|
|
|
def search(
|
|
collections: list[str],
|
|
vector: list[float],
|
|
limit: int = 8,
|
|
visibility_filter: Optional[dict] = None,
|
|
) -> list[dict]:
|
|
results: list[dict] = []
|
|
client = _get_qdrant()
|
|
if client is not None:
|
|
from qdrant_client.http import models as qm
|
|
|
|
for collection in collections:
|
|
ensure_collection(collection, len(vector))
|
|
flt = None
|
|
if visibility_filter:
|
|
must = [
|
|
qm.FieldCondition(key=k, match=qm.MatchValue(value=v))
|
|
for k, v in visibility_filter.items()
|
|
]
|
|
if must:
|
|
flt = qm.Filter(must=must)
|
|
hits = client.search(
|
|
collection_name=collection,
|
|
query_vector=vector,
|
|
limit=limit,
|
|
query_filter=flt,
|
|
)
|
|
for hit in hits:
|
|
payload = dict(hit.payload or {})
|
|
payload["score"] = hit.score
|
|
payload["collection"] = collection
|
|
results.append(payload)
|
|
else:
|
|
conn = _fallback_conn()
|
|
for collection in collections:
|
|
rows = conn.execute(
|
|
"SELECT vector,payload FROM vectors WHERE collection=?", (collection,)
|
|
).fetchall()
|
|
for vec_raw, payload_raw in rows:
|
|
payload = dict(json.loads(payload_raw))
|
|
if visibility_filter:
|
|
if any(payload.get(k) != v for k, v in visibility_filter.items()):
|
|
continue
|
|
score = _cosine(vector, json.loads(vec_raw))
|
|
payload["score"] = score
|
|
payload["collection"] = collection
|
|
results.append(payload)
|
|
results.sort(key=lambda x: x.get("score", 0), reverse=True)
|
|
return results[:limit]
|