How to Do RAG with Mastra and ZeroEntropy

npm install @mastra/core @mastra/rag @mastra/pg @ai-sdk/openai ai

POSTGRES_CONNECTION_STRING=postgresql://user:pass@localhost:5432/rag_db
OPENAI_API_KEY=sk-...
ZEROENTROPY_API_KEY=ze-...  # Sign up at zeroentropy.ai

import type { RelevanceScoreProvider } from '@mastra/core/relevance';
import ZeroEntropy from 'zeroentropy';
export class ZeroEntropyRelevanceScorer implements RelevanceScoreProvider {
  private client: ZeroEntropy;
  private model: string;
  constructor(model?: string, apiKey?: string) {
    this.client = new ZeroEntropy({
      apiKey: apiKey || process.env.ZEROENTROPY_API_KEY || '',
    });
    this.model = model || 'zerank-1';
  }
  async getRelevanceScore(query: string, text: string): Promise<number> {
    const response = await this.client.models.rerank({
      query,
      documents: [text],
      model: this.model,
      top_n: 1,
    });
    
    return response.results[0]?.relevance_score ?? 0;
  }
}

import { openai } from '@ai-sdk/openai';
import { Mastra } from '@mastra/core/mastra';
import { Agent } from '@mastra/core/agent';
import { PgVector } from '@mastra/pg';
import { MDocument, createVectorQueryTool } from '@mastra/rag';
import { embedMany } from 'ai';
import { ZeroEntropyRelevanceScorer } from './zeroentropy-scorer';

// Create vector query tool with ZeroEntropy reranker
const vectorQueryTool = createVectorQueryTool({
  vectorStoreName: 'pgVector',
  indexName: 'embeddings',
  model: openai.embedding('text-embedding-3-small'),
  reranker: {
    provider: new ZeroEntropyRelevanceScorer('zerank-1'),
  },
});

// Create RAG agent
export const ragAgent = new Agent({
  id: 'rag-agent',
  name: 'RAG Agent',
  instructions: `You are a helpful assistant that answers questions based on the provided context. Keep your answers concise and relevant.
    Important: When asked to answer a question, 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.`,
  model: openai('gpt-4o-mini'),
  tools: {
    vectorQueryTool,
  },
});

// Initialize Mastra with PgVector
const pgVector = new PgVector({ 
  connectionString: process.env.POSTGRES_CONNECTION_STRING! 
});

export const mastra = new Mastra({
  agents: { ragAgent },
  vectors: { pgVector },
});

// Prepare document with mixed content
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.
`);

// Chunk the document
const chunks = await doc1.chunk({
  strategy: 'recursive',
  maxSize: 150,
  overlap: 20,
  separator: '\n',
});

// Create embeddings
const { embeddings } = await embedMany({
  model: openai.embedding('text-embedding-3-small'),
  values: chunks.map(chunk => chunk.text),
});

// Initialize vector store
const vectorStore = mastra.getVector('pgVector');

await vectorStore.createIndex({
  indexName: 'embeddings',
  dimension: 1536,
});

// Store embeddings with metadata
await vectorStore.upsert({
  indexName: 'embeddings',
  vectors: embeddings,
  metadata: chunks?.map((chunk: any) => ({ text: chunk.text })),
});

// Get the agent
const agent = mastra.getAgent('ragAgent');

// Query 1: Technical trading
const queryOne = 'explain technical trading analysis';
const answerOne = await agent.generate(queryOne);
console.log('\nQuery:', queryOne);
console.log('Response:', answerOne.text);

// Query 2: Sports cards
const queryTwo = 'explain trading card valuation';
const answerTwo = await agent.generate(queryTwo);
console.log('\nQuery:', queryTwo);
console.log('Response:', answerTwo.text);

// Query 3: Market resistance
const queryThree = 'how do you analyze market resistance';
const answerThree = await agent.generate(queryThree);
console.log('\nQuery:', queryThree);
console.log('Response:', answerThree.text);