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 を認証付きで呼び出す構成は以下参考記事の構成にしています。本記事では説明を割愛します。

やってみたら面倒くさい、React アプリからの Amazon Cognito 認証付き AWS Lambda 関数 URL 呼出

Amazon Cognito 認証を施した React アプリから、認証済みユーザのみ AWS Lambda 関数を呼び出せるようにする React コードを紹介します。

blog.usize-tech.com

2024.01.15

AWS Lambda 関数コード

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

Lambda 関数 URL の作成と管理 - AWS Lambda

Lambda 関数の URL を設定して、他の AWS のサービスとの統合を必要とせずに、HTTP エンドポイントを Lambda 関数に割り当てます。

docs.aws.amazon.com

ストリームレスポンスに Lambda 関数を設定する - AWS Lambda

レスポンスをストリーミングする Lambda 関数を作成します。

docs.aws.amazon.com

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

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 のミリ秒数を調整することで、受け取ったレスポンス文章表示の流暢さを調整できます。

まとめ

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

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

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