Spaces:
Paused
Paused
| from dataclasses import dataclass | |
| from hashlib import sha256 | |
| import logging | |
| from app.config import Settings | |
| logger = logging.getLogger(__name__) | |
| class RetrievedDocument: | |
| title: str | |
| content: str | |
| source_url: str | None = None | |
| document_id: str | None = None | |
| section: str | None = None | |
| class CricketRetriever: | |
| def __init__(self, settings: Settings): | |
| self._settings = settings | |
| self.collection = settings.qdrant_collection | |
| self._client = None | |
| self._encoder = None | |
| if settings.qdrant_url and settings.qdrant_api_key: | |
| self.status = "configured" | |
| else: | |
| self.status = "not_configured" | |
| def _get_client(self): | |
| if self._client is None and self._settings.qdrant_url and self._settings.qdrant_api_key: | |
| try: | |
| from qdrant_client import QdrantClient | |
| self._client = QdrantClient( | |
| url=self._settings.qdrant_url, | |
| api_key=self._settings.qdrant_api_key, | |
| timeout=10, | |
| ) | |
| self.status = "connected" | |
| except Exception as exc: | |
| logger.warning(f"Could not connect to Qdrant: {exc}") | |
| self.status = f"unavailable: {type(exc).__name__}" | |
| self._client = None | |
| return self._client | |
| def _get_encoder(self): | |
| if self._encoder is None: | |
| try: | |
| import torch | |
| torch.set_num_threads(1) | |
| from sentence_transformers import SentenceTransformer | |
| self._encoder = SentenceTransformer(self._settings.embedding_model) | |
| except Exception as exc: | |
| logger.warning(f"Could not load SentenceTransformer: {exc}") | |
| self._encoder = None | |
| return self._encoder | |
| def _ensure_collection(self) -> None: | |
| client = self._get_client() | |
| encoder = self._get_encoder() | |
| if client and encoder: | |
| from qdrant_client.models import Distance, VectorParams | |
| size = encoder.get_sentence_embedding_dimension() | |
| if not client.collection_exists(self.collection): | |
| client.create_collection( | |
| collection_name=self.collection, | |
| vectors_config=VectorParams(size=size, distance=Distance.COSINE), | |
| ) | |
| def index(self, documents: list[RetrievedDocument]) -> int: | |
| client = self._get_client() | |
| encoder = self._get_encoder() | |
| if client is None or encoder is None: | |
| raise RuntimeError("Qdrant or embedding model is not configured") | |
| self._ensure_collection() | |
| from qdrant_client.models import PointStruct | |
| vectors = encoder.encode([document.content for document in documents]).tolist() | |
| points = [ | |
| PointStruct( | |
| id=int(sha256((document.document_id or str(index)).encode()).hexdigest()[:15], 16), | |
| vector=vector, | |
| payload={ | |
| "document_id": document.document_id, | |
| "title": document.title, | |
| "section": document.section, | |
| "content": document.content, | |
| "source_url": document.source_url, | |
| "document_type": "cricket_rule", | |
| }, | |
| ) | |
| for index, (document, vector) in enumerate(zip(documents, vectors)) | |
| ] | |
| client.upsert(collection_name=self.collection, points=points) | |
| return len(points) | |
| def retrieve(self, question: str, limit: int = 3) -> list[RetrievedDocument]: | |
| client = self._get_client() | |
| encoder = self._get_encoder() | |
| if client is None or encoder is None: | |
| return [] | |
| try: | |
| vector = encoder.encode(question).tolist() | |
| if hasattr(client, "query_points"): | |
| result = client.query_points( | |
| collection_name=self.collection, | |
| query=vector, | |
| limit=limit, | |
| with_payload=True, | |
| ) | |
| hits = result.points | |
| else: | |
| hits = client.search(collection_name=self.collection, query_vector=vector, limit=limit) | |
| return [ | |
| RetrievedDocument( | |
| title=str(hit.payload.get("title", "Cricket rules")), | |
| content=str(hit.payload.get("content", "")), | |
| source_url=hit.payload.get("source_url"), | |
| document_id=hit.payload.get("document_id"), | |
| section=hit.payload.get("section"), | |
| ) | |
| for hit in hits | |
| if hit.score >= 0.35 | |
| ] | |
| except Exception as exc: | |
| logger.warning(f"Error retrieving from Qdrant: {exc}") | |
| return [] | |