はじめに
こんにちわ、DROBE の都筑です。
OpenAI の提供する API を始めとして LLM をサービスで活用されている、もしくはこれから活用しようとしている方は多いと思います。 一方で、OpenAI が公式に提供している Library は Python と Node.js のみなので既存のサービスに直接インテグレーションする事が難しい場合も多いのではないでしょうか。
そういったケースに対応するために、 AWS Lambda で API の Wrapper 関数を作り公開し、サービスからはその API を叩くという構成のシステムをセットアップしたのでこの記事で解説します。
作るもの
以下のような構成のものを GitHub actions と terraform で構築していきます。
Lambda を採用した理由は以下になります。
- LLM を利用した機能を実装するサーバーになるが、入出力を揃えて API を叩くだけで重たい処理を行わない
- ECS などのサービスを利用して常時 API を立てておくよりも金額的なメリットがある
- 弊社が AWS を中心にインフラを構築しているので使いやすい
Lambda は python の container を利用します。Lambda の実行時に SSM から OpenAI の API Key を取得して API を叩く構成です。 クライアントからはカスタムドメインを貼った API Gateway 経由で Lambda にアクセスする構成としています。Log は Cloudwatch に、Metrics は Datadog に送り Lambda の実行回数などをモニタリングできるようにします。
解説
では、それぞれ解説していきます。
Lambda 関数の設計
まず Lambda 関数の設計ですが、以下のような要件を満たせるようにしたいと考えました。
- LLM を利用したい機能は複数出てくる事が想定されるので、新機能を作るたびにインフラを作らなくて良いようにしたい
- OpenAI 以外の API も気軽に試せるような構成にしたい
これらを踏まえて、以下のような Clean Architecture 風な File 構成で Lambda の内部を書きました。
├── main.py
└── src
├── domain // interface の定義
├── infra // interface の実装
│ └── openai // openai に依存した API の定義 (prompt もこの下で定義)
└── usecase
└── interactor.py
main.py には lambda の entry point となる関数を書きます。 main.py の主な役割は DI (interactor の準備) と request の validation と response の整形です。
domain に usecase で使う処理の interface を定義します。例として与えられた入力から特定の商品情報を抜き出す処理の interface は以下のような定義にしました。
RelatedItemExtractionResponse = TypedDict(
"RelatedItemExtractionResponse", {"items": List[str]}
)
class IFRelatedItemExtraction(metaclass=abc.ABCMeta):
@abc.abstractmethod
def handle(self, request_text: str) -> RelatedItemExtractionResponse:
pass
interactor 内部で business logic を呼び出す処理を書きます。ここでは LLM の呼び出し処理を行います。もちろん、このレイヤーは infra には直接依存しないレイヤーになります。
class Interactor: def __init__( self, repo: IFRelatedItemExtraction, ): self.repo = repo def extract(self, request_text: str) -> RelatedItemExtractionResponse: return self.repo.handle(request_text)
実際に API を呼び出す処理や prompt の定義は infra 以下に書きます。特定の商品情報を抜き出す処理の実装はこのような形です。
MODEL_NAME = "gpt-3.5-turbo-16k" # gpt-4 なども指定できる。Application に合わせて指定する。 SYSTEM_PROMPT = "" # システムプロンプト を書く PROMPT = """ {passage} """ class OpenAIRelatedItemExtraction: @timeout_decorator.timeout(60) # timeout させるための decoration def chat_complete(self, request_text: str): return openai.ChatCompletion.create( model=MODEL_NAME, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": PROMPT.format(passage=request_text)}, ], max_tokens=4098, temperature=0, ) def handle(self, request_text: str) -> RelatedItemExtractionResponse: output = self.chat_complete(request_text) content = output["choices"][0]["message"]["content"] return json.loads(content)
ポイントはモデルの指定やプロンプトをここに書いている事です。プロンプトなどは実装の詳細なので、ifra 以外のレイヤーでは意識しなくて良くなり、プロンプトの修正や呼び出し方の変更 (例えば Function calling を使うように修正するなど) をしても、infra レイヤー以外の修正は行わなくて良くなります。
また、こういった構成にしておくことで、例えば OpenAI をやめて別の LLM に変更するなどをする場合でも、domain で定義された interface に沿って infra を新規で作り main.py で interactor を作る時にどちらを使うかを選べるため、 比較コードの修正量を抑えながら AB テストなどを行う事ができます。
Lamda へのリクエストは以下のような形に統一しています。
{
"type": "function_name",
"input": "input_text"
}
type によって interactor のどの処理を呼び出すかを決定するようにします。
この構成により、新しく Lambda を Deploy しなくても Lambda のソースコードを変更することで新しいを機能を追加できます。
もちろんコンテナのサイズが肥大化していくという懸念はありますが、API の Wrapper であることを想定しているのである程度までは機能を増やしても問題ないと考えました。
CI/CD
コンテナのビルドと Lambda の更新は GitHub Actions で行います。
Actions は以下のような yaml で定義しました。
name: Build on: push: branches: - main env: CONTAINER_NAME: xxx # make の中で定義しているコンテナ名 REPO_NAME: yyy # ECR repo の名前 LAMBDA_FUNCTION_NAME: zzz # terraform で定義する Lambda 関数の名前 jobs: build: runs-on: ubuntu-latest steps: - name: Checkout uses: actions/checkout@v2 - name: Configure AWS credentials uses: aws-actions/configure-aws-credentials@v1 with: aws-access-key-id: ${{ secrets.ML_AWS_ACCESS_KEY_ID }} aws-secret-access-key: ${{ secrets.ML_AWS_SECRET_ACCESS_KEY }} aws-region: ap-northeast-1 - name: Login to Amazon ECR id: login-ecr uses: aws-actions/amazon-ecr-login@v1 - name: Build and push to ecr run: | make build # コンテナを build するためのコマンド。コンテナ名はこの内部で定義している。 docker tag ${CONTAINER_NAME}:latest ${REPO_NAME}:latest docker push <repo-name>:latest docker tag ${CONTAINER_NAME}:latest ${REPO_NAME}:${GITHUB_SHA} docker push ${REPO_NAME}:${GITHUB_SHA} - name: Deploy to lambda run: | aws lambda update-function-code --function-name ${LAMBDA_FUNCTION_NAME} --image-uri ${REPO_NAME}:latest
環境の構築
Lambda の実行環境は terraform で作ります。 それぞれのリソースに合わせて file を分割して作りました。 ここでは point となる lambda function と api gateway の部分だけピックアップして説明します。
まず lambda function の定義です。
resource "aws_lambda_function" "service_llm_lambda" { depends_on = [ aws_cloudwatch_log_group.service_llm_lambda, # cloudwatch の log group を定義しておきます ] function_name = "service-llm-lambda-${var.app_env}" package_type = "Image" image_uri = "${aws_ecr_repository.ecr_repo.repository_url}:latest" # ECR repository の名前です role = aws_iam_role.lambda_role.arn # Lambda の実行時の権限です。後述しますが SSM から機密情報を読み出せる権限を渡しました memory_size = 128 timeout = 60 # OpenAI の API は通常の API よりも実行時間が長い傾向があるので長めにしておきます lifecycle { ignore_changes = [ image_uri, last_modified ] } image_config { command = ["datadog_lambda.handler.handler"] # 後述しますが Datadog にメトリクスを送るためにコマンドを上書きしておきます。 } environment { variables = { DD_LAMBDA_HANDLER = "main.lambda_handler", # コマンドを上書しているので、実際のハンドラーはここに記載します。 DD_SITE = "datadoghq.com", DD_API_KEY = data.aws_ssm_parameter.dd_api_key.value, ## Datadog の API key を parameter store から渡します。この書き方だとコンソールから鍵が見えてしまうので実際はそうならないように工夫が必要です。 DD_TRACE_ENABLED = "true" } } }
api gateway はこんな感じです。HTTP モードでの設定とし、あらかじめ Route53 に作っておいたカスタムドメインと紐付けます。
resource "aws_apigatewayv2_api" "gw_api" { name = "gateway_api_service-llm-lambda-${var.app_env}" protocol_type = "HTTP" cors_configuration { allow_origins = ["*"] # FE から直接呼び出したいというニーズがある場合には CORS 設定をしておく。実際には `*` ではなくちゃんとドメインを指定しましょう allow_methods = ["POST", "GET", "OPTIONS"] allow_headers = ["*"] max_age = 300 } } resource "aws_apigatewayv2_stage" "gw_stage" { api_id = aws_apigatewayv2_api.gw_api.id name = "gateway_api_service-llm-lambda-${var.app_env}" auto_deploy = true access_log_settings { destination_arn = aws_cloudwatch_log_group.gw_lg.arn } resource "aws_apigatewayv2_integration" "gw_integration" { api_id = aws_apigatewayv2_api.gw_api.id integration_uri = aws_lambda_function.service_llm_lambda.invoke_arn integration_type = "AWS_PROXY" integration_method = "POST" } resource "aws_apigatewayv2_route" "gw_route" { api_id = aws_apigatewayv2_api.gw_api.id route_key = "POST /invoke" target = "integrations/${aws_apigatewayv2_integration.gw_integration.id}" } resource "aws_cloudwatch_log_group" "gw_lg" { name = "/aws/api_gw/service-llm-lambda" retention_in_days = 1 } resource "aws_lambda_permission" "api_gw" { statement_id = "AllowExecutionFromAPIGateway" action = "lambda:InvokeFunction" function_name = aws_lambda_function.service_llm_lambda.function_name principal = "apigateway.amazonaws.com" source_arn = "${aws_apigatewayv2_api.gw_api.execution_arn}/*/*" } resource "aws_apigatewayv2_domain_name" "api_gw_domain" { domain_name = "${var.app_env}.xxx" # あらかじめ route53 に作っておいた domain 名 domain_name_configuration { certificate_arn = var.domain_cert_arn # あらかじめ ACM で取得しておきた certificate の ARN を指定 endpoint_type = "REGIONAL" security_policy = "TLS_1_2" } } resource "aws_route53_record" "api_gw_domain_record" { # ここで Route53 の ZONE に API Gateway にリクエストを向けるための A record を追加 name = aws_apigatewayv2_domain_name.api_gw_domain.domain_name type = "A" zone_id = var.zone_id alias { evaluate_target_health = true name = aws_apigatewayv2_domain_name.api_gw_domain.domain_name_configuration[0].target_domain_name zone_id = aws_apigatewayv2_domain_name.api_gw_domain.domain_name_configuration[0].hosted_zone_id } } resource "aws_apigatewayv2_api_mapping" "api_gw_domain_mapping" { api_id = aws_apigatewayv2_api.gw_api.id stage = aws_apigatewayv2_stage.gw_stage.id domain_name = aws_apigatewayv2_domain_name.api_gw_domain.id }
設定のポイント
Timeout について
OpenAI の API を通常の API よりも実行時間が長い傾向があります。(公式の python sdk の設定 は 2023年8月現在、default の timeout はなんと 600sec です!)
そのため Lambda の実行時間を 60sec など長めにとっておくのが良いと思います。
また、コード自体でも timeout を適切に入れておくのが望ましいです。 今回は OpenAI の SDK に依存する形ではなく timeout-decorator という library を使って時前で timeout させるようにしました。
@timeout_decorator.timeout(60) def chat_complete(self, request_text: str): return openai.ChatCompletion.create( model=MODEL_NAME, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": PROMPT.format(passage=request_text)}, ], max_tokens=2048, temperature=0, )
API Key をどこから取り出すか
API Key は機密情報なので Lambda の内部で ssm を経由して取り出すようにしました。 Lambda の Execution role に SSM の該当リソースを取得してデコードする権限を渡しておきます。
data "aws_iam_policy_document" "lambda_execution_policy" { ... 省略 statement { effect = "Allow" actions = [ "kms:Decrypt", "ssm:GetParameter" ] resources = [ "arn:aws:ssm:*:*:parameter/xxx" # parameter 名 ] } }
python 側では Lambda 起動時に以下のような処理が走るような実装にしています。ただしこれは OpenAI のみで必要な処理なので infra レイヤー内部に実装を入れています。
import boto3 import openai ssm = boto3.client("ssm", region_name="ap-northeast-1") response = ssm.get_parameter( Name="xxx", WithDecryption=True ) openai.api_key = response["Parameter"]["Value"]
Datadog との接続
Datadog にメトリクスを送るには、以下の 2 点を行う必要があります。
- コンテナに Datadog に関連するライブラリを入れておく
- コンテナな起動設定でエントリポイントを Datadog のライブラリに上書きし、必要な環境変数を設定する
1 に関しては Dockerfile に public.ecr.aws/datadog/lambda-extension:latest から設定をコピーする設定を書きつつ、 datadog-lambda ライブラリをインストールします。今回は python のパッケージ管理を poetry で行っているので、 pyproject.toml に設定を追記します。
FROM public.ecr.aws/lambda/python:3.10 ... 省略 COPY --from=public.ecr.aws/datadog/lambda-extension:latest /opt/extensions/ /opt/extensions ... 省略 RUN poetry install --no-root
[tool.poetry.dependencies] python = "^3.10" openai = "^0.27.8" boto3 = "^1.28.22" timeout-decorator = "^0.5.0" datadog-lambda = "^4.78.0" # ここ
詳しい設定はこちらを参考にしてください。
設定が完了した上で Lambda を実行すると、コンソールの Menu で Infrastructure -> Serverless 内で図のようにメトリクスが表示されるようになります。
ハマりポイント
Lambda 関数の event について
api gateway で integration した場合には event の中身が変化するという事を理解できておらず、ややハマりました。
コンソールからのテストでは正常に動きましたが、postman などで api を試すと internal serveer error が返ってきてしまいました。
{"message":"Internal Server Error"}
調べていくと、公式に以下のような解説がありました。
ペイロード形式バージョンでは、API Gateway が Lambda 統合に送信するデータの形式と、API Gateway が Lambda からのレスポンスをどのように解釈するかを指定します。ペイロード形式バージョンを指定しない場合、AWS Management Console はデフォルトで最新バージョンを使用します。
Api Gateway からの呼び出しの際には event の構造が変化するので、それに対応するように event からパラメータを取り出したり、レスポンスの形式を揃える必要があります。
インフラ的に解決する道もあると思いますが、簡易的に以下の関数を main.py に定義して対応しました。
def get_request(event): if "type" not in event: event = json.loads(event["body"]) request_type = event["type"] request_text = event["input"] return request_type, request_text def create_response(response: dict): return { "statusCode": 200, "body": json.dumps(response), }
さいごに
OpenAI などの LLM を活用するための Lambda の構成について解説しました。 がっつり使う場合はまた別の構成があると思いますが、LLM をちょっとした機能で使う場合には出来る限りインフラを追加で作らずに機能を追加できるような構成にしておく事にはメリットがあると思います。
