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
- Uses existing embeddings to retrieve relevant context
- 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