Infer executable tool input type from input schema (#2070)

* Infer executable tool input type from input schema

* Infer executable tool input type from input schema

* tests

* npm install

* that is better

* fix that

Author: Kamil Sobol

.changeset/eleven-mails-scream.md

@@ -0,0 +1,6 @@
+---
+'@aws-amplify/ai-constructs': minor
+'@aws-amplify/backend-ai': minor
+---
+
+Infer executable tool input type from input schema

package-lock.json

@@ -9,8 +9,8 @@
 import * as bedrock from '@aws-sdk/client-bedrock-runtime';
 import { Construct } from 'constructs';
 import { FunctionResources } from '@aws-amplify/plugin-types';
+import * as jsonSchemaToTypeScript from 'json-schema-to-ts';
 import { ResourceProvider } from '@aws-amplify/plugin-types';
-import * as smithy from '@smithy/types';
 
 declare namespace __export__conversation {
     export {
@@ -25,10 +25,12 @@ declare namespace __export__conversation__runtime {
         ConversationMessage,
         ConversationMessageContentBlock,
         ConversationTurnEvent,
+        createExecutableTool,
         ExecutableTool,
+        FromJSONSchema,
+        JSONSchema,
         handleConversationTurnEvent,
         ToolDefinition,
-        ToolExecutionInput,
         ToolInputSchema,
         ToolResultContentBlock
     }
@@ -111,28 +113,36 @@ type ConversationTurnEvent = {
     };
 };
 
+// @public
+const createExecutableTool: <TJSONSchema extends JSONSchema = JSONSchema, TToolInput = FromJSONSchema<TJSONSchema>>(name: string, description: string, inputSchema: ToolInputSchema<TJSONSchema>, handler: (input: TToolInput) => Promise<bedrock.ToolResultContentBlock>) => ExecutableTool<TJSONSchema, TToolInput>;
+
 // @public (undocumented)
-type ExecutableTool = ToolDefinition & {
-    execute: (input: ToolExecutionInput | undefined) => Promise<ToolResultContentBlock>;
+type ExecutableTool<TJSONSchema extends JSONSchema = JSONSchema, TToolInput = FromJSONSchema<TJSONSchema>> = ToolDefinition<TJSONSchema> & {
+    execute: (input: TToolInput) => Promise<ToolResultContentBlock>;
 };
 
+// @public (undocumented)
+type FromJSONSchema<TJSONSchema extends JSONSchema> = jsonSchemaToTypeScript.FromSchema<TJSONSchema>;
+
 // @public
 const handleConversationTurnEvent: (event: ConversationTurnEvent, props?: {
-    tools?: Array<ExecutableTool>;
+    tools?: Array<ExecutableTool<JSONSchema, any>>;
 }) => Promise<void>;
 
 // @public (undocumented)
-type ToolDefinition = {
+type JSONSchema = jsonSchemaToTypeScript.JSONSchema;
+
+// @public (undocumented)
+type ToolDefinition<TJSONSchema extends JSONSchema = JSONSchema> = {
     name: string;
     description: string;
-    inputSchema: ToolInputSchema;
+    inputSchema: ToolInputSchema<TJSONSchema>;
 };
 
 // @public (undocumented)
-type ToolExecutionInput = smithy.DocumentType;
-
-// @public (undocumented)
-type ToolInputSchema = bedrock.ToolInputSchema;
+type ToolInputSchema<TJSONSchema extends JSONSchema> = {
+    json: TJSONSchema;
+};
 
 // @public (undocumented)
 type ToolResultContentBlock = bedrock.ToolResultContentBlock;

packages/ai-constructs/API.md

@@ -28,7 +28,11 @@
   "dependencies": {
     "@aws-amplify/plugin-types": "^1.0.1",
     "@aws-sdk/client-bedrock-runtime": "^3.622.0",
-    "@smithy/types": "^3.3.0"
+    "@smithy/types": "^3.3.0",
+    "json-schema-to-ts": "^3.1.1"
+  },
+  "devDependencies": {
+    "typescript": "^5.0.0"
   },
   "peerDependencies": {
     "aws-cdk-lib": "^2.152.0",

packages/ai-constructs/package.json

@@ -14,6 +14,7 @@ import {
   ConverseCommandOutput,
   Message,
   ToolConfiguration,
+  ToolInputSchema,
   ToolResultContentBlock,
 } from '@aws-sdk/client-bedrock-runtime';
 import { ConversationTurnEventToolsProvider } from './event-tools-provider';
@@ -252,14 +253,14 @@ void describe('Bedrock converse adapter', () => {
           toolSpec: {
             name: eventTool.name,
             description: eventTool.description,
-            inputSchema: eventTool.inputSchema,
+            inputSchema: eventTool.inputSchema as ToolInputSchema,
           },
         },
         {
           toolSpec: {
             name: additionalTool.name,
             description: additionalTool.description,
-            inputSchema: additionalTool.inputSchema,
+            inputSchema: additionalTool.inputSchema as ToolInputSchema,
           },
         },
       ],
@@ -689,14 +690,14 @@ void describe('Bedrock converse adapter', () => {
           toolSpec: {
             name: additionalTool.name,
             description: additionalTool.description,
-            inputSchema: additionalTool.inputSchema,
+            inputSchema: additionalTool.inputSchema as ToolInputSchema,
           },
         },
         {
           toolSpec: {
             name: clientTool.name,
             description: clientTool.description,
-            inputSchema: clientTool.inputSchema,
+            inputSchema: clientTool.inputSchema as ToolInputSchema,
           },
         },
       ],

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

@@ -7,6 +7,7 @@ import {
   Message,
   Tool,
   ToolConfiguration,
+  ToolInputSchema,
 } from '@aws-sdk/client-bedrock-runtime';
 import {
   ConversationTurnEvent,
@@ -171,7 +172,9 @@ export class BedrockConverseAdapter {
           toolSpec: {
             name: t.name,
             description: t.description,
-            inputSchema: t.inputSchema,
+            // We have to cast to bedrock type as we're using different types to describe JSON schema in our API.
+            // These types are runtime compatible.
+            inputSchema: t.inputSchema as ToolInputSchema,
           },
         };
       }),

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

@@ -1,5 +1,5 @@
 import { ConversationTurnResponseSender } from './conversation_turn_response_sender.js';
-import { ConversationTurnEvent, ExecutableTool } from './types.js';
+import { ConversationTurnEvent, ExecutableTool, JSONSchema } from './types.js';
 import { BedrockConverseAdapter } from './bedrock_converse_adapter.js';
 
 /**
@@ -54,7 +54,10 @@ export class ConversationTurnExecutor {
  */
 export const handleConversationTurnEvent = async (
   event: ConversationTurnEvent,
-  props?: { tools?: Array<ExecutableTool> }
+  // This is by design, so that tools with different input types can be added
+  // to single arrays. Downstream code doesn't use these types.
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  props?: { tools?: Array<ExecutableTool<JSONSchema, any>> }
 ): Promise<void> => {
   await new ConversationTurnExecutor(event, props?.tools ?? []).execute();
 };

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

@@ -36,14 +36,17 @@ void describe('events tool provider', () => {
       description: 'toolDescription1',
       inputSchema: {
         json: {
-          tool1: 'value1',
+          type: 'object',
+          properties: {
+            tool1Property: { type: 'string' },
+          },
         },
       },
       graphqlRequestInputDescriptor: {
         queryName: 'queryName1',
         selectionSet: 'selection1',
         propertyTypes: {
-          property1: 'type1',
+          tool1Property: 'type1',
         },
       },
     };
@@ -52,14 +55,17 @@ void describe('events tool provider', () => {
       description: 'toolDescription2',
       inputSchema: {
         json: {
-          tool1: 'value2',
+          type: 'object',
+          properties: {
+            tool2Property: { type: 'string' },
+          },
         },
       },
       graphqlRequestInputDescriptor: {
         queryName: 'queryName2',
         selectionSet: 'selection2',
         propertyTypes: {
-          property1: 'type2',
+          tool2Property: 'type2',
         },
       },
     };