"""Main retrieval pipeline combining dense, sparse, and reranking.""" import asyncio import logging from typing import Dict, Any, List, Optional, Tuple from datetime import datetime import uuid from config import PipelineConfig, SearchMode from document_store import DocumentStore from retrieval_client import RetrievalClient, SearchResult from reranker import Reranker, RerankResult from fusion import fuse logger = logging.getLogger(__name__) class RetrievalPipeline: """Main retrieval pipeline orchestrating dense, sparse, and reranking.""" def __init__(self, config: Optional[PipelineConfig] = None): self.config = config or PipelineConfig() # Initialize components self.document_store = DocumentStore() self.retrieval_client = RetrievalClient( dense_url=self.config.services.dense_service_url, sparse_url=self.config.services.sparse_service_url ) self.reranker = Reranker( model_name=self.config.reranker.model_name, device=self.config.reranker.device, use_fp16=self.config.reranker.use_fp16, max_length=self.config.reranker.max_length ) logger.info("Retrieval pipeline initialized") async def index_document(self, text: str, doc_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Index a document in both dense and sparse services. Args: text: Document text doc_id: Optional document ID (generated if not provided) metadata: Optional metadata Returns: Indexing result with status from both services """ # Generate doc_id if not provided if not doc_id: doc_id = f"doc_{uuid.uuid4().hex[:8]}" logger.info(f"Indexing document {doc_id}") # Store document locally self.document_store.add_document(doc_id, text, metadata) # Index in both services in parallel result = await self.retrieval_client.index_document(text, doc_id, metadata) # Add timestamp and document info result["timestamp"] = datetime.now().isoformat() result["text_length"] = len(text) if self.config.debug: logger.debug(f"Indexing result: dense={result['dense'].get('success')}, " f"sparse={result['sparse'].get('success')}") return result async def search(self, query: str, mode: SearchMode = SearchMode.HYBRID, top_k: Optional[int] = None, rerank_top_k: Optional[int] = None, skip_reranking: bool = False) -> Dict[str, Any]: """Search for documents using specified mode. Args: query: Search query mode: Search mode (dense, sparse, or hybrid) top_k: Number of candidates to retrieve from each service rerank_top_k: Number of results after reranking skip_reranking: Skip reranking step (for comparison) Returns: Search results with scores and rankings """ top_k = top_k or self.config.default_top_k rerank_top_k = rerank_top_k or self.config.rerank_top_k logger.info(f"Searching with mode={mode}, top_k={top_k}, rerank_top_k={rerank_top_k}") start_time = datetime.now() # Retrieve from services dense_results, sparse_results = await self.retrieval_client.search( query, top_k, mode.value ) retrieval_time = (datetime.now() - start_time).total_seconds() # Prepare response structure response = { "query": query, "mode": mode.value, "timestamp": start_time.isoformat(), "retrieval_time_ms": retrieval_time * 1000, "dense_results": [], "sparse_results": [], "combined_results": [], "reranked_results": [], "statistics": {} } # Process dense results if dense_results: response["dense_results"] = [ { "doc_id": r.doc_id, "score": r.score, "rank": r.rank or idx + 1, "text": r.text[:200] if r.text else None } for idx, r in enumerate(dense_results[:10]) # Show top 10 for educational purposes ] # Process sparse results if sparse_results: response["sparse_results"] = [ { "doc_id": r.doc_id, "score": r.score, "rank": r.rank or idx + 1, "text": r.text[:200] if r.text else None, "matched_terms": r.debug_info.get("matched_terms", []) if r.debug_info else [] } for idx, r in enumerate(sparse_results[:10]) # Show top 10 for educational purposes ] # Combine results for reranking combined_docs = self._combine_results(dense_results, sparse_results) if not combined_docs: logger.warning("No documents to rerank") return response response["combined_results"] = [ { "doc_id": doc["doc_id"], "dense_score": doc.get("dense_score"), "sparse_score": doc.get("sparse_score"), "dense_rank": doc.get("dense_rank"), "sparse_rank": doc.get("sparse_rank") } for doc in combined_docs[:20] # Show top 20 combined ] # Perform reranking if not skipped if not skip_reranking and combined_docs: rerank_start = datetime.now() reranked = self.reranker.rerank( query=query, documents=combined_docs, top_k=rerank_top_k, return_scores=self.config.show_scores ) rerank_time = (datetime.now() - rerank_start).total_seconds() response["rerank_time_ms"] = rerank_time * 1000 # Format reranked results response["reranked_results"] = self._format_reranked_results(reranked) # Calculate statistics response["statistics"] = self._calculate_statistics( dense_results, sparse_results, response.get("reranked_results", []) ) total_time = (datetime.now() - start_time).total_seconds() response["total_time_ms"] = total_time * 1000 return response def _combine_results(self, dense_results: List[SearchResult], sparse_results: List[SearchResult]) -> List[Dict[str, Any]]: """Combine dense and sparse results for reranking. Args: dense_results: Results from dense search sparse_results: Results from sparse search Returns: Combined document list with scores and ranks from both sources """ combined = {} # Add dense results for idx, result in enumerate(dense_results): doc_id = result.doc_id if doc_id not in combined: # Get full document from store doc = self.document_store.get_document(doc_id) combined[doc_id] = { "doc_id": doc_id, "text": doc["text"] if doc else result.text, "metadata": doc["metadata"] if doc else result.metadata, "dense_score": result.score, "dense_rank": idx + 1, "sparse_score": None, "sparse_rank": None } else: combined[doc_id]["dense_score"] = result.score combined[doc_id]["dense_rank"] = idx + 1 # Add sparse results for idx, result in enumerate(sparse_results): doc_id = result.doc_id if doc_id not in combined: # Get full document from store doc = self.document_store.get_document(doc_id) combined[doc_id] = { "doc_id": doc_id, "text": doc["text"] if doc else result.text, "metadata": doc["metadata"] if doc else result.metadata, "dense_score": None, "dense_rank": None, "sparse_score": result.score, "sparse_rank": idx + 1 } else: combined[doc_id]["sparse_score"] = result.score combined[doc_id]["sparse_rank"] = idx + 1 # Convert to list and keep the legacy average-rank field for reference. combined_list = list(combined.values()) for doc in combined_list: ranks = [] if doc["dense_rank"] is not None: ranks.append(doc["dense_rank"]) if doc["sparse_rank"] is not None: ranks.append(doc["sparse_rank"]) doc["avg_rank"] = sum(ranks) / len(ranks) if ranks else float('inf') # Fuse the two ranked lists into one unified candidate pool. This is the # dedicated fusion stage described in the book (RRF / weighted). The order # produced here is the candidate pool that neural reranking then refines. method = getattr(self.config, "fusion_method", "rrf") if method == "avg_rank": combined_list.sort(key=lambda x: x["avg_rank"]) return combined_list dense_ranked = [(r.doc_id, r.score) for r in dense_results] sparse_ranked = [(r.doc_id, r.score) for r in sparse_results] fused = fuse( {"dense": dense_ranked, "sparse": sparse_ranked}, method=method, k=getattr(self.config, "rrf_k", 60), ) fused_order = {doc_id: rank for rank, (doc_id, _) in enumerate(fused)} for doc in combined_list: doc["fusion_score"] = dict(fused).get(doc["doc_id"]) # Sort by fused rank; any doc not in the fused output (shouldn't happen) # falls back to the end, then to its average rank for stability. combined_list.sort( key=lambda x: (fused_order.get(x["doc_id"], len(fused)), x["avg_rank"]) ) return combined_list def _format_reranked_results(self, reranked: List[RerankResult]) -> List[Dict[str, Any]]: """Format reranked results for response. Args: reranked: List of reranked results Returns: Formatted results with educational information """ formatted = [] for idx, result in enumerate(reranked): item = { "rank": idx + 1, "doc_id": result.doc_id, "rerank_score": result.rerank_score, "text": result.text, "metadata": result.metadata } # Add educational information about rank changes if self.config.show_scores: item["original_scores"] = { "dense": result.original_dense_score, "sparse": result.original_sparse_score } item["original_ranks"] = { "dense": result.original_dense_rank, "sparse": result.original_sparse_rank } # Calculate rank changes rank_changes = [] if result.original_dense_rank: change = result.original_dense_rank - (idx + 1) rank_changes.append(f"dense: {change:+d}") if result.original_sparse_rank: change = result.original_sparse_rank - (idx + 1) rank_changes.append(f"sparse: {change:+d}") item["rank_changes"] = rank_changes formatted.append(item) return formatted def _calculate_statistics(self, dense_results: List[SearchResult], sparse_results: List[SearchResult], reranked_results: List[Dict[str, Any]]) -> Dict[str, Any]: """Calculate statistics for educational purposes. Args: dense_results: Dense search results sparse_results: Sparse search results reranked_results: Reranked results Returns: Statistics dictionary """ stats = { "dense_retrieved": len(dense_results), "sparse_retrieved": len(sparse_results), "total_unique_documents": 0, "reranked_count": len(reranked_results) } # Count unique documents unique_docs = set() for r in dense_results: unique_docs.add(r.doc_id) for r in sparse_results: unique_docs.add(r.doc_id) stats["total_unique_documents"] = len(unique_docs) # Calculate overlap if dense_results and sparse_results: dense_ids = {r.doc_id for r in dense_results} sparse_ids = {r.doc_id for r in sparse_results} overlap = dense_ids & sparse_ids stats["overlap_count"] = len(overlap) stats["overlap_percentage"] = (len(overlap) / len(unique_docs)) * 100 if unique_docs else 0 # Analyze rank changes if available if reranked_results and self.config.show_scores: avg_dense_change = [] avg_sparse_change = [] for r in reranked_results: if "original_ranks" in r: if r["original_ranks"]["dense"]: avg_dense_change.append(r["original_ranks"]["dense"] - r["rank"]) if r["original_ranks"]["sparse"]: avg_sparse_change.append(r["original_ranks"]["sparse"] - r["rank"]) if avg_dense_change: stats["avg_dense_rank_change"] = sum(avg_dense_change) / len(avg_dense_change) if avg_sparse_change: stats["avg_sparse_rank_change"] = sum(avg_sparse_change) / len(avg_sparse_change) return stats async def delete_document(self, doc_id: str) -> Dict[str, Any]: """Delete a document from all services. Args: doc_id: Document ID to delete Returns: Deletion result """ logger.info(f"Deleting document {doc_id}") # Delete from local store local_deleted = self.document_store.delete_document(doc_id) # Delete from remote services remote_result = await self.retrieval_client.delete_document(doc_id) return { "doc_id": doc_id, "local_deleted": local_deleted, "remote_result": remote_result, "success": local_deleted } def get_statistics(self) -> Dict[str, Any]: """Get pipeline statistics. Returns: Statistics dictionary """ return { "document_store": self.document_store.get_stats(), "pipeline_config": { "default_top_k": self.config.default_top_k, "rerank_top_k": self.config.rerank_top_k, "reranker_model": self.config.reranker.model_name, "device": self.config.reranker.device }, "services": { "dense_url": self.config.services.dense_service_url, "sparse_url": self.config.services.sparse_service_url } }