Basic RAG

This example demonstrates how to implement a Retrieval-Augmented Generation (RAG) system using Mastra, OpenAI embeddings, and PGVector for vector storage.

Overview

The system implements RAG using Mastra and OpenAI. Here’s what it does:

Sets up a Mastra agent with GPT-4o-mini for response generation
Creates a vector query tool to manage vector store interactions
Chunks text documents into smaller segments
Creates embeddings for these chunks
Stores them in a PostgreSQL vector database
Retrieves relevant chunks based on queries using vector query tool
Generates context-aware responses using the Mastra agent

Setup

Environment Setup

Make sure to set up your environment variables:

.env

POSTGRES_CONNECTION_STRING=your_connection_string_here

Dependencies

Then, import the necessary dependencies:

src/mastra/index.ts

import { Mastra, Agent, EmbedManyResult } from '@mastra/core';
import { embed, MDocument, PgVector, createVectorQueryTool } from '@mastra/rag';

Vectory Query Tool Creation

Using createVectorQueryTool imported from @mastra/rag, you can create a tool that can query the vector database.

src/mastra/index.ts

const vectorQueryTool = createVectorQueryTool({
  vectorStoreName: 'pgVector',
  indexName: 'embeddings',
  options: {
    provider: 'OPEN_AI',
    model: 'text-embedding-ada-002',
    maxRetries: 3,
  },
  topK: 3,
});

Agent Configuration

Set up the Mastra agent that will handle the responses:

src/mastra/index.ts

export const ragAgent = new Agent({
  name: 'RAG Agent',
  instructions:
    'You are a helpful assistant that answers questions based on the provided context. Keep your answers concise and relevant.',
  model: {
    provider: 'OPEN_AI',
    name: 'gpt-4o-mini',
  },
  tools: {
    vectorQueryTool,
  },
})

Instantiate PgVector and Mastra

Instantiate PgVector and Mastra with all components:

src/mastra/index.ts

const pgVector = new PgVector(process.env.POSTGRES_CONNECTION_STRING!);
 
export const mastra = new Mastra({
  agents: { ragAgent },
  vectors: { pgVector },
})
 
const agent = mastra.getAgent('ragAgent')

Document Processing

Create a document and process it into chunks:

src/mastra/index.ts

const doc = MDocument.fromText(`The Impact of Climate Change on Global Agriculture...`)
 
const chunks = await doc.chunk({
  strategy: 'recursive',
  size: 512,
  overlap: 50,
  separator: '\n',
})

Creating and Storing Embeddings

Generate embeddings for the chunks and store them in the vector database:

src/mastra/index.ts

const { embeddings } = await embed(chunks, {
  provider: "OPEN_AI",
  model: "text-embedding-ada-002",
  maxRetries: 3,
}) as EmbedManyResult<string>
 
const vectorStore = mastra.getVector('pgVector');
await vectorStore.createIndex("embeddings", 1536)
await vectorStore.upsert(
  "embeddings",
  embeddings,
  chunks?.map((chunk: any) => ({ text: chunk.text }))
)

Response Generation

Function to generate responses based on retrieved context:

src/mastra/index.ts

async function generateResponse(query: string) {
  const prompt = `
      Please answer the following question:
      ${query}
 
      Please base your answer only on the context provided in the tool. If the context doesn't 
      contain enough information to fully answer the question, please state that explicitly.
      `
 
  const completion = await agent.generate(prompt)
  return completion.text
}

Example Usage

src/mastra/index.ts

async function answerQueries(queries: string[]) {
  for (const query of queries) {
    try {
      const answer = await generateResponse(query)
      console.log('\nQuery:', query)
      console.log('Response:', answer)
    } catch (error) {
      console.error(`Error processing query "${query}":`, error)
    }
  }
}
 
const queries = [
  "What are the main points in the article?",
  "How does temperature affect crop yields?",
  "What solutions are farmers implementing?",
  "What are the future challenges mentioned in the text?",
];
 
await answerQueries(queries)

View Example on GitHub

Retrieve Results Optimizing Information Density