情報密度の最適化
この例では、Mastra、OpenAIの埋め込み、PGVectorによるベクトルストレージを用いたRetrieval-Augmented Generation(RAG)システムの実装方法を示します。 このシステムでは、エージェントが初期チャンクをクリーンアップし、情報密度を最適化しつつデータの重複を排除します。
概要
このシステムは、Mastra と OpenAI を用いて RAG を実装しており、今回は LLM ベースの処理によって情報密度の最適化を行っています。主な処理内容は以下の通りです。
- gpt-4o-mini を搭載した Mastra エージェントをセットアップし、クエリとドキュメントのクリーニングの両方に対応
- エージェントが利用するためのベクトルクエリおよびドキュメント分割ツールを作成
- 初期ドキュメントの処理:
- テキストドキュメントを小さなセグメントに分割
- 各セグメントの埋め込みを作成
- それらを PostgreSQL ベクトルデータベースに保存
- 初回クエリを実行し、ベースラインとなる応答品質を確認
- データの最適化:
- エージェントを使ってセグメントのクリーニングと重複排除を実施
- クリーニング後のセグメントに対して新たに埋め込みを作成
- 最適化されたデータでベクトルストアを更新
- 同じクエリを再度実行し、応答品質の向上を確認
セットアップ
環境のセットアップ
環境変数を必ず設定してください:
.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,
createDocumentChunkerTool,
} 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"),
});
ドキュメントチャンク化ツール
@mastra/rag からインポートした createDocumentChunkerTool を使用すると、ドキュメントをチャンク化し、そのチャンクをエージェントに送信するツールを作成できます。
index.ts
const doc = MDocument.fromText(yourText);
const documentChunkerTool = createDocumentChunkerTool({
doc,
params: {
strategy: "recursive",
size: 512,
overlap: 25,
separator: "\n",
},
});
エージェント設定
クエリとクリーニングの両方を処理できる単一のMastraエージェントをセットアップします。
index.ts
const ragAgent = new Agent({
name: "RAG Agent",
instructions: `You are a helpful assistant that handles both querying and cleaning documents.
When cleaning: Process, clean, and label data, remove irrelevant information and deduplicate content while preserving key facts.
When querying: Provide answers based on the available 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,
documentChunkerTool,
},
});
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 chunks = await doc.chunk({
strategy: "recursive",
size: 256,
overlap: 50,
separator: "\n",
});
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
// Generate response using the original embeddings
const query = "What are all the technologies mentioned for space exploration?";
const originalResponse = await agent.generate(query);
console.log("\nQuery:", query);
console.log("Response:", originalResponse.text);
データの最適化
初期結果を確認した後、データの品質を向上させるためにクリーニングを行うことができます:
index.ts
const chunkPrompt = `Use the tool provided to clean the chunks. Make sure to filter out irrelevant information that is not space related and remove duplicates.`;
const newChunks = await agent.generate(chunkPrompt);
const updatedDoc = MDocument.fromText(newChunks.text);
const updatedChunks = await updatedDoc.chunk({
strategy: "recursive",
size: 256,
overlap: 50,
separator: "\n",
});
const { embeddings: cleanedEmbeddings } = await embedMany({
model: openai.embedding("text-embedding-3-small"),
values: updatedChunks.map((chunk) => chunk.text),
});
// Update the vector store with cleaned embeddings
await vectorStore.deleteIndex({ indexName: "embeddings" });
await vectorStore.createIndex({
indexName: "embeddings",
dimension: 1536,
});
await vectorStore.upsert({
indexName: "embeddings",
indexName: "embeddings",
vectors: cleanedEmbeddings,
metadata: updatedChunks?.map((chunk: any) => ({ text: chunk.text })),
});
最適化されたクエリ
データをクリーンアップした後、再度クエリを実行してレスポンスに違いがあるか観察します。
index.ts
// Query again with cleaned embeddings
const cleanedResponse = await agent.generate(query);
console.log("\nQuery:", query);
console.log("Response:", cleanedResponse.text);
GitHubで例を見る