Bedrock Agent Functions

class BedrockAgentFunctionResolver {
  constructor(options?: { logger?: Pick<GenericLogger, 'debug' | 'warn' | 'error'> });
  tool<TParams extends Record<string, ParameterValue>>(fn: ToolFunction<TParams>, config: { name: string; description?: string }): void;
  resolve(event, context): Promise<BedrockAgentFunctionResponse>;
}

import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/bedrock-agent';
const app = new BedrockAgentFunctionResolver();
export const handler = async (event, context) => app.resolve(event, context);

app.tool(async ({ city }) => {
  const weather = await getWeather(city);
  return { city, temperature: weather.temp, condition: weather.condition };
}, { name: 'getWeatherForCity', description: 'Get current weather for a specific city' });

app.tool<{ city: string }>(async ({ city }) => await fetchWeather(city), {
  name: 'getWeather',
  description: 'Get weather for a city',
});

app.tool<{ query: string; limit: number }>(async ({ query, limit }) =>
  await searchDatabase(query, limit), { name: 'searchData', description: 'Search the database' }
);

app.tool<{ message: string }>(async ({ message }, { event, context }) => {
  console.log('Session ID:', event.sessionId);
  console.log('Remaining time:', context.getRemainingTimeInMillis());
  return await processMessage(message);
}, { name: 'processMessage' });

app.tool(async () => await getCurrentTime(), {
  name: 'getCurrentTime',
  description: 'Get the current server time',
});

class BedrockFunctionResponse {
  constructor(options: {
    body: string;                                      // Required: JSON string response body
    responseState?: 'FAILURE' | 'REPROMPT';           // Optional: response state
    sessionAttributes?: Record<string, JSONValue>;     // Optional: session attributes to persist
    promptSessionAttributes?: Record<string, JSONValue>; // Optional: prompt session attributes
    knowledgeBasesConfiguration?: Record<string, JSONValue>; // Optional: knowledge base config
  });
  build(options: { actionGroup: string; func: string }): BedrockAgentFunctionResponse;
}

import { BedrockFunctionResponse } from '@aws-lambda-powertools/event-handler/bedrock-agent';

app.tool<{ city: string }>(async ({ city }, { event }) => {
  const weather = await getWeather(city);
  return new BedrockFunctionResponse({
    body: JSON.stringify({ city, temperature: weather.temp, condition: weather.condition }),
    sessionAttributes: { ...event.sessionAttributes, lastCity: city, weatherChecked: true },
  });
}, { name: 'getWeather', description: 'Get weather for a city' });

app.tool<{ query: string }>(async ({ query }) => {
  const result = await search(query);
  if (!result) {
    return new BedrockFunctionResponse({
      body: JSON.stringify({ error: 'No results found' }),
      responseState: 'FAILURE',
    });
  }
  return result;
}, { name: 'search' });

app.tool<{ input: string }>(async ({ input }) => {
  if (!isValid(input)) {
    return new BedrockFunctionResponse({
      body: JSON.stringify({ message: 'Input is unclear. Please provide more details.' }),
      responseState: 'REPROMPT',
    });
  }
  return await processInput(input);
}, { name: 'processInput' });

app.tool<{ question: string }>(async ({ question }, { event }) => {
  return new BedrockFunctionResponse({
    body: JSON.stringify({ question }),
    knowledgeBasesConfiguration: event.knowledgeBasesConfiguration,
  });
}, { name: 'askKnowledgeBase' });

app.tool<{ item: string }>(async ({ item }, { event }) => {
  const cart = event.sessionAttributes.cart ? JSON.parse(event.sessionAttributes.cart) : [];
  cart.push(item);
  return new BedrockFunctionResponse({
    body: JSON.stringify({ message: `Added ${item} to cart`, cartSize: cart.length }),
    sessionAttributes: { ...event.sessionAttributes, cart: JSON.stringify(cart) },
  });
}, { name: 'addToCart', description: 'Add an item to the shopping cart' });

app.tool(async (params, { event }) => {
  const cart = event.sessionAttributes.cart ? JSON.parse(event.sessionAttributes.cart) : [];
  return { items: cart, total: cart.length };
}, { name: 'viewCart', description: 'View shopping cart contents' });

app.tool<{ userId: string }>(async ({ userId }) => {
  const user = await getUser(userId);
  if (!user) throw new Error(`User ${userId} not found`);
  return user;
}, { name: 'getUser' });

app.tool<{ userId: string }>(async ({ userId }) => {
  const user = await getUser(userId);
  if (!user) {
    return new BedrockFunctionResponse({
      body: JSON.stringify({ error: `User ${userId} not found` }),
      responseState: 'FAILURE',
    });
  }
  return user;
}, { name: 'getUser' });

app.tool<{
  name: string;      // Remains string
  age: number;       // Parsed from string to number
  active: boolean;   // Parsed from "true"/"false" to boolean
  tags: string[];    // Parsed from string array
}>(async ({ name, age, active, tags }) => {
  console.log(typeof name);    // "string"
  console.log(typeof age);     // "number"
  console.log(typeof active);  // "boolean"
  console.log(Array.isArray(tags)); // true
  return { name, age: age + 1, active: !active, tagCount: tags.length };
}, { name: 'processUser' });

// Lambda function
app.tool(async ({ city }) => getWeather(city), {
  name: 'getWeather',  // Must match agent config
  description: 'Get weather for a city',
});

// Bedrock Agent configuration:
// {
//   "functions": [{
//     "name": "getWeather",
//     "description": "Get weather for a city",
//     "parameters": {
//       "city": { "description": "The city name", "required": true, "type": "string" }
//     }
//   }]
// }

interface BedrockAgentFunctionEvent {
  actionGroup: string;
  function: string;
  messageVersion: string;
  agent: { name: string; id: string; alias: string; version: string };
  parameters: Array<{ name: string; type: 'string' | 'number' | 'integer' | 'boolean' | 'array'; value: string }>;
  inputText: string;
  sessionId: string;
  sessionAttributes: Record<string, string>;
  promptSessionAttributes: Record<string, string>;
  knowledgeBasesConfiguration?: Record<string, JSONValue>[];
}

interface BedrockAgentFunctionResponse {
  messageVersion: string;
  response: {
    actionGroup: string;
    function: string;
    functionResponse: {
      responseBody: { TEXT: { body: string } };
      responseState?: 'FAILURE' | 'REPROMPT';
    };
  };
  sessionAttributes?: Record<string, JSONValue>;
  promptSessionAttributes?: Record<string, JSONValue>;
  knowledgeBasesConfiguration?: Record<string, JSONValue>;
}

type ToolFunction<TParams extends Record<string, ParameterValue> = Record<string, ParameterValue>> = (
  params: TParams,
  options: { event: BedrockAgentFunctionEvent; context: Context }
) => Promise<JSONValue | BedrockFunctionResponse> | JSONValue | BedrockFunctionResponse;

type ParameterValue = string | number | boolean | Array<ParameterValue>;
type ResponseState = 'FAILURE' | 'REPROMPT';