再ランキング結果

この例では、Mastra、OpenAIの埋め込み、そしてベクトルストレージとしてPGVectorを使用し、再ランキングを取り入れたRetrieval-Augmented Generation（RAG）システムの実装方法を示します。

概要

このシステムは、Mastra と OpenAI を用いた再ランキング付き RAG を実装しています。主な処理内容は以下の通りです。

テキストドキュメントを小さなセグメントに分割し、それらから埋め込みを作成します
ベクトルを PostgreSQL データベースに保存します
初期のベクトル類似度検索を実行します
Mastra の rerank 関数を使い、ベクトル類似度・セマンティック関連性・位置スコアを組み合わせて結果を再ランキングします
初期結果と再ランキング後の結果を比較し、改善点を表示します

セットアップ

環境設定

環境変数を必ず設定してください：

.env


OPENAI_API_KEY=your_openai_api_key_here
POSTGRES_CONNECTION_STRING=your_connection_string_here

依存関係

次に、必要な依存関係をインポートします：

src/index.ts


import { openai } from "@ai-sdk/openai";
import { PgVector } from "@mastra/pg";
import { MDocument, rerank } from "@mastra/rag";
import { embedMany, embed } from "ai";

ドキュメント処理

ドキュメントを作成し、チャンクに分割します：

src/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.
`);
 
const chunks = await doc1.chunk({
  strategy: "recursive",
  size: 150,
  overlap: 20,
  separator: "\n",
});

エンベディングの作成と保存

チャンクのエンベディングを生成し、ベクトルデータベースに保存します：

src/index.ts


const { embeddings } = await embedMany({
  values: chunks.map((chunk) => chunk.text),
  model: openai.embedding("text-embedding-3-small"),
});
 
const pgVector = new PgVector({ connectionString: process.env.POSTGRES_CONNECTION_STRING! });
await pgVector.createIndex({
  indexName: "embeddings",
  dimension: 1536,
});
await pgVector.upsert({
  indexName: "embeddings",
  vectors: embeddings,
  metadata: chunks?.map((chunk: any) => ({ text: chunk.text })),
});

ベクター検索とリランキング

ベクター検索を実行し、結果をリランキングします：

src/index.ts


const query = "explain technical trading analysis";
 
// Get query embedding
const { embedding: queryEmbedding } = await embed({
  value: query,
  model: openai.embedding("text-embedding-3-small"),
});
 
// Get initial results
const initialResults = await pgVector.query({
  indexName: "embeddings",
  queryVector: queryEmbedding,
  topK: 3,
});
 
// Re-rank results
const rerankedResults = await rerank(
  initialResults,
  query,
  openai("gpt-4o-mini"),
  {
    weights: {
      semantic: 0.5, // How well the content matches the query semantically
      vector: 0.3, // Original vector similarity score
      position: 0.2, // Preserves original result ordering
    },
    topK: 3,
  },
);

重み付けによって、さまざまな要素が最終的なランキングにどのように影響するかを制御します：

semantic: 値が高いほど、クエリとの意味的な理解や関連性が優先されます
vector: 値が高いほど、元のベクター類似度スコアが重視されます
position: 値が高いほど、元の結果の順序が維持されやすくなります

結果の比較

初期結果と再ランク付け後の結果の両方を表示して、改善点を確認しましょう。

src/index.ts


console.log("Initial Results:");
initialResults.forEach((result, index) => {
  console.log(`Result ${index + 1}:`, {
    text: result.metadata.text,
    score: result.score,
  });
});
 
console.log("Re-ranked Results:");
rerankedResults.forEach(({ result, score, details }, index) => {
  console.log(`Result ${index + 1}:`, {
    text: result.metadata.text,
    score: score,
    semantic: details.semantic,
    vector: details.vector,
    position: details.position,
  });
});

再ランク付け後の結果は、ベクトル類似度とセマンティックな理解を組み合わせることで、検索品質がどのように向上するかを示しています。各結果には以下が含まれます：

すべての要素を組み合わせた総合スコア
言語モデルによるセマンティック関連度スコア
埋め込み比較によるベクトル類似度スコア
必要に応じて元の順序を維持するための位置ベースのスコア

GitHubで例を見る