ベクタークエリツールの使い方

この例では、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で例を見る