Re-ranking Results

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

Overview

The system implements RAG with re-ranking 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 with re-ranking capabilities
Chunks text documents into smaller segments
Creates embeddings for these chunks
Stores them in a PostgreSQL vector database
Retrieves and re-ranks relevant chunks based on queries
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';

Vector Query Tool Creation with Re-ranking

Using createVectorQueryTool imported from @mastra/rag, you can create a tool that can query the vector database and re-rank results:

src/mastra/index.ts

const vectorQueryTool = createVectorQueryTool({
  vectorStoreName: 'pgVector',
  indexName: 'embeddings',
  options: {
    provider: 'OPEN_AI',
    model: 'text-embedding-ada-002',
    maxRetries: 3,
  },
  topK: 5,
  rerankOptions: {
    semanticProvider: 'agent',
    agentProvider: {
      provider: 'OPEN_AI',
      name: 'gpt-4o-mini',
    },
  },
});

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 doc1 = MDocument.fromText(`
market data shows price resistance levels.
technical charts display moving averages.
support levels guide trading decisions.
breakout patterns signal entry points.
price action determines trade timing.
 
baseball cards show gradual value increase.
rookie cards command premium prices.
card condition affects resale value.
authentication prevents fake trading.
grading services verify card quality.
 
volume analysis confirms price trends.
sports cards track seasonal demand.
chart patterns predict movements.
mint condition doubles card worth.
resistance breaks trigger orders.
rare cards appreciate yearly.
`)
 
const chunks = await doc1.chunk({
  strategy: 'recursive',
  size: 150,
  overlap: 20,
  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 and re-ranked 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 = [
  "explain technical trading analysis",
  "explain trading card valuation",
  "how do you analyze market resistance",
];
 
await answerQueries(queries)

Metadata Filtering Chain of Thought Prompting