Initial commit: Pegas Memory (PG + Qdrant + Memgraph, MCP server, 14 tools)

pegas_memory/qdrant_store.py

@@ -0,0 +1,57 @@
+"""Qdrant — векторный срез (EmbeddingGemma 2, dim=768)."""
+from __future__ import annotations
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
+from . import config
+from .embedding import embed
+
+_client: QdrantClient | None = None
+
+def client() -> QdrantClient:
+    global _client
+    if _client is None:
+        _client = QdrantClient(url=config.QDRANT_URL)
+    return _client
+
+def ensure_collection():
+    c = client()
+    try:
+        c.get_collection(config.QDRANT_COLLECTION)
+    except Exception:
+        c.create_collection(
+            collection_name=config.QDRANT_COLLECTION,
+            vectors_config=VectorParams(size=config.EMBED_DIM, distance=Distance.COSINE),
+        )
+
+def upsert_atom(atom_id: int, text: str, payload: dict) -> None:
+    """text — что эмбеддим; payload — type/confidence/status/doc_id/claim."""
+    vec = embed(text)
+    c = client()
+    c.upsert(
+        collection_name=config.QDRANT_COLLECTION,
+        points=[PointStruct(id=atom_id, vector=vec, payload=payload)],
+    )
+
+def delete_atom(atom_id: int) -> None:
+    client().delete(
+        collection_name=config.QDRANT_COLLECTION,
+        points_selector=[atom_id],
+    )
+
+def search(query: str, top_k: int = 5, type_filter: str | None = None,
+           status_filter: str | None = None) -> list[dict]:
+    """Семантический поиск по атомам."""
+    qvec = embed(query)
+    must = []
+    if type_filter:
+        must.append(FieldCondition(key="type", match=MatchValue(value=type_filter)))
+    if status_filter:
+        must.append(FieldCondition(key="status", match=MatchValue(value=status_filter)))
+    qfilter = Filter(must=must) if must else None
+    hits = client().query_points(
+        collection_name=config.QDRANT_COLLECTION,
+        query=qvec,
+        limit=top_k,
+        query_filter=qfilter,
+    ).points
+    return [{"id": h.id, "score": h.score, **h.payload} for h in hits]