Optimizing Information Density

This example demonstrates how to implement a Retrieval-Augmented Generation (RAG) system using Mastra, OpenAI embeddings, and PGVector for vector storage. The system uses an agent to clean the initial chunks to optimize information density and deduplicate data.

Overview

The system implements RAG using Mastra and OpenAI, this time optimizing information density through LLM-based processing. Here’s what it does:

Sets up two Mastra agents with GPT-4o-mini for response generation
Sets up another Mastra agent to handle cleaning up chunk data before vector storage
Creates a vector query tool to manage vector store interactions
Create a document chunking tool for agent to use to get chunks
Chunks text documents into smaller segments
Takes those chunks and filters them to remove irrelevant or duplicate information
Creates embeddings for both the initial chunks and the updated chunks
Stores them both in a PostgreSQL vector database
Retrieves relevant chunks based on queries using vector query tool
Generates context-aware responses using the Mastra agents

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, EmbedResult, EmbedManyResult } from '@mastra/core'
import { embed, MDocument, PgVector, createVectorQueryTool, createDocumentChunker } from '@mastra/rag';

Tool Creation

Vectory Query Tool

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,
  },
});
 
const cleanedVectorQueryTool = createVectorQueryTool({
  vectorStoreName: 'pgVector',
  indexName: 'cleanedEmbeddings',
  options: {
    provider: 'OPEN_AI',
    model: 'text-embedding-ada-002',
    maxRetries: 3,
  },
});

Chunk Tool

Using createDocumentChunker imported from @mastra/rag, you can create a tool that chunk the document and send the chunks to your agent.

src/mastra/index.ts

const doc = MDocument.fromText(yourText);
 
const documentChunkerTool = createDocumentChunker({
  doc,
  params: {
    strategy: 'recursive',
    size: 256,
    overlap: 50,
    separator: '\n',
  },
});

Agent Configuration

Set up three Mastra agents:

src/mastra/index.ts

export const ragAgentOne = new Agent({
  name: 'RAG Agent One',
  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,
  },
});
 
export const ragAgentTwo = new Agent({
  name: 'RAG Agent Two',
  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: {
    cleanedVectorQueryTool,
  },
});
 
export const ragAgentThree = new Agent({
  name: 'RAG Agent Three',
  instructions: 'You are a helpful assistant that processes, cleans, and labels data before storage.',
  model: {
    provider: 'OPEN_AI',
    name: 'gpt-4o-mini',
  },
  tools: { documentChunkerTool },
});

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: { ragAgentOne, ragAgentTwo, ragAgentThree },
  vectors: { pgVector },
});
const dataAgentOne = mastra.getAgent('ragAgentOne');
const dataAgentTwo = mastra.getAgent('ragAgentTwo');
const processAgent = mastra.getAgent('ragAgentThree');

Document Processing

Chunk the initial document clean them using the processAgent.

src/mastra/index.ts

const chunks = await doc.chunk({
  strategy: 'recursive',
  size: 256,
  overlap: 50,
  separator: '\n',
});
 
const chunkPrompt = `Take the chunks returned from the tool and clean them up according to the instructions provided. Make sure to filter out irrelevant information and remove duplicates.`;
 
const newChunks = await processAgent.generate(chunkPrompt);
const updatedDoc = MDocument.fromText(newChunks.text);
 
const updatedChunks = await updatedDoc.chunk({
  strategy: 'recursive',
  size: 256,
  overlap: 50,
  separator: '\n',
});

Creating and Storing Embeddings

Generate and store both raw and cleaned embeddings:

src/mastra/index.ts

const { embeddings } = await embed(chunks, {
  provider: 'OPEN_AI',
  model: 'text-embedding-ada-002',
  maxRetries: 3,
}) as EmbedManyResult<string>;
 
const { embeddings: cleanedEmbeddings } = await embed(updatedChunks, {
  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.createIndex('cleanedEmbeddings', 1536);
 
await vectorStore.upsert(
  'embeddings',
  embeddings,
  chunks?.map((chunk: any) => ({ text: chunk.text })),
);
 
await vectorStore.upsert(
  'cleanedEmbeddings',
  cleanedEmbeddings,
  updatedChunks?.map((chunk: any) => ({ text: chunk.text })),
);

Response Generation

Function to generate responses with index selection:

src/mastra/index.ts

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

Example Usage

src/mastra/index.ts

async function answerQueries(queries: string[], agent: Agent) {
  for (const query of queries) {
    try {
      const answer = await generateResponse(query, agent);
      console.log('\nQuery:', query);
      console.log('Response:', answer);
    } catch (error) {
      console.error(`Error processing query "${query}":`, error);
    }
  }
}
 
const queries = [
  'What is the average temperature on Mars?',
  'What technologies are used in modern spacecraft?',
  'What are all the requirements for space settlements?',
  'What are all the dates mentioned related to space stations?',
  'What are all the mentions of sustainability in space settlements?',
];
 
// Compare responses between raw and cleaned embeddings
await answerQueries(queries, dataAgentOne);
await answerQueries(queries, dataAgentTwo);

View Example on GitHub

Basic RAG Metadata Filtering