Retrieving Relevant Chunks and Building LLM Prompts in JavaScript

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.

Before your LLM can generate a coherent, context-rich answer, you need to fetch the right information. Your vector database (for instance, using ChromaDB) will rank which document chunks are most relevant for a given query.

async function retrieveTopChunks(query, collection, topK = 2) {
    /**
     * Retrieves the topK chunks relevant to the given query from 'collection'.
     * Returns an array of retrieved chunks, each containing 'chunk' text,
     * 'docId', and 'distance'.
     */
    // Search for topK results matching the user's query
    const results = await collection.query({
        queryTexts: [query],
        nResults: topK
    });

    const retrievedChunks = [];

    // Safeguard in case no results are found
    if (!results.documents || !results.documents[0]) {
        return retrievedChunks;
    }

    // Gather each retrieved chunk, along with its distance score
    for (let i = 0; i < results.documents[0].length; i++) {
        retrievedChunks.push({
            chunk: results.documents[0][i],
            docId: results.ids[0][i],
            distance: results.distances[0][i]
        });
    }
    return retrievedChunks;
}

Let's break down this code in detail:

• Function Definition:

  • retrieveTopChunks takes three parameters:
    • query: The user's question or search term;
    • collection: The ChromaDB collection object containing our embedded documents;
    • topK: The number of most relevant chunks to retrieve (default is 2).

• Vector Search:

  • The collection.query() function performs a vector-based similarity search to pinpoint which chunks are most aligned with the query.
  • queryTexts: [query] passes the user's query as an array (ChromaDB's API expects an array).
  • nResults: topK specifies how many matching chunks to return.

• Results Structure:

  • The query returns an object with multiple keys:
    • documents: Contains the actual text chunks;
    • ids: Contains the document identifiers;
    • distances: Each result includes a distance, which indicates how semantically close a chunk is to your query — the lower the distance, the better the match.
  • Each of these keys maps to a nested array structure: [[item1, item2, ...]].

• Processing Results:

  • For each result, the function creates an object with three key pieces of information:
    • chunk: The actual text content from results.documents[0][i];
    • docId: The document identifier from results.ids[0][i];
    • distance: The similarity score from results.distances[0][i];
  • These objects are appended to the retrievedChunks array, which is then ultimately returned.