Feature request: Amazon Bedrock Agents Functions

[dreamorosi]

Use case

Bedrock Agents allows you to define action groups for your agents in two ways: OpenAPI schemas, and direct function integration. This issue focuses on the latter.

As a customer I can create Bedrock Agents that have tools at their disposal. These tools, or functions, can be defined as AWS Lambda functions. One Lambda function can hold one or more tools and when looked at together, they are what makes an action group.

When I build a Lambda function with multiple tools in it, I am responsible for parsing the payload sent by Bedrock and, based on certain fields, call the corresponding tool in my code (aka the tool use). The response of this tool use is then returned by my Lambda function handler according to a specific format that Bedrock expects.

This can result in some degree of boilerplate code that I have to repeat for each action group, specifically:

• parsing/validating the incoming Bedrock Agent request payload
• handling the event using the correct tool/function
• building the response according to the response Bedrock Agent payload schema

Click to see Action Group example

import type { Context } from 'aws-lambda';

// Define tools
function getCurrentTime(): string {
  // Implementation goes here
  return new Date().toISOString();
}

function greetUser(name: string): string {
  // Implementation goes here
  return `Hello, ${name}!`;
}

function simpleCalculator(a: number, b: number, operation: string): number | string {
  // Implementation goes here
  switch (operation) {
    case 'add':
      return a + b;
    case 'subtract':
      return a - b;
    case 'multiply':
      return a * b;
    case 'divide':
      return b !== 0 ? a / b : "Error: Division by zero";
    default:
      throw new Error(`Unknown operation: ${operation}`);
  }
}

// Main Lambda handler
export const handler = (event: unknown, context: Context) => {
  const agent = event.agent;
  const actionGroup = event.actionGroup;
  const functionName = event.function;
  const parameters = event.parameters || [];

  let result: string | number;
  
  // Execute the appropriate function based on the 'function' parameter
  if (functionName === 'get_current_time') {
    result = getCurrentTime();
  } else if (functionName === 'greet_user') {
    const name = parameters.length > 0 ? parameters[0].value : "User";
    result = greetUser(name);
  } else if (functionName === 'simple_calculator') {
    if (parameters.length >= 3) {
      const a = Number.parseFloat(parameters[0].value);
      const b = Number.parseFloat(parameters[1].value);
      const operation = parameters[2].value;
      result = simpleCalculator(a, b, operation);
    } else {
      result = "Error: Insufficient parameters for simple_calculator";
    }
  } else {
    result = `Unknown function: ${functionName}`;
  }

  const responseBody = {
    TEXT: {
      body: String(result)
    }
  };

  const actionResponse = {
    actionGroup: actionGroup,
    function: functionName,
    functionResponse: {
      responseBody: responseBody
    }
  };

  const functionResponse = {
    response: actionResponse, 
    messageVersion: event.messageVersion
  };
  
  console.log(`Response: ${JSON.stringify(functionResponse)}`);
  
  return functionResponse;
}

As a customer I would like to abstract all of that and instead focus primarily on building the tools for my agents, which are the only part of this that is specific to my business.

Solution/User Experience

When paired with a Lambda function via action group, Bedrock sends and expects payloads of known shapes.

Click to see request payload example

{
    "messageVersion": "1.0",
    "agent": {
        "alias": "PROD",
        "name": "hr-assistant-function-def",
        "version": "1",
        "id": "1234abcd-56ef-78gh-90ij-klmn12345678"
    },
    "sessionId": "87654321-abcd-efgh-ijkl-mnop12345678",
    "sessionAttributes": {
        "employeeId": "EMP123",
        "department": "Engineering"
    },
    "promptSessionAttributes": {
        "lastInteraction": "2024-02-01T15:30:00Z",
        "requestType": "vacation"
    },
    "inputText": "I want to request vacation from March 15 to March 20",
    "actionGroup": "VacationsActionGroup",
    "function": "submitVacationRequest",
    "parameters": [{
        "employeeId": "EMP123",
        "startDate": "2024-03-15",
        "endDate": "2024-03-20",
        "vacationType": "annual"
    }]
}

Documentation

Click to response payload example

{
   "response":{
      "actionGroup":"SmarterAgentActionGroup",
      "function":"submitVacationRequest",
      "functionResponse":{
         "responseBody":{
            "TEXT":{
               "body":"Your vacation was scheduled!"
            }
         }
      }
   },
   "messageVersion":"1.0"
}

Documentation

Since the input event includes both the function and parameters fields, we can abstract most/all the boilerplate and provide a more streamlined experience.

For example, borrowing heavily from the Event Handler RFC we concluded a few weeks ago, we could implement a BedrockAgentFunctionResolver resolver that provides a structured way to register functions, resolve function calls, and handle requests within Lambda.

For reference, the example shown above could be rewritten like:

import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/resolvers/bedrock-agents-fn';
import { z } from 'zod';

const app = new BedrockAgentFunctionResolver();

app.tool(() => {
  // Implementation goes here
  return new Date().toISOString();
}, {
  name: 'getCurrentTime',
  description: 'Gets the current UTC time'
  validation: {
    input: z.never(),
    output: z.string(),
  }
});

app.tool(({ name }) => {
  // Implementation goes here
  return `Hello, ${name}!`;
}, {
  name: 'greetUser',
  description: 'Greets the user using name',
  validation: {
    input: z.object({
      name: z.string(),
    }),
    output: z.string(),
  }
});

app.tool(({ a, b, operation } ) => {
  // Implementation goes here
  switch (operation) {
    case 'add':
      return a + b;
    case 'subtract':
      return a - b;
    case 'multiply':
      return a * b;
    case 'divide':
      return b !== 0 ? a / b : "Error: Division by zero";
    default:
      throw new Error(`Unknown operation: ${operation}`);
}, {
  name: 'simpleCalculator',
  description: 'Calculates operations between two numbers',
  validation: {
    input: z.object({
      a: z.number(),
      b: z.number(),
      operation: z.union([
        z.literal('add'),
        z.literal('subtract'),
        z.literal('multiply'),
        z.literal('divide')
      ])
    }),
    output: z.number()
  }
});

export const handler = async (event, context) =>
  app.resolve(event, context);

The implementation above allows customers to focus on defining and implementing the tools within the action group rather than the undifferentiated boilerplate required to parse an event, resolve which tool to use, call it, and build the response payload.

Similar to what we do with the API Event Handler, customers can use the pattern above as well as the class method decorator-based pattern, which would look like this (simplified for brevity):

import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/resolvers/bedrock-agents-fn';
import { z } from 'zod';

const app = new BedrockAgentFunctionResolver();

class Lambda {
  @app.tool({
    name: 'getCurrentTime',
    description: 'Gets the current UTC time'
    validation: {
      input: z.never(),
      output: z.string(),
    }
  })
  public getCurrentTime() {
    // Implementation goes here
    return new Date().toISOString();
  }
  
  // ... other tools

  public async handler(event, context) {
    return app.resolve(event, context);
  }
}

const lambda = new Lambda();
export const handler = lambda.handler.bind(lambda);

Contrary to the API Event Handler, validation and parsing via Zod or other Standard Schema-compliant libraries cannot be disabled. This is because we want to ensure both inputs and outputs are always compliant to the expected shapes, but also because we want to allow customers to automatically generate the schema for the action group to create the resource more easily.

Finally, customers can also access the current request event and context in the tool definitions. This is useful when the tools rely on information included in the request like the session, or they need to access the context. The event and context are passed to the tool handler as second argument, i.e.:

app.tool(async (params, {
  event,
  context,
}) => {
  return event.sessionId
}, {
  name: 'echo session',
  description: 'returns the session id',
  validation: {
    input: z.any(),
    output: z.string(),
  }
});

---

In terms of practical directions, the resolver and its logic will be reside under packages/event-handler and be part of the event handler package. With that said, the implementation will be completely independent form the rest of the event handler at least for its initial version.

Alternative solutions

See this discussion for considerations on alternative solutions: aws-powertools/powertools-lambda-python#6081

Acknowledgment

• This feature request meets Powertools for AWS Lambda (TypeScript) Tenets
• Should this be considered in other Powertools for AWS Lambda languages? i.e. Python, Java, and .NET

Future readers

Please react with 👍 and your use case to help us understand customer demand.

[ConnorKirk]

I'll take a look at this

[ConnorKirk]

There are two additional parts I think it would be useful to consider in the RFC.

1. Errors
2. Session Attributes, Prompt Session Attributes and KnowledgeBaseConfiguration

1. Errors

When there is an unhandled exception in the execution, the Agent returns an DependencyFailedException to the invoking client.

The lambda function can also signify two possible error states to the Bedrock Agent with a property in the response object:

1. FAILURE - Similar to unhandled exceptions, the invoking client receives a DependencyFailedException.
2. REPROMPT - Used when the function execution fails because of invalid input.
   The error state is signified by setting the response.functionResponse.responseState property to either FAILURE or REPROMPT in the response object.

I think it would be useful for the user to be able to emit both of these error states when using the function resolver.

Bedrock Agents - Response Object

2. Session Attributes, Prompt Session Attributes and KnowledgeBaseConfiguration

The lambda response object can contain three extra properties, sessionAttributes, promptSessionAttributes and bedrockKnowledgeBaseConfiguration. These exist on the input object, and can be added, amended or removed in the response object.

It would be useful to support setting these on the response object when using the function resolver.
I need to play with this further to see if there are any constraints.

• Session Attributes
  • Prompt Session Attributes
  • Knowledge Bases Configuration

[anafalcao]

@ConnorKirk ! One question regarding:

1. Errors

When there is an unhandled exception in the execution, the Agent returns an DependencyFailedException to the invoking client.

What's the user experience you're thinking here? It's that the functionResponse.responseState will be automatically passed with FAILURE/REPROMPT to the output event when the unhandled exception occurs?
Or is it more as an extra attribute in the BedrockResponse?