""" What this file does (list of functions): 1. Embed the query and check "paper_index" to see if it confidently matches one specific paper 2. Rewrite the query (with matched paper's summary context) using LLM for better retrieval 3. Use hybrid search (dense + BM25 sparse) on paper_chunk, fuse results via RRF 4. Reranks the top candidates with cross-encoder and returns top 3 5. Retrieval function which connect all this: route the query route_query() -> rewrite the query using resulted paper summary and name rewrite_query() -> use hybrid search to retrieve chunks hybrid_search() -> reranks the chunks and return top 3 rerank() """ import os from dotenv import load_dotenv from sentence_transformers import SentenceTransformer, CrossEncoder from qdrant_client import QdrantClient # Prefetch is used in qdrant as part of flow, like we first search top chunks in both sparse and dense with ranks assigned to them during fetching then we fuse them, pre is used because the final fetch isn't actually a fetch but the fusion of those ranked chunks is final fetch # FusionQuery is used to fuse the ranked chunks, qdrant pass the fetched chunks to this internally, we just need to specify the fusion method # Fusion is param of FusionQuery telling it which method to use to fuse the queries (RRF here) # Filter, FieldCondition, MatchValue are pydantic classes in qdrant for proper input structure from qdrant_client.models import Prefetch, FusionQuery, Fusion, Filter, FieldCondition, MatchValue, Document from google import genai from config import ( EMBED_MODEL, CROSS_ENCODER, LLM_MODEL, INDEX_COLLECTION, CHUNKS_COLLECTION ) from logger import get_logger load_dotenv() logger = get_logger("retrieve") QDRANT_URL = os.getenv("QDRANT_URL") QDRANT_API_KEY = os.getenv("QDRANT_API_KEY") GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") embed_model = SentenceTransformer(EMBED_MODEL) qdrant_client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY, timeout=60) reranker = CrossEncoder(CROSS_ENCODER) llm = genai.Client(api_key=GEMINI_API_KEY) # 1. Embed the query and check "paper_index" to see if it confidently matches one specific paper's summary, return the paper_name and its summary def route_query(query): """ Embeds the query and searches papers_index Returns the best matching paper_name and its_summary Falls back to None if confidence is low (score < 0.4) """ try: query_vector = embed_model.encode(query).tolist() # .query_points() finds the match (point) by similarity (while .scroll() matches exact keyword) result = qdrant_client.query_points( collection_name=INDEX_COLLECTION, query=query_vector, using="dense", limit=1 ) if not result.points: logger.info("route_query: no papers in index") return None, None top = result.points[0] if top.score < 0.4: logger.info("route_query: low confidence score %.3f, searching all papers", top.score) return None, None logger.info("route_query: routed to '%s' (score=%.3f)", top.payload["paper_name"], top.score) return top.payload["paper_name"], top.payload["summary"] except Exception as e: logger.warning("route_query failed, falling back to all-paper search: %s", e) return None, None # safe fallback: search all papers # 2. Rewrite the query (with matched paper's summary context) using LLM for better retrieval def rewrite_query(query, paper_name, paper_summary): """ Uses LLM to rewrite query to be more precise for retrieval Passes paper name + summary + query so the rewritten query has context of paper """ try: prompt = f"""You are helping retrieve relevant chunks from resarch papers vector database using rewritten query, you will be provided with research paper name, its summary and the original query. Rewrite the query into a precise search query that will retrieve most relevant chunks from the databse. Use ONLY the given summary to rewrite the query, DO NOT USE your own knowlege. Paper: {paper_name} Summary/Abstract: {paper_summary} Original Query: {query} Return ONLY the rewritten query, nothing else. Rewritten Query:""" response = llm.models.generate_content( model=LLM_MODEL, contents=prompt ) rewritten = response.text.strip() logger.info("rewrite_query: '%s' -> '%s'", query, rewritten) return rewritten except Exception as e: logger.warning("rewrite_query failed, using original query: %s", e) return query # safe fallback: original query still works # 3. Use hybrid search (dense + BM25 sparse) on paper_chunk, fuse results via RRF def hybrid_search(query, paper_name=None, limit=10): """ Runs dense + BM25 hybrid search IF paper_name is provided, filters to only that paper's chunks """ try: query_filter = None if paper_name: query_filter = Filter( must=[FieldCondition(key="paper_name", match=MatchValue(value=paper_name))] ) results = qdrant_client.query_points( collection_name=CHUNKS_COLLECTION, prefetch=[ # prefetch for dense vector Prefetch( query=embed_model.encode(query).tolist(), using="dense", limit=15, filter=query_filter ), # prefetch for sparse vector Prefetch( query=Document(text=query, model="Qdrant/bm25"), using="sparse", limit=15, filter=query_filter ) ], query = FusionQuery(fusion=Fusion.RRF), # fuse both resulted chunks limit=limit ) logger.info("hybrid_search: got %d candidates", len(results.points)) return results.points except Exception as e: logger.error("hybrid_search failed: %s", e) return [] # rerank() and answer() handle empty lists gracefully # 4. Reranks the top candidates with cross-encoder and returns top 3 def rerank(query, points, top_n=3): """ reranker scores each [query, chunk_text] pair, returns top_n chunks sorted by reranker score """ if not points: return [] try: scores = reranker.predict([[query, p.payload["text"]] for p in points]) # zip points with their scores, sort in descending using scores, take top_n ranked = sorted(zip(points, scores), key=lambda x:x[1], reverse=True) top = [point for point, score in ranked[:top_n]] logger.info("rerank: returning top %d chunks", len(top)) return top except Exception as e: logger.warning("rerank failed, returning unranked top chunks: %s", e) return points[:top_n] # skip reranking, still return something # 5. Retrieval function which connect all this: route the query route_query() -> rewrite the query using resulted paper summary and name rewrite_query() -> use hybrid search to retrieve chunks hybrid_search() -> reranks the chunks and return top 3 rerank() def retrieve(query): """ Full retrieval pipeline returns top 3 most relevant chunks for the query """ # route paper_name, paper_summary = route_query(query) if paper_name: logger.info("Routed to: %s", paper_name) # rewrite the query re_query = rewrite_query(query, paper_name, paper_summary) logger.info("Rewritten query: %s", re_query) else: logger.info("No confident paper match — searching all papers") re_query = query # hybrid search points = hybrid_search(re_query, paper_name) # rerank top_chunks = rerank(query, points) return top_chunks