Update CustomRetriever.py

CustomRetriever.py

@@ -7,44 +7,26 @@ from langchain.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain_core.documents import Document
 from langchain_core.runnables import chain
 
-class CustomRetriever():
-    def __init__(self, v_db, thold=0.7):
 
-        self.vdb=v_db
-
-        #self.retriever=RetrieverWithScores()
-
-        @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>thold:
-                    doc.metadata["score"] = score
-                    result.append(doc)
-            if len(result)==0:
-                result.append(Document(metadata={}, page_content='No data'))
-
-            return result #docs
-
-        class RetrieverWithScores(BaseRetriever):
-                #def __init__(self, vdb):
-                    #self.vdb=vdb
+class RetrieverWithScores(BaseRetriever):
+    def __init__(self, vdb, thold=0.7):
+        self.vdb=vdb
+        self.thold=thold
                 #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]:
-
-                    
-
-                    return retr_func.invoke(query) #(self.vdb, query)
-
-        self.retriever=RetrieverWithScores()
+    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun)-> List[Document]:                   
+
+                    return self.retr_func.invoke(query) #(self.vdb, query)
+
+    @chain
+    def retr_func(self, 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'))
         
-
-    
-    
-
-
-
-
+        return result #docs