Building a Document Processor

Welcome to the first lesson of our course on building a RAG-powered chatbot with Go! In this course, we'll create a complete Retrieval-Augmented Generation (RAG) system that can intelligently answer questions based on your documents.

At the heart of any RAG system is the document processor. This component is responsible for taking your raw documents, processing them into a format that can be efficiently searched, and retrieving the most relevant information when a query is made. Think of it as the librarian of your RAG system — organizing information and fetching exactly what you need when you ask for it.

The document processing pipeline we'll build today consists of several key steps:

1. Loading documents from files (such as PDFs)
2. Splitting these documents into smaller, manageable chunks
3. Creating vector embeddings for each chunk
4. Storing these embeddings in a vector database
5. Retrieving the most relevant chunks when a query is made

This document processor will serve as the foundation for our RAG chatbot. In later units, we'll build a chat engine that can maintain conversation history and then integrate both components into a complete RAG system. By the end of this course, you'll have a powerful chatbot that can answer questions based on your document collection with remarkable accuracy.

Let's start building our document processor!

First, we need to create a struct that will handle all our document processing needs. This struct will encapsulate the functionality for loading, processing, and retrieving information from documents using LangChain Go packages.

Let's start by setting up the basic structure of our DocumentProcessor struct:

package documentprocessor

import (
    "context"

    "github.com/tmc/langchaingo/embeddings"
    "github.com/tmc/langchaingo/schema"
    "github.com/tmc/langchaingo/textsplitter"

    "codesignal/memstore"
)

type DocumentProcessor struct {
    ChunkSize    int
    ChunkOverlap int
    Embedder     *embeddings.EmbedderImpl
    VectorStore  *memstore.Store
}

In this initialization, we're setting up several important parameters:

• ChunkSize: This determines how large each document chunk will be (measured in characters). A common starting point is around 1000 characters, which is a good balance between context size and specificity.
• ChunkOverlap: This specifies how much overlap there should be between consecutive chunks. Overlap helps maintain context across chunk boundaries.
• Embedder: This will hold our embedder instance from LangChain Go, which converts text into vector representations.
• VectorStore: This will hold our vector store from the memstore package, which efficiently stores and retrieves document embeddings.

These parameters can be adjusted based on your specific needs. For example, if you're working with technical documents where context is crucial, you might want to increase the chunk size and overlap.

Now that we have our struct set up, let's implement the methods for loading documents and splitting them into chunks. We'll use LangChain Go's document loaders and text splitters for this purpose.

First, we'll create a method to load documents from PDF files:

import (
    "os"

    "github.com/tmc/langchaingo/documentloaders"
)

func (dp *DocumentProcessor) LoadDocument(filePath string) ([]schema.Document, error) {
    // Open the PDF file
    file, err := os.Open(filePath)
    if err != nil {
        return nil, err
    }
    defer file.Close()

    // Create a PDF document loader
    loader := documentloaders.NewPDF(file, 0) // 0 = let the loader choose a default page batch size

    // Load the document
    docs, err := loader.Load(context.Background())
    if err != nil {
        return nil, err
    }

    return docs, nil
}

This method uses LangChain Go's documentloaders.NewPDF to load PDF files. The loader automatically extracts text from the PDF and returns it as a slice of schema.Document objects. You could easily extend this to support other file types by using different loaders like documentloaders.NewText for plain text files.

Next, let's implement the method that will process a document and add it to our vector store:

func (dp *DocumentProcessor) ProcessDocument(ctx context.Context, filePath string) error {
    // Load the document
    docs, err := dp.LoadDocument(filePath)
    if err != nil {
        return err
    }

    // Create a text splitter with our specified chunk size and overlap
    splitter := textsplitter.NewRecursiveCharacter(
        textsplitter.WithChunkSize(dp.ChunkSize),
        textsplitter.WithChunkOverlap(dp.ChunkOverlap),
    )

    // Split all loaded documents into chunks
    var allChunks []schema.Document
    for _, doc := range docs {
        // Split the text content of each document
        texts, err := splitter.SplitText(doc.PageContent)
        if err != nil {
            return err
        }
        
        // Create new documents from the text chunks
        for _, text := range texts {
            allChunks = append(allChunks, schema.Document{
                PageContent: text,
                Metadata:    doc.Metadata,
            })
        }
    }

    // Initialize vector store if it doesn't exist
    if dp.VectorStore == nil {
        store, err := memstore.New(
            memstore.WithEmbedder(dp.Embedder),
        )
        if err != nil {
            return err
        }
        dp.VectorStore = &store
    }

    // Add document chunks to the vector store
    _, err = dp.VectorStore.AddDocuments(ctx, allChunks)
    if err != nil {
        return err
    }

    return nil
}

This method demonstrates the complete document processing pipeline:

1. We load the document using our LoadDocument method
2. We create a RecursiveCharacter text splitter configured with our chunk size and overlap settings
3. We split all loaded documents into smaller chunks using SplitText, which returns the text content as strings
4. We reconstruct schema.Document objects from the split text chunks, preserving the original metadata
5. We initialize our vector store if it doesn't exist yet, passing in our embedder
6. We add all the document chunks to the vector store, which automatically generates and stores embeddings

The RecursiveCharacter splitter is particularly effective because it tries to split text at natural boundaries (like paragraphs and sentences) while maintaining our specified chunk size and overlap.

Now that we can process documents and store their embeddings, we need a way to retrieve relevant context when a query is made. This is where the "retrieval" part of RAG comes into play.

Let's implement a method to retrieve relevant document chunks for a given query:

func (dp *DocumentProcessor) RetrieveRelevantContext(ctx context.Context, query string, k int) ([]schema.Document, error) {
    // Check if vector store exists
    if dp.VectorStore == nil {
        return []schema.Document{}, nil
    }

    // Perform similarity search to find the k most relevant chunks
    docs, err := dp.VectorStore.SimilaritySearch(ctx, query, k)
    if err != nil {
        return nil, err
    }

    return docs, nil
}

This method takes a query string and a parameter k, which specifies how many chunks to retrieve. It then performs a similarity search in our vector store using the SimilaritySearch method from the memstore package. This method:

1. Converts the query into an embedding using our embedder
2. Computes similarity scores between the query embedding and all stored document embeddings
3. Returns the k most similar document chunks

If we haven't processed any documents yet, we simply return an empty slice.

Finally, let's add a utility method to reset our document processor:

func (dp *DocumentProcessor) Reset() {
    dp.VectorStore = nil
}

This method simply sets VectorStore to nil, effectively clearing our knowledge base. This can be useful if you want to start fresh with a new set of documents or when testing different document sets.

Now that we've built all the components of our document processor, let's see how to use it in a complete workflow. We'll create a simple example that:

1. Initializes our document processor with an embedder
2. Processes a PDF document
3. Retrieves relevant context for a query
4. Displays the retrieved chunks

Here's the complete example:

package main

import (
    "context"
    "fmt"
    "log"

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

    "codesignal/documentprocessor"
)

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

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

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

    // Initialize the document processor
    processor := &documentprocessor.DocumentProcessor{
        ChunkSize:    1000,
        ChunkOverlap: 100,
        Embedder:     embedder,
    }

    // Process a document
    filePath := "../data/a_scandal_in_bohemia.pdf"
    err = processor.ProcessDocument(ctx, filePath)
    if err != nil {
        log.Fatalf("error processing document: %v", err)
    }

    fmt.Println("Document processed successfully!")

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

    // Retrieve relevant context
    k := 3 // Number of chunks to retrieve
    relevantDocs, err := processor.RetrieveRelevantContext(ctx, query, k)
    if err != nil {
        log.Fatalf("error retrieving context: %v", err)
    }

    // Display the retrieved chunks
    fmt.Printf("\nQuery: %s\n", query)
    fmt.Printf("\nRetrieved %d relevant chunks:\n\n", len(relevantDocs))

    for i, doc := range relevantDocs {
        fmt.Printf("--- Chunk %d ---\n%s\n\n", i+1, doc.PageContent)
    }
}

This example demonstrates the complete document processing workflow:

1. We initialize an OpenAI LLM client configured for embeddings using the text-embedding-3-small model
2. We create an embedder by wrapping the LLM client
3. We initialize our DocumentProcessor with the embedder and our desired chunk settings
4. We process a PDF document, which loads it, splits it into chunks, and stores the embeddings
5. We define a query about the document
6. We retrieve the most relevant chunks for this query using similarity search
7. We display the retrieved chunks to see what context would be used for answering the query

When you run this code with the PDF of "A Scandal in Bohemia" (a Sherlock Holmes story), you might see output like:

Document processed successfully!

Query: What is the main mystery in the story?

Retrieved 3 relevant chunks:

--- Chunk 1 ---
The King of Bohemia has come to consult Sherlock Holmes about a delicate matter. He fears that a photograph in the possession of Irene Adler, an opera singer, could be used to prevent his upcoming marriage to a Scandinavian princess. The photograph shows the King with Irene Adler and could cause a scandal if revealed.

--- Chunk 2 ---
Holmes devises a plan to locate the photograph. He disguises himself and stages a fire alarm at Irene Adler's house, knowing that people instinctively save their most precious possessions during a fire. This trick reveals the location where she keeps the photograph.

--- Chunk 3 ---
However, when Holmes returns to retrieve the photograph, he discovers that Irene Adler has already left the country with it. She leaves behind a letter and a photograph of herself, explaining that she has no intention of using the compromising photograph but will keep it as protection against the King.

These retrieved chunks contain the most relevant information from the document related to our query. In later lessons, we'll use this retrieved context to generate intelligent responses with a chat model, but for now, we've successfully built a document processor that can load, process, and retrieve relevant information from documents!

In this lesson, we've built a powerful document processor for our RAG chatbot using Go and LangChain Go packages. We've learned how to:

• Create a DocumentProcessor struct that encapsulates document processing functionality
• Load documents from PDF files using LangChain Go's document loaders
• Split documents into manageable chunks with appropriate overlap using text splitters
• Generate embeddings and store them in a vector store using the memstore package
• Retrieve relevant context for user queries using similarity search

In the next unit, we'll build on this foundation by creating a chat engine that can maintain conversation history. This will allow our chatbot to have more natural, contextual conversations with users. Eventually, we'll integrate both components into a complete RAG system that can intelligently answer questions based on your documents while maintaining conversational context.

Get ready to practice what you've learned and take your RAG chatbot to the next level!