AWS Bedrock knowledgeとmetadataフィルターで効率的な検索をしよう！

## はじめに

AWS Amplify Gen2とAWSのナレッジベース機能を活用した検索システムの構築方法について解説します。本記事では、Lambdaを使用したナレッジベースの作成方法や、メタデータを活用した検索機能の実装方法について、実践的なコード例を交えながら詳しく説明していきます。

なお、この記事で紹介するシステムは、最近アップデートされたAmplify AIの機能を使用する前に作成されたものですので、Amplify AIの内容は含まれていません。

## 目次

1. ナレッジベースの基本設定
2. プロジェクト構成の概要
3. バックエンド実装の詳細
4. メタデータフィルターの実装
5. Bedrock統合の実装
6. データ同期の実装
7. まとめ

## ナレッジベースの基本設定

ナレッジベースの作成はAWSマネジメントコンソールから行います。本記事では、ナレッジベースIDとデータソースIDが既に取得済みであることを前提に進めていきます。ベクトルデータベースとしては、コスト効率の良いPineconeを採用しています。

以下参考記事です。

https://qiita.com/minorun365/items/ed680c2fc246f97985d2

## プロジェクト構成の概要

プロジェクトは以下のような構造で組織されています：

```
Root/
|_ amplify/
|  |_ functions/
|     |_ invoke-bedrock/      # Bedrock呼び出し用Lambda
|     |  |_ resource.ts
|     |  |_ handler.ts
|     |
|     |_ start-ingestion-job/ # データ取り込み用Lambda
|        |_ resource.ts
|        |_ handler.ts
|
|_ src/
|  |_ app/
|  |_ components/
|  |_ lib/
|
|_ public/
```

## バックエンド実装の詳細

Lambdaの実装には以下の記事を参考にしました。

https://qiita.com/moritalous/items/76a05676a564960ac974

### 環境変数の設定

まず、ナレッジベースの基本設定として、必要な環境変数を設定します。`amplify/invoke-bedrock/resource.ts`で、ナレッジベースIDとデータソースIDを環境変数として定義します。
secretsは以下を参考にしています。

https://docs.amplify.aws/react/deploy-and-host/fullstack-branching/secrets-and-vars/

```typescript:amplify/invoke-bedrock/resource.ts
import { defineFunction, secret } from '@aws-amplify/backend'

export const invokeBedrock = defineFunction({
  name: 'invoke-bedrock',
  entry: './handler.ts',
  timeoutSeconds: xxx,
  memoryMB: xxx,
  environment: {
    KN_ID: secret('KN_ID'),
    KN_DATA_SOURCE_ID: secret('KN_DATA_SOURCE_ID'),
  },
})
```

この設定により、Lambda関数内で環境変数を通じて安全にナレッジベースの設定にアクセスできます。

### メタデータフィルターの実装

メタデータは以下のような形式で登録されます：

```typescript:src/app/page.tsx
const metadataContent = JSON.stringify({
    metadataAttributes: {
        category: categoryId,
    },
})
```

検索時のフィルタリング機能は以下のように実装します：

```typescript:amplify/invoke-bedrock/handler.ts
const prompt = event.prompt
const selectedCategoryId = event.category

// カテゴリーに基づいてフィルターを作成
let filter
if (selectedCategoryId === 'all') {
  // 全カテゴリー検索の場合：データソースIDでフィルタリング
  filter = {
    equals: {
      key: 'x-amz-bedrock-kb-data-source-id',
      value: env.KN_DATA_SOURCE_ID,
    },
  }
} else {
  // 特定カテゴリー検索の場合：カテゴリーIDでフィルタリング
  filter = {
    equals: {
      key: 'category',
      value: selectedCategoryId,
    },
  }
}
```

このフィルター実装には以下のような特徴があります：
- カテゴリー別の検索が可能
- 全カテゴリー検索時はデータソースIDでフィルタリング
- 特定カテゴリー検索時はカテゴリーIDで絞り込み
- 完全一致検索による正確な結果の取得

### Bedrockの実装

Bedrockとの統合は以下のように実装します。以下参考です。

https://knowledge-oasis.net/aws/bedrock-knowledge-base-metadata-filtering/

```typescript:amplify/invoke-bedrock/handler.ts
const bedrockResponse = await bedrock.send(
  new RetrieveAndGenerateCommand({
    input: {
      text: prompt,  // ユーザーからの質問文
    },
    retrieveAndGenerateConfiguration: {
      type: 'KNOWLEDGE_BASE',
      knowledgeBaseConfiguration: {
        knowledgeBaseId: env.KN_ID,
        modelArn: modelId,
        retrievalConfiguration: {
          vectorSearchConfiguration: {
            numberOfResults: 5,     // 検索結果の最大数
            filter: filter,         // 上で作成したフィルター
          },
        },
        generationConfiguration: {
          inferenceConfig: {
            textInferenceConfig: {
              temperature: 0.7,     // 生成の多様性
              topP: 0.95,          // 出力のトークン選択の制御
              maxTokens: 1000,      // 最大生成トークン数
            },
          },
        },
      },
    },
  }),
)
```

このコードでは以下の設定を行っています：
- 検索結果の上限を5件に設定
- 生成時の温度パラメータを0.7に設定
- トークン数の上限を1000に設定
- フィルター条件の適用による検索結果の絞り込み

### レスポンスの処理

Bedrockからのレスポンスは3つの主要な部分から構成されます。

```typescript
const response = response.output.text // 生成された回答
const generatedResponsePart = response.citations[0].generatedResponsePart  // 引用部分
const retrievedReferences = response.citations[0].retrievedReferences    // 参照情報
```

これにより、生成された回答と、その根拠となった参照情報を取得できます。

<details><summary>amplify/invoke-bedrock/handler.ts</summary>

```typescript
import { env } from '$amplify/env/invoke-bedrock'
import {
  BedrockAgentRuntimeClient,
  RetrieveAndGenerateCommand,
} from '@aws-sdk/client-bedrock-agent-runtime'
import { Context, Handler } from 'aws-lambda'

// Lambdaが受け取るイベントの型定義
type eventType = {
  prompt: string    // ユーザーからの質問文
  category: string  // 検索対象のカテゴリー
}

// Bedrockのモデル指定
const modelId = 'anthropic.claude-3-5-sonnet-20240620-v1:0'
// Bedrockクライアントの初期化
const bedrock = new BedrockAgentRuntimeClient({ region: 'ap-northeast-1' })

export const handler: Handler = async (event: eventType, _context: Context) => {
  try {
    const prompt = event.prompt
    const selectedCategoryId = event.category

    // カテゴリーに基づいてフィルターを作成
    let filter
    if (selectedCategoryId === 'all') {
      // 全カテゴリー検索の場合：データソースIDでフィルタリング
      filter = {
        equals: {
          key: 'x-amz-bedrock-kb-data-source-id',
          value: env.KN_DATA_SOURCE_ID,
        },
      }
    } else {
      // 特定カテゴリー検索の場合：カテゴリーIDでフィルタリング
      filter = {
        equals: {
          key: 'category',
          value: selectedCategoryId,
        },
      }
    }

    // Bedrockに検索・生成リクエストを送信
    const bedrockResponse = await bedrock.send(
      new RetrieveAndGenerateCommand({
        input: {
          text: prompt,  // ユーザーからの質問文
        },
        retrieveAndGenerateConfiguration: {
          type: 'KNOWLEDGE_BASE',
          knowledgeBaseConfiguration: {
            knowledgeBaseId: env.KN_ID,
            modelArn: modelId,
            retrievalConfiguration: {
              vectorSearchConfiguration: {
                numberOfResults: 5,     // 検索結果の最大数
                filter: filter,         // 上で作成したフィルター
              },
            },
            generationConfiguration: {
              inferenceConfig: {
                textInferenceConfig: {
                  temperature: 0.7,     // 生成の多様性（0-1）
                  topP: 0.95,          // 出力のトークン選択の制御
                  maxTokens: 1000,      // 最大生成トークン数
                },
              },
            },
          },
        },
      }),
    )

    // Bedrockからのレスポンスを解析
    const generatedResponse = bedrockResponse.output.text                          // 生成された回答本文
    const generatedResponsePart = bedrockResponse.citations[0].generatedResponsePart  // 引用された部分
    const retrievedReferences = bedrockResponse.citations[0].retrievedReferences      // 参照元の情報

    // 成功時のレスポンスを返す
    return {
      statusCode: 200,
      body: JSON.stringify({
        response: generatedResponse,
        generatedResponsePart,
        retrievedReferences,
      }),
    }
  } catch (error) {
    // エラー発生時のログ記録とエラーレスポンスの返却
    console.error('Error generating response:', error)
    return {
      statusCode: 500,
      body: '何かしらのエラーが発生しました。',
    }
  }
}
```
</details>


## データ同期の実装

データの同期処理は以下のように実装します。同期ボタンで起動することを前提にしています。

```typescript:amplify/start-ingestion-job/resource.ts
import { defineFunction, secret } from '@aws-amplify/backend'

export const startIngestionJob = defineFunction({
  name: 'start-ingestion-job',
  entry: './handler.ts',
  timeoutSeconds: xxx,
  memoryMB: xxx,
  environment: {
    KN_ID: secret('KN_ID'),
    KN_DATA_SOURCE_ID: secret('KN_DATA_SOURCE_ID'),
  },
})
```

```typescript:amplify/start-ingestion-job/handler.ts
// 必要なライブラリとモジュールのインポート
import { env } from '$amplify/env/invoke-bedrock'
import { BedrockAgentClient, StartIngestionJobCommand } from '@aws-sdk/client-bedrock-agent'
import { Context, Handler } from 'aws-lambda'

const bedrock = new BedrockAgentClient({ region: 'ap-northeast-1' })

export const handler: Handler = async (event: eventType, _context: Context) => {
  try {
    // ナレッジベースとデータソースの同期に必要な入力パラメータを設定
    const input = {
      knowledgeBaseId: env.KN_ID,
      dataSourceId: env.KN_DATA_SOURCE_ID,
    }

    // 同期ジョブを開始するコマンドを作成
    const command = new StartIngestionJobCommand(input)

    // Bedrockに同期ジョブの開始リクエストを送信
    const response = await bedrock.send(command as any)

    // 成功時のレスポンスを返却
    return {
      statusCode: 200,
      body: response,
    }
  } catch (error) {
    // エラー発生時のログ記録とエラーレスポンスの返却
    console.error('Error generating response:', error)
    return {
      statusCode: 500,
      body: 'An error occurred while generating the response.',
    }
  }
}
```

## まとめ

本記事では、AWS Amplify Gen2を使用してナレッジベース検索システムを構築する方法について解説しました。

1. メタデータフィルターの実装
2. Bedrockとの適切な統合
3. ナレッジベースの同期機能

## その他の参考リンク



https://docs.amplify.aws/

https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/client/bedrock-agent-runtime/command/RetrieveAndGenerateCommand/

https://aws.amazon.com/jp/blogs/news/leveraging-pinecone-on-aws-marketplace-as-a-knowledge-base-for-amazon-bedrock/