Asking Questions with Retrieved Context and Templates in Go

Welcome to the final lesson of this course! In this lesson, we will integrate context retrieval with a chat model using Go. 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 templates 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 Go.

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, "Alice in Wonderland," and generate embeddings to create a vector store. This process is essential for effective context retrieval.

package main

import (
    "context"
    "fmt"
    "log"
    "os"

    "github.com/tmc/langchaingo/documentloaders"
    "github.com/tmc/langchaingo/embeddings"
    "github.com/tmc/langchaingo/llms/openai"
    "github.com/tmc/langchaingo/textsplitter"

    "codesignal/memstore"
)

func main() {
    ctx := context.Background()

    // Load the text file
    file, err := os.Open("data/alice_in_wonderland.txt")
    if err != nil {
        log.Fatalf("failed to open file: %v", err)
    }
    defer file.Close()

    loader := documentloaders.NewText(file)

    // Create a text splitter with specific chunk size and overlap
    splitter := textsplitter.NewRecursiveCharacter(
        textsplitter.WithChunkSize(500),
        textsplitter.WithChunkOverlap(100),
    )

    // Load and split the document in one step
    docs, err := loader.LoadAndSplit(ctx, splitter)
    if err != nil {
        log.Fatalf("failed to load and split document: %v", err)
    }

    // Initialize OpenAI LLM with embedding model
    llm, err := openai.New(
        openai.WithEmbeddingModel("text-embedding-3-small"),
    )
    if err != nil {
        log.Fatalf("failed to initialize OpenAI client: %v", err)
    }

    // Wrap the LLM in an Embedder
    embedder, err := embeddings.NewEmbedder(llm)
    if err != nil {
        log.Fatalf("failed to create embedder: %v", err)
    }

    // Create the in-memory vector store with the embedder
    store, err := memstore.New(
        memstore.WithEmbedder(embedder),
    )
    if err != nil {
        log.Fatalf("failed to create memstore: %v", err)
    }

    // Add all document chunks to the vector store
    _, err = store.AddDocuments(ctx, docs)
    if err != nil {
        log.Fatalf("failed to add documents to memstore: %v", err)
    }

    fmt.Println("Vector store created successfully!")
}

In this code, we load a document, split it into chunks, generate embeddings, and create a vector store, setting the stage for efficient context retrieval in our question-answering tasks.

Note: The RecursiveCharacter splitter uses character-based measurements for ChunkSize and ChunkOverlap. If you use a token-based splitter like TokenSplitter instead, these parameters would represent tokens rather than characters—so 500 tokens is not equivalent to 500 characters and would typically represent significantly more text.

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.

// Define a query
query := "What happens when Alice drinks the potion?"

// Retrieve relevant documents
k := 3 // Number of results to retrieve
retrievedDocs, err := store.SimilaritySearch(ctx, query, k)
if err != nil {
    log.Fatalf("similarity search failed: %v", err)
}

// Combine the content of retrieved documents
var context string
for _, doc := range retrievedDocs {
    context += doc.PageContent + "\n\n"
}

In this example, we define a query and retrieve the top three most relevant document chunks using the SimilaritySearch method. The content of these chunks is combined to form a context that will be used in the next step.

To effectively communicate with the chat model, we need to format our messages using templates. In Go, we can use the text/template package to create these templates. It allows us to define a structure for our message, ensuring that the chat model receives all the necessary information to provide an accurate response.

import (
    "text/template"
    "bytes"
)

// Create a prompt template for RAG
const promptTemplate = `Answer the following question based on the provided context.

Context:
{{.Context}}

Question: {{.Question}}
`

// Define a struct to hold template data
type TemplateData struct {
    Context  string
    Question string
}

// Format the prompt with our context and query
tmpl, err := template.New("prompt").Parse(promptTemplate)
if err != nil {
    log.Fatalf("failed to parse template: %v", err)
}

data := TemplateData{
    Context:  context,
    Question: query,
}

var formattedPrompt bytes.Buffer
err = tmpl.Execute(&formattedPrompt, data)
if err != nil {
    log.Fatalf("failed to execute template: %v", err)
}

In this code, we create a prompt template that includes placeholders for the context and the question. The text/template package helps us define this structure by allowing us to specify template variables using {{.FieldName}} syntax.

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

// Print the formatted prompt
fmt.Println(formattedPrompt.String())

This will output something like:

Answer the following question based on the provided context.

Context:
Alice opened the door and found that it led into a small
passage, not much larger than a rat-hole: she knelt down and
looked along the passage into the loveliest garden you ever saw.
How she longed to get out of that dark hall, and wander about
among those beds of bright flowers and those cool fountains...

Question: What happens when Alice drinks the potion?

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.

With our prompt ready, we can now move on to interacting with the chat model. We'll use the OpenAI LLM client from LangChain to get a response. Note that we create a new LLM instance configured for chat completion rather than embeddings.

import (
    "github.com/tmc/langchaingo/llms"
)

// Initialize the chat model (separate from embeddings)
chatLLM, err := openai.New(
    openai.WithModel("gpt-3.5-turbo"),
)
if err != nil {
    log.Fatalf("failed to initialize chat LLM: %v", err)
}

// Get the response from the model
response, err := llms.GenerateFromSinglePrompt(ctx, chatLLM, formattedPrompt.String())
if err != nil {
    log.Fatalf("failed to get response from chat model: %v", err)
}

// Print the question and the AI's answer
fmt.Printf("\nQuestion: %s\n", query)
fmt.Printf("Answer: %s\n", response)

In this section, we initialize a separate OpenAI LLM client configured for chat completion with the gpt-3.5-turbo model. We then use llms.GenerateFromSinglePrompt to send our formatted prompt to the model. 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.

The output will look something like this:

Question: What happens when Alice drinks the potion?
Answer: When Alice drinks the potion, she shrinks down to a very small size, allowing her to fit through the small passage and enter the beautiful garden she had been longing to explore.

This demonstrates how the chat model uses the context retrieved from the vector store to provide a relevant and accurate response to the question.

You've successfully completed this lesson, where you learned how to integrate context retrieval with a chat model using Go and LangChain. We explored the use of templates 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 Go.

Key takeaways:

• Templates help structure prompts with context and questions for chat models
• Go's text/template package provides powerful templating capabilities
• Retrieved context from similarity search enhances the quality of AI responses
• Separate LLM instances are used for embeddings and chat completion
• The complete RAG (Retrieval-Augmented Generation) pipeline combines all these components

As you continue your learning journey, consider experimenting with different queries and document types to deepen your understanding. Try adjusting the number of retrieved documents (k parameter) to see how it affects response quality. Stay tuned for the practice exercises, where you'll get hands-on experience implementing these concepts yourself!