Include padding from bedrock in AI chunked responses (#2453)

Author: sobolk

.changeset/hip-hairs-rescue.md

@@ -0,0 +1,5 @@
+---
+'@aws-amplify/ai-constructs': patch
+---
+
+Include padding from bedrock in AI chunked responses

package-lock.json

@@ -116,6 +116,7 @@ void describe('Bedrock converse adapter', () => {
               contentBlockText: 'b',
               contentBlockIndex: 0,
               contentBlockDeltaIndex: 0,
+              p: 'testPadding',
             },
             {
               accumulatedTurnContent: [
@@ -128,6 +129,7 @@ void describe('Bedrock converse adapter', () => {
               contentBlockText: 'lock1',
               contentBlockIndex: 0,
               contentBlockDeltaIndex: 1,
+              p: 'testPadding',
             },
             {
               accumulatedTurnContent: [
@@ -154,6 +156,7 @@ void describe('Bedrock converse adapter', () => {
               contentBlockText: 'b',
               contentBlockIndex: 1,
               contentBlockDeltaIndex: 0,
+              p: 'testPadding',
             },
             {
               accumulatedTurnContent: [
@@ -169,6 +172,7 @@ void describe('Bedrock converse adapter', () => {
               contentBlockText: 'lock2',
               contentBlockIndex: 1,
               contentBlockDeltaIndex: 1,
+              p: 'testPadding',
             },
             {
               accumulatedTurnContent: [
@@ -1202,6 +1206,8 @@ const mockConverseStreamCommandOutput = (
             // simulate chunked input
             text: input.substring(0, 1),
           },
+          // @ts-expect-error padding is sent by Bedrock but not in their API.
+          p: 'testPadding',
         },
       });
       if (input.length > 1) {
@@ -1212,6 +1218,8 @@ const mockConverseStreamCommandOutput = (
               // simulate chunked input
               text: input.substring(1),
             },
+            // @ts-expect-error padding is sent by Bedrock but not in their API.
+            p: 'testPadding',
           },
         });
       }

packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts

@@ -239,14 +239,22 @@ export class BedrockConverseAdapter {
               }
             } else if (chunk.contentBlockDelta.delta?.text) {
               text += chunk.contentBlockDelta.delta.text;
-              yield {
+              const amplifyChunk: StreamingResponseChunk = {
                 accumulatedTurnContent: [...accumulatedTurnContent, { text }],
                 conversationId: this.event.conversationId,
                 associatedUserMessageId: this.event.currentMessageId,
                 contentBlockText: chunk.contentBlockDelta.delta.text,
                 contentBlockIndex: blockIndex,
                 contentBlockDeltaIndex: blockDeltaIndex,
               };
+              // padding is sent from Bedrock but not included in the API.
+              if ('p' in chunk.contentBlockDelta) {
+                const bedrockPadding = chunk.contentBlockDelta.p;
+                if (typeof bedrockPadding === 'string') {
+                  amplifyChunk.p = bedrockPadding;
+                }
+              }
+              yield amplifyChunk;
               lastBlockDeltaIndex = blockDeltaIndex;
               blockDeltaIndex++;
             }

packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts

@@ -106,6 +106,11 @@ export type StreamingResponseChunk = {
   associatedUserMessageId: string;
   contentBlockIndex: number;
   accumulatedTurnContent: Array<bedrock.ContentBlock>;
+
+  // optional padding
+  // it is a random string included in text chunks,
+  // the value does not matter, but we need to pass it through to AppSync
+  p?: string;
 } & (
   | {
       // text chunk

packages/ai-constructs/src/conversation/runtime/types.ts

@@ -95,6 +95,7 @@ type ConversationTurnAppSyncResponseChunk = {
   contentBlockDoneAtIndex?: number;
   contentBlockToolUse?: string;
   stopReason?: string;
+  p?: string;
   errors?: Array<ConversationTurnError>;
 };
 
@@ -698,6 +699,7 @@ class ConversationHandlerTestProject extends TestProjectBase {
                     errorType
                     message
                   }
+                  p
                   id
                   owner
                   stopReason
@@ -747,6 +749,10 @@ class ConversationHandlerTestProject extends TestProjectBase {
 
       const content = chunks.reduce((accumulated, current) => {
         if (current.contentBlockText) {
+          assert.ok(
+            current.p && current.p.length > 0,
+            'Text chunks must include padding'
+          );
           accumulated += current.contentBlockText;
         }
         if (current.contentBlockToolUse) {

packages/integration-tests/src/test-project-setup/conversation_handler_project.ts

@@ -108,6 +108,9 @@ const schema = a.schema({
       contentBlockIndex: a.integer(),
       accumulatedTurnContent: a.ref('MockContentBlock').array(),
 
+      // padding
+      p: a.string(),
+
       // these describe chunks or end of block
       contentBlockText: a.string(),
       contentBlockToolUse: a.string(),