Spaces:
Paused
Paused
File size: 4,895 Bytes
29fd274 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | from dataclasses import dataclass
from hashlib import sha256
import logging
from app.config import Settings
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
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 []
|