Update processing.py

processing.py

@@ -1,13 +1,13 @@
-from langchain.schema import HumanMessage
+from langchain.schema import HumanMessage, BaseRetriever
 from output_parser import output_parser
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import FAISS
 from llm_loader import load_model
 from config import openai_api_key
 from langchain.chains import RetrievalQA, LLMChain
-from langchain.retrievers import MultiQueryRetriever, MultiVectorRetriever
-from langchain.schema import Document
+from langchain.retrievers import MultiQueryRetriever
 from langchain.prompts import PromptTemplate
+from typing import List
 import os
 import json
 
@@ -43,22 +43,27 @@ text_retriever = text_faiss_index.as_retriever()
 attachments_retriever = attachments_faiss_index.as_retriever()
 personalities_retriever = personalities_faiss_index.as_retriever()
 
-# Create a list of all retrievers
-all_retrievers = [text_retriever, attachments_retriever, personalities_retriever]
-
-# Create a function to combine results from all retrievers
-def combined_retriever_function(query):
-    combined_docs = []
-    for retriever in all_retrievers:
-        docs = retriever.get_relevant_documents(query)
-        combined_docs.extend(docs)
-    return combined_docs
-
-# Create a MultiVectorRetriever
-combined_retriever = MultiVectorRetriever(
-    retrievers=all_retrievers,
-    retriever_query=combined_retriever_function
-)
+# Create a custom combined retriever
+class CombinedRetriever(BaseRetriever):
+    def __init__(self, retrievers: List[BaseRetriever]):
+        self.retrievers = retrievers
+
+    def get_relevant_documents(self, query: str):
+        combined_docs = []
+        for retriever in self.retrievers:
+            docs = retriever.get_relevant_documents(query)
+            combined_docs.extend(docs)
+        return combined_docs
+
+    async def aget_relevant_documents(self, query: str):
+        combined_docs = []
+        for retriever in self.retrievers:
+            docs = await retriever.aget_relevant_documents(query)
+            combined_docs.extend(docs)
+        return combined_docs
+
+# Create an instance of the combined retriever
+combined_retriever = CombinedRetriever([text_retriever, attachments_retriever, personalities_retriever])
 
 # Create prompt template for query generation
 prompt_template = PromptTemplate(
@@ -76,7 +81,7 @@ multi_query_retriever = MultiQueryRetriever(
     parser_key="lines"  # Assuming the LLM outputs queries line by line
 )
 
-# Use the multi-query retriever in the QA chain
+# Create QA chain with multi-query retriever
 qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=multi_query_retriever)
 
 def load_text(file_path: str) -> str: