Iterative Retrieval for Enhanced RAG Pipelines

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

Before we dive into the code, let's clarify what we mean by "iterative retrieval" and why it matters. This section will set the stage for the practical implementation details that follow. 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. Below, we'll walk through the building blocks of an iterative retrieval system in detail.

Practical Example: Iterative Retrieval in Action

Now that you understand the concept, let's see how iterative retrieval works in practice. This example will help you visualize the process before we break it down into code. 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. With this example in mind, let's move on to the code that enables each step of this process.

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 in iterative retrieval is to fetch the most relevant chunk for a given query. In this section, we'll look at a function that does exactly that. We'll walk through the code and highlight the key parts with comments so you can see how the retrieval works under the hood.

Rust

pub async fn retrieve_best_chunk(
    collection: &ChromaCollection,
    embedder: &SentenceEmbedder,
    query: &str,
    n_results: usize,
) -> Result<Option<(String, f32, Value)>, Box<dyn Error>> {
    // Generate embeddings for the query text
    let query_embeddings = embedder.embed_texts(&[query])?;
    // Set up query options for the vector database
    let opts = QueryOptions {
        query_texts: None,
        query_embeddings: Some(query_embeddings),
        n_results: Some(n_results),
        where_metadata: None,
        where_document: None,
        include: Some(vec!["documents".into(), "distances".into(), "metadatas".into()]),
    };

    // Query the collection asynchronously
    let res = collection.query(opts, None).await?;

    // Safely extract the list of documents (chunks) from the response
    let empty_vec = Vec::new();
    let docs = res
        .documents
        .as_ref()
        .and_then(|groups| groups.get(0))
        .unwrap_or(&empty_vec);

    // If no documents are found, return None
    if docs.is_empty() {
        return Ok(None);
    }

    // Get the best matching chunk (first result)
    let text = docs[0].clone();
    // Calculate the similarity score (higher is better)
    let distance = res
        .distances
        .as_ref()
        .and_then(|groups| groups.get(0))
        .and_then(|row| row.get(0))
        .copied()
        .unwrap_or(0.0);

    let score = 1.0 / (1.0 + distance);

    // Extract metadata for the chunk, if available
    let metadata = res
        .metadatas
        .as_ref()
        .and_then(|groups| groups.get(0))
        .and_then(|row| row.get(0))
        .and_then(|m| serde_json::to_value(m).ok())
        .unwrap_or(Value::Null);

    // Return the chunk text, score, and metadata
    Ok(Some((text, score, metadata)))
}

pub async fn retrieve_best_chunk(
    collection: &ChromaCollection,
    embedder: &SentenceEmbedder,
    query: &str,
    n_results: usize,
) -> Result<Option<(String, f32, Value)>, Box<dyn Error>> {
    // Generate embeddings for the query text
    let query_embeddings = embedder.embed_texts(&[query])?;
    // Set up query options for the vector database
    let opts = QueryOptions {
        query_texts: None,
        query_embeddings: Some(query_embeddings),
        n_results: Some(n_results),
        where_metadata: None,
        where_document: None,
        include: Some(vec!["documents".into(), "distances".into(), "metadatas".into()]),
    };

    // Query the collection asynchronously
    let res = collection.query(opts, None).await?;

    // Safely extract the list of documents (chunks) from the response
    let empty_vec = Vec::new();
    let docs = res
        .documents
        .as_ref()
        .and_then(|groups| groups.get(0))
        .unwrap_or(&empty_vec);

    // If no documents are found, return None
    if docs.is_empty() {
        return Ok(None);
    }

    // Get the best matching chunk (first result)
    let text = docs[0].clone();
    // Calculate the similarity score (higher is better)
    let distance = res
        .distances
        .as_ref()
        .and_then(|groups| groups.get(0))
        .and_then(|row| row.get(0))
        .copied()
        .unwrap_or(0.0);

    let score = 1.0 / (1.0 + distance);

    // Extract metadata for the chunk, if available
    let metadata = res
        .metadatas
        .as_ref()
        .and_then(|groups| groups.get(0))
        .and_then(|row| row.get(0))
        .and_then(|m| serde_json::to_value(m).ok())
        .unwrap_or(Value::Null);

    // Return the chunk text, score, and metadata
    Ok(Some((text, score, metadata)))
}

This function takes a query, embeds it, and retrieves the best-matching chunk from your vector database. The comments above explain each step, from embedding the query to extracting the top result and its score.

Now that we can retrieve the best chunk, let's see how we can refine our queries to dig deeper.

Extracting And Refining Queries

After retrieving a chunk, the next step is to refine your query for the next iteration. This involves extracting a useful keyword from the retrieved chunk and appending it to your current query. The following code snippets show how to do this, with comments to clarify each part.

Rust

pub fn extract_refinement_keyword(chunk_text: &str, current_query: &str) -> String {
    // Convert both the chunk and query to lowercase for comparison
    let chunk_lower = chunk_text.to_lowercase();
    let query_lower = current_query.to_lowercase();

    // Split the chunk into words, remove punctuation, and collect as a vector
    let chunk_words: Vec<String> = chunk_lower
        .split_whitespace()
        .map(|w| w.chars().filter(|c| c.is_alphanumeric()).collect())
        .filter(|w: &String| !w.is_empty())
        .collect();

    // Build a set of words already present in the query
    let query_words: std::collections::HashSet<String> = query_lower
        .split_whitespace()
        .map(|w| w.chars().filter(|c| c.is_alphanumeric()).collect())
        .filter(|w: &String| !w.is_empty())
        .collect();

    // Find candidate words that are not in the query, not stopwords, and are long enough
    let candidates: Vec<String> = chunk_words.into_iter()
        .filter(|w| w.len() > 4 && !STOPWORDS.contains(&w.as_str()) && !query_words.contains(w))
        .collect();

    // Return the longest candidate word, or an empty string if none found
    candidates.into_iter()
        .max_by_key(|w| w.len())
        .unwrap_or_default()
}

// Appends the refinement keyword to the current query, if one was found
pub fn refine_query(current_query: &str, refine_word: &str) -> String {
    if refine_word.is_empty() {
        current_query.to_string()
    } else {
        format!("{} {}", current_query, refine_word)
    }
}

pub fn extract_refinement_keyword(chunk_text: &str, current_query: &str) -> String {
    // Convert both the chunk and query to lowercase for comparison
    let chunk_lower = chunk_text.to_lowercase();
    let query_lower = current_query.to_lowercase();

    // Split the chunk into words, remove punctuation, and collect as a vector
    let chunk_words: Vec<String> = chunk_lower
        .split_whitespace()
        .map(|w| w.chars().filter(|c| c.is_alphanumeric()).collect())
        .filter(|w: &String| !w.is_empty())
        .collect();

    // Build a set of words already present in the query
    let query_words: std::collections::HashSet<String> = query_lower
        .split_whitespace()
        .map(|w| w.chars().filter(|c| c.is_alphanumeric()).collect())
        .filter(|w: &String| !w.is_empty())
        .collect();

    // Find candidate words that are not in the query, not stopwords, and are long enough
    let candidates: Vec<String> = chunk_words.into_iter()
        .filter(|w| w.len() > 4 && !STOPWORDS.contains(&w.as_str()) && !query_words.contains(w))
        .collect();

    // Return the longest candidate word, or an empty string if none found
    candidates.into_iter()
        .max_by_key(|w| w.len())
        .unwrap_or_default()
}

// Appends the refinement keyword to the current query, if one was found
pub fn refine_query(current_query: &str, refine_word: &str) -> String {
    if refine_word.is_empty() {
        current_query.to_string()
    } else {
        format!("{} {}", current_query, refine_word)
    }
}

In these functions, we first extract a keyword from the retrieved chunk that isn't already in the query and isn't a common stopword. We then append this keyword to the query to make it more specific for the next retrieval step. The comments in the code walk you through the logic.

Also note that appending a keyword to the end of the query can sometimes create awkward or ambiguous phrases, which may reduce retrieval accuracy. When possible, consider rephrasing the query or inserting the keyword in a way that preserves the original meaning and grammar.

With the ability to retrieve and refine, let's see how to combine these steps into an iterative process.

Iterative Retrieval Loop

Now that you know how to retrieve the best chunk and refine your query, let's look at how to loop through these steps to perform iterative retrieval. The following function demonstrates the full process, with comments to explain each stage.

Rust

pub async fn iterative_retrieval(
    collection: &ChromaCollection,
    embedder: &SentenceEmbedder,
    initial_query: &str,
    steps: usize,
    improvement_threshold: f32,
) -> Result<Vec<IterationResult>, Box<dyn std::error::Error>> {
    let mut results = Vec::new();
    let mut current_query = initial_query.to_string();
    let mut best_score = 0.0;

    for step in 1..=steps {
        println!("Iteration {}, current query: '{}'", step, current_query);
        // Retrieve the best chunk for the current query
        let opt = retrieve_best_chunk(collection, embedder, &current_query, 1).await?;
        let (text, score, metadata) = match opt {
            Some(t) => t,
            None => {
                println!("No chunks found at this step. Ending.");
                break;
            }
        };

        println!("Best chunk (50 chars): '{}' | Score: {:.4}", &text[..text.len().min(50)], score);

        // Check if the improvement in score is significant enough to continue
        if score - best_score < improvement_threshold {
            println!("Improvement threshold not met. Stopping.");
            break;
        }
        best_score = score;

        // Store the result of this iteration
        results.push(IterationResult {
            step,
            query: current_query.clone(),
            retrieved_text: text.clone(),
            metadata,
            score,
        });

        // Extract a keyword from the retrieved chunk to refine the query
        let kw = extract_refinement_keyword(&text, &current_query);
        if kw.is_empty() {
            println!("No suitable keyword for further refinement.");
            break;
        }
        println!("Refining query with keyword: {}", kw);
        // Update the query for the next iteration
        current_query = refine_query(&current_query, &kw);
    }

    Ok(results)
}

pub async fn iterative_retrieval(
    collection: &ChromaCollection,
    embedder: &SentenceEmbedder,
    initial_query: &str,
    steps: usize,
    improvement_threshold: f32,
) -> Result<Vec<IterationResult>, Box<dyn std::error::Error>> {
    let mut results = Vec::new();
    let mut current_query = initial_query.to_string();
    let mut best_score = 0.0;

    for step in 1..=steps {
        println!("Iteration {}, current query: '{}'", step, current_query);
        // Retrieve the best chunk for the current query
        let opt = retrieve_best_chunk(collection, embedder, &current_query, 1).await?;
        let (text, score, metadata) = match opt {
            Some(t) => t,
            None => {
                println!("No chunks found at this step. Ending.");
                break;
            }
        };

        println!("Best chunk (50 chars): '{}' | Score: {:.4}", &text[..text.len().min(50)], score);

        // Check if the improvement in score is significant enough to continue
        if score - best_score < improvement_threshold {
            println!("Improvement threshold not met. Stopping.");
            break;
        }
        best_score = score;

        // Store the result of this iteration
        results.push(IterationResult {
            step,
            query: current_query.clone(),
            retrieved_text: text.clone(),
            metadata,
            score,
        });

        // Extract a keyword from the retrieved chunk to refine the query
        let kw = extract_refinement_keyword(&text, &current_query);
        if kw.is_empty() {
            println!("No suitable keyword for further refinement.");
            break;
        }
        println!("Refining query with keyword: {}", kw);
        // Update the query for the next iteration
        current_query = refine_query(&current_query, &kw);
    }

    Ok(results)
}

This function repeatedly retrieves the best chunk, checks if the retrieval is improving, and refines the query with a new keyword. The process continues until the improvement threshold is not met or the maximum number of steps is reached. The comments in the code clarify each part of the loop.

Tip: What counts as a "significant improvement" depends on your dataset. In dense datasets, even a small score increase (e.g., 0.01) might matter, while sparser datasets may require a larger jump. Always calibrate improvement_threshold to fit your data.

With the iterative retrieval loop in place, let's discuss some practical considerations and limitations you should be aware of.

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

Iterative Retrieval Loop

Practical Considerations

Conclusion And Next Steps