reranker added

app/retrieval/reranker.py

@@ -0,0 +1,48 @@
+from typing import List, Optional
+from langchain_core.documents import Document
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser 
+
+class Reranker:
+    """Reranker class to rerank retrieved documents based on relevance to the query."""
+    def __init__(self,llm, retrieved_docs:List[Document],query:str) ->List[Document]:
+        self.llm = llm
+        self.retrieved_docs = retrieved_docs
+        self.query = query
+
+    def rerank_documents(self) -> List[Document]:
+        """
+        Rerank the retrieved documents based on their relevance to the query.
+        
+        Args:
+            retrieved_docs: List of Document objects retrieved from the retriever.
+            query: The original user query string.
+        
+        Returns:
+            List of Document objects sorted by relevance to the query.
+        """
+        # Create a prompt template for scoring
+       # Prompt Template
+        prompt_template = PromptTemplate.from_template("""
+        You are a helpful assistant. Your task is to rank the following documents from most to least relevant to the user's question.
+
+        User Question: "{question}"
+
+        Documents:
+        {documents}
+
+        Instructions:
+        - Think about the relevance of each document to the user's question.
+        - Return a list of document indices in ranked order, starting from the most relevant.
+
+        Output format: comma-separated document indices (e.g., 2,1,3,0,...)
+        """)
+
+        chain=prompt_template | self.llm | StrOutputParser()
+        doc_texts = [f"{i+1}. {doc.page_content}" for i,doc in enumerate(self.retrieved_docs)]
+        response = chain.invoke({
+            "question": self.query,
+            "documents": "\n".join(doc_texts)
+        })
+        ranked_indices = [int(i) for i in response.split(",")]
+        return [self.retrieved_docs[i-1] for i in ranked_indices]

app/retrieval/retriever.py

@@ -1,5 +1,6 @@
 
 from langchain.retrievers import EnsembleRetriever
+from app.retrieval.reranker import Reranker
 
 class Retriever:
     def __init__(self, pinecone_index, query = None, metadata = None, namespace=None, vectore_store = None,sparse_retriever = None, llm = None):
@@ -66,4 +67,7 @@ class Retriever:
         results = self.hybrid_retriever.invoke(self.query)
         for doc in results:
             print(f"printing Doc content : {doc.page_content}")
+        if self.llm:
+            reranker = Reranker(self.llm, results, self.query)
+            results = reranker.rerank_documents()    
         return results