LlamaIndex のインデックスを読み込んで応答する Slack アプリを AWS Fargate で実行する

SlackアプリでLlamaIndexを動作させる実行環境として、AWS Fargateのコンテナ環境を利用する構成例をすぐに試せるCloudFormationテンプレート付きでご紹介します
2023.03.19

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

ども、大瀧です。

LlamaIndexはOpenAIのLLMに独自のデータを読み込ませる仕組みです。本ブログではSlackアプリでLlamaIndexを動作させる実行環境として、AWS Fargateのコンテナ環境を利用する例をご紹介します。

構成

AWSの構成は非常にシンプルです。Slackアプリのソケットモードはインターネットからのアクセスが不要なので、AWS FargateのサービスをELBなしで構成します。あらかじめ作成しておいたLlamaIndexのインデックスファイルをS3にアップロードし、コンテナではインデックスファイルを読み込んでSlack API、OpenAI APIとそれぞれ連携し動作します。

サンプルプログラム

LlamaIndexを動作させる以下のプログラムとDockerfileを用意しました。

llama_index_inference/
├── Dockerfile
├── app.py
└── requirements.txt

app.py

import boto3
import os
import re
from llama_index import GPTSimpleVectorIndex
from slack_bolt import App
from slack_bolt.adapter.socket_mode import SocketModeHandler

# S3にあるインデックスファイルを読み込んでインデックスを構築
s3 = boto3.resource('s3')
json = s3.Bucket(os.environ['INDEX_S3_BUCKET_NAME']).Object(os.environ['INDEX_S3_OBJECT_KEY']).get()['Body'].read()
index = GPTSimpleVectorIndex.load_from_string(json)

app = App(token=os.environ.get("SLACK_BOT_TOKEN"))

@app.event("app_mention")
def reply_app_mention(event, say):
    # Slackでメンションを受け取ったらメッセージからメンションを取り除く
    mention_text = re.sub(r'<@.*?> ', "", event["text"])
    # LlamaIndexにクエリし、レスポンスを取得
    response = index.query(mention_text)
    # メンションを受けたSlackメッセージのスレッドにレスポンスを返信
    channel = event["channel"]
    thread_ts = event["thread_ts"] if "thread_ts" in event else event["ts"]
    say(text=response.response, channel=channel, thread_ts=thread_ts)

if __name__ == "__main__":
    SocketModeHandler(app, os.environ["SLACK_APP_TOKEN"]).start()

ざっくり解説します。

  • 9-11行目: S3からLlamaIndexのインデックスファイルを読み込み、インデックスとしてロードします
  • 15-24行目: Slackアプリとしてメンションされたときにスレッドを作成し、LlamaIndex GPTのレスポンスをメッセージとして応答します
  • 27行目: Slackアプリをソケットモードで実行します

Slackアプリの実装は以下のブログとほぼ同様です。というか踏襲させてもらいました?
Slack APIの3秒ルールは動作確認した範囲では発動しなかったので、一旦特に対処はしていません。

Pythonの依存ライブラリ周りは特に難しいことはしていません。それぞれの処理に必要な最低限のライブラリを記載しました。

requiements.txt

boto3==1.26.94
llama-index==0.4.29
slack-bolt==1.16.4

Dockerfileも同様にSlack Bolt for Pythonの一般的な動作を記述しました。

Dockerfile

FROM python:3.11.0-slim

WORKDIR /app

COPY . .

RUN pip install -r requirements.txt && pip cache purge

CMD ["python", "app.py"]

コンテナのビルド

  • OS: Ubuntu Linux 20.04 Arm64
  • Docker: docker-ce 23.0.1

以下のコマンドでDockerイメージをビルド、ECR Publicにプッシュしました。(Privateリポジトリでも同様です)

$ sudo docker build -t takipone/llama_index_inference:latest .
$ aws ecr-public get-login-password --region us-east-1 | sudo docker login --username AWS --password-stdin public.ecr.aws/XXXXXXXX
$ sudo docker tag takipone/llama_index_inference:latest public.ecr.aws/XXXXXXXX/llama_index_inference:latest
$ sudo docker push public.ecr.aws/XXXXXXXX/llama_index_inference:latest

以下のイメージを公開しているので、任意のAWSアカウントから後述のCloudFormationスタックでビルドなしで利用できます。

※ コンテナイメージの提供を保証するものではありません。予告なくイメージの公開を終了する場合があります

インデックスファイルのアップロード

今回の記事にLlamaIndexのインデックスファイルの作成は含みません。以下の記事を参考に、インデックスを保存したJSONファイルをあらかじめ作成しておきます。前述のサンプルプログラムではGPTSimpleVectorIndex関数でロードしているので、インデックスファイルも同じベクターインデックスを作成します。

これをS3の任意のバケットにアップロードします。バケット名とフォルダ、ディレクトリ名はあとでCloudFormationのパラメータで指定するので、メモしておきます。

これでOKです。

Slackアプリの登録

Bolt(Slackアプリのフレームワーク)の入門ガイドにある以下を行います。

  1. アプリを作成する
  2. トークンとアプリのインストール
  3. イベントを設定する

途中のOAuth Scopeの追加で、メンションに返すために必要な以下のスコープを追加します。

  • app_mentions:read
  • calls:write

[トークンとアプリのインストール]までを完了すると、手元に以下2つのトークンを控えているはずです。このあとのCloudFormationのパラメータに入力します。

  • ボットトークン(Bot User OAuth Token)
  • アプリレベルトークン

AWS環境の構築

先ほどの構成図にあったFargate周りをCloudFormationで構築します。以下のテンプレートを実行しCloudFormationスタックを作成します。

AWSTemplateFormatVersion: 2010-09-09
Parameters:
  StackSuffix:
    Description: Stack Suffix use to resource name suffix
    Type: String
  SlackAppToken:
    Description: Slack app token
    Type: String
    NoEcho: true
  SlackBotToken:
    Description: Slack bot boken
    Type: String
    NoEcho: true
  OpenAIAPIKey:
    Description: OpenAI API key
    Type: String
    NoEcho: true
  IndexS3BucketName:
    Description: S3 bucket name where the llama index index file stored
    Type: String
  IndexS3ObjectKey:
    Description: S3 object key name  where the llama index index file stored
    Type: String
    Default: index.json
  DesiredCount:
    Description: Number of tasks to launch in your service.
    Type: Number
    Default: 1
  VPCSubnet:
    Type: AWS::EC2::Subnet::Id
  VPCSecurityGroup:
    Type: AWS::EC2::SecurityGroup::Id
Resources:
  ECSCluster:
    Type: AWS::ECS::Cluster
    Properties:
      ClusterName: !Sub 'llama-index-inference-${StackSuffix}'
  ECSTaskDefinition:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: !Sub 'llama-index-inference-${StackSuffix}'
      TaskRoleArn: !Ref ECSTaskRole
      ExecutionRoleArn: !Ref ECSTaskExecutionRole
      RequiresCompatibilities:
        - 'FARGATE'
      RuntimePlatform:
        CpuArchitecture: 'ARM64'
        OperatingSystemFamily: 'LINUX'
      Cpu: 256
      Memory: 512
      NetworkMode: 'awsvpc'
      ContainerDefinitions:
        - Name: !Sub 'llama-index-inference-${StackSuffix}'
          Essential: 'true'
          Image: 'public.ecr.aws/f4a3l1g8/llama_index_inference:latest'
          MemoryReservation: 128
          Secrets:
            - Name: 'OPENAI_API_KEY'
              ValueFrom: !Sub '${SecretOpenAiApiKey}:OPENAI_API_KEY::'
            - Name: 'SLACK_APP_TOKEN'
              ValueFrom: !Sub '${SecretSlackTokens}:SLACK_APP_TOKEN::'
            - Name: 'SLACK_BOT_TOKEN'
              ValueFrom: !Sub '${SecretSlackTokens}:SLACK_BOT_TOKEN::'
            - Name: 'INDEX_S3_BUCKET_NAME'
              ValueFrom: !Ref ParameterIndexS3BucketName
            - Name: 'INDEX_S3_OBJECT_KEY'
              ValueFrom: !Ref ParameterIndexS3ObjectKey
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref CWLogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: ecs
  ECSService:
    Type: AWS::ECS::Service
    Properties:
      ServiceName: !Sub 'llama-index-inference-${StackSuffix}'
      LaunchType: 'FARGATE'
      Cluster: !Ref ECSCluster
      TaskDefinition: !Ref ECSTaskDefinition
      DesiredCount: !Ref DesiredCount
      NetworkConfiguration:
        AwsvpcConfiguration:
          AssignPublicIp: "ENABLED"
          SecurityGroups:
            - !Ref VPCSecurityGroup
          Subnets:
            - !Ref VPCSubnet
  ECSTaskRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "llama-index-inference-${StackSuffix}-ECSTaskRole"
      Path: /
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: ecs-tasks.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: GetObjectPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: 's3:GetObject'
                Resource: !Join
                  - ''
                  - - 'arn:aws:s3:::'
                    - !GetAtt ParameterIndexS3BucketName.Value
                    - '/*' 
  ECSTaskExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "llama-index-inference-${StackSuffix}-ECSTaskExecutionRole"
      Path: /
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: ecs-tasks.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy
      Policies:
        - PolicyName: GetSecretsPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: 
                - 'ssm:GetParameters'
                - 'secretsmanager:GetSecretValue'
                Resource:
                - !Ref SecretOpenAiApiKey
                - !Ref SecretSlackTokens
                - !Sub 'arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter/IndexS3*'
  CWLogGroup:
    Type: AWS::Logs::LogGroup
    Properties: 
      LogGroupName: !Sub '/ecs/llama-index-inference-${StackSuffix}'              
  SecretOpenAiApiKey:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: !Sub 'OpenAiApiKey-${StackSuffix}'
      Description: 'OpenAI API Key'
      SecretString: !Sub '{"OPENAI_API_KEY": "${OpenAIAPIKey}"}'
  SecretSlackTokens:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: !Sub 'SlackTokens-${StackSuffix}'
      Description: 'Slack App and Bot tokens'
      SecretString: !Sub '{"SLACK_APP_TOKEN": "${SlackAppToken}","SLACK_BOT_TOKEN": "${SlackBotToken}"}'
  ParameterIndexS3BucketName:
    Type: AWS::SSM::Parameter
    Properties:
      Name: !Sub 'IndexS3BucketName-${StackSuffix}'
      Type: String
      Value: !Ref IndexS3BucketName
      Description: "Index S3 Bucket Name"
  ParameterIndexS3ObjectKey:
    Type: AWS::SSM::Parameter
    Properties:
      Name: !Sub 'IndexS3ObjectKey-${StackSuffix}'
      Type: String
      Value: !Ref IndexS3ObjectKey
      Description: "Index S3 Object Key"
Outputs:
  Cluster:
    Value: !Ref ECSCluster
  TaskDefinition:
    Value: !Ref ECSTaskDefinition
  ECSService:
    Value: !Ref ECSService

作成時のパラメータは以下です。

  • DesiredCount: ECSサービスのタスク数です。2以上にするとSlackのリプライが重複して返ってくる気がするので、1しか試したことはありません。
  • IndexS3BucketName: インデックスファイルをアップロードしたS3バケット名
  • IndexS3ObjectKey: インデックスファイルのフォルダ名とファイル名(バケット直下に置いた場合はファイル名のみでOK)
  • OpenAIAPIKey: OpenAI PlatformのAPIキー sk-XXXXの文字列
  • SlackAppToken: Slackのアプリレベルトークン xapp-XXXXの文字列
  • SlackBotToken: Slackのボットトークン xoxb-XXXXの文字列
  • StackSuffix: 他のスタックをかぶらない一意な任意の文字列
  • VPCSecurityGroup: VPCのセキュリティグループの選択 アウトバウンド通信のみなのでdefaultで構わない
  • VPCSubnet: 任意のPublicサブネット(タスク定義でAssignPublicIpを有効にしているため)

スタック作成が完了すると、ECSサービスのログに以下のようなメッセージが表示されれば成功です。

ではおもむろにSlackアプリにメンションしてみましょう。今回はLlamaIndexのチュートリアルに出てくるポール・グラハムのエッセイの一節をインデックスに読み込んでいるので名前を聞いてみます。

インデックスの情報から返答している様子がわかりますね!

まとめ

SlackアプリでLlamaIndexを動作させる実行環境として、AWS Fargateのコンテナ環境を利用する例をご紹介しました。インデックスファイルはコンテナ起動時に読み込むますので、インデックスファイルを更新したらECSサービスのタスク(コンテナ)を終了し新しいタスクを実行すれば新しいインデックスを利用するようになります。様々なドキュメントをインデックス化して、Slackアプリで手軽にGPTを活用していきましょう!