React アプリに Amazon Bedrock への問い合わせ画面を組み込む [レスポンスストリーミング対応]

こんにちは、広野です。

React アプリに Amazon Bedrock に問い合わせる画面をつくってみたのでコードを紹介します。よくある生成系 AI チャットボットのように、回答文が作成されながら表示されるやつです。↓

アーキテクチャ

  • 生成系 AI の基盤モデルには、Amazon Bedrock の Anthropic Claude 2.1 を使用しています。
  • Amazon Bedrock に問い合わせるのは AWS Lambda 関数 (Node.js 20) です。
  • AWS Lambda 関数は関数 URL を持たせており、React アプリから直接呼び出します。
  • AWS Lambda 関数からのレスポンスはレスポンスストリーミング対応にしています。それにより、Amazon Bedrock が回答を作成しつつ、レスポンスを画面に表示できるようにしています。Amazon Bedrock を呼び出す API はレスポンスストリーミング対応のInvokeModelWithResponseStreamCommand を使用しています。

補足ですが、React アプリから AWS Lambda 関数 URL を認証付きで呼び出す構成は以下参考記事の構成にしています。本記事では説明を割愛します。

AWS Lambda 関数コード

AWS Lambda 関数に関数 URL を持たせる設定や、レスポンスストリーミングを有効にする設定は以下 AWS 公式ドキュメントを参照ください。この点において難しいことはありません。

難しいのは関数コードです。

invokeModelWithStreamResponseCommand の API が Amazon Bedrock からのレスポンスを chunk 分割された状態で受け取るので、chunk ごとに React アプリへのレスポンス用ストリーム responseStream に放り込んで (write) います。レスポンスは Blob なので、途中で人が読み取れるテキストデータに変換する処理が入っています。

const { BedrockRuntimeClient, InvokeModelWithResponseStreamCommand } = require("@aws-sdk/client-bedrock-runtime");
const bedrock = new BedrockRuntimeClient({region: "ap-northeast-1"});
exports.handler = awslambda.streamifyResponse(async (event, responseStream, _context) => {
  try {
    console.log((JSON.parse(event.body)).prompt);
    const prompt = (JSON.parse(event.body)).prompt;
    if (prompt == '') {
      responseStream.write("No prompt");
      responseStream.end();
    }
    const enclosed_prompt = "\n\nHuman: " + prompt + "\n\nAssistant:";
    const body = {
      "prompt": enclosed_prompt,
      "max_tokens_to_sample": 3000,
      "temperature": 0.5,
      "top_k": 250,
      "top_p": 1,
      "stop_sequences": ["\n\nHuman:"],
      "anthropic_version": "bedrock-2023-05-31"
    };
    const input = {
      modelId: 'anthropic.claude-v2:1',
      accept: 'application/json',
      contentType: 'application/json',
      body: JSON.stringify(body)
    };
    const command = new InvokeModelWithResponseStreamCommand(input);
    const res = await bedrock.send(command);
    const actualStream = res.body.options.messageStream;
    const chunks = [];
    for await (const value of actualStream) {
      const jsonString = new TextDecoder().decode(value.body);
      const base64encoded = JSON.parse(jsonString).bytes;
      const decodedString = Buffer.from(base64encoded,'base64').toString();
      try {
        const streamingCompletion = JSON.parse(decodedString).completion;
        chunks.push(streamingCompletion);
        responseStream.write(streamingCompletion);
      } catch (error) {
        console.error(error);
        responseStream.write(null);
        responseStream.end();
      }
    }
    //log entire chunks 結合されたレスポンス。今後履歴データとしてDynamodBに放り込む予定
    console.log("stream ended: ", chunks.join(""));
    responseStream.end();
  } catch (error) {
    console.error(error);
    responseStream.write("error");
    responseStream.end();
  }
});

Python を採用していない理由は、AWS 公式ドキュメントにもありますが AWS Lambda 関数がレスポンスストリーミングをネイティブにサポートしているランタイムが Node.js のみだからです。※執筆時点

Node.js のコードを ES2015 (ES6) 構文にしていない理由は、私が AWS Lambda 関数を AWS CloudFormation でデプロイしていることに起因しています。AWS CloudFormation で Node.js の Lambda 関数をデプロイすると Lambda 関数コードファイルの拡張子が .js になってしまうため、ES2015 を使用できない (CommonJS 一択になる) のです。まあ他の方法でデプロイしろよって話なのですが、CFn テンプレート一撃で全アプリ環境を作りたいのでそこにこだわっております・・・。

React コード

React アプリ側は、Lambda 関数 URL を呼び出すために fetch を使用します。レスポンスストリーミング対応にするには fetch が良いのだとか。

AWS Lambda 関数 URL から chunk 分割されたレスポンスが送られてくるので、それを順次、直前の chunk と繋げて state に突っ込むだけの、非常に原始的なループのコードです。もっと簡単かつ効率的な処理になる SDK を AWS が開発してくれることを期待しております。SSR 対応フレームワーク限定になってしまうのかもですが。

//レスポンス表示用state
const [response, setResponse] = useState("");
//送信したいプロンプト
const body = {
  prompt: prompt
};
//Lambda関数URL呼出
try {
  const res = await window.fetch(
    "https://xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx.lambda-url.ap-northeast-1.on.aws",
    {
      method: 'POST',
      body: JSON.stringify(body),
      headers: signedRequest.headers //認証用の署名(説明割愛)
    }
  );
  const reader = res.body.getReader();
  const decoder = new TextDecoder();
  const sleep = ms => new Promise(resolve => setTimeout(resolve, ms));
  while (true) {
    const { value, done } = await reader.read();
    if (done) break;
    const chunk = decoder.decode(value, { stream: true });
    if (chunk) {
      setResponse(prevData => prevData + chunk);
    }
    await sleep(100); //100msごとに届いたchunkを処理(=画面更新)
  }
} catch (error) {
  console.error('Error fetching:', error);
}

sleep のミリ秒数を調整することで、受け取ったレスポンス文章表示の流暢さを調整できます。

まとめ

いかがでしたでしょうか。

レスポンスストリーミングの仕組みが未だ根本的に理解できていませんが、経験を積んで理解を深めたいと思います。今回は参考になるドキュメントが少なくて苦労しました。

本記事が皆様のお役に立てれば幸いです。

著者について
広野 祐司

AWS サーバーレスアーキテクチャを駆使して社内クラウド人材育成アプリとコンテンツづくりに勤しんでいます。React で SPA を書き始めたら快適すぎて、他の言語には戻れなくなりました。サーバーレス & React 仲間を増やしたいです。AWSは好きですが、それよりもAWSすげー!って気持ちの方が強いです。
取得資格:AWS 認定は12資格、ITサービスマネージャ、ITIL v3 Expert 等
2020 - 2023 Japan AWS Top Engineer 受賞
2022 - 2023 Japan AWS Ambassador 受賞
2023 当社初代フルスタックエンジニア認定
好きなAWSサービス:AWS Amplify / AWS AppSync / Amazon Cognito / AWS Step Functions / AWS CloudFormation

広野 祐司をフォローする
クラウドに強いによるエンジニアブログです。
SCSKは専門性と豊富な実績を活かしたクラウドサービス USiZE(ユーサイズ)を提供しています。
USiZEサービスサイトでは、お客様のDX推進をワンストップで支援するサービスの詳細や導入事例を紹介しています。
AI・MLAWSアプリケーション開発クラウドソリューション
シェアする
TechHarmony
タイトルとURLをコピーしました