Iterative Retrieval: Enhancing RAG Systems with JavaScript

Introduction

Welcome back to this second lesson in the " Beyond Basic RAG: Improving our Pipeline " course! In the previous lesson, we explored ways to ensure that your language model stays grounded by responding only with information taken from retrieved context. That approach helps us avoid hallucinations and keeps the output reliable. In this lesson, we'll improve the pipeline even further by making the retrieval process more iterative . Instead of collecting chunks of text just once before moving on to generation, we'll refine our queries step by step. This multi-stage retrieval can pinpoint the most relevant information and produce a more helpful final context.

The Concept of Iterative Retrieval

Imagine a scenario where a user asks: "Tell me about the regulations for staff members." The question might be too broad. A typical retrieval step might find chunks containing some relevant information, but you might also want to narrow in on "internal policies" or "mandatory forms" for more precision. Iterative retrieval does exactly that: Retrieve an initial chunk based on the user's query. Refine that query with a new keyword from the retrieved chunk (e.g., "internal" or "policies"). Repeat until you've gathered a set of chunks that thoroughly answers the question—or until improvements level off. This multi-pass approach can drastically improve the depth and breadth of the retrieved information, making your final context more complete. The process continues until you've gathered a set of chunks that thoroughly answers the question—or until improvements level off. “Leveling off” refers to a point where the retrieval quality no longer improves meaningfully across iterations. This is typically implemented using a numeric threshold on the similarity score—if the new score improves less than a set amount (for example, 0.02) over the previous best score, the system exits early. This prevents unnecessary retrieval steps and helps optimize performance. Below, we'll walk through the building blocks of an iterative retrieval system in detail.

Practical Example: Iterative Retrieval in Action

Imagine a user asks, "Tell me about the regulations for staff members." Our DB may include chunks like: Chunk 1: "Our company requires that all staff members adhere to internal policies such as punctuality, dress code, and ethical behavior..." Chunk 2: "Regulations for staff emphasize adherence to both internal policies and government standards, covering conduct, reporting, ..." Iteration 1: Query: "Tell me about the regulations for staff members" Best match: Chunk 1 (score: 0.87) Extracted keyword: "internal" Iteration 2: Updated Query: "Tell me about the regulations for staff members internal" Best match: Chunk 2 (score: 0.93) Since further refinement doesn't significantly improve the score, the process stops. The system then uses these accumulated chunks to generate a grounded and comprehensive answer.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Retrieving the Best Chunk

The first step is to define a function that fetches the best matching chunk given a query. Take a look at the example below, with additional comments to make it clear:

JavaScript

async function retrieveBestChunk(queryText, collection, nResults = 1) {
    /**
     * Retrieve the best matching chunk from the collection based on the given query.
     * Returns:
     *   bestChunkText, bestChunkScore, bestChunkMetadata
     *   (or null, null, null if retrieval fails)
     */
    // Perform a similarity search for the provided query
    const retrieval = await collection.query({ queryTexts: [queryText], nResults: nResults });
    
    // If nothing is found, return null
    if (!retrieval.documents[0]) {
        return [null, null, null];
    }

    // Extract the best match from the results
    const bestChunkText = retrieval.documents[0][0];
    const bestDistance = retrieval.distances[0][0];
    
    // Convert 'distance' to a simple similarity score
    const bestChunkScore = 1 / (1 + bestDistance);
    const bestChunkMetadata = retrieval.metadatas[0][0];
    return [bestChunkText, bestChunkScore, bestChunkMetadata];
}

async function retrieveBestChunk(queryText, collection, nResults = 1) {
    /**
     * Retrieve the best matching chunk from the collection based on the given query.
     * Returns:
     *   bestChunkText, bestChunkScore, bestChunkMetadata
     *   (or null, null, null if retrieval fails)
     */
    // Perform a similarity search for the provided query
    const retrieval = await collection.query({ queryTexts: [queryText], nResults: nResults });
    
    // If nothing is found, return null
    if (!retrieval.documents[0]) {
        return [null, null, null];
    }

    // Extract the best match from the results
    const bestChunkText = retrieval.documents[0][0];
    const bestDistance = retrieval.distances[0][0];
    
    // Convert 'distance' to a simple similarity score
    const bestChunkScore = 1 / (1 + bestDistance);
    const bestChunkMetadata = retrieval.metadatas[0][0];
    return [bestChunkText, bestChunkScore, bestChunkMetadata];
}

What's happening here?

We query our collection (which contains chunks of text stored in a vector database) for the user's queryText.
If no chunks are returned, we gracefully exit with null.
Otherwise, we pick the top chunk (the first one) and calculate a simple inverted distance as a similarity score. This score is computed as 1 / (1 + bestDistance): note how this formula transforms the distance metric into a similarity score, where a smaller distance results in a higher similarity score. The addition of 1 prevents division by zero.
We return the chunk text and its metadata for further use.

Extracting And Refining Queries

Once you've retrieved a chunk, you often want to refine your query before the next retrieval pass. For instance, maybe the returned chunk contains a useful new keyword—like "internal" or "procedures"—that can make the next query more specific. Below is a snippet showing how you can extract a refinement keyword and append it to your query:

JavaScript

function extractRefinementKeyword(chunkText, currentQuery) {
    /**
     * Extract a single keyword from the chunk that is not already in the current query.
     * - Ignores stopwords and short words.
     * - Picks the longest remaining candidate.
     */
    const stopwords = new Set(["the", "and", "is", "in", "at", "of", "a", "to"]);
    const chunkWords = chunkText.toLowerCase().match(/\b\w+\b/g);
    const queryWords = new Set(currentQuery.toLowerCase().match(/\b\w+\b/g));

    // Filter out stopwords, words that already appear in the query, or very short words
    const candidateWords = chunkWords.filter(w => 
        !stopwords.has(w) && !queryWords.has(w) && w.length > 4
    );

    // If no candidates remain, return an empty string
    if (candidateWords.length === 0) {
        return "";
    }

    // Pick the longest candidate word among those left
    const refineWord = candidateWords.reduce((a, b) => a.length > b.length ? a : b);
    return refineWord;
}

function refineQuery(currentQuery, refineWord) {
    /**
     * Append the chosen refineWord to the current query if it exists.
     */
    // If there's no refinement word, do nothing
    if (!refineWord) {
        return currentQuery;
    }
    // Otherwise, enrich the query with the new keyword
    return `${currentQuery} ${refineWord}`;
}

Why is this useful?

By filtering out stopwords (common words like “the,” “and,” “is” that add little meaning), you reduce noise.
By skipping words already in the current query, you make sure you're not just repeating the same request.
By capturing only longer words, you often get more meaningful terms (“policy” vs. “at”).
Finally, refineQuery simply appends the chosen keyword, shaping a more specific query for the next retrieval cycle.

Putting It All Together: Iterative Retrieval

Below is a condensed example showing how we can loop through multiple retrieval steps, each time refining the query and checking whether we're truly improving:

JavaScript

async function iterativeRetrieval(query, collection, steps = 3) {
    /**
     * Multi-step retrieval with a simple query refinement approach:
     *   1) Retrieve the best chunk for the current query.
     *   2) Extract one new keyword from that chunk and add it to the query.
     *   3) Repeat until no improvements or keywords are found,
     *      or we've reached 'steps' iterations.
     */
    const accumulatedChunks = [];
    let currentQuery = query;
    let bestScoreSoFar = 0.0;
    const IMPROVEMENT_THRESHOLD = 0.02;

    for (let step = 0; step < steps; step++) {
        // Get the best chunk for the current query
        const [bestChunkText, bestChunkScore, metadata] = await retrieveBestChunk(currentQuery, collection);
        
        // End if nothing relevant is found or if new chunk score hasn't improved
        if (!bestChunkText || bestChunkScore <= bestScoreSoFar + IMPROVEMENT_THRESHOLD) {
            break;
        }

        // Update the best score tracker
        bestScoreSoFar = bestChunkScore;

        // Record the current step result
        accumulatedChunks.push({
            step: step + 1,
            query: currentQuery,
            retrievedChunk: bestChunkText,
            score: bestChunkScore
        });

        // Pick a new keyword from the chunk
        const refineWord = extractRefinementKeyword(bestChunkText, currentQuery);

        // Refine the query based on this new keyword
        currentQuery = refineQuery(currentQuery, refineWord);
    }
    
    return accumulatedChunks;
}

async function iterativeRetrieval(query, collection, steps = 3) {
    /**
     * Multi-step retrieval with a simple query refinement approach:
     *   1) Retrieve the best chunk for the current query.
     *   2) Extract one new keyword from that chunk and add it to the query.
     *   3) Repeat until no improvements or keywords are found,
     *      or we've reached 'steps' iterations.
     */
    const accumulatedChunks = [];
    let currentQuery = query;
    let bestScoreSoFar = 0.0;
    const IMPROVEMENT_THRESHOLD = 0.02;

    for (let step = 0; step < steps; step++) {
        // Get the best chunk for the current query
        const [bestChunkText, bestChunkScore, metadata] = await retrieveBestChunk(currentQuery, collection);
        
        // End if nothing relevant is found or if new chunk score hasn't improved
        if (!bestChunkText || bestChunkScore <= bestScoreSoFar + IMPROVEMENT_THRESHOLD) {
            break;
        }

        // Update the best score tracker
        bestScoreSoFar = bestChunkScore;

        // Record the current step result
        accumulatedChunks.push({
            step: step + 1,
            query: currentQuery,
            retrievedChunk: bestChunkText,
            score: bestChunkScore
        });

        // Pick a new keyword from the chunk
        const refineWord = extractRefinementKeyword(bestChunkText, currentQuery);

        // Refine the query based on this new keyword
        currentQuery = refineQuery(currentQuery, refineWord);
    }
    
    return accumulatedChunks;
}

Notes on what's happening:

We limit to a certain number of steps (e.g., 3 iterations).
In each iteration, we retrieve the best matching chunk.
We compare its score to the best score we've seen so far. If there's no improvement beyond a small threshold (like 0.02), we assume that refining further might not be worthwhile.
If we continue, we extract a fresh keyword from the chunk and add it to our query.
By the end, the list accumulatedChunks holds the chunks from each iteration—often providing a richer context than a single retrieval could.

Introduction

The Concept of Iterative Retrieval

Practical Example: Iterative Retrieval in Action

Join the 1M+ learners on CodeSignal

Retrieving the Best Chunk

Extracting And Refining Queries

Putting It All Together: Iterative Retrieval

Practical Considerations

Conclusion And Next Steps