Update CustomRetriever.py

CustomRetriever.py

@@ -1,47 +1,50 @@
-from langchain.schema.retriever import BaseRetriever
-from langchain_core.documents import Document
-from typing import List
-
-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.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]:
-
-                    @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()
-        
-
-    
-    
-
-
-
-
+from langchain.schema.retriever import BaseRetriever
+from langchain_core.documents import Document
+from typing import List
+
+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
+                #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()
+        
+
+    
+    
+
+
+
+