-
Notifications
You must be signed in to change notification settings - Fork 155
feat(event-handler): add Amazon Bedrock Agents Functions Resolver #3957
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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<string, Tool> = new Map<string, Tool>(); | ||
readonly #envService: EnvironmentVariablesService; | ||
readonly #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); | ||
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]) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. One of the consequences of not having vaildation using the Zod parser is that all these arguments will be strings because they are coming from the |
||
); | ||
|
||
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, | ||
}); | ||
} | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
export { BedrockAgentFunctionResolver } from './BedrockAgentFunctionResolver.js'; |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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' | ||
); | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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>> = { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. As mentiioned, technically this function with only take strings because we are not using Zod to parse but in the future we will need the type to be open to accepting any arguments. Edit: this type is too lose even if we had parsing. According to the docs the only supported types are the following:
|
||
// 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>; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We use the type system here to only allow values to be returned that can be serialised to JSON. |
||
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, | ||
}; |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 }; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Take particular notice of this part, I've never worked with decorators before so any feedback here is particularly welcome.