Amazon Bedrockのモデル実行結果をStreamで受け取る方法

LLMを使用したアプリの開発では、ユーザーの待ち時間内にインタラクションを作成するために、Streamを使ってレスポンスを処理するのが一般的です。Amazon Bedrockでホストされているモデルを呼び出す際に、Streamでレスポンスを受け取る方法について調査しました。Bedrock Runtime ClientのInvokeModelWithResponseStreamを使用することで、Streamでの受け取りが可能です。ただし、必須項目としてmax_tokens_to_sampleを指定する必要があります。また、Claudeを使用する場合には、「Assistant:」で終わらせる必要があります。

LLMを使ったアプリを作る際、ユーザーがレスポンスを待てる範囲でインタラクションを作るため、Streamでレスポンスを処理することが一般的です。Amazon Bedrockでホストしているモデルを、AWS SDK経由で呼び出した場合に、Streamでレスポンスを返す方法について調べてみました。

Bedrock Runtime ClientのInvokeModelWithResponseStreamを利用する

Streamで返答を受けるには、InvokeModelWithResponseStreamCommandを使います。プロンプトなどを渡すBody部分がオブジェクト形式で、max_tokens_to_sampleが必須値なことに注意しましょう。

また、Claudeを利用する場合「Assistant:」で終わらせないとエラーになることがあります。

    const client = new BedrockRuntimeClient({
        region: env.aws.region,
        credentials: {
            accessKeyId: env.aws.credentials.accessKeyId,
            secretAccessKey: env.aws.credentials.secretAccessKey
        }
    });
    const input = { // InvokeModelWithResponseStreamRequest
        body: JSON.stringify({
            prompt: 'Human: Hello\nAssistant:',
            max_tokens_to_sample: 500
        }), // required
        contentType: "application/json",
        //accept: "STRING_VALUE",
        modelId: "anthropic.claude-v2", // required
    };
    const command = new InvokeModelWithResponseStreamCommand(input);
    const response = await client.send(command);
    if (!response.body) {
        return c.text('N/A')
    }

ちなみにこの時のresponseをそのままJSONで返そうとすると、このようなデータが返ってきます。「返事が返ってこない！」となった場合、このようなStreamを含むデータが返ってきていないかを確認することで、気づかないうちにStreamの実装が必要なAPIを叩いていないかをみることができそうです。

{
  "$metadata": {
    "httpStatusCode": 200,
    "requestId": "0e8cc468-2c65-4394-965d-41756c88a5a3",
    "attempts": 1,
    "totalRetryDelay": 0
  },
  "contentType": "application/json",
  "body": {
    "options": {
      "messageStream": {
        "options": {
          "inputStream": {},
          "decoder": {
            "headerMarshaller": {},
            "messageBuffer": [],
            "isEndOfStream": false
          }
        }
      }
    }
  }
}

Honoを使ってStreamでレスポンスを返す

ExpressやNest.jsでもよいのですが、手が馴れているHonoで進めます。Streamで処理する際のポイントをまとめてくださっていた記事がありましたので、参考にしながら処理を追加しました。

    return c.streamText(async stream => {
        if (!response.body) {
            return
        }
        for await (const data of response.body) {
            if (!data.chunk) {
                if (
                    data.internalServerException ||
                    data.modelStreamErrorException || 
                    data.modelTimeoutException || 
                    data.throttlingException ||
                    data.validationException
                ) {
                    console.log(`Error: ${JSON.stringify(data)}`)
                    await stream.writeln('Internal error')
                }
                break
            }
            if (!data.chunk.bytes) {
                break
            }
            const item = JSON.parse(Buffer.from(data.chunk.bytes).toString('utf-8'))
            await stream.write(item.completion)
        }
    })

エラー系のデータが含まれているのがありがたいですね。とはいえどんなエラーが返ってくるかまでは見ていないので、プロダクションで使うには検証が必要かもしれません。

実行してみる

用意ができたので動かしてみます。

$ curl http://localhost:3000/bedrock/llm -sN
 Hello! I'm Claude, an AI assistant created by Anthropic

curl http://localhost:3000/bedrock/llm -sN
$ Hello!

ベタ書きしているテキストが短いので、返事もシンプルですね。

質問を変えてみた

もう少し長めの、時間がかかりそうな質問を入れてみましょう。

    const input = { // InvokeModelWithResponseStreamRequest
        body: JSON.stringify({
            prompt: 'Human: ClaudeとLlamaやOpenAIとの違いを教えてください。\nAssistant:',
            max_tokens_to_sample: 500
        }), // required
        contentType: "application/json",
        //accept: "STRING_VALUE",
        modelId: "anthropic.claude-v2", // required
    };

長めのテキストが生成されました。curl -sNで呼びだしていれば、文字が少しずつ出てくる様子が見れると思います。

 はい、ClaudeとLlama、OpenAIという3つのAIシステムには以下のような違いがあります。

- ClaudeはAnthropic社が開発した対話AIシステムで、安全性と有用性に重点を置いています。人間に有益な会話をすることを目指しています。

- LlamaはMeta社が開発した大規模言語モデルで、広範な知識と会話能力を備えています。効率的な学習とスケーラビリティに重点を置いています。

- OpenAIのGPTモデルはOpenAI社が開発した大規模言語生成モデルで、 GPT-3などが有名です。テキスト生成能力に優れています。 

- Claudeは相手の意図を理解し安全性を重視するのに対し、Llamaは知識量と会話の自然さを追求し、GPTはテキスト生成能力そのものの向上に重点があります。

- ClaudeとLlamaは会話AIとして設計されていますが、GPTはそうした使い方を直接の目的とはしていません。

- データセットやアーキテクチャー設計の違いなど技術的な違いもありますが、目的とアプローチの方向性に大きな違いがあると言えます。

やってみての感想

「そもそもAWS SDKのAPIでStreamかそうじゃないかの判別ってどうやるんだ？」という疑問から始まったのですが、専用のAPIがあるのは流石です。この辺りの実装を抽象化したい場合には、LangChainを使うこともできますし、容量やv1が出ていないライブラリを使うことへの懸念がある方は、AWS SDKでの作り方を深堀していくのがよいかなと思います。