WebSocketを用いてChatGPT APIのレスポンスをストリーム送信する

# はじめに
2023年3月にChatGPTのAPIが公開されて以降、様々なWebサービスが誕生しています。
私も何か作ってみたくなり、[dAIary](https://daiary-w.web.app/)という日記作成サービスを作ってみました。
午前と午後のできごとを簡単に入力するだけで、AI(ChatGPT)が良い感じに日記を書いてくれるサービスです。結構面白いので試してみてください！

https://daiary.patchwork-technology.com/

さて、このサービスのバックエンドはAWSを利用し、API Gateway + AWS Lambdaのサーバーレス構成にしていました。クライアントからREST APIリクエストを受け取り、bodyの内容を用いてChatGPTのAPI(gpt-3.5-turbo)を呼び出し、そのレスポンスを返すという動作です。サーバーレスにすることでリクエスト時のみの課金となるため、リーズナブルにサービスを運用できます。
<img alt="before-architecture.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/61029/46e17682-54f7-c3f6-7441-aafbac785f5b.png">
しかしながらChatGPTのAPIは全ての文章の生成が完了してからレスポンスが返るため、時間帯にもよりますが速い時は5秒、遅い時は10秒程度の時間がかかります。そのため下記のように、ボタンを押してLambda起動後しばらく待たないと、結果が返ってこない状況でした。
<img width="40%" alt="before.gif" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/61029/ae68cc44-f75e-e99b-ff90-12a4f2627f31.gif">
もう少し速くしたいと思い、方法を調査してみたところ[ChatCompletion API のstream mode](https://platform.openai.com/docs/api-reference/chat/create)が使えそうでした。streamを有効にすることで、[Server-Sent Event](https://developer.mozilla.org/ja/docs/Web/API/Server-sent_events)を用いてメッセージがストリーミング配信され、ChatGPTのWeb UIのように生成文字列の断片を逐次レスポンスとして受け取れます。

ただしREST APIのサーバーレス構成だと、単にstreamを有効化しただけではレスポンスをストリーム送信できません。そこで今回は[WebSocket](https://triple-underscore.github.io/RFC6455-ja.html)を用いて、クライアント・サーバ間のコネクションを維持することで、サーバからクライアントに向けてストリーム送信できるようにしました。
改修を加えた結果が下記で、逐次表示することで体感速度が速くなりました。
<img width="40%" alt="after.gif" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/61029/42c2f717-56ae-dc80-3142-aeab35930222.gif">
以降では、その実現方法を記載します。なお本記事で紹介するソースコードは下記GitHubリポジトリで公開しています。AWS SAMを用いたリソースの立ち上げもできるため、興味のある方は参考にしてください。

https://github.com/hama1080/ServerlessChatGPTStream

# 構成
元のものと同じくAPI Gateway+Lambdaの構成です。
<img alt="after-architecture.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/61029/ead410f4-5cea-e868-a3eb-0223eb3d5810.png">
ただし今回はAPI GatewayのAPIタイプとして、[WebSocket API](https://docs.aws.amazon.com/ja_jp/apigateway/latest/developerguide/apigateway-websocket-api-overview.html)を指定しています。WebSocket APIにすることで、API GatewayがWebSocket接続を維持し、クライアントから受信したメッセージの内容に従い特定のLambdaを起動するといったことが可能になります。また[コスト](https://aws.amazon.com/jp/api-gateway/pricing/)も安いです。なおAPI GatewayでWebSocketを扱う方法は、下記のチャットアプリの実装が参考になります。

https://github.com/aws-samples/simple-websockets-chat-app

API GatewayのWebSocket APIでは、ルート選択式の指定が必須になります。ルート選択式はAPI Gatewayがクライアントからのメッセージを受信したときに、実施する動作を決定するのに使われます。今回はルート選択式に`request.body.action`を指定し、`request.body.action = sendmessage`の時にLambdaを呼び出す設定にしました。なお定義済みのルートとして、接続時に利用される`$connect`、切断時に利用される`$disconnect`、一致するルートが無い時に利用される`$default`が用意されていますが、今回は利用しないためそれらのルートの動作は未定義です。
このような構築を行い、クライアントがAPI GatewayとのWebSocket接続確立後に下記のようなメッセージを送信すると、sendmessageと紐づけたLambdaに`"data":"こんにちは"`が渡ります。
```
{"action":"sendmessage", "data":"こんにちは"}
```

# ChatGPT APIの呼び出し結果をストリーム送信するLambdaの実装
今回作成するLambdaでは、ChatGPT APIの`stream = True`でのコールと、そのレスポンスをクライアントにストリーム送信する、といったことを行います。
クライアントへのストリーム送信は、[ApiGatewayManagementApiのpost_to_connection](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/apigatewaymanagementapi/client/post_to_connection.html)を使うことで実施できます。API GatewayはConnectionIdで個々のクライアントを識別しますが、post_to_connectionでは引数にそのConnectionIdを指定することで、対応するconnectionに対するメッセージ送信ができます。lambda_handlerに渡されるeventにはクライアントのConnectionIdが格納されているため、そのConnectionIdを利用することでメッセージ送信元クライアントに対してメッセージを返信することができます。
以上を踏まえて作成したコードが下記になります。LambdaのRuntimeはPython3.9で、ChatCompletions APIを叩くためにopenaiのライブラリを利用しています。OpenAIのAPI KEYは環境変数で渡しています。
```python
import json
import boto3
import openai
import os


def lambda_handler(event, context):
    # OpenAI APIキーを環境変数から取得
    openai.api_key = os.environ["OPENAI_API_KEY"]

    # 入力メッセージおよびconnection識別のための情報を取得
    data = json.loads(event.get('body', '{}')).get('data')
    domain_name = event.get('requestContext', {}).get('domainName')
    stage = event.get('requestContext', {}).get('stage')
    connectionId = event.get('requestContext', {}).get('connectionId')
    apigw_management = boto3.client(
        'apigatewaymanagementapi', endpoint_url=F"https://{domain_name}/{stage}")

    # Request ChatGPT API
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "user", "content": data},
        ],
        stream=True
    )

    # メッセージ送信元クライアントに逐次メッセージ送信
    for partial_message in response:
        content = partial_message['choices'][0]['delta'].get('content')
        if content:
            apigw_management.post_to_connection(
                ConnectionId=connectionId, Data=content)

    return {
        "statusCode": 200,
        "body": json.dumps({
            "message": "data sent.",
        }),
    }
```

# 接続テスト
以上の内容で環境を構築することで、ChatGPT APIのレスポンスのストリーム送信が可能になります。
実際に構築を行うと、`wss://hhyd8fkhtl.execute-api.ap-northeast-1.amazonaws.com/Prod`のようなAPI Gatewayのエンドポイントが得られます。REST APIのテストにはcurlコマンド等が使えますが、今回はWebSocket接続であるため[wscat](https://github.com/websockets/wscat)を用いてテストします。nodeがインストールされている環境であれば簡単にインストールできます。

wscatインストール後、下記のようにコマンドを実行することで、ChatGPTに"こんにちは"と送信するようなテストが行えます。
```bash
$ wscat -c wss://hhyd8fkhtl.execute-api.ap-northeast-1.amazonaws.com/Prod
Connected (press CTRL+C to quit)
> {"action":"sendmessage", "data":"こんにちは"}
```
実際の動作の様子は下記のような感じで、生成した文章がサーバ側からストリーム送信されていることが確認できます。
<img width="80%" alt="ChatGPTStream.gif" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/61029/7a30cfea-08e4-3380-d599-de7ae4a1bfe1.gif">

# おわりに
AWS API GatewayのWebSocket API + Lambdaを用いることでサーバーレス環境でChatGPT APIのストリーム送信を可能にしました。API GatewayにWebSocketの管理を任せることで、比較的容易に環境の構築ができたと思います。ChatGPTのWeb UIのような見た目を、サーバーレスで実現する際には参考にしてみてください。