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
aws-amplify · Oct 3, 2024 · 300a72d · 300a72d
1 parent 27202b4
commit 300a72d
Show file tree

Hide file tree

Showing 17 changed files with 375 additions and 92 deletions.
diff --git a/.changeset/eleven-mails-scream.md b/.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
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/ai-constructs/API.md b/packages/ai-constructs/API.md
@@ -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;

diff --git a/packages/ai-constructs/package.json b/packages/ai-constructs/package.json
@@ -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",

diff --git a/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts b/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.test.ts
@@ -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,
  },
  },
  ],

diff --git a/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts b/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.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,
  },
  };
  }),

diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_turn_executor.ts b/packages/ai-constructs/src/conversation/runtime/conversation_turn_executor.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();
 };
diff --git a/.../ai-constructs/src/conversation/runtime/event-tools-provider/event_tools_provider.test.ts b/.../ai-constructs/src/conversation/runtime/event-tools-provider/event_tools_provider.test.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',
  },
  },
  };