llm/CustomRetriever.py

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