Retrieving and Prompt Building in RAG Systems

Welcome to the third lesson in our journey through "Scaling Up RAG with Vector Databases"! Well done, you're halfway through this course. In the previous lesson, you learned how to split or chunk your text data and store those chunks in a vector database collection. Now, we'll delve into retrieving the most relevant chunks for any given query and building an LLM prompt to produce more accurate, context-driven answers.

Let's start by understanding how to fetch the most relevant information from your vector database, which is the foundation for effective retrieval-augmented generation.

Before your LLM can generate a coherent, context-rich answer, you need to fetch the right information. Your vector database (for instance, using Chroma) will rank which document chunks are most relevant for a given query. Let's explore how this is achieved with the retrieve_top_chunks function.

pub async fn retrieve_top_chunks(
    collection: &ChromaCollection,
    query: &str,
    top_k: usize,
    embedder: &SentenceEmbedder,
) -> Result<Vec<RetrievedChunk>, Box<dyn std::error::Error>> {
    // Create embeddings for the query text.
    let query_embeddings = embedder.embed_texts(&[query])?;

In this initial part of the function, we use the SentenceEmbedder to convert the query text into a dense vector representation. This embedding is crucial as it allows us to perform a similarity search in the vector space, comparing the query against the stored document embeddings.

    let query_options = QueryOptions {
        query_texts: None,
        query_embeddings: Some(query_embeddings),
        n_results: Some(top_k),
        where_metadata: None,
        where_document: None,
        include: Some(vec!["documents", "distances"]),
    };

    let result = collection.query(query_options, None).await?;
    let mut chunks = Vec::new();

    // Early return if no results
    if result.documents.is_none() || result.documents.as_ref().unwrap().is_empty() {
        return Ok(chunks);
    }

    let documents = &result.documents.as_ref().unwrap()[0];
    let distances = result.distances.as_ref().map(|d| d[0].clone()).unwrap_or_default();

Here, we define the QueryOptions struct, specifying that we want to use the query embeddings for the search. We also set n_results to top_k, indicating the number of top results we wish to retrieve. The include field specifies that we want both the document texts and their distances (similarity scores) in the results. The query method of the collection is then called with these options, returning the most relevant document chunks. We also handle the case where no documents are returned by checking if the documents field is None or empty, allowing for an early return.

    for i in 0..documents.len() {
        chunks.push(RetrievedChunk {
            chunk: documents[i].clone(),
            doc_id: i, // Use the index as the document ID
            distance: distances.get(i).copied().unwrap_or(0.0),
        });
    }

    Ok(chunks)
}

In this final section, we process the query results. We iterate over the documents and their corresponding distances, creating a RetrievedChunk struct for each document containing the document text, an index-based document ID, and the distance. These chunks are collected into a vector, which is returned as the function's result.

This function is essential for fetching the most relevant information for your query, ensuring that the LLM has the right context to generate accurate and context-driven answers.

Now that you know how to retrieve the best-matching chunks, let's see how to use them to build a prompt that guides your LLM to generate focused and reliable answers.