ツールを使った再ランキング結果
この例では、Mastra のベクトルクエリツールを使用して、OpenAI の埋め込みと PGVector をベクトルストレージとして活用し、再ランキングを行う Retrieval-Augmented Generation(RAG)システムを実装する方法を示します。
概要
このシステムは、Mastra と OpenAI を用いた再ランク付き RAG を実装しています。主な機能は以下の通りです。
- 応答生成のために gpt-4o-mini を使った Mastra エージェントをセットアップします
- 再ランク機能を備えたベクトルクエリツールを作成します
- テキストドキュメントを小さなセグメントに分割し、それらから埋め込みを作成します
- それらを PostgreSQL ベクトルデータベースに保存します
- クエリに基づいて関連するチャンクを取得し、再ランクします
- Mastra エージェントを使ってコンテキストに応じた応答を生成します
セットアップ
環境設定
環境変数を設定してください:
.env
OPENAI_API_KEY=your_openai_api_key_here
POSTGRES_CONNECTION_STRING=your_connection_string_here
依存関係
次に、必要な依存関係をインポートします:
index.ts
import { openai } from "@ai-sdk/openai";
import { Mastra } from "@mastra/core";
import { Agent } from "@mastra/core/agent";
import { PgVector } from "@mastra/pg";
import { MDocument, createVectorQueryTool } from "@mastra/rag";
import { embedMany } from "ai";
リランキング付きベクタークエリツールの作成
@mastra/rag からインポートした createVectorQueryTool を使用すると、ベクターデータベースにクエリを実行し、結果をリランキングするツールを作成できます。
index.ts
const vectorQueryTool = createVectorQueryTool({
vectorStoreName: "pgVector",
indexName: "embeddings",
model: openai.embedding("text-embedding-3-small"),
reranker: {
model: openai("gpt-4o-mini"),
},
});
エージェント設定
応答を処理する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.
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,
},
});
PgVectorとMastraのインスタンス化
コンポーネントを使用してPgVectorとMastraをインスタンス化します:
index.ts
const pgVector = new PgVector({ connectionString: process.env.POSTGRES_CONNECTION_STRING! });
export const mastra = new Mastra({
agents: { ragAgent },
vectors: { pgVector },
});
const agent = mastra.getAgent("ragAgent");
ドキュメント処理
ドキュメントを作成し、チャンクに分割します:
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",
});
埋め込みの作成と保存
チャンクの埋め込みを生成し、ベクターデータベースに保存します。
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({
indexName: "embeddings",
dimension: 1536,
});
await vectorStore.upsert({
indexName: "embeddings",
vectors: embeddings,
metadata: chunks?.map((chunk: any) => ({ text: chunk.text })),
});
リランキングによるクエリ
リランキングが結果にどのように影響するか、さまざまなクエリを試してみましょう。
index.ts
const queryOne = "explain technical trading analysis";
const answerOne = await agent.generate(queryOne);
console.log("\nQuery:", queryOne);
console.log("Response:", answerOne.text);
const queryTwo = "explain trading card valuation";
const answerTwo = await agent.generate(queryTwo);
console.log("\nQuery:", queryTwo);
console.log("Response:", answerTwo.text);
const queryThree = "how do you analyze market resistance";
const answerThree = await agent.generate(queryThree);
console.log("\nQuery:", queryThree);
console.log("Response:", answerThree.text);
GitHubで例を見る