Update models/langOpen.py

models/langOpen.py

@@ -9,6 +9,10 @@ from langchain.embeddings import OpenAIEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain_pinecone import PineconeVectorStore
 
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CohereRerank
+from langchain_community.llms import Cohere
+
 prompt_template = """Answer the question using the given context to the best of your ability. 
 If you don't know, answer I don't know.
 Context: {context}
@@ -33,7 +37,13 @@ class LangOpen:
     def get_response(self, query_str):
         print("query_str: ", query_str)
         print("model_name: ", self.llm.model_name)
-        docs = self.index.similarity_search(query_str, k=4)
+        #docs = self.index.similarity_search(query_str, k=4)
+        vectorstore_retriever = self.index.as_retriever(search_type="similarity", search_kwargs={"k": 10})
+        compressor = CohereRerank()
+        compression_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor, base_retriever=vectorstore_retriever
+        )
+        docs = compression_retriever.get_relevant_documents(query_str)
         inputs = [{"context": doc.page_content, "topic": query_str} for doc in docs]
         result = self.chain.apply(inputs)[0]["text"]
         return result