feat: add reranking with keyword extractor

sorbobotapp/app.py

@@ -7,7 +7,6 @@ from chain import get_chain
 from chat_history import insert_chat_history, insert_chat_history_articles
 from connection import connect
 from css import load_css
-from keyword_extraction import KeywordExtractor
 from langchain.callbacks import get_openai_callback
 from message import Message
 
@@ -27,8 +26,6 @@ def initialize_session_state():
         st.session_state.token_count = 0
     if "conversation" not in st.session_state:
         st.session_state.conversation = get_chain(conn)
-    if "keyword_extractor" not in st.session_state:
-        st.session_state.keyword_extractor = KeywordExtractor()
 
 
 def send_message_callback():

sorbobotapp/conversation_retrieval_chain.py

@@ -1,12 +1,17 @@
 import inspect
+import json
 from typing import Any, Dict, Optional
 
+from keyword_extraction import KeywordExtractor
 from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.conversational_retrieval.base import (
     ConversationalRetrievalChain, _get_chat_history)
+from langchain.schema import Document
 
 
 class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
+    keyword_extractor: KeywordExtractor = KeywordExtractor()
+
     def _handle_docs(self, docs):
         if len(docs) == 0:
             return False, "No documents found. Can you rephrase ?"
@@ -16,6 +21,33 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
             return False, "Too many documents found. Can you specify your request ?"
         return True, ""
 
+    def rerank_documents(self, question: str, docs: list[Document]) -> list[Document]:
+        """Rerank documents based on the number of similar keywords
+
+        Args:
+            question (str): Orinal question
+            docs (list[Document]): List of documents
+
+        Returns:
+            list[Document]: List of documents sorted by the number of similar keywords
+        """
+        keywords = self.keyword_extractor(question)
+
+        for doc in docs:
+            doc.metadata["similar_keyword"] = 0
+            doc_keywords = json.loads(doc.page_content)["keywords"]
+            if doc_keywords is None:
+                continue
+            doc_keywords = doc_keywords.lower().split(",")
+
+            for kw in keywords:
+                if kw.lower() in doc_keywords:
+                    doc.metadata["similar_keyword"] += 1
+                    print("similar keyword : ", kw)
+
+        docs = sorted(docs, key=lambda x: x.metadata["similar_keyword"])
+        return docs
+
     def _call(
         self,
         inputs: Dict[str, Any],
@@ -40,6 +72,7 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
             docs = self._get_docs(new_question, inputs, run_manager=_run_manager)
         else:
             docs = self._get_docs(new_question, inputs)  # type: ignore[call-arg]
+
         valid_docs, message = self._handle_docs(docs)
         if not valid_docs:
             return {
@@ -47,6 +80,9 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
                 "source_documents": docs,
             }
 
+        # Add reranking
+        docs = self.rerank_documents(new_question, docs)
+
         new_inputs = inputs.copy()
         if self.rephrase_question:
             new_inputs["question"] = new_question