retriever.py

def initialize_document_retriever(top_k_documents, vector_database):
    """
    Initialize a document retriever using a Chroma vector database.

    This function initializes a document retriever that can be used to find and retrieve the most relevant documents
    for a specified search query. The number of documents to retrieve is determined by the top_k_documents parameter.
    
    Args:
        top_k_documents (int): The number of top relevant documents to retrieve.
        vector_database (obj): The Chroma vector database to use for retrieving documents.

    Returns:
        document_retriever (obj): The initialized document retriever.
    """
    
    # Initialize the document retriever with the Chroma vector database and the number of documents to retrieve.
    document_retriever = vector_database.as_retriever(
        search_kwargs = {"k": top_k_documents}
    )

    return document_retriever

def retrieve_relevant_documents(search_query, document_retriever):
    """
    Retrieve the most relevant documents for a given query.

    This function uses an initialized document retriever to find and retrieve the most relevant documents
    for a specified search query.
    
    Args:
        search_query (str): The search query for which to find and retrieve relevant documents.
        document_retriever (obj): The initialized document retriever.

    Returns:
        relevant_documents (list): The list of most relevant documents for the search query.
    """
    
    # Retrieve the most relevant documents for the search query using the document retriever.
    relevant_documents = document_retriever.get_relevant_documents(search_query)

    return relevant_documents