ExamplesRAGBasic RAG

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:

  1. Sets up a Mastra agent with gpt-4o-mini for response generation
  2. Creates a vector query tool to manage vector store interactions
  3. Uses existing embeddings to retrieve relevant context
  4. Generates context-aware responses using the Mastra agent

Note: To learn how to create and store embeddings in PGVector, see the Insert Embeddings in PGVector guide.

Setup

Environment Setup

Make sure to set up your environment variables:

.env
OPENAI_API_KEY=your_openai_api_key_here
POSTGRES_CONNECTION_STRING=your_connection_string_here

Dependencies

Import the necessary dependencies:

src/index.ts
import { openai } from '@ai-sdk/openai';
import { Mastra } from '@mastra/core';
import { Agent } from '@mastra/core/agent';
import { createVectorQueryTool } from '@mastra/rag';
import { PgVector } from '@mastra/pg';

Vector Query Tool Creation

Create a tool that can query the vector database:

src/index.ts
const vectorQueryTool = createVectorQueryTool({
  vectorStoreName: 'pgVector',
  indexName: 'embeddings',
  model: openai.embedding('text-embedding-3-small'),
});

Agent Configuration

Set up the Mastra agent that will handle the responses:

src/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: openai('gpt-4o-mini'),
  tools: {
    vectorQueryTool,
  },
});

Instantiate PgVector and Mastra

Instantiate PgVector and Mastra with all components:

src/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');

Example Usage

src/index.ts
const prompt = `
[Insert query based on document here]
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);
console.log(completion.text);





View Example on GitHub