MastraのRAGシステムを試してみる

Mastraで実装したサービス上でRAG（検索拡張生成）という技術を試してみました。

mastra.ai

RAGとは？

RAG（検索拡張生成）はRetrieval-Augmented Generationの略でLLMに学習データに加え社内文書などの外部データからの検索結果を組み合わせて回答を生成させることをできるようにする技術です。

参考

RAG（検索拡張生成）とは？意味・定義 | IT用語集 | docomo business Watch | ドコモビジネス | NTTコミュニケーションズ法人のお客さま

Mastra上でのRAGシステム構築の流れ（ざっくり

テキストをベクトル化（数値の配列（ベクトル）に変換する処理）します
ベクトル化した情報を任意のベクトルデータベースに保存します
Mastraにベクトルデータベースから情報を取得するツールを実装します
実装したツールを使うエージェントを登録して完了です

今回実装したもの

ローカルでPostgresqlベクトルデータベースを起動するDockerfile
Mastra、Next.js、Postgresqlを一括で起動できるdocker compose
テキストをベクトル化するスクリプト
ベクトルデータベースから情報を取得するMastra Toolとそれを使うエージェント

Dockerfileとcomposeの追加

MastraのサーバとMastraに対する独自フロントエンドとしてのNext.js、加えて今回ベクトルデータベースにPostgresqlを使用するのでこれらの開発用サーバをまとめて起動するためのDocker環境を用意します。Next.jsは必須ではないですが今回独自にフロントエンド実装したかったので用意しています。

version: '3.8'

services:
  nextjs:
    build:
      context: .
      dockerfile: docker/nextjs/Dockerfile
    image: nextjs-app
    container_name: nextjs-app
    ports:
      - '3000:3000'
    volumes:
      - ./:/app
      - node_modules:/app/node_modules
      - /.mastra # .mastraディレクトリをマウントから除外。除外しないとホスト側の設定が上書きされて落ちる
  mastra:
    build:
      context: .
      dockerfile: docker/mastra/Dockerfile
    image: mastra-app
    container_name: mastra-app
    ports:
      - '4222:4111' # 4111ポートローカルの別PJで使用されていたので、4222に変更
    environment:
      # コンテナ間の接続にはホスト名（localhost）ではなく、サービス名（postgresql）を使用
      - POSTGRES_CONNECTION_STRING=postgresql://user:password@postgresql:5432/vectordb
    depends_on:
      - postgresql
    volumes:
      - ./:/app
      - node_modules:/app/node_modules
  postgresql:
    build:
      context: .
      dockerfile: docker/postgresql/Dockerfile
    image: pgvector-db
    container_name: pgvector-db
    environment:
      POSTGRES_USER: user
      POSTGRES_PASSWORD: password
      POSTGRES_DB: vectordb
    ports:
      - '5432:5432'
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:
  node_modules:

Mastra、Next.jsのDockerfileにはnodeのイメージ、PostgresqlのDockerfileにはankane/pgvectorイメージを使用しています。

このcomposeから起動するコンテナ上で開発を進めていきます。

FAQテキストをベクトル化するスクリプト

@mastra/ragパッケージのMDocumentを使ってテキストの内容を読み込みます。

今回はfromTextメソッドを使ってテキストとして読み込んでいますがHTML、マークダウン、JSON形式で読み込むこともできます。

架空のWebサービスのFAQテキストとして用意します。

  const doc = MDocument.fromText(`
1. アカウント・ログインについて
Q1-1. パスワードを忘れてしまいました。どうすればいいですか？
A: ログイン画面の「パスワードを忘れた場合」リンクをクリックし、登録済みのメールアドレスを入力してください。パスワードリセット用のメールが届きます。届かない場合は迷惑メールフォルダもご確認ください。

Q1-2. メールアドレスを変更したいのですが？
A: ログイン後、「アカウント設定」から変更可能です。認証メールを確認して手続きを完了してください。

Q1-3. アカウントがロックされてしまいました。解除方法は？
A: 一時的にロックされた場合は24時間で自動解除されます。お急ぎの場合はサポートまでご連絡ください。

お問い合わせフォーム：管理画面右上の「ヘルプ」からアクセス

電話受付時間：平日10:00〜18:00

チャット対応：24時間以内に返信（土日祝を除く）
  `)

上記のテキストをチャンク化(数値の配列に変換)します

const chunks = await doc.chunk()

チャンク化したデータを元に埋め込み表現（Embeddings）を作ります。

Embeddingとは文字情報をAIモデルが処理しやすい数値ベクトル表現に変換する技術です。

const { embeddings } = await embedMany({
  values: chunks.map(chunk => chunk.text),
  model: openai.embedding('text-embedding-3-small')
})

console.logで標準出力にembeddingsの中身を表示すると以下のような値が確認できます。

✅ faq-embeddings: [
  [
      0.021157619,   0.031170301,    0.01990326,   0.038740866,  0.029238809,
     -0.012177287, -0.0013098631,   0.026619082,    0.00601649,   0.00705439,
     0.0049591637,   -0.04213763,  0.0024934576,   -0.01223279,  0.016473193,
     0.0116777625,  -0.052039307, -0.0013272077,  -0.029993646,  0.012210588,
      0.022822699,   0.026108459,   0.020436084,   0.023599736, -0.043092277,
      
      以下略

次にembeddingsの中身をベクトルデータベースに登録します。

ベクトルデータベースにはMongoDB、OpenSearch、CloudFlareなどさまざまなストレージサービスを指定できますが今回はPostgreSQLの拡張機能、PgVectorを使って実装します。

@mastra/pgパッケージにPgVectorのクライアントが用意されているのでそれを使います。

const pgVector = new PgVector({
  connectionString: process.env.POSTGRES_CONNECTION_STRING as string
})

初期化したクライアントからindexを作成します。indexはPostgresqlを使う場合はテーブルとして定義されます。dimensionにはモデルごとに最適な数字を定義します。

  await pgVector.createIndex({
    indexName: 'embeddings',
    dimension: 1536
  })

ベクトルデータベースに数値ベクトル表現（embeddings）とチャンク化したテキストを合わせて登録します。

await pgVector.upsert({
  indexName: 'embeddings',
  vectors: embeddings,
  metadata: chunks.map(chunk => ({
    text: chunk.text
  }))
})

ここまでの処理をmain関数にまとめてts-nodeから実行できるようにして完了です。

スクリプト全文（FAQテキストは一部のみ記載）

import 'dotenv/config'
import { embedMany } from 'ai'
import { openai } from '@ai-sdk/openai'
import { PgVector } from '@mastra/pg'
import { MDocument } from '@mastra/rag'

async function main() {
  const doc = MDocument.fromText(`
　
  `)

  const chunks = await doc.chunk()

  const { embeddings } = await embedMany({
    values: chunks.map(chunk => chunk.text),
    model: openai.embedding('text-embedding-3-small')
  })

  const pgVector = new PgVector({
    connectionString: process.env.POSTGRES_CONNECTION_STRING as string
  })

  console.log('✅ embeddings:', embeddings)

  try {
    await pgVector.createIndex({
      indexName: 'embeddings',
      dimension: 1536
    })

    await pgVector.upsert({
      indexName: 'embeddings',
      vectors: embeddings,
      metadata: chunks.map(chunk => ({
        text: chunk.text
      }))
    })

    console.log('✅ Embeddings inserted successfully')
  } finally {
    process.exit(0)
  }
}

main().catch(err => {
  console.error('❌ Error:', err)
  process.exit(1)
})

実行するとPgVectorに内容が登録されて検索可能になります。

ベクトルデータベースから検索結果を取得するMastra Toolの実装

上記のスクリプトで登録したベクトル値を使ってレスポンスを生成するToolを実装します。

ベクトルデータベースから値を取得するためにまずクライアントを初期化します。

const pgVector = new PgVector({
  connectionString: process.env.POSTGRES_CONNECTION_STRING as string
})

次に問い合わせの内容をembeddingに変換します。context.queryに問い合わせ内容の文字列が入っています。

const { embedding } = await embed({
  model: openai.embedding('text-embedding-3-small'),
  value: context.query
})

embeddingに変換したクエリ文字列を使って検索結果を取得します。

const results = await pgVector.query({
  indexName: 'embeddings',
  queryVector: embedding,
  topK: 3
})

この内容をToolとして登録すると以下のような実装になります。

export const faqQueryRagTool = createTool({
  id: 'Get Faq Information',
  inputSchema: z.object({
    query: z.string()
  }),
  description: `Fetches the FAQ information from a file`,
  execute: async ({ context }) => {
    const pgVector = new PgVector({
      connectionString: process.env.POSTGRES_CONNECTION_STRING as string
    })
    console.log('✅ connectionString:', process.env.POSTGRES_CONNECTION_STRING)
    const { embedding } = await embed({
      model: openai.embedding('text-embedding-3-small'),
      value: context.query
    })

    const results = await pgVector.query({
      indexName: 'embeddings',
      queryVector: embedding,
      topK: 3
    })

    return {
      contents: results ? results : 'No faq found.'
    }
  }
})

Agentの実装

以下のようにinstructionを設定してToolを指定したAgentを実装します

export const faqQueryRagAgent = new Agent({
  name: 'FAQ Vector Query RAG エージェント',
  instructions: `
    あなたはFAQデータベースから関連する情報を検索し、ユーザーの質問に答えるエージェントです。
    ユーザーからの質問に対して、以下のツールを使用して回答を生成してください。
    - faqQueryRagTool: ベクトルデータベースから質問に関連するFAQ情報を検索します。
    
    ユーザーの質問に対して、できるだけ具体的で明確な回答を提供してください。
  `,
  tools: { faqQueryRagTool },
  model: openai('gpt-4o-mini')
})

index.tsに追加します

export const mastra = new Mastra({
  agents: {
    faqQueryRagAgent
  }

これでlocalhost:4111にアクセスして追加したAgentを使ってMastraの対話UIからFAQの内容を問い合わせることが確認できます。

yikegaya’s blog

仕事関連（Webエンジニア）と資産運用について書いてます