Amazon Bedrock+PineconeでRAGを構築し､IPAの"安全なウェブサイトの作り方"に詳しいSlackチャットボットを作成する

Last updated at 2024-11-10Posted at 2024-04-26

はじめに

前回の記事｢Amazon Bedrock+Anthropic Claude 3 Sonnetで会話履歴を保持するSlackチャットボットを作成する｣では､会話履歴を保持したチャットボットを作成しました｡

本記事では､Knowledge Bases for Amazon BedrockとPinecone Serverlessを使って構築したRAGと､前回の記事で作成した会話履歴を保持するチャットボットを組み合わせてみます｡
RAGで使用するドキュメントは､独立行政法人情報処理推進機構(IPA)の"安全なウェブサイトの作り方"で公開されているPDFとExcelファイルです｡セキュリティに関する問い合わせに回答するチャットボットを作成します｡
言語モデルは､Anthropic Claude 3 Sonnetを使用します｡

PineconeとKnowledge Baseを使ったRAGについては､こちらの記事もご覧ください｡

デモ画面

2回目の回答は､トークン数の制約の影響か途中で切れてしまいました｡以下の履歴では省略していますが､回答のあとに参照リンクを出力します｡ここは､ConversationalRetrievalChainのreturn_source_documents=Trueで出力される内容を利用しS3バケット内のPDFやExcelファイルへのリンクを出しています｡
s3://で始まるアドレスをそのまま使っているので､実運用の際にはもう一工夫必要です｡

チャットボットとの会話履歴

このように､追加の質問に対してチャットボットは会話履歴を踏まえた回答を生成しました｡

tsuyoshi: 
ウェブページのある部分をクリックしたところ､そのページとは全く関係ないサイトに遷移しました｡これは､ウェブページにどのような攻撃が仕組まれていたのでしょうか?

bedrock-slack-backlog-rag-app:
クリックした際に意図しないサイトへ遷移したということは、「クリックジャッキング」という攻撃を受けた可能性があります。
クリックジャッキングは、ウェブサイト上に不正なコードを仕掛けることで、利用者の操作をそのサイト上で意図せず別のサイトや機能を実行させる攻撃手法です。
具体的には、攻撃者が細工したページの上に、正規のサイトの重要な機能画面(例えば金融サービスの送金画面など)を透明な状態で重ねて表示させます。利用者がページ上の何らかの箇所をクリックすると、実際には重なった正規サイト上の機能を実行してしまう仕組みです。
クリックジャッキング攻撃を防ぐには、HTTPレスポンスヘッダにX-Frame-Optionsを設定して、他サイトからのフレーム表示を制限する対策が有効とされています。つまり、説明の状況からすると、そのウェブサイトがクリックジャッキングの脆弱性を放置していた可能性があります。

tsuyoshi: 
この脆弱性の被害に遭わないために､ウェブサイトの利用者､構築担当者､運用担当者それぞれどのような点に気をつけるべきでしょうか｡そして､どのような対策が必要でしょうか｡

bedrock-slack-backlog-rag-app: 
この脆弱性の被害に遭わないために、ウェブサイトの利用者、構築担当者、運用担当者それぞれが気をつける点と必要な対策は以下のとおりです。
ウェブサイトの利用者側:
- 信頼できるサイトのみにアクセスする。不審なリンクはクリックしない。
- ブラウザのポップアップブロック機能を有効にする。
- 最新のブラウザやセキュリティソフトを使う。
構築担当者側:
- X-Frame-Optionsヘッダを適切に設定する。
- コンテンツセキュリティポリシー(CSP)を導入する。
- フレーム埋め込み機能を無効化する。
- サードパーティ製のスクリプトの利用は最小限に抑える。
運用担当者側:
- ウェブサイトのコードレビューを定期的に実施する。  
- 脆弱性情報に注意を払い、必要に応じてソフトウェアを更新する。
- ウェブサイトのアクセスログを監視する。
このように、利用者、開発者、運用担当者全員がそれぞれの立場で注意を払う必要があります。特に開発者側では、フレーム埋め込み対策やCSP導入などの技術的対策が重要になります。運用側も定期的な監視とタイムリーな対応が求められます。この脆弱性からウェブサイトを守るには、以下の点に注意する必要があります。
ウェブサイトの利用者:
- 不審なリンクやメールに簡単にアクセスしない
- 信頼できないウェブサイトではクッキーを削除するなどしてセッション情報を保護する
構築担当者:
- すべての外部入力データに対してエスケープ処理を適切に行う
- HTTPレスポンスヘッダにX-XSS-ProtectionやContent Security Policyを設定する
- フレーム機構を使って埋め込みを許可するかを適切に制御する  
運用担当者:
- 利用しているソフトウェアの最新の脆弱性情報を入手し、必要に応じて修正プログラムを適用する
- ウェブアプリケーションファイアウォールなどを導入して不正スクリプトの実行を防ぐ
- セキュリティ診断ツールを用いて定期的にウェブサイトの脆弱性チェックを行う
具体的な対策としては、次のようなものが有効です。
1. 外部からの入力データに対する適切なエスケープ処理
2. X-XSS-ProtectionやContent Security Policyなどのブラウザ機能の有効化
3. クリックジャッキング対策としてX-Frame-Optionsヘッダの設定
4. 最新の脆弱性情報の入手と修正プログラムの適用
5. ウェブアプリケーションファイアウォールの導入
6. 定期的なセキュリティ診断の

構成図

参考情報

環境構築

前提条件

以下のリソースを作成･削除･変更できる権限をもつAWSユーザーを利用すること
- AWS IAM
- AWS Lambda
- AWS CloudFormation
- AWS Secrets Manager
- Amazon API Gateway
- Amazon S3
- Amazon CloudWatch Logs
- Amazon Bedrock
  - Anthropic Claude 3 Sonnetが利用可能な状態
- Amazon DynamoDB
使用するAWSリージョンは、us-east-1
Slack Appを作成するためのアカウントや権限を持っている
Pineconeにサインイン可能なアカウント(Google, Github, Microsoftのいずれか)を持っている

PineconeセットアップとIndexの作成

Pineconeのサイトにアクセスし､画面右上の Sign Up Freeからアカウント作成画面に進みます。すでにアカウントがある場は､Log inからログインします｡
もしくは､AWS MarketplaceでPinecone serverlessをサブスクライブすることもできます｡すでにPineconeにアカウントを持っている場合は､AWS MarketplaceでサブスクライブすることによりAWSアカウントに紐付けることができるようです｡

AWS MarketplaceでPinecone serverlessをサブスクライブしナレッジベースをデータストアと同期するまでの手順が公式ドキュメントに詳しく書かれています｡ナレッジベースの画面が最新のものとことなりますが､問題無く進められると思います｡

以下は､Pinecone Serverlessに新規Indexを作成し､ナレッジベースをデータストアと同期するまでの手順です｡

Serverlesの画面を開き､Create IndexをクリックしてIndexを作成します｡

入力/選択項目	内容
Name	任意のIndex名｡使用可能な文字は小文字､数字､ハイフンのみ
Dimensions	1536
Metric	cosine
Capacity mode	Serverless
Cloud provider	aws
Region	us-east-1

後ほど作成するknowledge baseの埋め込みモデルTitan Embeddings G1 - Text v1.2が1536次元のため､Indexを1536次元で作成します｡ナレッジベースはus-east-1で構築するため､Regionはus-east-1を選択します｡

Indexの作成が完了すると､このようにIndexes画面にHOSTアドレスが､API Keys画面にAPIキーが表示されます｡

PineconeのAPI KeyをSecrets Managerに登録

AWS Secrets Managerを開き､シークレットの登録を行います｡

入力/選択項目	内容
シークレットのタイプ	その他のシークレットのタイプ
キー	apiKey
値	PineconeのAPIキー

シークレットの名前と説明を入力します｡

ローテーションの設定はすべてデフォルトのままにします｡

シークレットの作成が完了すると､Secret ARNが発行されます｡
Secret ARNは後ほど作成するknowledge baseの設定で使用します｡

Bedrock環境構築

Knowledge Bases for Amazon Bedrock

ファイルをS3バケットにアップロード

任意のS3バケットにファイルをアップロードします｡knowledge baseをus-east-1に作成するためリージョン間の通信費用を考慮し､S3バケットのリージョンもus-east-1としたほうが良いと思います｡
knowledge basesがサポートするファイルフォーマットは､Set up your data for ingestionに記載されています｡
引用すると以下のとおりです｡また､1ファイルあたりの最大ファイルサイズは50MBです｡

Plain text (.txt)
Markdown (.md)
HyperText Markup Language (.html)
Microsoft Word document (.doc/.docx)
Comma-separated values (.csv)
Microsoft Excel spreadsheet (.xls/.xlsx)
Portable Document Format (.pdf)

ここでは､､独立行政法人情報処理推進機構(IPA)の"安全なウェブサイトの作り方"で公開されているPDFやExcelファイルのなかから以下のファイルをひとつのS3バケットにアップロードします｡

安全なウェブサイトの作り方

「安全なウェブサイトの作り方」: 安全なウェブサイトの作り方（全115ページ）(PDF:2.2 MB); セキュリティ実装チェックリスト(Excel:18 KB)>
別冊：「安全なSQLの呼び出し方」: 安全なSQLの呼び出し方（全40ページ） (PDF:714 KB)
別冊：「ウェブ健康診断仕様」: ウェブ健康診断仕様（全30ページ）(PDF:771 KB)

安全なウェブサイトの運用管理に向けての20ヶ条〜セキュリティ対策のチェックポイント〜

チェックリスト: ウェブサイトのセキュリティ対策のチェックポイント20ヶ条チェックリスト(Excel:14 KB)
参考資料: ウェブサイト運営者のための脆弱性対応ガイド(PDF:1.2 MB); ウェブサイト構築事業者のための脆弱性対応ガイド(PDF:1.3 MB); セキュリティ担当者のための脆弱性対応ガイド(PDF:1.4 MB); ウェブサイト運営のファーストステップ～ウェブサイト運営者がまず知っておくべき脅威と責任～(PDF:1.7 MB)

Knowledge baseを作成

knowledge baseを開き､knowledge baseを作成します｡

入力/選択項目	内容
Knowledge base name	任意のknowledge base名
Knowkedge base description	knowledge baseの説明文
Runtime role	Create and use a new service role
Service role name	任意のロール名

続いて､データソースとしてファイルをアップロードしたS3バケットを指定します｡
S3バケットのURIは､直接URIを入力するか､Browse S3をクリックしてバケット一覧からバケットを選択することで入力できます｡
Chunking strategyは､チャンク分割の設定を選択します｡knowledge base作成後に設定を変更することができないため､この時点で選択が必要です｡それぞれどのように分割するのかはプルダウンメニュー内に記載されています｡また､Set up your data for ingestionにも記載があります｡

入力/選択項目	内容
Data source name	任意のデータソース名
Data source location	データソースの場所がこのAWSアカウントもしくは､別のAWSアカウントかを選択
S3 URI	S3バケットのURI
Chunking strategy	fixed size chunking
Max Tokens	512
Overlap percentage between chunks	25
Data deletion policy	Retain

チャンク分割の設定値はこれが正解というものはなく､様々な研究があるようです｡今回は､この記事を参考に設定をしました｡

つづいて､埋め込みモデルとVector storeを設定します｡
埋め込みモデルは､Titan Embeddings G1 - Text v1.2を使用します｡PineconeのIndexを作成する際に設定した Dimensions: 1536 はここに明記されています｡

入力/選択項目	内容
Embeddings model	Titan Embeddings G1 - Text v1.2
Select how you want to create your vector store	Choose a vector store you have created
Select your vector store	Pinecone
By selecting "Pinecone"...	✓
Connection String	PineconeのIndexのHOSTアドレス
Credentials secret ARN	PineconeのAPIキーを登録したSecretsのARN
Text field	Vector databaseのテキストフィールドの任意の名前｡ `text`など
Bedrock-managed metadata field	Vector databaseのメタデータフィールドの任意の名前｡ `metadata`など

By selecting "Pinecone"... の項目は､AWSがユーザーに代わってサードバーティーのソースにアクセスすることを許可することに同意する旨が書かれています｡

最後に入力内容を確認し､設定完了です｡

Knowledge baseの作成中､以下のエラーが発生しデータソースの作成に失敗しました｡そのため､実際はKnowledge base作成後にデータソースを作成しなおしました｡

failed to add datasource {データソース名} to the knowledge base. string_value can not be converted to an integer

knowledge baseは作成されましたが､まだS3バケットのデータはPineconeに登録されていません｡作成したknowledge baseのData Sourceの項目に移動し､Syncボタンをクリックします｡S3バケット内のデータに追加/更新があった場合も同様にSyncを実行します｡

Statusが Ready となれば完了です｡PineconeのIndex画面にアクセスすると､このようにデータが登録されていることが分かります｡

開発環境構築

作業環境のOSバージョン

Windows 11上のWSLでUbuntu 23.04を動かしています。

$ cat /etc/os-release | grep PRETTY_NAME
PRETTY_NAME="Ubuntu 23.04"

Python環境

$ python3 --version
Python 3.12.0

$ python3 -m venv .venv
$ source .venv/bin/activate

$ pip3 install --upgrade pip
$ pip3 --version
pip 24.0 from /home/xxx/.venv/lib/python3.12/site-packages/pip (python 3.12)

AWS環境構築

aws configureコマンドでデフォルトのリージョンやクレデンシャルを設定するか､もしくは~/.aws/configや~/.aws/credentialsを用意します。

AWS SAM CLIインストール

AWS上でサーバーレスアプリケーションを構築､実行するAWS SAMを使用します｡

Installing the AWS SAM CLI の手順に従い、AWS SAM CLIをインストールします。今回はx86_64環境でLinux OSを使用するため、x86_64 - command line installerの手順を実行します｡

$ sam --version
SAM CLI, version 1.113.0

Slack Appの作成

Slac APIを開き､From scratchからSlack Appを作成します｡ここでは、App Nameをbedrock-slack-backlog-rag-appとします｡

Basic Information画面のApp Credentialsに表示されているクレデンシャルはSlackSigningSecretとして後述のSecret Managerのシークレット登録に使用します。

OAuth & Permissions画面のOAuth Tokens for Your WorkspaceにあるBot User OAuth Tokenは、SlackBotTokenとして後述のSecret Managerのシークレット登録に使用します。

OAuth & Permissions画面のBot Token Scopesにapp_mentions:readとchat:writeを追加します｡

シークレット情報をSecret Managerに登録

あらたにシークレットを作成し、ここまでの手順で作成した以下のシークレット情報を登録します。

シークレットキー	値
SlackSigningSecret	前述のSlackのSigning Secret
SlackBotToken	前述のSlackのBot User OAuth Token

アプリケーションの構築

ディレクトリ構造は以下のとおりです｡

.
├── bedrock-slack-backlog-rag-app
│   ├── __init__.py
│   ├── app.py
│   └── requirements.txt
├── samconfig.toml
└── template.yaml

__init__.pyは空のファイルです｡
bedrock-slack-backlog-rag-app/requirements.txtは以下のとおりです｡boto3やrequestsも必要ですが､それらはLambdaレイヤーで追加するようtemplate.yamlに記述します｡

slack-bolt
slack-sdk
langchain

template.yamlの構成

template.yaml (長いので折りたたんでいます｡クリックして展開)

template.yaml

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: Slack Bedrock Assitant.

Resources:
  # Lambda function for Bedrock
  BedrockRAGFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: bedrock-slack-backlog-rag-app/
      Handler: app.lambda_handler
      Runtime: python3.12
      Role: !GetAtt LambdaRole.Arn
      Timeout: 300
      MemorySize: 512
      Architectures:
        - arm64
      Policies:
        - DynamoDBCrudPolicy:
            TableName: !Ref DynamoDBTable
      Environment:
        Variables:
          SECRET_NAME: 'Bedrock-sam-secrets-backlog-rag' # Name of the secret in Secrets Manager
          REGION_NAME: 'us-east-1' # Region of the secret in Secrets Manager
          DYNAMODB_TABLE_NAME: !Ref DynamoDBTable

      Events:
        Slack:
          Type: Api
          Properties:
            Method: POST
            Path: /slack/events
      Layers:
        # Layer for AWS Parameter Store and Secrets Manager
        # https://docs.aws.amazon.com/systems-manager/latest/userguide/ps-integration-lambda-extensions.html#ps-integration-lambda-extensions-add
        - arn:aws:lambda:us-east-1:177933569100:layer:AWS-Parameters-and-Secrets-Lambda-Extension-Arm64:11
        # Layer for boto3
        # https://github.com/keithrozario/Klayers?tab=readme-ov-file#list-of-arns
        - arn:aws:lambda:us-east-1:770693421928:layer:Klayers-p312-arm64-boto3:1

  # DynamoDB Table for storing chat history
  DynamoDBTable:
    Type: AWS::DynamoDB::Table
    Properties:
      TableName: 'bedrock-slack-backlog-rag-app-chat-history'
      AttributeDefinitions:
        - AttributeName: 'SessionId'
          AttributeType: 'S'
      KeySchema:
        - AttributeName: 'SessionId'
          KeyType: 'HASH'
      BillingMode: PAY_PER_REQUEST

  # IAM Role for lambda.
  LambdaRole:
    Type: "AWS::IAM::Role"
    Properties:
      RoleName: bedrock-slack-backlog-rag-app-role
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: allow-lambda-invocation
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                  - lambda:InvokeAsync
                Resource: "*"
        - PolicyName: SecretsManagerPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: 'secretsmanager:GetSecretValue' # Required for Lambda to retrieve the secret
                Resource: "*"
        - PolicyName: allow-bedrock-agent-access
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - bedrock:InvokeAgent
                  - bedrock:InvokeModel
                  - bedrock:Retrieve
                  - bedrock:InvokeModelWithResponseStream
                Resource: "*"
        - PolicyName: DynamoDBCrudPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - dynamodb:PutItem
                  - dynamodb:GetItem
                  - dynamodb:UpdateItem
                  - dynamodb:DeleteItem
                Resource: "*"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  BedrockRAGFunctionLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${BedrockRAGFunction}
      RetentionInDays: 14 # Optional. Default retention is 30 days.

Outputs:
  BedrockAssitantApi:
    Description: "The URL of Slack Event Subscriptions"
    Value: !Sub "https://${ServerlessRestApi}.execute-api.${AWS::Region}.amazonaws.com/Prod/slack/events"
  BedrockRAGFunction:
    Description: "Bedrock Rag Lambda Function ARN"
    Value: !GetAtt BedrockRAGFunction.Arn
  BedrockRAGFunctionIamRole:
    Description: "Implicit IAM Role created for Bedrock Assistant function"
    Value: !GetAtt LambdaRole.Arn

AWS SAM テンプレートファイル(template.yaml)に､作成するAWSリソースを定義します｡
Lambda関数用ロールやポリシー､Lambdaの環境変数などを記述します｡その他に､以下のレイヤーやリソースベースポリシーが含まれます｡

Lambda関数からSecrets ManagerにアクセスするためのAWS-Parameters-and-Secrets-Lambda-Extensionレイヤー
Lambda関数内からimportするためのboto3をパッケージにしたレイヤー
BedrockkからLambda関数を扱うためのリソースベースポリシー
DynamoDBの操作を許可するポリシー
DynamoDBテーブルの作成とSessionIdをプライマリキーに設定

DynamoDBテーブル名はbedrock-slack-app-chat-historyとしています｡そのほかに､テーブル操作を許可するポリシーを以下のようにtemplate.yamlに記述しています｡

Resources:
  BedrockAssitantFunction:
    Type: AWS::Serverless::Function
    Properties:
(途中省略)
      Policies:
        - DynamoDBCrudPolicy:
            TableName: !Ref DynamoDBTable

(途中省略)
  DynamoDBTable:
    Type: AWS::DynamoDB::Table
    Properties:
      TableName: 'bedrock-slack-rag-app-chat-history'
      AttributeDefinitions:
        - AttributeName: 'SessionId'
          AttributeType: 'S'
      KeySchema:
        - AttributeName: 'SessionId'
          KeyType: 'HASH'
      BillingMode: PAY_PER_REQUEST

(途中省略)
  LambdaRole:
    Type: "AWS::IAM::Role"
    Properties:
      Policies:
(途中省略)
        - PolicyName: DynamoDBCrudPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - dynamodb:PutItem
                  - dynamodb:GetItem
                  - dynamodb:UpdateItem
                  - dynamodb:DeleteItem
                Resource: "*"

Lambdaレイヤーは以下のようにtemplate.yamlに記述しています｡

      Layers:
        # Layer for AWS Parameter Store and Secrets Manager
        # https://docs.aws.amazon.com/systems-manager/latest/userguide/ps-integration-lambda-extensions.html#ps-integration-lambda-extensions-add
        - arn:aws:lambda:us-east-1:177933569100:layer:AWS-Parameters-and-Secrets-Lambda-Extension-Arm64:11
        # Layer for boto3
        # https://github.com/keithrozario/Klayers?tab=readme-ov-file#list-of-arns
        - arn:aws:lambda:us-east-1:770693421928:layer:Klayers-p312-boto3:4

適用すると､Lambda関数のLayersに以下のように表示されます｡

リソースベースポリシーは､以下のようにtemplate.yamlに記述しています｡

  BacklogSearchFunction:
    Type: AWS::Serverless::Function
    Properties:

(途中省略)

  # Resouse based policy for lambda.
  PermissionForBacklogSearchToInvokeLambda:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt BacklogSearchFunction.Arn
      Action: lambda:InvokeFunction
      Principal: bedrock.amazonaws.com

適用すると､Lambda関数の設定のResource-based policy statementsに以下のように表示されます｡

Knowledge baseが完了すると､このようにKnowledge base IDが表示されます｡

template.yaml内のEnvironmentにあるSECRET_NAMEとREGION_NAMEには､それぞれ先ほど作成したSecrets Managerのシークレットの名前とリージョンを設定します｡

samconfig.tomlの構成

samconfig.toml (長いので折りたたんでいます｡クリックして展開)

samconfig.toml

# More information about the configuration file can be found here:
# https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-sam-cli-config.html
version = 0.1

[default]
[default.global.parameters]
stack_name = "bedrock-slack-backlog-rag-app"

[default.build.parameters]
cached = true
parallel = true

[default.validate.parameters]
lint = true

[default.deploy.parameters]
capabilities = "CAPABILITY_NAMED_IAM"
confirm_changeset = true
resolve_s3 = true
region = "us-east-1"

[default.package.parameters]
resolve_s3 = true

[default.sync.parameters]
watch = true

[default.local_start_api.parameters]
warm_containers = "EAGER"

[default.local_start_lambda.parameters]
warm_containers = "EAGER"

SAM CLIの実行設定ファイル(samconfig.toml)に､SAM CLIを実行する際の設定を定義します｡AWS SAMのチュートリアル: Hello World アプリケーションのデプロイを実行した際に作成されるsamconfig.tomlをもとにしています｡今回の例では､以下の点を変更しています｡

[default.global.parameters]セクションのstack_nameを"sam-app"から"bedrock-slack-backlog-rag-app"に変更
[default.deploy.parameters]セクションにregion指定を追加
[default.deploy.parameters]セクションのcapabilitiesを"CAPABILITY_IAM"から"CAPABILITY_NAMED_IAM"に変更

bedrock-slack-backlog-rag-app/app.pyの構成

bedrock-slack-app/app.p (長いので折りたたんでいます｡クリックして展開)

bedrock-slack-backlog-rag-app/app.py

import ast
import logging
import os
import re
import time
from typing import Any

import boto3
from botocore.exceptions import ClientError
from langchain.callbacks.base import BaseCallbackHandler
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import (
    PromptTemplate,
)
from langchain.retrievers import AmazonKnowledgeBasesRetriever
from langchain.schema import LLMResult
from langchain_community.chat_message_histories import DynamoDBChatMessageHistory
from langchain_community.chat_models import BedrockChat
from slack_bolt import App
from slack_bolt.adapter.aws_lambda import SlackRequestHandler

CHAT_UPDATE_INTERVAL_SEC = 1

SlackRequestHandler.clear_all_log_handlers()
logging.basicConfig(
    format="%(asctime)s [%(levelname)s] %(message)s",
    level=logging.DEBUG
)

logger = logging.getLogger(__name__)

REGION_NAME = "us-east-1"
MODEL_ID = "anthropic.claude-3-sonnet-20240229-v1:0"
KNOWLEDGE_BASE_ID = ""
DYNAMODB_TABLE_NAME = "bedrock-slack-rag-app-chat-history"


class SecretsManager:
    """
    Class to retrieve secrets from Secrets Manager

    Attributes:
        secret_name (str): The name of the secret
        region_name (str): The name of the region
        client (boto3.client): The client for Secrets Manager
    """

    def __init__(self, secret_name, region_name):
        self.secret_name = secret_name
        self.region_name = region_name
        self.client = boto3.client(
            service_name='secretsmanager',
            region_name=region_name
        )

    def get_secret(self, key):
        """
        Retrieves the value of a secret based on the provided key.

        Args:
            key (str): The key of the secret to retrieve.

        Returns:
            str: The value of the secret.

        Raises:
            ClientError: If there is an error retrieving the secret.
        """
        try:
            get_secret_value_response = self.client.get_secret_value(
                SecretId=self.secret_name
            )
        except ClientError as e:
            raise e

        secret_data = get_secret_value_response['SecretString']
        secret = ast.literal_eval(secret_data)

        return secret[key]


secrets_manager = SecretsManager(
    secret_name=os.environ.get("SECRET_NAME"),
    region_name=os.environ.get("REGION_NAME")
)

app = App(
    signing_secret=secrets_manager.get_secret("SlackSigningSecret"),
    token=secrets_manager.get_secret("SlackBotToken"),
    process_before_response=True,
)


class SlackStreamingCallbackHandler(BaseCallbackHandler):
    """
    A callback handler for handling events during Slack streaming.

    Attributes:
        last_send_time (float): The timestamp of the last message sent.
        message (str): The accumulated message to be sent.

    Args:
        channel (str): The Slack channel to send messages to.
        ts (str): The timestamp of the message to be updated.
    """

    last_send_time = time.time()
    message = ""

    def __init__(self, userid, channel, ts):
        self.userid = userid
        self.channel = channel
        self.ts = ts
        self.interval = CHAT_UPDATE_INTERVAL_SEC
        self.update_count = 0

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        """
        Event handler for a new token received.

        Args:
            token (str): The new token received.
            **kwargs: Additional keyword arguments.
        """

        self.message += token

        now = time.time()
        if now - self.last_send_time > self.interval:

            # mention_message = f"<@{self.userid}> {self.message}"
            # message_blocks = create_message_blocks(mention_message)

            app.client.chat_update(
                channel=self.channel,
                ts=self.ts,
                text=f"<@{self.userid}> {self.message}",
                # blocks=message_blocks
            )
            self.last_send_time = now
            self.update_count += 1

            if self.update_count / 10 > self.interval:
                self.interval = self.interval * 2

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
        """
        Event handler for the end of Slack streaming.

        Args:
            response (LLMResult): The result of the Slack streaming.
            **kwargs: Additional keyword arguments.

        Returns:
            Any: The result of the event handling.
        """

        mention_message = f"<@{self.userid}> {self.message}"
        message_blocks = create_message_blocks(mention_message)

        app.client.chat_update(
            channel=self.channel,
            ts=self.ts,
            text=self.message,
            blocks=message_blocks
        )


def create_message_blocks(text):
    """
    Creates the message blocks for updating the Slack message.

    Args:
        text (str): The updated text for the Slack message.

    Returns:
        list: The message blocks for updating the Slack message.
    """

    message_context = "Claude 3 Sonnetで生成される情報は不正確な場合があります。"
    message_blocks = [
        {
            "type": "section",
            "text":
                {
                    "type": "mrkdwn",
                    "text": text
                }
        },
        {
            "type": "divider"
        },
        {
            "type": "context",
            "elements": [
                {
                    "type": "mrkdwn",
                    "text": message_context
                }
            ]
        },
    ]

    return message_blocks


def prompt_template():
    """
    Returns the prompt template for the chat.

    Returns:
        str: The prompt template.
    """

    chat_template = """
    Let's think step by step.
    Take a deep breath.

    Answer the question based on the context below.
    And also, follow the rules below.

    This is rules for chat:
    Answer in Japanese if the question is asked in Japanese.
    If you cannot answer a question due to lack of specificity, please advise on how to ask the question.


    This is your context:
    {context}

    Question: {question}
    Answer:
    """

    return PromptTemplate(
        input_variables=["context", "question"],
        template=chat_template
    )


def system_instruction_template():
    """
    Returns the system instruction template for the chat.

    Returns:
        str: The system instruction template.
    """

    # Define your system instruction
    system_instruction = "The assistant should provide detailed explanations."

    # Define your template with the system instruction
    template = (
        f"{system_instruction} "
        "Combine the chat history and follow up question into "
        "a standalone question. Chat History: {chat_history}"
        "Follow up question: {question}"
    )

    # Create the prompt template
    return PromptTemplate.from_template(template)


def get_bedrock_knowledge_base(knowledge_base_id, region_name):

    return AmazonKnowledgeBasesRetriever(
            knowledge_base_id=knowledge_base_id,
            region_name=region_name,
            retrieval_config={
                "vectorSearchConfiguration": {
                    "numberOfResults": 5
                }
            }
        )


def get_bedrock_llm(model_id, region_name, callback: SlackStreamingCallbackHandler):

    return BedrockChat(
            model_id=model_id,
            region_name=region_name,
            streaming=True,
            callbacks=[callback],
            model_kwargs={
                "max_tokens": 500,
                "temperature": 0.99,
                "top_p": 0.999
            },
            verbose=True
        )


def get_chain_bedrock_knowledge_base(llm, memory, knowledge_base_id, region_name):
    retriever = get_bedrock_knowledge_base(
                knowledge_base_id=knowledge_base_id,
                region_name=region_name
            )

    return ConversationalRetrievalChain.from_llm(
            llm=llm,
            retriever=retriever,
            chain_type="stuff",  # Or "refine" | "map_reduce"
            memory=memory,
            return_source_documents=True,
            # Prompt template for generated question .
            condense_question_prompt=system_instruction_template(),
            # Prompt template for combining documents.
            combine_docs_chain_kwargs={'prompt': prompt_template()},
            get_chat_history=lambda h: h,
            verbose=True,
            # Rephrase the question before asking the knowledge base.
            rephrase_question=False
        )


def handle_app_mentions(event, say):
    """
    Handle app mentions in Slack.

    Args:
        event (dict): The event data containing information about the mention.
        say (function): The function used to send a message in Slack.

    Returns:
        None
    """

    channel = event["channel"]
    thread_ts = event["ts"]
    input_text = re.sub("<@.*>", "", event["text"])
    userid = event["user"]

    # セッションIDとして、thread_tsを使用
    # 初回はevent["ts"]を使用､以降はevent["thread_ts"]を使用
    id_ts = event["ts"]
    if "thread_ts" in event:
        id_ts = event["thread_ts"]

    result = say("\n\nお待ちください...", thread_ts=thread_ts)
    ts = result["ts"]

    history = DynamoDBChatMessageHistory(
            table_name=DYNAMODB_TABLE_NAME,
            session_id=id_ts,
            ttl=3600
        )

    callback = SlackStreamingCallbackHandler(
                    userid=userid,
                    channel=channel,
                    ts=ts
                )

    llm = get_bedrock_llm(
            model_id=MODEL_ID,
            region_name=REGION_NAME,
            callback=callback
            )

    memory = ConversationBufferMemory(
                chat_memory=history,
                input_key="question",
                memory_key="chat_history",
                output_key="answer",
                # Return messages in the memory as list.
                return_messages=True,
                human_prefix="H",
                assistant_prefix="A"
            )

    chain = get_chain_bedrock_knowledge_base(
                llm=llm,
                memory=memory,
                knowledge_base_id=KNOWLEDGE_BASE_ID,
                region_name=REGION_NAME
            )

    result = chain.invoke(
                {
                    "question": input_text,
                    "chat_history": memory.chat_memory.messages
                }
            )

    source_documents = result.get('source_documents')
    uri, score, references = "", "", ""
    for i, refs in enumerate(source_documents):
        count = i + 1
        uri = refs.metadata['location']['s3Location']['uri']
        score = round(refs.metadata['score'] * 100, 2)
        text = re.sub(r"[\n\s]+", "", refs.page_content[:40])

        references += f'[{count}] <{uri}|{text}...>' + "  " + f"(関連度: {score}%)\n"

    say("[参照情報]\n\n" + references, thread_ts=thread_ts)


def respond_to_slack_within_3_seconds(ack):
    """
    Responds to a Slack message within 3 seconds.

    Parameters:
    - ack: A function to acknowledge the Slack message.

    Returns:
    None
    """
    ack()


app.event("app_mention")(
    ack=respond_to_slack_within_3_seconds,
    lazy=[handle_app_mentions]
)


def lambda_handler(event, context):
    """
    Lambda function handler for processing Slack events.

    Args:
        event (dict): The event data passed to the Lambda function.
        context (object): The runtime information of the Lambda function.

    Returns:
        dict: The response data to be returned by the Lambda function.
    """
    print(event)

    retry_counts = event.get("multiValueHeaders", {}).get("X-Slack-Retry-Num", [0])

    if retry_counts[0] != 0:
        logging.info("Skip slack retrying(%s).", retry_counts)
        return {}

    slack_handler = SlackRequestHandler(app=app)
    return slack_handler.handle(event, context)

bedrock-slack-backlog-rag-app/app.pyはメンションされたメッセージを取得し､Amazon Bedrockを使用して生成したテキストをストリーミングで返すアプリケーションです｡また､DynamoDBChatMessageHistoryを使用して､会話履歴をDynamoDBに保存します｡

LangChainのConversationalRetrievalChainを利用して質問応答を実現しています｡ConversationalRetrievalChainについては､以下の記事に使い方の解説があり､参考になります｡

参考

変更が必要な点

35行目のKNOWLEDGE_BASE_IDに､構築したKnowledge baseのIDを指定します｡
if now - self.last_send_time > self.interval:の箇所で以下のようにコメントアウトしています｡ストリーミングで回答を生成している最中もSection blockを使って出力するする場合はこのコメントアウトを外してください｡

        if now - self.last_send_time > self.interval:

            # mention_message = f"<@{self.userid}> {self.message}"
            # message_blocks = create_message_blocks(mention_message)

            app.client.chat_update(
                channel=self.channel,
                ts=self.ts,
                text=f"<@{self.userid}> {self.message}",
                # blocks=message_blocks
            )

ビルド

template.yamlがあるディレクトリで､ビルドコマンドを実行します｡

$ sam build

ビルドに成功すると､以下のようなメッセージが表示されます｡

Starting Build use cache
Manifest is not changed for (BedrockAssitantFunction), running incremental build
Building codeuri: /home/xxx/aws-sam-bedrock-slack-rag-app/bedrock_slack_rag_app runtime: python3.11 metadata: {} architecture:
x86_64 functions: BedrockAssitantFunction
 Running PythonPipBuilder:CopySource
 Running PythonPipBuilder:CopySource

Build Succeeded

Built Artifacts  : .aws-sam/build
Built Template   : .aws-sam/build/template.yaml

Commands you can use next
=========================
[*] Validate SAM template: sam validate
[*] Invoke Function: sam local invoke
[*] Test Function in the Cloud: sam sync --stack-name {{stack-name}} --watch
[*] Deploy: sam deploy --guided

デプロイ

ビルドでエラーがなければsam deployコマンドを実行し､デプロイを行います。

$ sam deploy

デプロイが成功すると､以下のような情報がコンソールに出力されます｡

CloudFormation outputs from deployed stack
-------------------------------------------------------------------------------------------------------------------------------------------------------------
Outputs
-------------------------------------------------------------------------------------------------------------------------------------------------------------
Key                 BedrockAssitantApi
Description         The URL of Slack Event Subscriptions
Value               https://xxxxxxxxxx.execute-api.us-east-1.amazonaws.com/Prod/slack/events

Key                 BedrockAssitantFunction
Description         Bedrock Assistant Lambda Function ARN
Value               arn:aws:lambda:us-east-1:xxxxxxxxxxxx:function:bedrock-slack-rag-app-BedrockAssitantFunction-xxxxxxxxxxxx

Key                 BedrockAssitantFunctionIamRole
Description         Implicit IAM Role created for Bedrock Assistant function
Value               arn:aws:iam::xxxxxxxxxxxx:role/bedrock-slack-rag-app-lambda-role
-------------------------------------------------------------------------------------------------------------------------------------------------------------

Slack Appの設定

メンションイベントに応答するため､Event Subscriptions画面のSubscribe to bot eventsにapp_mentionを追加します｡

Event Subscriptions画面のEnable EventsをOnにし､Request URLへさきほど出力されたURL https://xxxxxxxxxx.execute-api.us-east-1.amazonaws.com/Prod/slack/events を入力します｡
Verified✓ と表示されれば､正しいURLが入力されたことになります｡レスポンスが得られるようになるまで時間がかかる場合もあります。正しいURLを入力しているにもかかわらずVerifiedとならない場合は、時間をおいて再試行します。

設定を追加後､画面最下部にあるSave Changesをクリックし内容を保存します｡

Slack Workspaseにアプリをインストール

Install App画面のInstall to Workspaceをクリックし､Slack AppをWorkspaceにインストールします｡
インストールが成功すると､Thank you!画面が表示されます｡Slackアプリをインストールしている場合はclick hereのリンク､Webブラウザを使用している場合は､this linkをクリックしてSlackを開きます｡

動作確認

任意のチャンネンルに､@bedrock-slack-backlog-rag-app (Slack Appの作成時に設定した名前)を招待し､メンション形式で依頼をポストします｡動作しない場合のログやOpenAI APIが返すレスポンスは CloudWatch Logs に出力されているログが参考になります｡

作成したリソースの削除

最後に､作成したアプリケーションを削除する手順です｡リソースを削除するには sam delete コマンドを実行します｡

$ sam delete

まとめ

今回は､RAGと組み合わせたチャットボットを作成してみました｡ DynamoDBに会話履歴を持たせることで､ウェブサイトのセキュリティについて気軽に尋ねられるチャットボットが実現しました｡
IPAの資料は情報量が豊富で専門的な知識のかたまりなので､RAGに向いていると実感しました｡

たとえば、このようなチャットボットを社内に導入することで､

ちょっと気になったことを質問する
セキュリティ対策を行う際に要点整理の参考にする

などといった使い方ができそうです｡
ウェブサイトを設計､構築､運用する際に､様々なセキュリティを考慮する必要があります｡その際に､IPAの資料について回答するチャットボットがあれば効率的に情報整理ができそうです｡

今後の発展

Knowledge baseがCloudFormation によるデプロイをサポートしたそうなので､今後はKnowledge base構築をコード化したいと思います｡

LangChainのRetrieveAndGenerateを使用するとSessionIDのみで会話履歴を保持できるので、これを使って構成やコードをシンプルにできればと考えています。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up