29
17

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

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

Last updated at Posted at 2023-04-02

はじめに

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

さて、このサービスのバックエンドはAWSを利用し、API Gateway + AWS Lambdaのサーバーレス構成にしていました。クライアントからREST APIリクエストを受け取り、bodyの内容を用いてChatGPTのAPI(gpt-3.5-turbo)を呼び出し、そのレスポンスを返すという動作です。サーバーレスにすることでリクエスト時のみの課金となるため、リーズナブルにサービスを運用できます。
before-architecture.png
しかしながらChatGPTのAPIは全ての文章の生成が完了してからレスポンスが返るため、時間帯にもよりますが速い時は5秒、遅い時は10秒程度の時間がかかります。そのため下記のように、ボタンを押してLambda起動後しばらく待たないと、結果が返ってこない状況でした。
before.gif
もう少し速くしたいと思い、方法を調査してみたところChatCompletion API のstream modeが使えそうでした。streamを有効にすることで、Server-Sent Eventを用いてメッセージがストリーミング配信され、ChatGPTのWeb UIのように生成文字列の断片を逐次レスポンスとして受け取れます。

ただしREST APIのサーバーレス構成だと、単にstreamを有効化しただけではレスポンスをストリーム送信できません。そこで今回はWebSocketを用いて、クライアント・サーバ間のコネクションを維持することで、サーバからクライアントに向けてストリーム送信できるようにしました。
改修を加えた結果が下記で、逐次表示することで体感速度が速くなりました。
after.gif
以降では、その実現方法を記載します。なお本記事で紹介するソースコードは下記GitHubリポジトリで公開しています。AWS SAMを用いたリソースの立ち上げもできるため、興味のある方は参考にしてください。

構成

元のものと同じくAPI Gateway+Lambdaの構成です。
after-architecture.png
ただし今回はAPI GatewayのAPIタイプとして、WebSocket APIを指定しています。WebSocket APIにすることで、API GatewayがWebSocket接続を維持し、クライアントから受信したメッセージの内容に従い特定のLambdaを起動するといったことが可能になります。またコストも安いです。なおAPI GatewayでWebSocketを扱う方法は、下記のチャットアプリの実装が参考になります。

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を使うことで実施できます。API GatewayはConnectionIdで個々のクライアントを識別しますが、post_to_connectionでは引数にそのConnectionIdを指定することで、対応するconnectionに対するメッセージ送信ができます。lambda_handlerに渡されるeventにはクライアントのConnectionIdが格納されているため、そのConnectionIdを利用することでメッセージ送信元クライアントに対してメッセージを返信することができます。
以上を踏まえて作成したコードが下記になります。LambdaのRuntimeはPython3.9で、ChatCompletions APIを叩くためにopenaiのライブラリを利用しています。OpenAIのAPI KEYは環境変数で渡しています。

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を用いてテストします。nodeがインストールされている環境であれば簡単にインストールできます。

wscatインストール後、下記のようにコマンドを実行することで、ChatGPTに"こんにちは"と送信するようなテストが行えます。

$ wscat -c wss://hhyd8fkhtl.execute-api.ap-northeast-1.amazonaws.com/Prod
Connected (press CTRL+C to quit)
> {"action":"sendmessage", "data":"こんにちは"}

実際の動作の様子は下記のような感じで、生成した文章がサーバ側からストリーム送信されていることが確認できます。
ChatGPTStream.gif

おわりに

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

29
17
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
29
17

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?