celsius_csrd_chatbot/chains/retriever.py

def make_retriever_node(vectorstore, k=10):
    def retrieve_documents(state):
        sources = state["esrs_type"]
        query = state["query"]
        if sources == "none" or sources == "no_intent":
            docs_retrieved = vectorstore.similarity_search_with_score(query=query, k=k)
        else:
            filters = {"ESRS_filter": {"$in": sources}}
            docs_retrieved = vectorstore.similarity_search_with_score(
                query=query, filter=filters, k=k
            )
        docs = []
        for doc in docs_retrieved:
            doc_append = doc[0]
            doc_append.metadata["similarity_score"] = doc[1]
            docs.append(doc_append)

        docs = sorted(docs, key=lambda x: x.metadata["similarity_score"], reverse=True)
        new_state = {"documents": docs}

        return new_state

    return retrieve_documents