Update CustomRetriever.py

CustomRetriever.py

@@ -7,27 +7,41 @@ from langchain.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain_core.documents import Document
 from langchain_core.runnables import chain
 
-
-class RetrieverWithScores(BaseRetriever):
-    def __init__(self, vdb, thold=0.7):
-        self.vdb=vdb
-        self.thold=thold
-
-    @chain
-    def retr_func(query: str)-> List[Document]: #(vdb, query: str)-> List[Document]:
-        docs, scores = zip(*self.vdb.similarity_search_with_relevance_scores(query))#similarity_search_with_score(query))
-        result=[]
-        for doc, score in zip(docs, scores):
-            if score>self.thold:
-                doc.metadata["score"] = score
-                result.append(doc)
-        if len(result)==0:
-            result.append(Document(metadata={}, page_content='No data'))
+class CustomRetriever():
+    def __init__(self, v_db, thold=0.7):
         
-        return result #docs
+
+        #self.retriever=RetrieverWithScores()
+
+        class RetrieverWithScores(BaseRetriever):
+                #def __init__(self, vdb):
+                    #self.vdb=vdb
                 #def __init__(self, retriever: BaseRetriever):  # Add an __init__ to store the existing retriever
                     #super().__init__(retriever=retriever)
-    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun)-> List[Document]:                   
+                def _get_relevant_documents(
+                    self, query: str, *, run_manager: CallbackManagerForRetrieverRun)-> List[Document]:
+
+                    @chain
+                    def retr_func(query: str)-> List[Document]: #(vdb, query: str)-> List[Document]:
+                        docs, scores = zip(*v_db.similarity_search_with_relevance_scores(query))#similarity_search_with_score(query))
+                        result=[]
+                        for doc, score in zip(docs, scores):
+                            if score>thold:
+                                doc.metadata["score"] = score
+                                result.append(doc)
+                        if len(result)==0:
+                            result.append(Document(metadata={}, page_content='No data'))
+
+                        return result #docs
+
+                    return retr_func.invoke(query) #(self.vdb, query)
+
+        self.retriever=RetrieverWithScores()
+        
+
+    
+    
+
+
 
-        return self.retr_func.invoke(query) #(self.vdb, query)