Asking Questions with Retrieved Context and Templates

Welcome to the final lesson of this course! In this lesson, we will integrate context retrieval with a chat model using LangChain4j. This builds on the skills you've developed in previous lessons, where you learned about document embeddings and similarity search. Today, we'll focus on using Java's string manipulation techniques to format messages with extra context, enabling you to ask questions and receive answers based on the retrieved document content. This lesson will bring together all the skills you've learned so far, culminating in a comprehensive understanding of document processing and retrieval with Java.

Quick Reminder: Preparing Documents and Creating a Vector Store

Let's quickly recap what we've learned in previous lessons about preparing documents and creating a vector store. We'll load and prepare our document, "The Adventure of the Blue Carbuncle," and generate embeddings to create a vector store. This process is essential for effective context retrieval.

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import java.util.List;
import java.util.ArrayList;
import java.time.Duration;

public class Main {
    public static void main(String[] args) {
        // Define the file path
        String filePath = "data/the_adventure_of_the_blue_carbuncle.pdf";

        // Load and split the document
        Document document = FileSystemDocumentLoader.loadDocument(filePath, new ApachePdfBoxDocumentParser());
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(1000, 100);
        List<TextSegment> chunks = splitter.split(document);

        // Initialize the embedding model
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
            .apiKey(System.getenv("OPENAI_API_KEY"))
            .modelName("text-embedding-3-small")
            .dimensions(1536)
            .timeout(Duration.ofSeconds(60))
            .build();

        // Create a vector store
        InMemoryEmbeddingStore<TextSegment> vectorStore = new InMemoryEmbeddingStore<>();

        // Generate embeddings for all chunks and store in a list
        List<Embedding> embeddings = new ArrayList<>();
        for (TextSegment chunk : chunks) {
            Embedding embedding = embeddingModel.embed(chunk.text()).content();
            embeddings.add(embedding);
        }

        // Add all embeddings and chunks to the store at once
        vectorStore.addAll(embeddings, chunks);
    }
}

In this code, we load a document using LangChain4j's document loaders, split it into chunks with a paragraph splitter, generate embeddings using an OpenAI model, and create an in-memory vector store. This sets the stage for efficient context retrieval in our question-answering tasks.

Combining Retrieved Context

Now that we have our vector store, we can integrate context retrieval with a chat model. First, we'll define a query and perform a similarity search to retrieve relevant documents based on the query. This will allow us to combine the retrieved document content to form a context for our question.

public class Main {
    public static void main(String[] args) {
        // Previous code for loading and embedding the document...

        // Define a query
        String query = "What is the main mystery in the story?";

        // Generate embedding for the query
        Embedding queryEmbedding = embeddingModel.embed(query).content();

        // Create a search request for top 3 matches
        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
            .queryEmbedding(queryEmbedding)
            .maxResults(3)
            .build();

        // Perform the search
        EmbeddingSearchResult<TextSegment> matches = vectorStore.search(searchRequest);

        // Combine the content of retrieved documents
        StringBuilder contextBuilder = new StringBuilder();
        for (EmbeddingMatch<TextSegment> match : matches.matches()) {
            TextSegment segment = match.embedded();
            contextBuilder.append(segment.text()).append("\n\n");
        }
        String context = contextBuilder.toString();
    }
}

In this example, we define a query about the main mystery in the story, generate an embedding for it, and retrieve the top three most relevant document chunks. The content of these chunks is combined to form a context that will be used in the next step.

Formatting Messages with Templates

To effectively communicate with the chat model, we need to format our messages using Java's string manipulation techniques. Think of a template as a fill-in-the-blank form where you can insert specific pieces of information. In our case, we want to insert the context we retrieved from the similarity search and the question we want to ask.

public class Main {
    public static void main(String[] args) {
        // Previous code for loading, embedding, and searching...

        // Create a prompt template for RAG
        String promptTemplate = "Answer the following question based on the provided context.\n\n"
                + "Context:\n%s\n\n"
                + "Question: %s";

        // Format the prompt with our context and query
        String prompt = String.format(promptTemplate, context, query);
    }
}

In this code, we create a prompt template that includes placeholders for the context and the question. Java's String.format method helps us define this structure.

Exploring the Formatted Prompt

Now, let's print the formatted template to see how the context and question are structured in the message.

public class Main {
    public static void main(String[] args) {
        // Previous code for loading, embedding, searching, and formatting...

        // Print the formatted prompt
        System.out.println("Formatted Prompt:");
        System.out.println(prompt);
    }
}

This will output something like:

Formatted Prompt:
Answer the following question based on the provided context.

Context:
The little man stood glancing from one to the
other of us with half-frightened, half-hopeful eyes,
as one who is not sure whether he is on the verge
of a windfall or of a catastrophe...

Question: What is the main mystery in the story?

By printing the formatted template, we can verify that the context and question are correctly inserted into the template. This ensures that the chat model receives a well-structured message, allowing it to generate a relevant and accurate response.

Asking a Question with Retrieved Context to a Chat Model

With our prompt ready, we can now move on to interacting with the chat model. We'll instantiate the chat model and use our formatted template to get a response.

public class Main {
    public static void main(String[] args) {
        // Previous code for loading, embedding, searching, and formatting...

        // Initialize the chat model
        ChatLanguageModel chatModel = OpenAiChatModel.builder()
            .apiKey(System.getenv("OPENAI_API_KEY"))
            .modelName("gpt-3.5-turbo")
            .temperature(0.7)
            .maxTokens(500)
            .timeout(Duration.ofSeconds(60))
            .build();

        // Create a system message
        SystemMessage systemMessage = new SystemMessage(
                "You are an assistant that answers questions based on the provided context. " +
                "If no context is given, state that you don't have enough information.");

        // Invoke the model with context
        String response = chatModel.chat(List.of(
                systemMessage,
                new UserMessage(prompt)
        )).aiMessage().text();

        // Print the question and the AI's answer
        System.out.println("Question: " + query);
        System.out.println("Answer: " + response);
    }
}

In this section, we initialize a chat model using OpenAI's API and invoke it with our formatted prompt. The model processes the input and generates an answer based on the context and question provided. By printing both the question and the AI's answer, we can see how effectively the integration of context retrieval and templates enhances the interaction.

Summary and Next Steps

You've successfully completed this lesson, where you learned how to integrate context retrieval with a chat model using LangChain4j. We explored the use of Java's string manipulation techniques to format messages with extra context, allowing you to ask questions and receive answers based on the retrieved document content. This lesson consolidated all the skills you've acquired so far, equipping you with a solid understanding of document processing and retrieval with Java. As you continue your learning journey, consider experimenting with different queries and document types to deepen your understanding. Stay tuned for the next course, where we'll build on these concepts and explore even more advanced techniques.

Sign up
Join the 1M+ learners on CodeSignal
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal