wip(event-handler): add Amazon Bedrock Agents Functions Resolver

Author: svozza

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"

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 {
+  #tools: Map<string, Tool> = new Map<string, Tool>();
+  #envService: EnvironmentVariablesService;
+  #logger: Pick<GenericLogger, 'debug' | 'warn' | 'error'>;
+
+  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 as Configuration);
+      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<BedrockAgentFunctionResponse> {
+    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<string, string> = 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,
+      });
+    }
+  }
+}

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

@@ -0,0 +1 @@
+export { BedrockAgentFunctionResolver } from './BedrockAgentFunctionResolver.js';

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'
+  );
+}

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
 

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

@@ -0,0 +1,102 @@
+import type { JSONValue } from '@aws-lambda-powertools/commons/types';
+import type { GenericLogger } from '../types/common.js';
+
+type Configuration = {
+  name: string;
+  definition: string;
+};
+
+// biome-ignore lint/suspicious/noExplicitAny: this is a generic type that is intentionally open
+type Tool<TParams = Record<string, any>> = {
+  // 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>;
+  config: Configuration;
+};
+
+type ToolFunction = Tool['function'];
+
+type Parameter = {
+  name: string;
+  type: string;
+  value: string;
+};
+
+type Attributes = Record<string, string>;
+
+type FunctionIdentifier = {
+  actionGroup: string;
+  function: string;
+};
+
+type FunctionInvocation = FunctionIdentifier & {
+  parameters?: Array<Parameter>;
+};
+
+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,
+};

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 };