From 9cf6c7d83d9edeb7ecee7c907b9f6cc3765fc0ac Mon Sep 17 00:00:00 2001 From: Connor Kirkpatrick Date: Thu, 24 Apr 2025 15:45:10 +0100 Subject: [PATCH 1/3] Implement basic resolver with tool and resolve methods --- .../BedrockAgentFunctionResolver.ts | 82 ++++++++++++++++++ packages/event-handler/src/types/Tools.ts | 79 +++++++++++++++++ .../unit/BedrockAgentFunctionResolver.test.ts | 86 +++++++++++++++++++ 3 files changed, 247 insertions(+) create mode 100644 packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts create mode 100644 packages/event-handler/src/types/Tools.ts create mode 100644 packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts diff --git a/packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts b/packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts new file mode 100644 index 0000000000..84cb27373b --- /dev/null +++ b/packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts @@ -0,0 +1,82 @@ +import type { Context } from 'aws-lambda'; +import type { + ToolConfig, + ToolRegistry, + ToolDefinition, + BedrockAgentFunctionRequest, + ToolFunction, + BedrockAgentFunctionResponse, + ResponseOpts, +} from '../types/Tools'; + +export class BedrockAgentFunctionResolver { + public constructor() { + this.registry = new Map(); + } + + protected registry: ToolRegistry; + + public tool(fn: ToolFunction, config: ToolConfig) { + this.registry.set(config.name, { function: fn, config }); + } + + public resolve( + event: BedrockAgentFunctionRequest, + context: Context + ): BedrockAgentFunctionResponse { + const { function: toolName, parameters, actionGroup } = event; + + const tool = this.registry.get(toolName); + + if (tool === undefined) { + console.error(`Cant find tool ${tool}`); + return this.response({ + actionGroup, + function: toolName, + responseBody: 'error', + }); + } + + const parameterObject = parameters.reduce((acc, curr) => { + acc[curr.name] = curr.value; + return acc; + }, {}); + + console.debug(`Callin tool ${tool.config.name}`); + const response = tool.function(parameterObject); + + return this.response({ + actionGroup, + function: toolName, + responseBody: response, + }); + } + + private response(opts: ResponseOpts): BedrockAgentFunctionResponse { + const { + actionGroup, + function: fn, + responseBody, + errorType, + sessionAttributes, + promptSessionAttributes, + } = opts; + return { + messageVersion: '1.0', + response: { + actionGroup, + function: fn, + functionResponse: { + responseState: errorType, + responseBody: { + TEXT: { + body: responseBody, + }, + }, + }, + }, + sessionAttributes, + promptSessionAttributes, + }; + } +} diff --git a/packages/event-handler/src/types/Tools.ts b/packages/event-handler/src/types/Tools.ts new file mode 100644 index 0000000000..460c541e8a --- /dev/null +++ b/packages/event-handler/src/types/Tools.ts @@ -0,0 +1,79 @@ +import { number } from 'zod'; + +type ToolConfig = { + name: string; + definition: string; + validation: { input: object; output: object }; + requireConfirmation: boolean | undefined; +}; + +type ToolDefinition = { + function: ToolFunction; + config: ToolConfig; +}; + +type ToolFunction = Function; + +type ToolRegistry = Map; + +type Parameter = { + name: string; + type: string; + value: string; +}; + +type BedrockAgentFunctionRequest = { + messageVersion: string; + agent: { + name: string; + id: string; + alias: string; + version: string; + }; + inputText: string; + sessionId: string; + actionGroup: string; + function: string; + parameters: Array; + sessionAttributes: Attributes; + promptSessionAttributes: Attributes; +}; + +type Attributes = Map | undefined; + +type BedrockAgentFunctionResponse = { + messageVersion: string; + response: { + actionGroup: string; + function: string; + functionResponse: { + responseState?: 'ERROR' | 'REPROMPT'; + responseBody: { + TEXT: { + body: string; + }; + }; + }; + }; + sessionAttributes?: Attributes; + promptSessionAttributes?: Attributes; +}; + +type ResponseOpts = { + actionGroup: string; + function: string; + responseBody: string; + errorType?: 'ERROR' | 'REPROMPT'; + sessionAttributes?: Attributes; + promptSessionAttributes?: Attributes; +}; + +export type { + BedrockAgentFunctionRequest, + ToolConfig, + ToolRegistry, + ToolDefinition, + ToolFunction, + BedrockAgentFunctionResponse, + ResponseOpts, +}; diff --git a/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts b/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts new file mode 100644 index 0000000000..77c78f16ee --- /dev/null +++ b/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts @@ -0,0 +1,86 @@ +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import { BedrockAgentFunctionResolver } from '../../src/bedrockAgentFunction/BedrockAgentFunctionResolver.js'; +import context from '@aws-lambda-powertools/testing-utils/context'; + +const baseBedrockAgentFunctionRequest = { + messageVersion: '1.0', + agent: { + name: '', + id: 'string', + alias: 'string', + version: 'string', + }, + inputText: 'string', + sessionId: 'string', + actionGroup: 'string', + function: 'string', + parameters: [], + sessionAttributes: undefined, + promptSessionAttributes: undefined, +}; + +describe('BedrockAgentFunctionResolver', () => { + const ENVIRONMENT_VARIABLES = process.env; + + beforeEach(() => { + vi.clearAllMocks(); + vi.resetModules(); + process.env = { ...ENVIRONMENT_VARIABLES }; + }); + + afterAll(() => { + process.env = ENVIRONMENT_VARIABLES; + }); + + class WrappedResolver extends BedrockAgentFunctionResolver { + public getRegistry() { + return this.registry; + } + } + + it('registers tools', () => { + // Arrange + const resolver = new WrappedResolver(); + // Act + resolver.tool(() => {}, { + name: 'noop', + definition: 'Does nothing', + validation: { input: {}, output: {} }, + requireConfirmation: false, + }); + + // Assess + expect(resolver.getRegistry().get('noop')).toEqual( + expect.objectContaining({ + config: { + name: 'noop', + requireConfirmation: false, + definition: 'Does nothing', + validation: { input: {}, output: {} }, + }, + function: expect.any(Function), + }) + ); + }); + + it('resolves events to the correct tool', () => { + // Arrange + const resolver = new WrappedResolver(); + const noop = vi.fn(); + + resolver.tool(noop, { + name: 'noop', + definition: 'Does nothing', + validation: { input: {}, output: {} }, + requireConfirmation: false, + }); + + // Act + resolver.resolve( + { ...baseBedrockAgentFunctionRequest, function: 'noop' }, + context + ); + + expect(noop).toBeCalled(); + }); +}); From 2f2fe83eaf26ba3f81b8cc999c48bc36de7f402a Mon Sep 17 00:00:00 2001 From: Connor Kirkpatrick Date: Thu, 24 Apr 2025 16:05:02 +0100 Subject: [PATCH 2/3] Tidy up --- .../bedrockAgentFunction/BedrockAgentFunctionResolver.ts | 8 ++------ packages/event-handler/src/types/Tools.ts | 4 +--- .../tests/unit/BedrockAgentFunctionResolver.test.ts | 4 ++-- 3 files changed, 5 insertions(+), 11 deletions(-) diff --git a/packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts b/packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts index 84cb27373b..4732006631 100644 --- a/packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts +++ b/packages/event-handler/src/bedrockAgentFunction/BedrockAgentFunctionResolver.ts @@ -10,11 +10,7 @@ import type { } from '../types/Tools'; export class BedrockAgentFunctionResolver { - public constructor() { - this.registry = new Map(); - } - - protected registry: ToolRegistry; + protected registry: ToolRegistry = new Map(); public tool(fn: ToolFunction, config: ToolConfig) { this.registry.set(config.name, { function: fn, config }); @@ -42,7 +38,7 @@ export class BedrockAgentFunctionResolver { return acc; }, {}); - console.debug(`Callin tool ${tool.config.name}`); + console.debug(`Calling tool ${tool.config.name}`); const response = tool.function(parameterObject); return this.response({ diff --git a/packages/event-handler/src/types/Tools.ts b/packages/event-handler/src/types/Tools.ts index 460c541e8a..06bac68dbd 100644 --- a/packages/event-handler/src/types/Tools.ts +++ b/packages/event-handler/src/types/Tools.ts @@ -1,5 +1,3 @@ -import { number } from 'zod'; - type ToolConfig = { name: string; definition: string; @@ -39,7 +37,7 @@ type BedrockAgentFunctionRequest = { promptSessionAttributes: Attributes; }; -type Attributes = Map | undefined; +type Attributes = Record; type BedrockAgentFunctionResponse = { messageVersion: string; diff --git a/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts b/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts index 77c78f16ee..e224e1ecc6 100644 --- a/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts +++ b/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts @@ -15,8 +15,8 @@ const baseBedrockAgentFunctionRequest = { actionGroup: 'string', function: 'string', parameters: [], - sessionAttributes: undefined, - promptSessionAttributes: undefined, + sessionAttributes: {}, + promptSessionAttributes: {}, }; describe('BedrockAgentFunctionResolver', () => { From b7d3e71e2abdc211d2bc0a06f643cae643ec66ef Mon Sep 17 00:00:00 2001 From: Connor Kirkpatrick Date: Thu, 24 Apr 2025 16:21:51 +0100 Subject: [PATCH 3/3] Verify output --- .../unit/BedrockAgentFunctionResolver.test.ts | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts b/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts index e224e1ecc6..f7ba9188d7 100644 --- a/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts +++ b/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts @@ -83,4 +83,43 @@ describe('BedrockAgentFunctionResolver', () => { expect(noop).toBeCalled(); }); + + it('responds with the correct response structure when a tool is successfully invoked', () => { + const resolver = new WrappedResolver(); + const uppercaser = ({ str }): string => str.toUpperCase(); + + resolver.tool(uppercaser, { + name: 'uppercaser', + definition: 'Converts a string to uppercase', + validation: { input: {}, output: {} }, + requireConfirmation: false, + }); + + // Act + const response = resolver.resolve( + { + ...baseBedrockAgentFunctionRequest, + function: 'uppercaser', + parameters: [{ name: 'str', value: 'hello world', type: 'string' }], + }, + context + ); + + expect(response).toEqual( + expect.objectContaining({ + messageVersion: '1.0', + response: { + actionGroup: baseBedrockAgentFunctionRequest.actionGroup, + function: 'uppercaser', + functionResponse: { + responseBody: { + TEXT: { + body: 'HELLO WORLD', + }, + }, + }, + }, + }) + ); + }); });