Integrating the RAG Chatbot

Welcome to the third lesson of our course on building a RAG-powered chatbot! In the previous lessons, we've built two essential components: a document processor that handles the retrieval of relevant information and a chat engine that manages conversations with users. Now, it's time to bring these components together to create a complete Retrieval-Augmented Generation (RAG) system.

In this lesson, we'll integrate our document processor and chat engine into a unified RAGChatbot struct. This integration will create a seamless experience where users can upload documents, ask questions about them, and receive informed responses based on the document content. By the end of this lesson, you'll have a fully functional RAG chatbot that can answer questions about any documents you provide. This represents the culmination of our work so far, bringing together retrieval and generation in a practical, user-friendly system.

Let's start building our integrated RAG chatbot!

The first step in our integration is to create a new struct that will serve as the main interface for our RAG chatbot. This struct will coordinate between the document processor and chat engine components we've already built.

Let's create our RAGChatbot struct in a new package:

package ragchatbot

import (
    "context"
    "fmt"

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

    "codesignal/chatengine"
    "codesignal/documentprocessor"
)

type RAGChatbot struct {
    DocumentProcessor *documentprocessor.DocumentProcessor
    ChatEngine        *chatengine.ChatEngine
}

func NewRAGChatbot(ctx context.Context) (*RAGChatbot, error) {
    // Initialize OpenAI LLM for embeddings
    embeddingLLM, err := openai.New(
        openai.WithEmbeddingModel("text-embedding-3-small"),
    )
    if err != nil {
        return nil, fmt.Errorf("failed to initialize embedding LLM: %w", err)
    }

    // Create an embedder
    embedder, err := embeddings.NewEmbedder(embeddingLLM)
    if err != nil {
        return nil, fmt.Errorf("failed to create embedder: %w", err)
    }

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

    // Initialize chat engine
    chatEngine, err := chatengine.NewChatEngine(ctx)
    if err != nil {
        return nil, fmt.Errorf("failed to initialize chat engine: %w", err)
    }

    return &RAGChatbot{
        DocumentProcessor: docProcessor,
        ChatEngine:        chatEngine,
    }, nil
}

This initialization is straightforward but powerful. We're creating instances of both our DocumentProcessor and ChatEngine structs, which we developed in the previous lessons. This struct will serve as the coordinator between these components, handling the flow of information from document processing to context retrieval to conversation management.

This design follows the principle of separation of concerns, where each component has a specific responsibility:

• The DocumentProcessor handles document loading, chunking, embedding, and retrieval.
• The ChatEngine manages the conversation flow and language model interactions.
• The RAGChatbot coordinates between these components and provides a unified interface.

This architecture makes our system modular and maintainable. If we want to improve our document processing or chat capabilities in the future, we can update the respective components without affecting the overall system.

Now that we have our basic struct structure, let's implement the document management functionality. The first method we'll add is UploadDocument, which will handle document processing:

func (r *RAGChatbot) UploadDocument(ctx context.Context, filePath string) (string, error) {
    err := r.DocumentProcessor.ProcessDocument(ctx, filePath)
    if err != nil {
        return "", fmt.Errorf("failed to process document: %w", err)
    }
    return "Document successfully processed.", nil
}

This method serves as a wrapper around our document processor's ProcessDocument method, but with added error handling and proper context passing. If the document processor encounters an issue (such as an unsupported file format, a corrupted file, or a file that is too large to process), it will return an error. Our UploadDocument method wraps this error with additional context and returns it so it can be handled by the caller.

The heart of our RAG chatbot is the message processing pipeline, which connects user queries to document retrieval and response generation. Let's implement the SendMessage method:

func (r *RAGChatbot) SendMessage(ctx context.Context, message string) (string, error) {
    // Retrieve relevant document chunks based on the user's query
    k := 3 // Number of chunks to retrieve
    relevantDocs, err := r.DocumentProcessor.RetrieveRelevantContext(ctx, message, k)
    if err != nil {
        return "", fmt.Errorf("failed to retrieve context: %w", err)
    }

    // Initialize an empty string for the context
    context := ""

    // Loop through each relevant document
    for _, doc := range relevantDocs {
        // Extract the source from metadata, defaulting to "unknown" if not available
        source := "unknown"
        if val, ok := doc.Metadata["source"]; ok {
            // Safely assert the type to string
            if sourceStr, ok := val.(string); ok {
                source = sourceStr
            }
        }
        // Extract the content of the document
        content := doc.PageContent
        // Append the source and content to the context string
        context += "Source: " + source + "\n" + content + "\n\n"
    }

    // Send the user's message along with the context to the chat engine
    response, err := r.ChatEngine.SendMessage(ctx, message, context)
    if err != nil {
        return "", fmt.Errorf("failed to get response: %w", err)
    }
    return response, nil
}

This method implements the core RAG workflow:

1. It takes a user message and context as input.
2. It uses the document processor to retrieve the top 3 most relevant document chunks based on the message.
3. It builds a context string that includes both the content of each document chunk and its source.
4. It sends the original message and the retrieved context to the chat engine.
5. It returns the response from the chat engine.

The magic of RAG happens in this method. When a user asks a question, the system automatically searches through all processed documents to find relevant information. This relevant context, along with source attribution, is then provided to the language model along with the user's question, allowing it to generate an informed response based on the document content.

Including the source information enhances transparency by allowing the model to reference where the information came from. If no relevant documents are found, an empty context is provided. In this case, our chat engine (as we designed it in the previous lesson) will inform the user that it doesn't have enough information to answer the question.

To complete our RAG chatbot, let's add some system management features that will help users control the state of the chatbot.

First, let’s add a method to reset the document knowledge:

func (r *RAGChatbot) ResetDocuments() string {
    r.DocumentProcessor.Reset()
    return "Document knowledge has been reset."
}

This method calls the Reset method of our document processor, which clears the vector store. This is useful when users want to start fresh with a new set of documents or when they want to remove previously processed documents from the chatbot's knowledge.

We also need a way to reset the conversation history:

func (r *RAGChatbot) ResetConversation() string {
    r.ChatEngine.ResetConversation()
    return "Conversation history has been reset."
}

This method simply calls the ResetConversation method of our chat engine, which clears the conversation history while preserving the system message. This is useful when users want to start a new conversation without affecting the document knowledge.

Finally, let's add a method to reset both the conversation history and document knowledge:

func (r *RAGChatbot) ResetAll() string {
    r.ResetConversation()
    r.ResetDocuments()
    return "Both conversation history and document knowledge have been reset."
}

This method provides a convenient way to completely reset the chatbot's state. It calls both ResetConversation and ResetDocuments, effectively returning the chatbot to its initial state.

Now that we've built our integrated RAG chatbot, let's test it by uploading a document and asking a question about it:

package main

import (
    "context"
    "fmt"
    "log"

    "codesignal/ragchatbot"
)

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

    // Initialize the RAG chatbot
    chatbot, err := ragchatbot.NewRAGChatbot(ctx)
    if err != nil {
        log.Fatalf("failed to initialize RAG chatbot: %v", err)
    }

    // Upload a document
    result, err := chatbot.UploadDocument(ctx, "../data/a_scandal_in_bohemia.pdf")
    if err != nil {
        log.Fatalf("error uploading document: %v", err)
    }
    fmt.Println(result)

    // Send a message about the document
    query := "What is the main mystery in the story?"
    response, err := chatbot.SendMessage(ctx, query)
    if err != nil {
        log.Fatalf("error sending message: %v", err)
    }
    fmt.Printf("\nQuestion: %s\n", query)
    fmt.Printf("Answer: %s\n", response)
}

When you run this code, you'll see output similar to:

Document successfully processed.

Question: What is the main mystery in the story?
Answer: The main mystery in the story revolves around the King of Bohemia's concern about a photograph in the possession of Irene Adler. He fears that this compromising photograph could be used to prevent his upcoming marriage to a Scandinavian princess if revealed. The King hires Sherlock Holmes to retrieve this photograph to protect his reputation and ensure his marriage proceeds without scandal.

This demonstrates how our RAG system successfully retrieves relevant context from the document and uses it to inform the language model's response. The chatbot has processed the PDF, extracted relevant chunks about the story's mystery, and generated an accurate answer based on the document content.

To verify that our system management features work correctly, let's test what happens when we reset the chatbot and try to ask about documents that are no longer in its knowledge base:

package main

import (
    "context"
    "fmt"
    "log"

    "codesignal/ragchatbot"
)

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

    // Initialize the RAG chatbot
    chatbot, err := ragchatbot.NewRAGChatbot(ctx)
    if err != nil {
        log.Fatalf("failed to initialize RAG chatbot: %v", err)
    }

    // Upload a document
    result, err := chatbot.UploadDocument(ctx, "../data/a_scandal_in_bohemia.pdf")
    if err != nil {
        log.Fatalf("error uploading document: %v", err)
    }
    fmt.Println(result)

    // Send a message about the document
    query := "What is the main mystery in the story?"
    response, err := chatbot.SendMessage(ctx, query)
    if err != nil {
        log.Fatalf("error sending message: %v", err)
    }
    fmt.Printf("\nQuestion: %s\n", query)
    fmt.Printf("Answer: %s\n\n", response)

    // Reset everything
    result = chatbot.ResetAll()
    fmt.Println(result)

    // Try asking about Sherlock Holmes after reset
    finalQuery := "Tell me about Sherlock Holmes."
    response, err = chatbot.SendMessage(ctx, finalQuery)
    if err != nil {
        log.Fatalf("error sending message: %v", err)
    }
    fmt.Printf("\nQuestion: %s\n", finalQuery)
    fmt.Printf("Answer: %s\n", response)
}

When you run this code, you'll see output similar to:

Document successfully processed.

Question: What is the main mystery in the story?
Answer: The main mystery in the story revolves around the King of Bohemia's concern about a photograph in the possession of Irene Adler. He fears that this compromising photograph could be used to prevent his upcoming marriage to a Scandinavian princess if revealed. The King hires Sherlock Holmes to retrieve this photograph to protect his reputation and ensure his marriage proceeds without scandal.

Both conversation history and document knowledge have been reset.

Question: Tell me about Sherlock Holmes.
Answer: I'm sorry, but I don't have any context provided to answer your question about Sherlock Holmes. Without relevant information in the context, I cannot provide an accurate answer. If you could provide additional context, I'd be happy to help.

This confirms our reset functionality works as expected, clearing both conversation history and document knowledge. The chatbot has returned to its initial state, ready to process new documents and start fresh conversations. The assistant correctly refuses to answer the question about Sherlock Holmes because, after the reset, there is no document knowledge available.

In this lesson, we've successfully integrated our document processor and chat engine to create a complete RAG chatbot system. We've built a RAGChatbot struct that coordinates between these components, providing a unified interface for document upload, message processing, and system management.

Our integrated RAG chatbot can:

• Upload and process documents in supported formats (PDF)
• Retrieve relevant context from documents based on user queries using similarity search
• Generate informed responses using the retrieved context and the language model
• Maintain conversation history for display or logging purposes
• Reset conversation history or document knowledge as needed
• Properly handle errors throughout the pipeline

Get ready to put your knowledge into practice and take your RAG chatbot to the next level!