Amazon Bedrock를 이용한 텍스트 생성을 lambda stream mode로 올리기

2024. 5. 23. 14:55개발/AWS

다음은 amazon bedrock의 claude3을 이용해서 코드를 lambda stream모드로 올리고 값을 가져오는 예제입니다.

우선 handler.js에 claude3 sonnet을 이용하는 코드를 작성합니다.
대부분의 코드는 Amazon Bedrock에서 Anthropic Claude 3을 호출하여 텍스트 생성을 참고하여 작성했습니다.

"use strict";

const {BedrockRuntimeClient, InvokeModelWithResponseStreamCommand} = require("@aws-sdk/client-bedrock-runtime");

exports.hello = async (event) => {
    const client = new BedrockRuntimeClient({
        region:'us-east-1',
        credentials:{
            accessKeyId:'YOUR_ACCESS_KEY',
            secretAccessKey:'YOUR_SECRET_KEY'
        }
    });

    const modelId = "anthropic.claude-3-sonnet-20240229-v1:0";
    const inputText = "한국 여행은 어떻게 하는거죠?";
    const maxTokens = 4000;

    const body = {
        anthropic_version: "bedrock-2023-05-31",
        max_tokens: maxTokens,
        messages: [
            {
                role: "user",
                content: [
                    { type: "text", text: inputText }
                ]
            }
        ]
    };

    const command = new InvokeModelWithResponseStreamCommand({
        body: JSON.stringify(body),
        modelId,
        contentType: "application/json",
    });

    const response = await client.send(command);

    let completeMessage = "";

    for await (const item of response.body) {
        const chunk = JSON.parse(new TextDecoder().decode(item.chunk.bytes));

        const chunk_type = chunk.type;

        if (chunk_type === "content_block_delta") {
            const text = chunk.delta.text;
            completeMessage = completeMessage + text;
            process.stdout.write(text);
        }
    }

    return completeMessage;
};

 

코드를 작성후 로컬에서 테스트 하기 위해서는 sls invoke 명령어를 사용하여 테스트를 할 수 있습니다.

serverless invoke local --function hello

 

lambda에 STREAM 모드로 올리기 위해 handler.js의 변경이 필요합니다.
streamifyResponse를 이용하여 결과 값을 responseStream로 출력합니다.

변경된 handler.js는 아래와 같습니다.

"use strict";

const {BedrockRuntimeClient, InvokeModelWithResponseStreamCommand} = require("@aws-sdk/client-bedrock-runtime");

exports.hello = awslambda.streamifyResponse(async (event, responseStream) => {
    const client = new BedrockRuntimeClient({
        region:'us-east-1',
        credentials:{
            accessKeyId:'YOUR_ACCESS_KEY',
            secretAccessKey:'YOUR_SECRET_KEY'
        }
    });

    const modelId = "anthropic.claude-3-sonnet-20240229-v1:0";
    const inputText = "한국 여행은 어떻게 하는거죠?";
    const maxTokens = 4000;

    const body = {
        anthropic_version: "bedrock-2023-05-31",
        max_tokens: maxTokens,
        messages: [
            {
                role: "user",
                content: [
                    { type: "text", text: inputText }
                ]
            }
        ]
    };

    const command = new InvokeModelWithResponseStreamCommand({
        body: JSON.stringify(body),
        modelId,
        contentType: "application/json",
    });

    const response = await client.send(command);

    let completeMessage = "";

    for await (const item of response.body) {
        // Decode each chunk
        const chunk = JSON.parse(new TextDecoder().decode(item.chunk.bytes));

        const chunk_type = chunk.type;

        if (chunk_type === "content_block_delta") {
            const text = chunk.delta.text;
            responseStream.write(text);
        }
    }

    responseStream.end();
});

 

스트림 모드를 올리기 위해서는 serverless.yml에 invokeMode를 추가합니다.

functions:
  hello:
    handler: handler.hello
    events:
      - httpApi:
          path: /
          method: get
    url:
      invokeMode: RESPONSE_STREAM

 

그 후 배포를 하면 endpoints에 새로운 hello가 생긴 것을 확인할 수 있습니다.

 

curl 명령어를 통해 잘 작동하는지 확인할 수 있습니다.

 

참고로 lambda의 timeout 시간에 맞춰서 출력이 끊길 수 있습니다.
timeout 시간을 넉넉히 잡는 것을 추천합니다.

html 페이지를 하나 만들어서 잘 되는지 테스트를 해보겠습니다.

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Real-time Data from Lambda</title>
</head>
<body>
<div id="data"></div>

<script>
    const lambdaUrl = "AWS_ENDPOINT";

    const request = new Request(lambdaUrl, { method: 'GET' });

    const decoder = new TextDecoder('utf-8');

    fetch(request)
        .then(response => {
            const reader = response.body.getReader();
            return new ReadableStream({
                start(controller) {
                    function push() {
                        reader.read().then(({ done, value }) => {
                            if (done) {
                                controller.close();
                                return;
                            }
                            const decodedString = decoder.decode(value);
                            document.getElementById("data").innerText += decodedString;
                            push();
                        });
                    }
                    push();
                }
            });
        })
        .catch(error => {
            console.error('Error:', error);
        });
</script>
</body>
</html>

 

위 코드를 실행하면 stream으로 출력이 되는걸 확인할 수 있습니다.