سیستم بازیابی تقویتشده پیوندی (Hybrid RAG) با Sparse BM25 و Dense Embeddings

from typing import List, Dict, Any
from qdrant_client import QdrantClient, models
from sentence_transformers import CrossEncoder

class ProductionHybridRetriever:
    """
    Two-Stage Hybrid RAG Pipeline:
    Stage 1: Native Qdrant Sparse (BM25/SPLADE) + Dense HNSW Retrieval with Server-Side RRF
    Stage 2: Deep Cross-Encoder Reranking for Top-K Precision
    """
    def __init__(
        self,
        qdrant_url: str,
        collection_name: str,
        reranker_model_name: str = 'BAAI/bge-reranker-v2-m3',
        device: str = 'cuda'
    ):
        self.client = QdrantClient(url=qdrant_url, prefer_grpc=True)
        self.collection_name = collection_name
        self.reranker = CrossEncoder(reranker_model_name, max_length=1024, device=device)

    def create_hybrid_collection(self, dense_dim: int = 1024) -> None:
        self.client.recreate_collection(
            collection_name=self.collection_name,
            vectors_config={
                'dense_bge': models.VectorParams(
                    size=dense_dim,
                    distance=models.Distance.COSINE,
                    hnsw_config=models.HnswConfigDiff(m=16, ef_construct=200)
                )
            },
            sparse_vectors_config={
                'sparse_bm25': models.SparseVectorParams(
                    index=models.SparseIndexParams(on_disk=False)
                )
            }
        )

    def retrieve_and_rerank(
        self,
        query_text: str,
        dense_query_vec: List[float],
        sparse_indices: List[int],
        sparse_values: List[float],
        prefetch_limit: int = 40,
        final_top_k: int = 5
    ) -> List[Dict[str, Any]]:
        # Stage 1: Execute parallel Dense + Sparse prefetch and fuse via RRF inside Qdrant
        fused_response = self.client.query_points(
            collection_name=self.collection_name,
            prefetch=[
                models.Prefetch(
                    query=dense_query_vec,
                    using='dense_bge',
                    limit=prefetch_limit
                ),
                models.Prefetch(
                    query=models.SparseVector(indices=sparse_indices, values=sparse_values),
                    using='sparse_bm25',
                    limit=prefetch_limit
                ),
            ],
            query=models.FusionQuery(fusion=models.Fusion.RRF),
            limit=prefetch_limit,
            with_payload=True
        )

        candidates = fused_response.points
        if not candidates:
            return []

        # Stage 2: Cross-Encoder full self-attention scoring over (Query, Passage) pairs
        sentence_pairs = [(query_text, pt.payload.get('text', '')) for pt in candidates]
        cross_scores = self.reranker.predict(sentence_pairs, batch_size=16, show_progress_bar=False)

        reranked_results = []
        for point, score in zip(candidates, cross_scores):
            reranked_results.append({
                'doc_id': point.id,
                'rrf_score': point.score,
                'rerank_score': float(score),
                'text': point.payload.get('text', ''),
                'metadata': point.payload.get('metadata', {})
            })

        reranked_results.sort(key=lambda x: x['rerank_score'], reverse=True)
        return reranked_results[:final_top_k]