Graph RAG

この例では、Mastra、OpenAIの埋め込み、そしてベクトルストレージとしてPGVectorを使用して、検索拡張生成（RAG）システムを実装する方法を示します。

概要

このシステムは、Mastra と OpenAI を用いて Graph RAG を実装しています。主な機能は以下の通りです。

応答生成のために gpt-4o-mini を使用した Mastra エージェントをセットアップ
ベクトルストアの操作やナレッジグラフの作成・トラバーサルを管理する GraphRAG ツールを作成
テキストドキュメントを小さなセグメントに分割
これらのチャンクに対して埋め込みを作成
PostgreSQL ベクトルデータベースに保存
GraphRAG ツールを使ってクエリに基づき関連するチャンクのナレッジグラフを作成
- ツールはベクトルストアから結果を返し、ナレッジグラフを作成
- クエリを使ってナレッジグラフをトラバース
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, createGraphRAGTool } from "@mastra/rag";
import { embedMany } from "ai";

GraphRAGツールの作成

@mastra/ragからインポートしたcreateGraphRAGToolを使用すると、ベクターデータベースにクエリを実行し、その結果をナレッジグラフに変換するツールを作成できます。

index.ts


const graphRagTool = createGraphRAGTool({
  vectorStoreName: "pgVector",
  indexName: "embeddings",
  model: openai.embedding("text-embedding-3-small"),
  graphOptions: {
    dimension: 1536,
    threshold: 0.7,
  },
});

エージェント設定

応答を処理するMastraエージェントを設定します：

index.ts


const ragAgent = new Agent({
  name: "GraphRAG Agent",
  instructions: `You are a helpful assistant that answers questions based on the provided context. Format your answers as follows:
 
1. DIRECT FACTS: List only the directly stated facts from the text relevant to the question (2-3 bullet points)
2. CONNECTIONS MADE: List the relationships you found between different parts of the text (2-3 bullet points)
3. CONCLUSION: One sentence summary that ties everything together
 
Keep each section brief and focus on the most important points.
 
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: {
    graphRagTool,
  },
});

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 doc = MDocument.fromText(`
# Riverdale Heights: Community Development Study
// ... text content ...
`);
 
const chunks = await doc.chunk({
  strategy: "recursive",
  size: 512,
  overlap: 50,
  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 =
  "What are the direct and indirect effects of early railway decisions on Riverdale Heights' current state?";
const answerOne = await ragAgent.generate(queryOne);
console.log("\nQuery:", queryOne);
console.log("Response:", answerOne.text);
 
const queryTwo =
  "How have changes in transportation infrastructure affected different generations of local businesses and community spaces?";
const answerTwo = await ragAgent.generate(queryTwo);
console.log("\nQuery:", queryTwo);
console.log("Response:", answerTwo.text);
 
const queryThree =
  "Compare how the Rossi family business and Thompson Steel Works responded to major infrastructure changes, and how their responses affected the community.";
const answerThree = await ragAgent.generate(queryThree);
console.log("\nQuery:", queryThree);
console.log("Response:", answerThree.text);
 
const queryFour =
  "Trace how the transformation of the Thompson Steel Works site has influenced surrounding businesses and cultural spaces from 1932 to present.";
const answerFour = await ragAgent.generate(queryFour);
console.log("\nQuery:", queryFour);
console.log("Response:", answerFour.text);

GitHubで例を見る