Add document support in AI kit (#2566)

* Add document support in AI kit

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

Co-authored-by: Kamil Sobol <[email protected]>

* adding e2e tests

* adding docx to dictionary

* fixing e2e test

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

Co-authored-by: Kamil Sobol <[email protected]>

---------

Co-authored-by: Kamil Sobol <[email protected]>

Author: dbanksdesign

.changeset/famous-cougars-fold.md

@@ -0,0 +1,5 @@
+---
+'@aws-amplify/ai-constructs': minor
+---
+
+Adding document support for ai conversation routes

.eslint_dictionary.json

@@ -49,6 +49,7 @@
   "deprecator",
   "deserializer",
   "disambiguator",
+  "docx",
   "downlevel",
   "durations",
   "dynamodb",

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

@@ -824,6 +824,72 @@ void describe('Bedrock converse adapter', () => {
           },
         ]);
       });
+
+      void it('decodes base64 encoded documents', async () => {
+        const event: ConversationTurnEvent = {
+          ...commonEvent,
+        };
+
+        const fakeDocumentPayload = randomBytes(32);
+
+        messageHistoryRetrieverMockGetEventMessages.mock.mockImplementationOnce(
+          () => {
+            return Promise.resolve([
+              {
+                id: '',
+                conversationId: '',
+                role: 'user',
+                content: [
+                  {
+                    document: {
+                      name: 'test',
+                      format: 'doc',
+                      source: {
+                        bytes: fakeDocumentPayload.toString('base64'),
+                      },
+                    },
+                  },
+                ],
+              },
+            ]);
+          }
+        );
+
+        const bedrockClient = new BedrockRuntimeClient();
+        const content = [{ text: 'block1' }, { text: 'block2' }];
+        const bedrockResponse = mockBedrockResponse(content, streamResponse);
+        const bedrockClientSendMock = mock.method(bedrockClient, 'send', () =>
+          Promise.resolve(bedrockResponse)
+        );
+
+        await new BedrockConverseAdapter(
+          event,
+          [],
+          bedrockClient,
+          undefined,
+          messageHistoryRetriever
+        ).askBedrock();
+
+        assert.strictEqual(bedrockClientSendMock.mock.calls.length, 1);
+        const bedrockRequest = bedrockClientSendMock.mock.calls[0]
+          .arguments[0] as unknown as ConverseCommand;
+        assert.deepStrictEqual(bedrockRequest.input.messages, [
+          {
+            role: 'user',
+            content: [
+              {
+                document: {
+                  format: 'doc',
+                  name: 'test',
+                  source: {
+                    bytes: fakeDocumentPayload,
+                  },
+                },
+              },
+            ],
+          },
+        ]);
+      });
     });
   });
 

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

@@ -385,6 +385,18 @@ export class BedrockConverseAdapter {
               },
             },
           });
+        } else if (typeof contentElement.document?.source?.bytes === 'string') {
+          messageContent.push({
+            document: {
+              ...contentElement.document,
+              source: {
+                bytes: Buffer.from(
+                  contentElement.document.source.bytes,
+                  'base64'
+                ),
+              },
+            },
+          });
         } else {
           // Otherwise type conforms to Bedrock's type and it's safe to cast.
           messageContent.push(contentElement as ContentBlock);

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

@@ -21,6 +21,20 @@ export type ConversationMessage = {
 
 export type ConversationMessageContentBlock =
   | bedrock.ContentBlock
+  | {
+      image?: never;
+      // These are needed so that union with other content block types works.
+      // See https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/Package/-aws-sdk-client-bedrock-runtime/TypeAlias/ContentBlock/.
+      text?: never;
+      document: Omit<bedrock.DocumentBlock, 'source'> & {
+        // Upstream (Appsync) may send documents in a form of Base64 encoded strings
+        source: { bytes: string };
+      };
+      toolUse?: never;
+      toolResult?: never;
+      guardContent?: never;
+      $unknown?: never;
+    }
   | {
       image: Omit<bedrock.ImageBlock, 'source'> & {
         // Upstream (Appsync) may send images in a form of Base64 encoded strings

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

@@ -60,6 +60,20 @@ type ConversationMessage = {
 
 type ConversationMessageContentBlock =
   | bedrock.ContentBlock
+  | {
+      image?: never;
+      // These are needed so that union with other content block types works.
+      // See https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/Package/-aws-sdk-client-bedrock-runtime/TypeAlias/ContentBlock/.
+      text?: never;
+      document: Omit<bedrock.DocumentBlock, 'source'> & {
+        // Upstream (Appsync) may send images in a form of Base64 encoded strings
+        source: { bytes: string };
+      };
+      toolUse?: never;
+      toolResult?: never;
+      guardContent?: never;
+      $unknown?: never;
+    }
   | {
       image: Omit<bedrock.ImageBlock, 'source'> & {
         // Upstream (Appsync) may send images in a form of Base64 encoded strings
@@ -352,6 +366,26 @@ class ConversationHandlerTestProject extends TestProjectBase {
       )
     );
 
+    await this.executeWithRetry(() =>
+      this.assertDefaultConversationHandlerCanExecuteTurnWithDocument(
+        backendId,
+        authenticatedUserCredentials.accessToken,
+        dataUrl,
+        apolloClient,
+        false
+      )
+    );
+
+    await this.executeWithRetry(() =>
+      this.assertDefaultConversationHandlerCanExecuteTurnWithDocument(
+        backendId,
+        authenticatedUserCredentials.accessToken,
+        dataUrl,
+        apolloClient,
+        true
+      )
+    );
+
     await this.executeWithRetry((attempt) =>
       this.assertDefaultConversationHandlerCanPropagateError(
         backendId,
@@ -496,6 +530,74 @@ class ConversationHandlerTestProject extends TestProjectBase {
     assert.match(response.content, /(aws)|(AWS)|(Amazon Web Services)/);
   };
 
+  private assertDefaultConversationHandlerCanExecuteTurnWithDocument = async (
+    backendId: BackendIdentifier,
+    accessToken: string,
+    graphqlApiEndpoint: string,
+    apolloClient: ApolloClient<NormalizedCacheObject>,
+    streamResponse: boolean
+  ): Promise<void> => {
+    const defaultConversationHandlerFunction = (
+      await this.resourceFinder.findByBackendIdentifier(
+        backendId,
+        'AWS::Lambda::Function',
+        (name) => name.includes('default')
+      )
+    )[0];
+
+    const documentPath = resolve(
+      fileURLToPath(import.meta.url),
+      '..',
+      '..',
+      '..',
+      'src',
+      'test-projects',
+      'conversation-handler',
+      'resources',
+      'sample-document.docx'
+    );
+
+    const documentSource = await fs.readFile(documentPath, 'base64');
+
+    const message: CreateConversationMessageChatInput = {
+      id: randomUUID().toString(),
+      conversationId: randomUUID().toString(),
+      role: 'user',
+      content: [
+        {
+          text: 'What is in the attached document?',
+        },
+        {
+          document: {
+            format: 'docx',
+            name: 'sample-document',
+            source: { bytes: documentSource },
+          },
+        },
+      ],
+    };
+
+    // send event
+    const event: ConversationTurnEvent = {
+      conversationId: message.conversationId,
+      currentMessageId: message.id,
+      graphqlApiEndpoint: graphqlApiEndpoint,
+      request: {
+        headers: { authorization: accessToken },
+      },
+      ...this.getCommonEventProperties(streamResponse),
+    };
+    await this.insertMessage(apolloClient, message);
+    const response = await this.executeConversationTurn(
+      event,
+      defaultConversationHandlerFunction,
+      apolloClient
+    );
+    // The document contains a hello world string. Responses may vary, but they should always contain statements below.
+    assert.match(response.content, /document/);
+    assert.match(response.content, /(H|h)ello (W|w)orld/);
+  };
+
   private assertDefaultConversationHandlerCanExecuteTurnWithDataTool = async (
     backendId: BackendIdentifier,
     accessToken: string,