File size: 356 Bytes
bf6dbfa
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
from rag.vectorstore import get_vectorstore

def retrieve_documents(query: str, k: int = 3):
    """
    Retrieves the top k relevant documents for the given query.
    """
    vectorstore = get_vectorstore()
    retriever = vectorstore.as_retriever(search_kwargs={"k": k})
    docs = retriever.invoke(query)

    return [doc.page_content for doc in docs]