diff --git a/packages/event-handler/package.json b/packages/event-handler/package.json index 09ca6be37..2fc2a7f90 100644 --- a/packages/event-handler/package.json +++ b/packages/event-handler/package.json @@ -39,6 +39,16 @@ "default": "./lib/esm/appsync-events/index.js" } }, + "./bedrock-agent-function": { + "require": { + "types": "./lib/cjs/bedrock-agent-function/index.d.ts", + "default": "./lib/cjs/bedrock-agent-function/index.js" + }, + "import": { + "types": "./lib/esm/bedrock-agent-function/index.d.ts", + "default": "./lib/esm/bedrock-agent-function/index.js" + } + }, "./types": { "require": { "types": "./lib/cjs/types/index.d.ts", @@ -56,6 +66,10 @@ "./lib/cjs/appsync-events/index.d.ts", "./lib/esm/appsync-events/index.d.ts" ], + "bedrock-agent-function": [ + "./lib/cjs/bedrock-agent-function/index.d.ts", + "./lib/esm/bedrock-agent-function/index.d.ts" + ], "types": [ "./lib/cjs/types/index.d.ts", "./lib/esm/types/index.d.ts" diff --git a/packages/event-handler/src/bedrock-agent-function/BedrockAgentFunctionResolver.ts b/packages/event-handler/src/bedrock-agent-function/BedrockAgentFunctionResolver.ts new file mode 100644 index 000000000..32927dcba --- /dev/null +++ b/packages/event-handler/src/bedrock-agent-function/BedrockAgentFunctionResolver.ts @@ -0,0 +1,196 @@ +import { EnvironmentVariablesService } from '@aws-lambda-powertools/commons'; +import type { Context } from 'aws-lambda'; +import type { + BedrockAgentFunctionEvent, + BedrockAgentFunctionResponse, + Configuration, + GenericLogger, + ResolverOptions, + ResponseOptions, + Tool, + ToolFunction, +} from '../types/index.js'; +import { isPrimitive } from './utils.js'; + +export class BedrockAgentFunctionResolver { + readonly #tools: Map = new Map(); + readonly #envService: EnvironmentVariablesService; + readonly #logger: Pick; + + constructor(options?: ResolverOptions) { + this.#envService = new EnvironmentVariablesService(); + const alcLogLevel = this.#envService.get('AWS_LAMBDA_LOG_LEVEL'); + this.#logger = options?.logger ?? { + debug: alcLogLevel === 'DEBUG' ? console.debug : () => {}, + error: console.error, + warn: console.warn, + }; + } + + /** + * Register a tool function for the Bedrock Agent. + * + * This method registers a function that can be invoked by a Bedrock Agent. + * + * @example + * ```ts + * import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/bedrock-agent-function'; + * + * const app = new BedrockAgentFunctionResolver(); + * + * app.tool(async (params) => { + * const { name } = params; + * return `Hello, ${name}!`; + * }, { + * name: 'greeting', + * definition: 'Greets a person by name', + * }); + * + * export const handler = async (event, context) => + * app.resolve(event, context); + * ``` + * + * The method also works as a class method decorator: + * + * @example + * ```ts + * import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/bedrock-agent-function'; + * + * const app = new BedrockAgentFunctionResolver(); + * + * class Lambda { + * @app.tool({ name: 'greeting', definition: 'Greets a person by name' }) + * async greeting(params) { + * const { name } = params; + * return `Hello, ${name}!`; + * } + * + * async handler(event, context) { + * return app.resolve(event, context); + * } + * } + * + * const lambda = new Lambda(); + * export const handler = lambda.handler.bind(lambda); + * ``` + * + * @param fn - The tool function + * @param config - The configuration object for the tool + */ + public tool(fn: ToolFunction, config: Configuration): void; + public tool(config: Configuration): MethodDecorator; + public tool( + fnOrConfig: ToolFunction | Configuration, + config?: Configuration + ): MethodDecorator | void { + // When used as a method (not a decorator) + if (typeof fnOrConfig === 'function') { + return this.#registerTool(fnOrConfig, config as Configuration); + } + + // When used as a decorator + return (_target, _propertyKey, descriptor: PropertyDescriptor) => { + const toolFn = descriptor.value as ToolFunction; + this.#registerTool(toolFn, fnOrConfig); + return descriptor; + }; + } + + #registerTool(fn: ToolFunction, config: Configuration): void { + const { name } = config; + + if (this.#tools.size >= 5) { + this.#logger.warn( + `The maximum number of tools that can be registered is 5. Tool ${name} will not be registered.` + ); + return; + } + + if (this.#tools.has(name)) { + this.#logger.warn( + `Tool ${name} already registered. Overwriting with new definition.` + ); + } + + this.#tools.set(name, { function: fn, config }); + this.#logger.debug(`Tool ${name} has been registered.`); + } + + #buildResponse(options: ResponseOptions): BedrockAgentFunctionResponse { + const { + actionGroup, + function: func, + body, + errorType, + sessionAttributes, + promptSessionAttributes, + } = options; + + return { + messageVersion: '1.0', + response: { + actionGroup, + function: func, + functionResponse: { + responseState: errorType, + responseBody: { + TEXT: { + body, + }, + }, + }, + }, + sessionAttributes, + promptSessionAttributes, + }; + } + + async resolve( + event: BedrockAgentFunctionEvent, + context: Context + ): Promise { + const { + function: toolName, + parameters = [], + actionGroup, + sessionAttributes, + promptSessionAttributes, + } = event; + + const tool = this.#tools.get(toolName); + + if (tool == null) { + this.#logger.error(`Tool ${toolName} has not been registered.`); + return this.#buildResponse({ + actionGroup, + function: toolName, + body: 'Error: tool has not been registered in handler.', + }); + } + + const parameterObject: Record = Object.fromEntries( + parameters.map((param) => [param.name, param.value]) + ); + + try { + const res = (await tool.function(parameterObject)) ?? ''; + const body = isPrimitive(res) ? String(res) : JSON.stringify(res); + return this.#buildResponse({ + actionGroup, + function: toolName, + body, + sessionAttributes, + promptSessionAttributes, + }); + } catch (error) { + this.#logger.error(`An error occurred in tool ${toolName}.`, error); + return this.#buildResponse({ + actionGroup, + function: toolName, + body: `Error when invoking tool: ${error}`, + sessionAttributes, + promptSessionAttributes, + }); + } + } +} diff --git a/packages/event-handler/src/bedrock-agent-function/index.ts b/packages/event-handler/src/bedrock-agent-function/index.ts new file mode 100644 index 000000000..a18e9dd72 --- /dev/null +++ b/packages/event-handler/src/bedrock-agent-function/index.ts @@ -0,0 +1 @@ +export { BedrockAgentFunctionResolver } from './BedrockAgentFunctionResolver.js'; diff --git a/packages/event-handler/src/bedrock-agent-function/utils.ts b/packages/event-handler/src/bedrock-agent-function/utils.ts new file mode 100644 index 000000000..525c9f89f --- /dev/null +++ b/packages/event-handler/src/bedrock-agent-function/utils.ts @@ -0,0 +1,14 @@ +import type { + JSONPrimitive, + JSONValue, +} from '@aws-lambda-powertools/commons/types'; + +export function isPrimitive(value: JSONValue): value is JSONPrimitive { + return ( + value === null || + value === undefined || + typeof value === 'string' || + typeof value === 'number' || + typeof value === 'boolean' + ); +} diff --git a/packages/event-handler/src/types/appsync-events.ts b/packages/event-handler/src/types/appsync-events.ts index 7c391f331..1367cacd9 100644 --- a/packages/event-handler/src/types/appsync-events.ts +++ b/packages/event-handler/src/types/appsync-events.ts @@ -1,22 +1,7 @@ import type { Context } from 'aws-lambda'; import type { RouteHandlerRegistry } from '../appsync-events/RouteHandlerRegistry.js'; import type { Router } from '../appsync-events/Router.js'; - -// #region Shared - -// biome-ignore lint/suspicious/noExplicitAny: We intentionally use `any` here to represent any type of data and keep the logger is as flexible as possible. -type Anything = any; - -/** - * Interface for a generic logger object. - */ -type GenericLogger = { - trace?: (...content: Anything[]) => void; - debug: (...content: Anything[]) => void; - info?: (...content: Anything[]) => void; - warn: (...content: Anything[]) => void; - error: (...content: Anything[]) => void; -}; +import type { Anything, GenericLogger } from './common.js'; // #region OnPublish fn diff --git a/packages/event-handler/src/types/bedrock-agent-function.ts b/packages/event-handler/src/types/bedrock-agent-function.ts new file mode 100644 index 000000000..b3f1037fd --- /dev/null +++ b/packages/event-handler/src/types/bedrock-agent-function.ts @@ -0,0 +1,105 @@ +import type { JSONValue } from '@aws-lambda-powertools/commons/types'; +import type { GenericLogger } from '../types/common.js'; + +type Configuration = { + name: string; + definition: string; +}; + +type Parameter = { + name: string; + type: 'string' | 'number' | 'integer' | 'boolean' | 'array'; + value: string; +}; + +type ParameterPrimitives = string | number | boolean + +type ParameterValue = ParameterPrimitives | Array + +type Tool> = { + // biome-ignore lint/suspicious/noConfusingVoidType: we need to support async functions that don't have an explicit return value + function: (params: TParams) => Promise; + config: Configuration; +}; + +type ToolFunction = Tool['function']; + +type Attributes = Record; + +type FunctionIdentifier = { + actionGroup: string; + function: string; +}; + +type FunctionInvocation = FunctionIdentifier & { + parameters?: Array; +}; + +type BedrockAgentFunctionEvent = FunctionInvocation & { + messageVersion: string; + agent: { + name: string; + id: string; + alias: string; + version: string; + }; + inputText: string; + sessionId: string; + sessionAttributes: Attributes; + promptSessionAttributes: Attributes; +}; + +type ResponseState = 'ERROR' | 'REPROMPT'; + +type TextResponseBody = { + TEXT: { + body: string; + }; +}; + +type SessionData = { + sessionAttributes?: Attributes; + promptSessionAttributes?: Attributes; +}; + +type BedrockAgentFunctionResponse = SessionData & { + messageVersion: string; + response: FunctionIdentifier & { + functionResponse: { + responseState?: ResponseState; + responseBody: TextResponseBody; + }; + }; +}; + +type ResponseOptions = FunctionIdentifier & + SessionData & { + body: string; + errorType?: ResponseState; + }; + +/** + * Options for the {@link BedrockAgentFunctionResolver} class + */ +type ResolverOptions = { + /** + * A logger instance to be used for logging debug, warning, and error messages. + * + * When no logger is provided, we'll only log warnings and errors using the global `console` object. + */ + logger?: GenericLogger; +}; + +export type { + Configuration, + Tool, + ToolFunction, + Parameter, + Attributes, + FunctionIdentifier, + FunctionInvocation, + BedrockAgentFunctionEvent, + BedrockAgentFunctionResponse, + ResponseOptions, + ResolverOptions, +}; diff --git a/packages/event-handler/src/types/common.ts b/packages/event-handler/src/types/common.ts new file mode 100644 index 000000000..a3f71b397 --- /dev/null +++ b/packages/event-handler/src/types/common.ts @@ -0,0 +1,15 @@ +// biome-ignore lint/suspicious/noExplicitAny: We intentionally use `any` here to represent any type of data and keep the logger is as flexible as possible. +type Anything = any; + +/** + * Interface for a generic logger object. + */ +type GenericLogger = { + trace?: (...content: Anything[]) => void; + debug: (...content: Anything[]) => void; + info?: (...content: Anything[]) => void; + warn: (...content: Anything[]) => void; + error: (...content: Anything[]) => void; +}; + +export type { Anything, GenericLogger }; diff --git a/packages/event-handler/src/types/index.ts b/packages/event-handler/src/types/index.ts index 424189e05..c0fe3d44e 100644 --- a/packages/event-handler/src/types/index.ts +++ b/packages/event-handler/src/types/index.ts @@ -8,3 +8,22 @@ export type { RouteOptions, RouterOptions, } from './appsync-events.js'; + +export type { + Configuration, + Tool, + ToolFunction, + Parameter, + Attributes, + FunctionIdentifier, + FunctionInvocation, + BedrockAgentFunctionEvent, + BedrockAgentFunctionResponse, + ResponseOptions, + ResolverOptions, +} from './bedrock-agent-function.js'; + +export type { + GenericLogger, + Anything, +} from './common.js'; 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 000000000..8a7386bf0 --- /dev/null +++ b/packages/event-handler/tests/unit/BedrockAgentFunctionResolver.test.ts @@ -0,0 +1,457 @@ +import context from '@aws-lambda-powertools/testing-utils/context'; +import type { Context } from 'aws-lambda'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { BedrockAgentFunctionEvent, Parameter } from '../../src/types'; +import { BedrockAgentFunctionResolver } from '../../src/bedrock-agent-function'; + +function createEvent(functionName: string, parameters?: Parameter[]) { + return { + messageVersion: '1.0', + agent: { + name: 'agentName', + id: 'agentId', + alias: 'agentAlias', + version: '1', + }, + sessionId: 'sessionId', + inputText: 'inputText', + function: functionName, + parameters, + actionGroup: 'myActionGroup', + sessionAttributes: {}, + promptSessionAttributes: {}, + }; +} + +describe('Class: BedrockAgentFunctionResolver', () => { + beforeEach(() => { + vi.unstubAllEnvs(); + }); + + it('uses a default logger with only warnings if none is provided', () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + app.tool(async (_params) => {}, { + name: 'noop', + definition: 'Does nothing', + }); + + // Assess + expect(console.debug).not.toHaveBeenCalled(); + }); + + it('emits debug message when AWS_LAMBDA_LOG_LEVEL is set to DEBUG', () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_LOG_LEVEL', 'DEBUG'); + const app = new BedrockAgentFunctionResolver(); + + app.tool(async (_params) => {}, { + name: 'noop', + definition: 'Does nothing', + }); + + // Assess + expect(console.debug).toHaveBeenCalled(); + }); + + it('accepts custom logger', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_LOG_LEVEL', 'DEBUG'); + + const logger = { + debug: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + const app = new BedrockAgentFunctionResolver({ logger }); + + app.tool(async (_params) => {}, { + name: 'noop', + definition: 'Does nothing', + }); + + await app.resolve(createEvent('noop', []), context); + + // Assess + expect(logger.debug).toHaveBeenCalled(); + expect(logger.debug).toHaveBeenCalled(); + expect(logger.debug).toHaveBeenCalled(); + }); + + it('only allows five tools to be registered', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + for (const num of [1, 2, 3, 4, 5]) { + app.tool(async (_params) => {}, { + name: `noop${num}`, + definition: 'Does nothing', + }); + } + + app.tool(async (params) => { + return Number(params.a) + Number(params); + }, { + name: 'mult', + definition: 'Multiplies two numbers', + }); + + expect(console.warn).toHaveBeenLastCalledWith( + 'The maximum number of tools that can be registered is 5. Tool mult will not be registered.' + ); + + const event = createEvent('mult', [ + { + name: 'a', + type: 'number', + value: '1', + }, + { + name: 'b', + type: 'number', + value: '2', + }, + ]); + + const actual = await app.resolve(event, context); + + expect(actual.response.function).toEqual('mult'); + expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual( + 'Error: tool has not been registered in handler.' + ); + }); + + it('overwrites tools with the same name and uses the latest definition', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + const event = createEvent('math', [ + { + name: 'a', + type: 'number', + value: '10', + }, + { + name: 'b', + type: 'number', + value: '2', + }, + ]); + + app.tool( + async (params) => { + return Number(params.a) + Number(params.b); + }, + { + name: 'math', + definition: 'Adds two numbers', + } + ); + + const addResult = await app.resolve(event, context); + expect(addResult.response.function).toEqual('math'); + expect(addResult.response.functionResponse.responseBody.TEXT.body).toEqual( + '12' + ); + + app.tool( + async (params) => { + return Number(params.a) * Number(params.b); + }, + { + name: 'math', + definition: 'Multiplies two numbers', + } + ); + + const multiplyResult = await app.resolve(event, context); + expect(multiplyResult.response.function).toEqual('math'); + expect( + multiplyResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('20'); + + app.tool( + async (params) => { + return Number(params.a) / Number(params.b); + }, + { + name: 'math', + definition: 'Divides two numbers', + } + ); + + const divideResult = await app.resolve(event, context); + expect(divideResult.response.function).toEqual('math'); + expect( + divideResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('5'); + + expect(console.warn).toHaveBeenCalledTimes(3); + expect(console.warn).toHaveBeenCalledWith( + 'Tool math already registered. Overwriting with new definition.' + ); + }); + + it('can be invoked using the decorator pattern', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + class Lambda { + @app.tool({ name: 'hello', definition: 'Says hello' }) + async helloWorld() { + return 'Hello, world!'; + } + + @app.tool({ name: 'add', definition: 'Adds two numbers' }) + async add(params: { a: string; b: string }) { + const { a, b } = params; + return Number.parseInt(a) + Number.parseInt(b); + } + + public async handler(event: BedrockAgentFunctionEvent, context: Context) { + return app.resolve(event, context); + } + } + + const lambda = new Lambda(); + + const helloEvent = createEvent('hello'); + + const helloResult = await lambda.handler(helloEvent, context); + expect(helloResult.response.function).toEqual('hello'); + expect( + helloResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('Hello, world!'); + + const addEvent = createEvent('add', [ + { + name: 'a', + type: 'number', + value: '1', + }, + { + name: 'b', + type: 'number', + value: '2', + }, + ]); + + const addResult = await lambda.handler(addEvent, context); + expect(addResult.response.function).toEqual('add'); + expect(addResult.response.functionResponse.responseBody.TEXT.body).toEqual( + '3' + ); + }); + + it('handles functions that return different primitive types', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + app.tool(async () => 'Hello, world!', { + name: 'string-tool', + definition: 'Returns a string', + }); + + app.tool(async () => 42, { + name: 'number-tool', + definition: 'Returns a number', + }); + + app.tool(async () => true, { + name: 'boolean-tool', + definition: 'Returns a boolean', + }); + + const stringResult = await app.resolve(createEvent('string-tool'), context); + expect(stringResult.response.function).toEqual('string-tool'); + expect( + stringResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('Hello, world!'); + + const numberResult = await app.resolve(createEvent('number-tool'), context); + expect(numberResult.response.function).toEqual('number-tool'); + expect( + numberResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('42'); + + const booleanResult = await app.resolve( + createEvent('boolean-tool'), + context + ); + expect(booleanResult.response.function).toEqual('boolean-tool'); + expect( + booleanResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('true'); + }); + + it('handles functions that return complex types (array and object)', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + app.tool(async () => [1, 'two', false, null], { + name: 'array-tool', + definition: 'Returns an array', + }); + + app.tool( + async () => ({ + name: 'John Doe', + age: 30, + isActive: true, + address: { + street: '123 Main St', + city: 'Anytown', + }, + }), + { + name: 'object-tool', + definition: 'Returns an object', + } + ); + + const arrayResult = await app.resolve(createEvent('array-tool'), context); + expect(arrayResult.response.function).toEqual('array-tool'); + expect( + arrayResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('[1,"two",false,null]'); + + const objectResult = await app.resolve(createEvent('object-tool'), context); + expect(objectResult.response.function).toEqual('object-tool'); + expect( + objectResult.response.functionResponse.responseBody.TEXT.body + ).toEqual( + '{"name":"John Doe","age":30,"isActive":true,"address":{"street":"123 Main St","city":"Anytown"}}' + ); + }); + + it('handles functions that return null or undefined by returning a string', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + app.tool(async () => null, { + name: 'null-tool', + definition: 'Returns null', + }); + + app.tool(async () => undefined, { + name: 'undefined-tool', + definition: 'Returns undefined', + }); + + app.tool(async () => {}, { + name: 'no-return-tool', + definition: 'Has no return statement', + }); + + const nullResult = await app.resolve(createEvent('null-tool'), context); + expect(nullResult.response.function).toEqual('null-tool'); + expect(nullResult.response.functionResponse.responseBody.TEXT.body).toEqual( + '' + ); + + const undefinedResult = await app.resolve( + createEvent('undefined-tool'), + context + ); + expect(undefinedResult.response.function).toEqual('undefined-tool'); + expect( + undefinedResult.response.functionResponse.responseBody.TEXT.body + ).toEqual(''); + + const noReturnResult = await app.resolve( + createEvent('no-return-tool'), + context + ); + expect(noReturnResult.response.function).toEqual('no-return-tool'); + expect( + noReturnResult.response.functionResponse.responseBody.TEXT.body + ).toEqual(''); + }); + + it('handles functions that throw errors', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + app.tool( + async () => { + throw new Error('Something went wrong'); + }, + { + name: 'error-tool', + definition: 'Throws an error', + } + ); + + const errorResult = await app.resolve( + createEvent('error-tool', []), + context + ); + expect(errorResult.response.function).toEqual('error-tool'); + expect( + errorResult.response.functionResponse.responseBody.TEXT.body + ).toEqual('Error when invoking tool: Error: Something went wrong'); + expect(console.error).toHaveBeenCalledWith( + 'An error occurred in tool error-tool.', + new Error('Something went wrong') + ); + }); + + it('returns a fully structured BedrockAgentFunctionResponse', async () => { + // Prepare + const app = new BedrockAgentFunctionResolver(); + + // Register a tool that returns a simple value + app.tool( + async (params) => { + return `Hello, ${params.name}!`; + }, + { + name: 'greeting', + definition: 'Greets a person by name', + } + ); + + // Define custom session attributes and parameters + const customSessionAttrs = { + sessionAttr: '12345', + }; + + const customPromptAttrs = { + promptAttr: 'promptAttr', + }; + + // Create a custom event with session attributes + const customEvent = { + ...createEvent('greeting', [ + { + name: 'name', + type: 'string', + value: 'John', + }, + ]), + actionGroup: 'actionGroup', + sessionAttributes: customSessionAttrs, + promptSessionAttributes: customPromptAttrs, + }; + + const result = await app.resolve(customEvent, context); + + expect(result).toEqual({ + messageVersion: '1.0', + response: { + actionGroup: 'actionGroup', + function: 'greeting', + functionResponse: { + responseBody: { + TEXT: { + body: 'Hello, John!', + }, + }, + }, + }, + sessionAttributes: customSessionAttrs, + promptSessionAttributes: customPromptAttrs, + }); + }); +});