58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
"""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]
|