AWS Lambdaのレスポンスストリーミングを使ってChatGPTっぽいUIにする

sam init --name lambda-response-stream-sample \
  --app-template hello-world-typescript --runtime nodejs18.x 

npm install openai @aws-sdk/client-ssm

import * as AWS from '@aws-sdk/client-ssm';
import { Configuration, OpenAIApi } from 'openai';
import { IncomingMessage } from 'http';

import { pipeline as pipelineSync, Transform, TransformOptions } from 'stream';
import { promisify } from 'util';
const pipeline = promisify(pipelineSync);

// OpenAIのストリームレスポンスをChatメッセージに変換するカスタムTransformer
class ChunkTransformer extends Transform {
  constructor(options?: TransformOptions) {
    super(options);
  }
  _transform(chunk: Buffer, encoding: BufferEncoding, callback: (error?: Error | null, data?: any) => void): void {
    console.log(chunk.toString());
    const payloads = chunk.toString().split('\n\n');
    for (const payload of payloads) {
      if (payload.includes('[DONE]')) {
        break;
      }
      if (payload.startsWith('data:')) {
        const data = payload.replaceAll(/(\n)?^data:\s*/g, '');
        try {
          const delta = JSON.parse(data.trim());
          this.push(delta.choices[0].delta?.content || '');
        } catch (error) {
          console.log(error);
        }
      }
    }
    callback();
  }
}

// SSMパラメータストアからAPIキーを取得
const ssmClient = new AWS.SSM({});
const ssmParam = ssmClient.getParameter({ Name: '/openai/api-key', WithDecryption: true });

// Lambdaイベントハンドラー本体
export const lambdaHandler = awslambda.streamifyResponse(async (event, responseStream, context) => {
  const json = JSON.parse(event.body || '{}');
  if (!json.message) {
    responseStream.write('no message');
    responseStream.end();
    return;
  }

  try {
    // OpenAI APIの実行
    const apiKey = (await ssmParam).Parameter?.Value || '';
    const configuration = new Configuration({
      apiKey
    });
    const openai = new OpenAIApi(configuration);
    const response = await openai.createChatCompletion({
      model: 'gpt-3.5-turbo',
      messages: [{
        role: 'user',
        content: json.message
      }],
      max_tokens: 1024,
      temperature: 0.7,
      stream: true // enable stream
    }, {
      responseType: 'stream' // axios option(required)
    });

    // Node.jsのPipelineに連携
    const stream = response.data as unknown as IncomingMessage;
    await pipeline(stream, new ChunkTransformer(), responseStream);
  } catch (err) {
    console.log(err);
    responseStream.write('エラーが発生しました！');
    responseStream.end();
  }
});

import { APIGatewayProxyEvent, Context, Handler } from 'aws-lambda';

declare global {
  namespace awslambda {
    function streamifyResponse(f: (event: APIGatewayProxyEvent,
                                   responseStream: NodeJS.WritableStream,
                                   context: Context) => Promise<void>): Handler;
  }
}

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: lambda-response-stream-sample

Resources:
  # Lambdaの実行ロール
  ChatFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: SSMOpenAIApiKeyRead
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "ssm:GetParameter"
                Resource:
                  - !Sub "arn:aws:ssm:*:${AWS::AccountId}:parameter/openai/api-key"
      Path: "/"

  # Lambda関数本体
  ChatFunction:
    Type: AWS::Serverless::Function
    Properties:
      Role: !GetAtt ChatFunctionRole.Arn
      CodeUri: chat/dist
      Handler: app.lambdaHandler
      Runtime: nodejs18.x
      Timeout: 60
      Architectures:
        - x86_64
      FunctionUrlConfig:
        AuthType: NONE
        InvokeMode: RESPONSE_STREAM # LambdaレスポンスストリームON
        Cors:
          AllowCredentials: true
          AllowMethods: ["POST"]
          AllowOrigins: ["*"]
          AllowHeaders: ["*"]

Outputs:
  ChatApi:
    Description: "Chat Lambda Function URL"
    Value: !GetAtt ChatFunctionUrl.FunctionUrl
  ChatFunction:
    Description: "Chat Lambda Function ARN"
    Value: !GetAtt ChatFunction.Arn

require('esbuild').build({
  entryPoints: ['app.ts'],
  bundle: true,
  minify: true,
  sourcemap: true,
  outfile: 'dist/app.js',
  target: 'node18',
  platform: 'node'
}).catch((e) => console.log(e) && process.exit(1));

# Lambda関数ソースコードディレクトリに移動
pushd chat
# esbuildでバンドル
node build.js
# プロジェクトルートに戻る
popd
# SAM CLIでビルド
sam build

sam deploy

curl https://xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx.lambda-url.ap-northeast-1.on.aws/ \
  -H 'Content-Type:application/json' -d '{"message": "こんにちわ！あなたは誰ですか？"}'
> こんにちは！私はAIアシスタントのGPT-3です。あなたのサポートになります。

npx nuxi init nuxt-app
cd nuxt-app
npm install

<script setup lang="ts">
const message = ref('');
const response = ref('');
const ask = async () => {
  response.value = '';
  const config = useRuntimeConfig()
  const stream = await $fetch(config.public.API_ENDPOINT, {
    method: 'POST',
    body: { message: message.value },
    responseType: 'stream'
  });
  const reader = (stream as ReadableStream).getReader();
  let done = false;
  const decoder = new TextDecoder();

  while (!done) {
    const { value, done: isDone } = await reader.read();
    if (isDone) {
      done = true;
      break;
    }
    const newData = new Uint8Array(value);
    response.value += decoder.decode(newData);
  }
}
</script>

<template>
  <div>
    <textarea v-model="message" cols="50" rows="5" style="display: block;margin-bottom: 5px"></textarea>
    <button @click="ask">送信!!</button>
    <div style="margin-top: 5px">{{ response }}</div>
  </div>
</template>

export default defineNuxtConfig({
  runtimeConfig: {
    public: {
      API_ENDPOINT: 'https://xxxxxxxxxxxxxxxxxxxx.lambda-url.ap-northeast-1.on.aws'
    }
  }
})

npm run dev