Upload CustomRetriever.py

CustomRetriever.py

@@ -0,0 +1,47 @@
+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()
+        
+
+    
+    
+
+
+
+