Understanding RAG: The Theory
Large Language Models (LLMs) face a fundamental challenge: they can only work with information they were trained on. This creates limitations when dealing with new, private, or domain-specific information. Additionally, LLMs sometimes generate plausible but incorrect information – a problem known as hallucination.
Retrieval-Augmented Generation (RAG) addresses these limitations by connecting LLMs to external knowledge sources. Think of RAG as giving an LLM the ability to "look up" information before responding, similar to how a human might consult reference materials to answer a question accurately.
How RAG Works
The RAG process mirrors how humans research and answer questions. When a query comes in, the system first searches its knowledge base for relevant information using vector embeddings – numerical representations that capture the meaning of text. This retrieved information is then formatted into a prompt for the LLM, which combines this specific knowledge with its general capabilities to generate an accurate, contextual response.
Here's a visual representation of how RAG works:
Mastra has built-in support for RAG, and allows developers to easily integrate a RAG workflow into their applications.
Now, let's build a RAG workflow using Mastra, including reranking capabilities for improved accuracy.
Part 1: Document Ingestion and Chunking
Let’s start by ingesting a document and chunking it. We want to split the document into bite-sized pieces for our search; this is called “chunking.”
The goal is to split at natural boundaries like topic transitions and new sections. This is called “semantic coherence.”
import { Mastra, EmbedManyResult, EmbedResult } from "@mastra/core";
import { embed, MDocument, PgVector, Reranker } from "@mastra/rag";
const doc = MDocument.fromText(`Your text content here...`);
const chunks = await doc.chunk({
strategy: "recursive",
size: 512,
overlap: 50,
separator: "\n",
});
Part 2: Embedding Generation
After chunking, we'll need to embed our data – transform it into a vector, or an array of 1536 values between 0 and 1, representing the meaning of the text.
We do this with LLMs, because they make the embeddings much more accurate. OpenAI has an API for this.
const { embeddings } = (await embed(chunks, {
provider: "OPEN_AI",
model: "text-embedding-3-small",
maxRetries: 3,
})) as EmbedManyResult<string>;
Part 3: Vector Storage and Management
We need to use a vector DB which can store these vectors and do the math to search on them. We'll use pgvector, which comes out of the box with Postgres.
Once we pick a vector DB, we need to set up an index to store our document chunks, represented as vector embeddings.
const pgVector = new PgVector(process.env.POSTGRES_CONNECTION_STRING!);
// add to the Mastra object to get logging
export const mastra = new Mastra({
vectors: { pgVector },
});
await pgVector.createIndex("embeddings", 1536);
await pgVector.upsert(
"embeddings",
embeddings,
chunks?.map((chunk: any) => ({ text: chunk.text }))
);
Part 4: Query Processing and Response Generation
Let's set up the LLM we'll use for generating responses and define our query. We'll then generate an embedding for the query using the same embedding API we used for the document chunks.
const llm = mastra.LLM({
provider: "OPEN_AI",
name: "gpt-4o-mini",
});
const query = "insert query here";
const { embedding } = await embed(query, {
provider: "OPEN_AI",
model: "text-embedding-3-small",
maxRetries: 3,
}) as EmbedResult<string>;
Okay, after that setup, we can now query the database!
Under the hood, Mastra is running an algorithm that compares our query string to all the chunks in the database and returning the most similar ones.
The actual algorithm is called “cosine similarity”. The implementation is similar to geo queries searching latitude/longitude, except the search goes over 1536 dimensions instead of two. We can use other algorithms as well.
Now, we can construct a prompt using the results as context.
// Perform vector similarity search
const results = await pgVector.query("embeddings", embedding);
// Extract and combine relevant chunks
const relevantChunks = results.map((result) => result?.metadata?.text);
const relevantContext = relevantChunks.join("\n\n");
const prompt = `
Please answer the following question:
${query}
Please base your answer only on this context ${relevantContext}.
If the context doesn't contain enough information to fully answer the question, please state that explicitly.
`;
We can now pass this prompt into an LLM to generate a response.
const completion = await llm.generate(prompt);
console.log(completion.text);Part 5: Enhanced Retrieval with Reranking
Optionally, after querying and getting our results, we can use a reranker. Reranking is basically a more computationally expensive way of searching the dataset.
It would take too long to run it over the entire database, but we can run it over our results to improve the ordering.
const reranker = new Reranker({
semanticProvider: "agent",
agentProvider: {
provider: "OPEN_AI",
name: "gpt-4o-mini",
},
});
const rerankedResults = await reranker.rerank({
query: query,
vectorStoreResults: results,
topK: 3,
});
// Process reranked results
const rerankedChunks = rerankedResults.map(
({ result }) => result?.metadata?.text
);
// Combine reranked chunks into a context
const rerankedContext = rerankedChunks.join("\n\n");
const rerankedPrompt = `
Please answer the following question:
${query}
Please base your answer only on this context ${rerankedContext}.
If the context doesn't contain enough information to fully answer the question, please state that explicitly.
`;
Finally, we can generate a response using the reranked prompt:
const rerankedCompletion = await llm.generate(rerankedPrompt);
console.log(rerankedCompletion.text);Common Patterns and Best Practices
When working with this RAG implementation, keep these principles in mind:
The chunk size and overlap settings in the document processing significantly impact retrieval quality. Using a recursive strategy with moderate overlap (like 50 tokens) often provides good results while maintaining context.
The choice of embedding model affects both the quality of retrieval and the cost of operation. While OpenAI's embeddings provide excellent results, there are open-source alternatives available for different needs.
The number of chunks retrieved (topK in the vector query tool) is a balance between providing enough context and staying within the LLM's context window. Start with 3-4 chunks and adjust based on need.
Consider the tradeoff between response time and quality when deciding whether to use reranking. While it can improve accuracy, it adds an additional processing step to the workflow.