Upload CustomRetriever.py

llm/CustomRetriever.py

@@ -0,0 +1,29 @@
+from langchain.vectorstores import VectorStore
+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(BaseRetriever):
+    vectorstore:VectorStore    
+    thold:float
+
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[Document]:
+
+        docs, scores = zip(*self.vectorstore.similarity_search_with_relevance_scores(query, callbacks=run_manager.get_child()))#get_relevant_documents(query, callbacks=run_manager.get_child())
+        
+        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