ベクタークエリツールの使い方
この例では、RAGシステムにおけるセマンティック検索のための createVectorQueryTool
の実装と使用方法を示します。ツールの設定方法、ベクターストレージの管理、関連するコンテキストの効果的な取得方法について説明します。
概要
このシステムは、Mastra と OpenAI を使って RAG を実装しています。主な機能は以下の通りです。
- 応答生成のために gpt-4o-mini を使用した Mastra エージェントをセットアップします
- ベクトルストアとのやり取りを管理するベクトルクエリツールを作成します
- 既存の埋め込みを使って関連するコンテキストを取得します
- Mastra エージェントを使ってコンテキストに応じた応答を生成します
注意: 埋め込みの作成と保存方法については、Upsert Embeddings ガイドをご覧ください。
セットアップ
環境設定
環境変数を設定してください:
.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 { Mastra } from "@mastra/core";
import { Agent } from "@mastra/core/agent";
import { createVectorQueryTool } from "@mastra/rag";
import { PgVector } from "@mastra/pg";
ベクタークエリツールの作成
ベクターデータベースをクエリできるツールを作成します:
src/index.ts
const vectorQueryTool = createVectorQueryTool({
vectorStoreName: "pgVector",
indexName: "embeddings",
model: openai.embedding("text-embedding-3-small"),
});
エージェントの設定
応答を処理するMastraエージェントをセットアップします:
src/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: openai("gpt-4o-mini"),
tools: {
vectorQueryTool,
},
});
PgVector と Mastra のインスタンス化
すべてのコンポーネントを使って PgVector と Mastra をインスタンス化します。
src/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");
使用例
src/index.ts
const prompt = `
[Insert query based on document here]
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);
console.log(completion.text);
GitHubで例を見る