Fix a case when Bedrock throws validation error if tool input is not persisted in history (#2230)

* Fix a case when Bedrock throws validation error if tool input is not persisted in history

* some prompt engineering

Author: Kamil Sobol

.changeset/clever-emus-hope.md

@@ -0,0 +1,5 @@
+---
+'@aws-amplify/ai-constructs': patch
+---
+
+Fix a case when Bedrock throws validation error if tool input is not persisted in history

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

@@ -901,8 +901,8 @@ void describe('Bedrock converse adapter', () => {
     assert.strictEqual(responseText, 'finalResponse');
 
     assert.strictEqual(toolExecuteMock.mock.calls.length, 2);
-    assert.strictEqual(toolExecuteMock.mock.calls[0].arguments[0], undefined);
-    assert.strictEqual(toolExecuteMock.mock.calls[1].arguments[0], undefined);
+    assert.deepStrictEqual(toolExecuteMock.mock.calls[0].arguments[0], {});
+    assert.deepStrictEqual(toolExecuteMock.mock.calls[1].arguments[0], {});
   });
 
   void it('throws if tool is duplicated', () => {

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

@@ -252,6 +252,10 @@ export class BedrockConverseAdapter {
           if (toolUseBlock) {
             if (toolUseInput) {
               toolUseBlock.toolUse.input = JSON.parse(toolUseInput);
+            } else {
+              // Bedrock API requires tool input to be non-null in message history.
+              // Therefore, falling back to empty object.
+              toolUseBlock.toolUse.input = {};
             }
             accumulatedAssistantMessage.content?.push(toolUseBlock);
             if (

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

@@ -26,6 +26,7 @@ import assert from 'assert';
 import { NormalizedCacheObject } from '@apollo/client';
 import {
   bedrockModelId,
+  expectedRandomNumber,
   expectedTemperatureInDataToolScenario,
   expectedTemperaturesInProgrammaticToolScenario,
 } from '../test-projects/conversation-handler/amplify/constants.js';
@@ -278,6 +279,16 @@ class ConversationHandlerTestProject extends TestProjectBase {
       )
     );
 
+    await this.executeWithRetry(() =>
+      this.assertCustomConversationHandlerCanExecuteTurnWithParameterLessTool(
+        backendId,
+        authenticatedUserCredentials.accessToken,
+        dataUrl,
+        apolloClient,
+        true
+      )
+    );
+
     await this.executeWithRetry(() =>
       this.assertDefaultConversationHandlerCanExecuteTurnWithDataTool(
         backendId,
@@ -853,6 +864,56 @@ class ConversationHandlerTestProject extends TestProjectBase {
     );
   };
 
+  private assertCustomConversationHandlerCanExecuteTurnWithParameterLessTool =
+    async (
+      backendId: BackendIdentifier,
+      accessToken: string,
+      graphqlApiEndpoint: string,
+      apolloClient: ApolloClient<NormalizedCacheObject>,
+      streamResponse: boolean
+    ): Promise<void> => {
+      const customConversationHandlerFunction = (
+        await this.resourceFinder.findByBackendIdentifier(
+          backendId,
+          'AWS::Lambda::Function',
+          (name) => name.includes('custom')
+        )
+      )[0];
+
+      const message: CreateConversationMessageChatInput = {
+        conversationId: randomUUID().toString(),
+        id: randomUUID().toString(),
+        role: 'user',
+        content: [
+          {
+            text: 'Give me a random number',
+          },
+        ],
+      };
+      await this.insertMessage(apolloClient, message);
+
+      // send event
+      const event: ConversationTurnEvent = {
+        conversationId: message.conversationId,
+        currentMessageId: message.id,
+        graphqlApiEndpoint: graphqlApiEndpoint,
+        request: {
+          headers: { authorization: accessToken },
+        },
+        ...this.getCommonEventProperties(streamResponse),
+      };
+      const response = await this.executeConversationTurn(
+        event,
+        customConversationHandlerFunction,
+        apolloClient
+      );
+      // Assert that tool was used. I.e. LLM used value provided by the tool.
+      assert.match(
+        response.content,
+        new RegExp(expectedRandomNumber.toString())
+      );
+    };
+
   private assertDefaultConversationHandlerCanPropagateError = async (
     backendId: BackendIdentifier,
     accessToken: string,

packages/integration-tests/src/test-projects/conversation-handler/amplify/constants.ts

@@ -14,3 +14,5 @@ export const expectedTemperaturesInProgrammaticToolScenario = {
 };
 
 export const expectedTemperatureInDataToolScenario = 85;
+
+export const expectedRandomNumber = 42;

packages/integration-tests/src/test-projects/conversation-handler/amplify/custom-conversation-handler/custom_handler.ts

@@ -3,7 +3,10 @@ import {
   createExecutableTool,
   handleConversationTurnEvent,
 } from '@aws-amplify/backend-ai/conversation/runtime';
-import { expectedTemperaturesInProgrammaticToolScenario } from '../constants.js';
+import {
+  expectedRandomNumber,
+  expectedTemperaturesInProgrammaticToolScenario,
+} from '../constants.js';
 
 const thermometerInputSchema = {
   type: 'object',
@@ -33,11 +36,32 @@ const thermometer = createExecutableTool(
   }
 );
 
+// Parameter-less tool.
+const randomNumberGeneratorInputSchema = {
+  type: 'object',
+  properties: {},
+  required: [],
+} as const;
+
+const randomNumberGenerator = createExecutableTool(
+  'randomNumberGenerator',
+  'Returns a random number',
+  {
+    json: randomNumberGeneratorInputSchema,
+  },
+  () => {
+    return Promise.resolve({
+      // We use this value in test assertion.
+      text: `${expectedRandomNumber}`,
+    });
+  }
+);
+
 /**
  * Handler with simple tool.
  */
 export const handler = async (event: ConversationTurnEvent) => {
   await handleConversationTurnEvent(event, {
-    tools: [thermometer],
+    tools: [randomNumberGenerator, thermometer],
   });
 };