make tool method generic for type Record<string, ParameterValue>>

Author: svozza

packages/event-handler/src/bedrock-agent-function/BedrockAgentFunctionResolver.ts

@@ -4,6 +4,7 @@ import type {
   BedrockAgentFunctionEvent,
   BedrockAgentFunctionResponse,
   Configuration,
+  ParameterValue,
   ResolverOptions,
   ResponseOptions,
   Tool,
@@ -13,7 +14,8 @@ import type { GenericLogger } from '../types/common.js';
 import { isPrimitive } from './utils.js';
 
 export class BedrockAgentFunctionResolver {
-  readonly #tools: Map<string, Tool> = new Map<string, Tool>();
+  readonly #tools: Map<string, Tool<Record<string, ParameterValue>>> =
+    new Map();
   readonly #envService: EnvironmentVariablesService;
   readonly #logger: Pick<GenericLogger, 'debug' | 'warn' | 'error'>;
 
@@ -77,10 +79,15 @@ export class BedrockAgentFunctionResolver {
    * @param fn - The tool function
    * @param config - The configuration object for the tool
    */
-  public tool(fn: ToolFunction, config: Configuration): undefined;
-  public tool(config: Configuration): MethodDecorator;
-  public tool(
-    fnOrConfig: ToolFunction | Configuration,
+  public tool<TParams extends Record<string, ParameterValue>>(
+    fn: ToolFunction<TParams>,
+    config: Configuration
+  ): undefined;
+  public tool<TParams extends Record<string, ParameterValue>>(
+    config: Configuration
+  ): MethodDecorator;
+  public tool<TParams extends Record<string, ParameterValue>>(
+    fnOrConfig: ToolFunction<TParams> | Configuration,
     config?: Configuration
   ): MethodDecorator | undefined {
     // When used as a method (not a decorator)
@@ -97,7 +104,10 @@ export class BedrockAgentFunctionResolver {
     };
   }
 
-  #registerTool(fn: ToolFunction, config: Configuration): void {
+  #registerTool<TParams extends Record<string, ParameterValue>>(
+    handler: ToolFunction<TParams>,
+    config: Configuration
+  ): void {
     const { name } = config;
 
     if (this.#tools.size >= 5) {
@@ -113,7 +123,10 @@ export class BedrockAgentFunctionResolver {
       );
     }
 
-    this.#tools.set(name, { function: fn, config });
+    this.#tools.set(name, {
+      handler: handler as ToolFunction<Record<string, ParameterValue>>,
+      config,
+    });
     this.#logger.debug(`Tool ${name} has been registered.`);
   }
 
@@ -169,12 +182,29 @@ export class BedrockAgentFunctionResolver {
       });
     }
 
-    const parameterObject: Record<string, string> = Object.fromEntries(
-      parameters.map((param) => [param.name, param.value])
-    );
+    const toolParams: Record<string, ParameterValue> = {};
+    for (const param of parameters) {
+      switch (param.type) {
+        case 'boolean': {
+          toolParams[param.name] = param.value === 'true';
+          break;
+        }
+        case 'number':
+        case 'integer': {
+          toolParams[param.name] = Number(param.value);
+          break;
+        }
+        // this default will also catch array types but we leave them as strings
+        // because we cannot reliably parse them
+        default: {
+          toolParams[param.name] = param.value;
+          break;
+        }
+      }
+    }
 
     try {
-      const res = (await tool.function(parameterObject)) ?? '';
+      const res = (await tool.handler(toolParams, event, context)) ?? '';
       const body = isPrimitive(res) ? String(res) : JSON.stringify(res);
       return this.#buildResponse({
         actionGroup,

packages/event-handler/src/types/bedrock-agent-function.ts

@@ -1,4 +1,5 @@
 import type { JSONValue } from '@aws-lambda-powertools/commons/types';
+import type { Context } from 'aws-lambda';
 import type { GenericLogger } from '../types/common.js';
 
 type Configuration = {
@@ -16,14 +17,18 @@ type ParameterPrimitives = string | number | boolean;
 
 type ParameterValue = ParameterPrimitives | Array<ParameterValue>;
 
-type Tool<TParams = Record<string, ParameterValue>> = {
+type ToolFunction<TParams = Record<string, ParameterValue>> = (
+  params: TParams,
+  event: BedrockAgentFunctionEvent,
+  context: Context
   // biome-ignore lint/suspicious/noConfusingVoidType: we need to support async functions that don't have an explicit return value
-  function: (params: TParams) => Promise<JSONValue | void>;
+) => Promise<JSONValue | void>;
+
+type Tool<TParams = Record<string, ParameterValue>> = {
+  handler: ToolFunction<TParams>;
   config: Configuration;
 };
 
-type ToolFunction = Tool['function'];
-
 type Attributes = Record<string, string>;
 
 type FunctionIdentifier = {
@@ -96,6 +101,7 @@ export type {
   ToolFunction,
   Parameter,
   Attributes,
+  ParameterValue,
   FunctionIdentifier,
   FunctionInvocation,
   BedrockAgentFunctionEvent,

packages/event-handler/tests/unit/bedrock-agent/BedrockAgentFunctionResolver.test.ts

@@ -2,7 +2,13 @@ import context from '@aws-lambda-powertools/testing-utils/context';
 import type { Context } from 'aws-lambda';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { BedrockAgentFunctionResolver } from '../../../src/bedrock-agent-function/index.js';
-import type { BedrockAgentFunctionEvent, Parameter } from '../../../src/types';
+import type {
+  BedrockAgentFunctionEvent,
+  Configuration,
+  Parameter,
+  ParameterValue,
+  ToolFunction,
+} from '../../../src/types/bedrock-agent-function.js';
 
 function createEvent(functionName: string, parameters?: Parameter[]) {
   return {
@@ -178,6 +184,32 @@ describe('Class: BedrockAgentFunctionResolver', () => {
     ).toEqual('20');
   });
 
+  it('tool function has access to the event variable', async () => {
+    // Prepare
+    const app = new BedrockAgentFunctionResolver();
+
+    app.tool(
+      async (_params, event) => {
+        return event;
+      },
+      {
+        name: 'event-accessor',
+        description: 'Accesses the event object',
+      }
+    );
+
+    const event = createEvent('event-accessor');
+
+    // Act
+    const result = await app.resolve(event, context);
+
+    // Assess
+    expect(result.response.function).toEqual('event-accessor');
+    expect(result.response.functionResponse.responseBody.TEXT.body).toEqual(
+      JSON.stringify(event)
+    );
+  });
+
   it('can be invoked using the decorator pattern', async () => {
     // Prepare
     const app = new BedrockAgentFunctionResolver();
@@ -327,6 +359,92 @@ describe('Class: BedrockAgentFunctionResolver', () => {
     }
   );
 
+  it.each<{
+    toolFunction: ToolFunction<Record<string, ParameterValue>>;
+    toolParams: Configuration;
+    parameters: Parameter[];
+    expected: string;
+  }>([
+    {
+      toolFunction: async (params: Record<string, ParameterValue>) =>
+        params.arg,
+      toolParams: {
+        name: 'boolean-handler',
+        description: 'Handles boolean parameters',
+      },
+      parameters: [
+        { name: 'arg', type: 'boolean', value: 'true' },
+      ] as Parameter[],
+      expected: 'true',
+    },
+    {
+      toolFunction: async (params: Record<string, ParameterValue>) =>
+        (params.arg as number) + 10,
+      toolParams: {
+        name: 'number-handler',
+        description: 'Handles number parameters',
+      },
+      parameters: [{ name: 'arg', type: 'number', value: '42' }] as Parameter[],
+      expected: '52',
+    },
+    {
+      toolFunction: async (params: Record<string, ParameterValue>) =>
+        (params.arg as number) + 10,
+      toolParams: {
+        name: 'integer-handler',
+        description: 'Handles integer parameters',
+      },
+      parameters: [
+        { name: 'arg', type: 'integer', value: '37' },
+      ] as Parameter[],
+      expected: '47',
+    },
+    {
+      toolFunction: async (params: Record<string, ParameterValue>) =>
+        `String: ${params.arg}`,
+      toolParams: {
+        name: 'string-handler',
+        description: 'Handles string parameters',
+      },
+      parameters: [
+        { name: 'arg', type: 'string', value: 'hello world' },
+      ] as Parameter[],
+      expected: 'String: hello world',
+    },
+    {
+      toolFunction: async (params: Record<string, ParameterValue>) =>
+        `Array as string: ${params.arg}`,
+      toolParams: {
+        name: 'array-handler',
+        description: 'Handles array parameters (as string)',
+      },
+      parameters: [
+        { name: 'arg', type: 'array', value: '[1,2,3]' },
+      ] as Parameter[],
+      expected: 'Array as string: [1,2,3]',
+    },
+  ])(
+    'correctly parses $toolParams.name parameters',
+    async ({ toolFunction, toolParams, parameters, expected }) => {
+      // Prepare
+      const app = new BedrockAgentFunctionResolver();
+
+      app.tool(toolFunction, toolParams);
+
+      // Act
+      const actual = await app.resolve(
+        createEvent(toolParams.name, parameters),
+        context
+      );
+
+      // Assert
+      expect(actual.response.function).toEqual(toolParams.name);
+      expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
+        expected
+      );
+    }
+  );
+
   it('handles functions that throw errors', async () => {
     // Prepare
     const app = new BedrockAgentFunctionResolver();