Retrieving and Utilizing Relevant Chunks with Java in RAG Systems

Welcome to the last lesson in our journey through Scaling Up RAG with Vector Databases! Well done, you're at the end of 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.

Chroma supports filtering by metadata and document contents using the where filter. This filter allows you to specify conditions that the metadata must meet for a document to be included in the results of a query. The where filter is structured as follows:

{
    "metadata_field": {
        <Operator>: <Value>
    }
}

• $eq: Equal to (string, int, float)
• $ne: Not equal to (string, int, float)
• $gt: Greater than (int, float)
• $gte: Greater than or equal to (int, float)
• $lt: Less than (int, float)
• $lte: Less than or equal to (int, float)
• $in: A value is in a predefined list (string, int, float, bool)
• $nin: A value is not in a predefined list (string, int, float, bool)

You can combine multiple filters using logical operators $and and $or.

• $and: Returns results that match all of the filters in the list.
• $or: Returns results that match any of the filters in the list.

{
    "$and": [
        {
            "category": {
                "$eq": "food"
            }
        },
        {
            "author": {
                "$in": ["Robert", "John", "Daniel"]
            }
        }
    ]
}

For example, this map will select all the documents with the category "food" and written by one of the authors of the list ["Robert", "John", "Daniel"]

To filter based on document contents, you need to provide a where_document filter dictionary in your query. This dictionary supports two filtering keys: $contains and $not_contains. The structure of the dictionary is as follows:

{
    "$contains": "technology",
    "$not_contains": "AI"
}

For example, this map will select all the documents that contain the word "technology" but not "AI"

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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import tech.amikos.chromadb.Collection;
import tech.amikos.chromadb.Collection.QueryResponse;
import tech.amikos.chromadb.handler.ApiException;

public class ChunkRetriever {

    public static List<Map<String, Object>> retrieveTopChunks(String query, Collection collection, int topK) 
            throws ApiException {
        QueryResponse results = collection.query(Arrays.asList(query), topK, null, null, null);

        List<Map<String, Object>> retrievedChunks = new ArrayList<>();

        // Safeguard for empty results
        if (results.getDocuments().isEmpty() || results.getDocuments().get(0).isEmpty()) {
            return retrievedChunks;
        }

        for (int i = 0; i < results.getDocuments().get(0).size(); i++) {
            Map<String, Object> chunk = new HashMap<>();
            chunk.put("chunk", results.getDocuments().get(0).get(i));
            chunk.put("doc_id", i);
            chunk.put("distance", results.getDistances().get(0).get(i));
            retrievedChunks.add(chunk);
        }

        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 vector database object containing our embedded documents;
    • topK: The number of most relevant chunks to retrieve.

• Vector Search:

  • The collection.query() function performs a vector-based similarity search to pinpoint which chunks are most aligned with the query.
  • The method returns a QueryResponse object, containing the documents and their distance scores.

• Results Structure:

  • Each result includes:
    • chunk: Contains the actual text chunk;
    • doc_id: Contains the document identifier;
    • distance: Indicates how semantically close a chunk is to your query — the lower the distance, the better the match.
  • For each result, the function creates a map with that information and appends it to the retrievedChunks list, which is then ultimately returned.

Once you have your relevant chunks, the next step is constructing a prompt that ensures the LLM focuses on precisely those chunks. This helps maintain factual accuracy and a tight context.

public class PromptBuilder {

    public static String buildPrompt(String query, List<Map<String, Object>> retrievedChunks) {
        StringBuilder prompt = new StringBuilder("Question: " + query + "\nAnswer using only the following context:\n");
        for (Map<String, Object> rc : retrievedChunks) {
            prompt.append("- ").append(rc.get("chunk")).append("\n");
        }
        prompt.append("Answer:");
        return prompt.toString();
    }
}

Why is this important?

1. Controlled Context: By explicitly instructing the LLM to focus on the given context, you reduce the probability of hallucinations.
2. Flexibility: You can modify the prompt format — like adding bullet points or rewording instructions — to direct the LLM's style or depth of response.
3. Clarity: Including the question upfront reminds the model of the exact query it must address.

We'll be seeing an actual prompt example later in the lesson!

With your collection in place, it's time to retrieve the most relevant chunks and put them to use in your prompt. The snippet below ties everything together: from forming the query, to constructing the prompt, and finally getting the answer from your Large Language Model.

public class QueryProcessor {

    public static void main(String[] args) {
        try {
            String query = "What are some recent technological breakthroughs?";
            List<Map<String, Object>> retrievedDocs = ChunkRetriever.retrieveTopChunks(query, collection, 5);
            String finalPrompt = PromptBuilder.buildPrompt(query, retrievedDocs);
            String answer = LLM.get_llm_response(finalPrompt);

            System.out.println("Prompt:\n");
            System.out.println(finalPrompt);
            System.out.println("\nLLM Answer: " + answer);

        } catch (ApiException e) {
            System.err.println("Error: " + e.getMessage());
            e.printStackTrace();
        }
    }
}

Here's what's happening step by step:

1. Formulating the Query: We define a query string that reflects the user's question or information request.
2. Retrieving Chunks: Using retrieveTopChunks, you get the top five chunks that closely match the query based on semantic similarity.
3. Prompt Construction: The function buildPrompt takes the user's question and the retrieved chunks to assemble a cohesive prompt.
4. LLM Response: Finally, get_llm_response is called with the constructed prompt, prompting the model to generate a context-informed answer.

By printing both the prompt and the answer, you can debug, refine, and further tailor your approach to retrieval and prompt design.

Below is an example of the system's final output after retrieving the most relevant chunks and assembling them into a prompt:

Prompt:

Question: What are some recent technological breakthroughs?
Answer using only the following context:
- The Industrial Revolution brought significant technological and social changes. It reshaped economies and altered the fabric of society. Scholars examine its impact on labor, innovation, and modern industrial practices.
- Breakthroughs in renewable energy technologies are reducing global dependence on fossil fuels. Solar and wind systems are becoming more efficient and affordable. These innovations are crucial to combating climate change and ensuring a sustainable future.
- The digital revolution is transforming how we approach health and wellness. Technological innovations, from fitness trackers to health apps, are empowering individuals to manage their well-being. This integration of technology and lifestyle is reshaping daily habits for a healthier future.
- Advances in medical technology are revolutionizing patient care through new diagnostic and treatment methods. Breakthroughs in imaging and robotics are enhancing the precision of medical procedures. Healthcare professionals are optimistic about the potential for improved outcomes.
- Scientists are developing renewable materials that could replace traditional plastics. Innovations in biopolymers are leading to sustainable manufacturing practices. These breakthroughs promise to reduce environmental waste and support a circular economy.
Answer:

LLM Answer: Recent technological breakthroughs include advancements in renewable energy technologies, which are making solar and wind systems more efficient and affordable, thereby reducing global dependence on fossil fuels and aiding in the fight against climate change. Additionally, the digital revolution is enhancing health and wellness through innovations like fitness trackers and health apps, empowering individuals to better manage their well-being. In the medical field, new diagnostic and treatment methods, along with improvements in imaging and robotics, are revolutionizing patient care and enhancing the precision of medical procedures. Furthermore, scientists are developing renewable materials, such as biopolymers, to replace traditional plastics, promoting sustainable manufacturing practices and supporting a circular economy.

In this snippet, the prompt clearly instructs the LLM to focus on the listed chunks. By doing so, the final LLM Answer highlights the key points about recent breakthroughs in renewable energy, healthcare innovations, and sustainable materials, reflecting the relevance of the context. Interestingly, the chunk referencing the Industrial Revolution is not directly invoked in the final answer, showcasing the LLM's ability to select and incorporate only the most suitable context. Notice how each retrieved chunk contributes to a coherent, context-based response, demonstrating how RAG systems help reduce hallucinations and maintain factual alignment.

In this lesson, you discovered how to:

• Retrieve the most relevant text chunks in your vector database through semantic similarity.
• Construct a well-structured prompt so the LLM stays true to the provided text.

These steps are central to building a robust Retrieval-Augmented Generation pipeline. By creating focused, context-driven prompts, your LLM's responses tend to be more accurate and trustworthy.

Next, you'll have the opportunity to practice and solidify this knowledge. Look for the exercises that follow to test retrieving chunks with different queries, adjusting the prompt format, and experimenting with how the LLM responds. Keep pushing those boundaries — your mastery of RAG systems is well underway!