Agent-Driven Metadata Filtering

This example demonstrates how to implement a Retrieval-Augmented Generation (RAG) system using Mastra, OpenAI embeddings, and PGVector for vector storage. This system uses an agent to construct metadata filters from a user’s query to search for relevant chunks in the vector store, reducing the amount of results returned.

Overview

The implementation consists of the following components:

Sets up a Mastra agent with gpt-4o-mini to understand queries and identify filter requirements
Creates a vector query tool to handle metadata filtering and semantic search
Processes documents into chunks with metadata and embeddings
Stores both vectors and metadata in PGVector for efficient retrieval
Processes queries by combining metadata filters with semantic search

When a user asks a question:

The agent analyzes the query to understand the intent
Constructs appropriate metadata filters (e.g., by topic, date, category)
Uses the vector query tool to find the most relevant information
Generates a contextual response based on the filtered results

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

Then, import the necessary dependencies:

src/mastra/index.ts

import { openai } from '@ai-sdk/openai';
import { Mastra } from '@mastra/core';
import { Agent } from '@mastra/core/agent';
import { OpenAI } from '@mastra/core/llm/openai';
import { createVectorQueryTool, MDocument, PGVECTOR_PROMPT } from '@mastra/rag';
import { PgVector } from '@mastra/pg';
import { embedMany } from 'ai';

Vector Query Tool Creation

Using createVectorQueryTool imported from @mastra/rag, you can create a tool that enables the agent to query the vector database with filtering capabilities.

src/mastra/index.ts

const vectorQueryTool = createVectorQueryTool({
  id: 'vectorQueryTool',
  vectorStoreName: "pgVector",
  indexName: "embeddings",
  model: openai.embedding('text-embedding-3-small'),
  enableFilter: true,
});

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',
  extract: {
    keywords: true,  // Extracts keywords from each chunk
  },
});

Transform Chunks into Metadata

Transform chunks into metadata that can be filtered:

src/mastra/index.ts

const chunkMetadata = chunks?.map((chunk: any, index: number) => ({
  text: chunk.text,
  ...chunk.metadata,
  nested: {
    keywords: chunk.metadata.excerptKeywords
      .replace('KEYWORDS:', '')
      .split(',')
      .map(k => k.trim()),
    id: index,
  },
}));

Agent Configuration

The agent is configured to understand user queries and translate them into appropriate metadata filters.

The agent requires both the vector query tool and a system prompt containing:

Metadata structure for available filter fields
Vector store prompt for filter operations and syntax

src/mastra/index.ts

export const ragAgent = new Agent({
  name: 'RAG Agent',
  model: openai('gpt-4o-mini'),
  instructions: `
  You are a helpful assistant that answers questions based on the provided context. Keep your answers concise and relevant.
 
  Use the vectorQueryTool to filter the context by searching the metadata.
  
  The metadata is structured as follows:
 
  {
    text: string,
    excerptKeywords: string,
    nested: {
      keywords: string[],
      id: number,
    },
  }
 
  ${PGVECTOR_PROMPT}
  `,
  tools: { vectorQueryTool },
});

The agent’s instructions are designed to:

Process user queries to identify filter requirements
Use the metadata structure to find relevant information
Apply appropriate filters through the vectorQueryTool and the provided vector store prompt
Generate responses based on the filtered context

Note: Different vector stores have specific prompts available. See Vector Store Prompts for details.

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

Creating and Storing Embeddings

Generate embeddings for the chunks and store them in the vector database along with their metadata:

src/mastra/index.ts

const { embeddings } = await embedMany({
  model: openai.embedding('text-embedding-3-small'),
  values: chunks.map(chunk => chunk.text),
});
 
const vectorStore = mastra.getVector('pgVector');
await vectorStore.createIndex('embeddings', 1536);
 
// Store both embeddings and metadata together
await vectorStore.upsert('embeddings', embeddings, chunkMetadata);

The upsert operation stores both the vector embeddings and their associated metadata, enabling combined semantic search and metadata filtering capabilities.

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.
      `;
 
  // Call the agent to generate a response
  const completion = await agent.generate(prompt);
 
  return completion.text;
}

Example Usage

src/mastra/index.ts

const queries = [
  "What adaptation strategies are mentioned? Use regex to search for the word 'adaptation' in the 'nested.keywords' field.",
  "Show me recent sections. Check the 'nested.id' field and return values that are greater than 2.",
  "Search the 'text' field using regex operator to find sections containing 'temperature'.",
];
 
async function answerQueries() {
  for (const query of queries) {
    try {
      // Generate and log the response
      const answer = await generateResponse(query);
      console.log('\nQuery:', query);
      console.log('Response:', answer);
    } catch (error) {
      console.error(`Error processing query "${query}":`, error);
    }
  }
}
 
await answerQueries();

View Example on GitHub

Optimizing Information Density Re-ranking Results