Spaces:
Runtime error
Runtime error
| """Shared singletons: embedding model, re-ranker and the vector DB (all thread-safe).""" | |
| import logging | |
| import threading | |
| from functools import lru_cache | |
| from pathlib import Path | |
| from qdrant_client import QdrantClient | |
| from qdrant_client.models import ( | |
| Distance, | |
| FieldCondition, | |
| Filter, | |
| FilterSelector, | |
| MatchAny, | |
| PointStruct, | |
| VectorParams, | |
| ) | |
| from .config import settings | |
| log = logging.getLogger("rag.store") | |
| __all__ = ["PointStruct"] | |
| _qlock = threading.RLock() # qdrant local mode is not thread-safe, so every call goes through this | |
| _model_lock = threading.Lock() | |
| _embedder = None | |
| _reranker = None | |
| _reranker_failed = False | |
| # ---------------------------------------------------------------- embeddings | |
| def set_embedder(obj) -> None: | |
| """Swap the embedder (used by tests with a fake one).""" | |
| global _embedder | |
| _embedder = obj | |
| def get_embedder(): | |
| global _embedder | |
| if _embedder is None: | |
| with _model_lock: | |
| if _embedder is None: | |
| from sentence_transformers import SentenceTransformer # lazy: heavy import | |
| log.info("Loading embedding model %s", settings.embed_model) | |
| _embedder = SentenceTransformer(settings.embed_model) | |
| return _embedder | |
| def embed(texts: list[str]) -> list[list[float]]: | |
| vecs = get_embedder().encode(texts, normalize_embeddings=True, batch_size=32, show_progress_bar=False) | |
| return vecs.tolist() | |
| def embedding_dim() -> int: | |
| return int(get_embedder().get_sentence_embedding_dimension()) | |
| # ------------------------------------------------------------------ reranker | |
| def get_reranker(): | |
| """Cross-encoder used to re-rank candidates. Returns None if disabled or unavailable.""" | |
| global _reranker, _reranker_failed | |
| if not settings.use_reranker or _reranker_failed: | |
| return None | |
| if _reranker is None: | |
| with _model_lock: | |
| if _reranker is None and not _reranker_failed: | |
| try: | |
| from sentence_transformers import CrossEncoder # lazy | |
| log.info("Loading re-ranker %s", settings.rerank_model) | |
| _reranker = CrossEncoder(settings.rerank_model, max_length=512) | |
| except Exception as e: # no internet, blocked DLL, etc. -> degrade gracefully | |
| log.warning("Re-ranker unavailable, continuing without it: %s", e) | |
| _reranker_failed = True | |
| return _reranker | |
| # -------------------------------------------------------------------- qdrant | |
| def _ensure_collection(client: QdrantClient) -> None: | |
| existing = {c.name for c in client.get_collections().collections} | |
| if settings.collection not in existing: | |
| client.create_collection( | |
| collection_name=settings.collection, | |
| vectors_config=VectorParams(size=embedding_dim(), distance=Distance.COSINE), | |
| ) | |
| def get_qdrant() -> QdrantClient: | |
| path = Path(settings.data_dir) / "qdrant" | |
| path.mkdir(parents=True, exist_ok=True) | |
| with _qlock: | |
| client = QdrantClient(path=str(path)) | |
| _ensure_collection(client) | |
| return client | |
| def make_filter(doc_ids=None, pages=None) -> Filter | None: | |
| must = [] | |
| if doc_ids: | |
| must.append(FieldCondition(key="doc_id", match=MatchAny(any=list(doc_ids)))) | |
| if pages: | |
| must.append(FieldCondition(key="page", match=MatchAny(any=[int(p) for p in pages]))) | |
| return Filter(must=must) if must else None | |
| def _hit(p) -> dict: | |
| return {"id": str(p.id), "score": float(getattr(p, "score", 0.0) or 0.0), **(p.payload or {})} | |
| def upsert(points: list[PointStruct]) -> None: | |
| with _qlock: | |
| get_qdrant().upsert(collection_name=settings.collection, points=points) | |
| def search(vector: list[float], flt: Filter | None, limit: int) -> list[dict]: | |
| with _qlock: | |
| res = get_qdrant().query_points( | |
| collection_name=settings.collection, | |
| query=vector, | |
| query_filter=flt, | |
| limit=limit, | |
| with_payload=True, | |
| ) | |
| return [_hit(p) for p in res.points] | |
| def scroll_all(flt: Filter | None = None) -> list[dict]: | |
| """Every stored chunk (payload only) that matches the filter.""" | |
| out: list[dict] = [] | |
| offset = None | |
| while True: | |
| with _qlock: | |
| points, offset = get_qdrant().scroll( | |
| collection_name=settings.collection, | |
| scroll_filter=flt, | |
| limit=512, | |
| offset=offset, | |
| with_payload=True, | |
| with_vectors=False, | |
| ) | |
| out.extend(_hit(p) for p in points) | |
| if offset is None: | |
| break | |
| return out | |
| def delete_doc(doc_id: str) -> None: | |
| with _qlock: | |
| get_qdrant().delete( | |
| collection_name=settings.collection, | |
| points_selector=FilterSelector(filter=make_filter([doc_id])), | |
| ) | |
| def clear_all() -> None: | |
| with _qlock: | |
| client = get_qdrant() | |
| client.delete_collection(settings.collection) | |
| _ensure_collection(client) | |