Merge pull request #2858 from asmovian/fix-lambda-bedrock-streaming

Fix lambda bedrock streaming

Author: julianwood

lambda-streaming-ttfb-write-sam-with-bedrock-streaming/README.md

@@ -1,6 +1,6 @@
-# Lambda Response streaming: time-to-first-byte using write() method
+# AWS Lambda Response streaming: time-to-first-byte using write() method
 
-This pattern shows how to use Lambda response streaming ability and bedrock LLM streaming inference api to improve time-to-first byte using the write() method and `InvokeModelWithResponseStreamCommand`. For more information on the feature, see the [launch blog post](https://aws.amazon.com/blogs/compute/introducing-aws-lambda-response-streaming/).
+This pattern shows how to use AWS Lambda response streaming ability and Amazon Bedrock LLM streaming inference api to improve time-to-first byte using the write() method and `InvokeModelWithResponseStreamCommand`. For more information on the feature, see the [launch blog post](https://aws.amazon.com/blogs/compute/introducing-aws-lambda-response-streaming/).
 
 Learn more about this pattern at Serverless Land Patterns: [lambda-streaming-ttfb-write-sam](https://github.com/aws-samples/serverless-patterns/tree/main/lambda-streaming-ttfb-write-sam)
 
@@ -69,4 +69,4 @@ You can see the gradual display of the streamed response.
 ----
 Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-SPDX-License-Identifier: MIT-0
+SPDX-License-Identifier: MIT-0

lambda-streaming-ttfb-write-sam-with-bedrock-streaming/src/index.js

@@ -6,6 +6,19 @@ const {
 exports.handler = awslambda.streamifyResponse(
     async (event, responseStream, context) => {
         const lambdaRequestBody = JSON.parse(event.body);
+        try {
+            // Handle base64 encoded body from AWS SigV4 requests
+            const body = event.isBase64Encoded ? 
+                Buffer.from(event.body, 'base64').toString('utf-8') : 
+                event.body;
+            lambdaRequestBody = JSON.parse(body);
+        } catch (error) {
+            console.error('Error parsing request body:', error);
+            responseStream.write("Invalid JSON in request body.");
+            responseStream.end();
+            return;
+        }
+
         const prompt = lambdaRequestBody.prompt || '';
         if (prompt == '') {
             responseStream.write("No prompt provided.");
@@ -20,18 +33,15 @@ exports.handler = awslambda.streamifyResponse(
             }
         };
 
-        const
+         const
             maxTokens = 2500,
             temperature = .7,
             topP = 1,
-            stopSequences = ["\n\nHuman:"],
-            anthropicVersion = "bedrock-2023-05-31",
-            modelId = 'anthropic.claude-v2',
+            anthropicVersion = "bedrock-2023-05-31"
+            modelId = 'anthropic.claude-3-haiku-20240307-v1:0',
             contentType = 'application/json',
             accept = '*/*';
 
-        const formattedPrompt = `Human: ${prompt}\n\nAssistant:`
-
         try {
             responseStream = awslambda.HttpResponseStream.from(responseStream, httpResponseMetadata);
 
@@ -45,12 +55,16 @@ exports.handler = awslambda.streamifyResponse(
             });
 
             const llmRequestBody = {
-                prompt: formattedPrompt,
-                max_tokens_to_sample: maxTokens,
+                max_tokens: maxTokens,
+                messages: [
+                    {
+                        role: "user",
+                        content: prompt
+                    }
+                ],
                 temperature,
-                top_p: topP,
-                stop_sequences: stopSequences,
-                anthropic_version: anthropicVersion
+                top_p: topP
+                anthropic_version: anthropicVersion,
             };
 
             const params = {
@@ -68,19 +82,21 @@ exports.handler = awslambda.streamifyResponse(
             responseStream = awslambda.HttpResponseStream.from(responseStream, httpResponseMetadata);
             const chunks = []
             for await (const value of actualStream) {
-                const jsonString = new TextDecoder().decode(value.body); // body is a Uint8Array. jsonString->'{"bytes":"eyJjb21wbGV0aW9uIjoiIEkiLCJzdG9wX3JlYXNvbiI6bnVsbH0="}'
-                const base64encoded = JSON.parse(jsonString).bytes; // base64 encoded string. 
-                const decodedString = Buffer.from(base64encoded, "base64").toString(
-                    "utf-8"
-                  );
+                const jsonString = new TextDecoder().decode(value.body);
+                const base64encoded = JSON.parse(jsonString).bytes;
+                const decodedString = Buffer.from(base64encoded, "base64").toString("utf-8");
+                
                 try {
-                    const streamingCompletion = JSON.parse(decodedString).completion;
-                    chunks.push(streamingCompletion)
-                    responseStream.write(streamingCompletion)
+                    const chunk = JSON.parse(decodedString);
+                    if (chunk.type === 'content_block_delta' && chunk.delta?.text) {
+                        const text = chunk.delta.text;
+                        chunks.push(text);
+                        responseStream.write(text);
+                    }
                 } catch (error) {
                     console.error(error);
                     responseStream.write(null);
-                    responseStream.end()
+                    responseStream.end();
                 }
             }
             console.log("stream ended: ", chunks.join(''))