@lobehub/chat/src/libs/model-runtime/utils/streams/bedrock/common.ts

Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.

import {
  InvokeModelWithResponseStreamResponse,
  ResponseStream,
} from '@aws-sdk/client-bedrock-runtime';

import { readableFromAsyncIterable } from '../protocol';

const chatStreamable = async function* (stream: AsyncIterable<ResponseStream>) {
  for await (const response of stream) {
    if (response.chunk) {
      const decoder = new TextDecoder();

      const value = decoder.decode(response.chunk.bytes, { stream: true });
      try {
        const chunk = JSON.parse(value);

        yield chunk;
      } catch (e) {
        console.log('bedrock stream parser error:', e);

        yield value;
      }
    } else {
      yield response;
    }
  }
};

/**
 * covert the bedrock response to a readable stream
 */
export const createBedrockStream = (res: InvokeModelWithResponseStreamResponse) =>
  readableFromAsyncIterable(chatStreamable(res.body!));