Update processing.py

processing.py

@@ -7,8 +7,9 @@ from config import openai_api_key
 from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
 from langchain_core.runnables import RunnablePassthrough, RunnableLambda
-from typing import List
+from typing import List, Any, Optional
 from pydantic import Field
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
 import os
 import json
 
@@ -51,20 +52,21 @@ class CombinedRetriever(BaseRetriever):
     class Config:
         arbitrary_types_allowed = True
 
-    def get_relevant_documents(self, query: str) -> List[Document]:
-        return self.invoke(query)
-
-    async def aget_relevant_documents(self, query: str) -> List[Document]:
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None
+    ) -> List[Document]:
         combined_docs = []
         for retriever in self.retrievers:
-            docs = await retriever.aget_relevant_documents(query)
+            docs = retriever.get_relevant_documents(query, run_manager=run_manager)
             combined_docs.extend(docs)
         return combined_docs
 
-    def invoke(self, query: str) -> List[Document]:
+    async def _aget_relevant_documents(
+        self, query: str, *, run_manager: Optional[CallbackManagerForRetrieverRun] = None
+    ) -> List[Document]:
         combined_docs = []
         for retriever in self.retrievers:
-            docs = retriever.get_relevant_documents(query)
+            docs = await retriever.aget_relevant_documents(query, run_manager=run_manager)
             combined_docs.extend(docs)
         return combined_docs