File size: 863 Bytes
7bfa7e6
 
 
 
8ca00e0
 
7bfa7e6
8ca00e0
 
 
 
7bfa7e6
d708cb9
 
 
 
7bfa7e6
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
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