Skip to content

Commit ecc2470

Browse files
dreamorosidreamorosi
committed
chore(event-handler): align implementation
1 parent 7e23468 commit ecc2470

File tree

4 files changed

+302
-165
lines changed

4 files changed

+302
-165
lines changed

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

Lines changed: 96 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,34 @@
1-
import { EnvironmentVariablesService } from '@aws-lambda-powertools/commons';
1+
import { isNullOrUndefined } from '@aws-lambda-powertools/commons/typeutils';
2+
import { getStringFromEnv } from '@aws-lambda-powertools/commons/utils/env';
23
import type { Context } from 'aws-lambda';
34
import type {
45
BedrockAgentFunctionResponse,
56
Configuration,
67
ParameterValue,
78
ResolverOptions,
8-
ResponseOptions,
99
Tool,
1010
ToolFunction,
1111
} from '../types/bedrock-agent.js';
1212
import type { GenericLogger } from '../types/common.js';
13+
import { BedrockFunctionResponse } from './BedrockFunctionResponse.js';
1314
import { assertBedrockAgentFunctionEvent } from './utils.js';
1415

16+
/**
17+
* Resolver for AWS Bedrock Agent Function invocations.
18+
*
19+
* This resolver is designed to handle function invocations from Bedrock Agents.
20+
*
21+
*
22+
*/
1523
export class BedrockAgentFunctionResolver {
1624
readonly #tools: Map<string, Tool> = new Map();
17-
readonly #envService: EnvironmentVariablesService;
1825
readonly #logger: Pick<GenericLogger, 'debug' | 'warn' | 'error'>;
1926

2027
constructor(options?: ResolverOptions) {
21-
this.#envService = new EnvironmentVariablesService();
22-
const alcLogLevel = this.#envService.get('AWS_LAMBDA_LOG_LEVEL');
28+
const alcLogLevel = getStringFromEnv({
29+
key: 'AWS_LAMBDA_LOG_LEVEL',
30+
defaultValue: '',
31+
});
2332
this.#logger = options?.logger ?? {
2433
debug: alcLogLevel === 'DEBUG' ? console.debug : () => {},
2534
error: console.error,
@@ -34,7 +43,9 @@ export class BedrockAgentFunctionResolver {
3443
*
3544
* @example
3645
* ```ts
37-
* import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/bedrock-agent-function';
46+
* import {
47+
* BedrockAgentFunctionResolver
48+
* } from '@aws-lambda-powertools/event-handler/bedrock-agent';
3849
*
3950
* const app = new BedrockAgentFunctionResolver();
4051
*
@@ -50,11 +61,36 @@ export class BedrockAgentFunctionResolver {
5061
* app.resolve(event, context);
5162
* ```
5263
*
64+
* If you know the function signature, you can also use a type parameter to specify the parameters of the tool function:
65+
*
66+
* @example
67+
* ```ts
68+
* import {
69+
* BedrockAgentFunctionResolver,
70+
* } from '@aws-lambda-powertools/event-handler/bedrock-agent';
71+
*
72+
* const app = new BedrockAgentFunctionResolver();
73+
*
74+
* app.tool<{ name: string }>(async (params) => {
75+
* const { name } = params;
76+
* // ^ name: string
77+
* return `Hello, ${name}!`;
78+
* }, {
79+
* name: 'greeting',
80+
* description: 'Greets a person by name',
81+
* });
82+
*
83+
* export const handler = async (event, context) =>
84+
* app.resolve(event, context);
85+
* ```
86+
*
5387
* The method also works as a class method decorator:
5488
*
5589
* @example
5690
* ```ts
57-
* import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/bedrock-agent-function';
91+
* import {
92+
* BedrockAgentFunctionResolver
93+
* } from '@aws-lambda-powertools/event-handler/bedrock-agent';
5894
*
5995
* const app = new BedrockAgentFunctionResolver();
6096
*
@@ -74,8 +110,38 @@ export class BedrockAgentFunctionResolver {
74110
* export const handler = lambda.handler.bind(lambda);
75111
* ```
76112
*
113+
* When defining a tool, you can also access the original `event` and `context` objects from the Bedrock Agent function invocation.
114+
* This is useful if you need to access the session attributes or other context-specific information.
115+
*
116+
* @example
117+
* ```ts
118+
* import {
119+
* BedrockAgentFunctionResolver
120+
* } from '@aws-lambda-powertools/event-handler/bedrock-agent';
121+
*
122+
* const app = new BedrockAgentFunctionResolver();
123+
*
124+
* app.tool(async (params, { event, context }) => {
125+
* const { name } = params;
126+
* // Access session attributes from the event
127+
* const sessionAttributes = event.sessionAttributes || {};
128+
* // You can also access the context if needed
129+
* sessionAttributes.requestId = context.awsRequestId;
130+
*
131+
* return `Hello, ${name}!`;
132+
* }, {
133+
* name: 'greetingWithContext',
134+
* description: 'Greets a person by name',
135+
* });
136+
*
137+
* export const handler = async (event, context) =>
138+
* app.resolve(event, context);
139+
* ```
140+
*
77141
* @param fn - The tool function
78142
* @param config - The configuration object for the tool
143+
* @param config.name - The name of the tool, which must be unique across all registered tools.
144+
* @param config.description - A description of the tool, which is optional but highly recommended.
79145
*/
80146
public tool<TParams extends Record<string, ParameterValue>>(
81147
fn: ToolFunction<TParams>,
@@ -107,54 +173,17 @@ export class BedrockAgentFunctionResolver {
107173
config: Configuration
108174
): void {
109175
const { name } = config;
110-
111-
if (this.#tools.size >= 5) {
112-
this.#logger.warn(
113-
`The maximum number of tools that can be registered is 5. Tool ${name} will not be registered.`
114-
);
115-
return;
116-
}
117-
118176
if (this.#tools.has(name)) {
119177
this.#logger.warn(
120-
`Tool ${name} already registered. Overwriting with new definition.`
178+
`Tool "${name}" already registered. Overwriting with new definition.`
121179
);
122180
}
123181

124182
this.#tools.set(name, {
125183
handler: handler as ToolFunction,
126184
config,
127185
});
128-
this.#logger.debug(`Tool ${name} has been registered.`);
129-
}
130-
131-
#buildResponse(options: ResponseOptions): BedrockAgentFunctionResponse {
132-
const {
133-
actionGroup,
134-
function: func,
135-
body,
136-
errorType,
137-
sessionAttributes,
138-
promptSessionAttributes,
139-
} = options;
140-
141-
return {
142-
messageVersion: '1.0',
143-
response: {
144-
actionGroup,
145-
function: func,
146-
functionResponse: {
147-
responseState: errorType,
148-
responseBody: {
149-
TEXT: {
150-
body,
151-
},
152-
},
153-
},
154-
},
155-
sessionAttributes,
156-
promptSessionAttributes,
157-
};
186+
this.#logger.debug(`Tool "${name}" has been registered.`);
158187
}
159188

160189
async resolve(
@@ -175,11 +204,13 @@ export class BedrockAgentFunctionResolver {
175204

176205
if (tool == null) {
177206
this.#logger.error(`Tool ${toolName} has not been registered.`);
178-
return this.#buildResponse({
207+
return new BedrockFunctionResponse({
179208
actionGroup,
180-
function: toolName,
181-
body: 'Error: tool has not been registered in handler.',
182-
});
209+
func: toolName,
210+
body: `Error: tool ${toolName} has not been registered.`,
211+
sessionAttributes,
212+
promptSessionAttributes,
213+
}).build();
183214
}
184215

185216
const toolParams: Record<string, ParameterValue> = {};
@@ -204,24 +235,30 @@ export class BedrockAgentFunctionResolver {
204235
}
205236

206237
try {
207-
const res = await tool.handler(toolParams, { event, context });
208-
const body = res == null ? '' : JSON.stringify(res);
209-
return this.#buildResponse({
238+
const response = await tool.handler(toolParams, { event, context });
239+
if (response instanceof BedrockFunctionResponse) {
240+
return response.build();
241+
}
242+
const body =
243+
isNullOrUndefined(response) || response === ''
244+
? ''
245+
: JSON.stringify(response);
246+
return new BedrockFunctionResponse({
210247
actionGroup,
211-
function: toolName,
248+
func: toolName,
212249
body,
213250
sessionAttributes,
214251
promptSessionAttributes,
215-
});
252+
}).build();
216253
} catch (error) {
217254
this.#logger.error(`An error occurred in tool ${toolName}.`, error);
218-
return this.#buildResponse({
255+
return new BedrockFunctionResponse({
219256
actionGroup,
220-
function: toolName,
221-
body: `Error when invoking tool: ${error}`,
257+
func: toolName,
258+
body: `Unable to complete tool execution due to ${error instanceof Error ? `${error.name} - ${error.message}` : String(error)}`,
222259
sessionAttributes,
223260
promptSessionAttributes,
224-
});
261+
}).build();
225262
}
226263
}
227264
}

‎packages/event-handler/src/bedrock-agent/BedrockFunctionResponse.ts

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import type { BedrockAgentFunctionResolver } from './BedrockAgentFunctionResolver.js';
2+
3+
/**
4+
* Class representing a response from a Bedrock agent function.
5+
*
6+
* You can use this class to customize the response sent back to the Bedrock agent with additional fields like:
7+
* - session attributes
8+
* - prompt session attributes
9+
* - response state (`FAILURE` or `REPROMPT`)
10+
*
11+
* When working with the {@link BedrockAgentFunctionResolver} class, this is built automatically
12+
* when you return anything from your function handler other than an instance of this class.
13+
*/
14+
class BedrockFunctionResponse {
15+
/**
16+
* The name of the action group, this comes from the `event.actionGroup` field
17+
* in the Bedrock agent function event.
18+
*/
19+
readonly actionGroup: string;
20+
/**
21+
* The name of the function returning the response, this comes from the `event.function` field
22+
* in the Bedrock agent function event.
23+
*/
24+
readonly func: string;
25+
/**
26+
* The response object that defines the response from execution of the function.
27+
*/
28+
readonly body: string;
29+
/**
30+
* Optional field to indicate the whether the response is a failure or a reprompt.
31+
* If not provided, the default is undefined, which means no specific response state is set.
32+
*
33+
* - `FAILURE`: The agent throws a `DependencyFailedException` for the current session.
34+
* - `REPROMPT`: The agent passes a response string to the model to reprompt it.
35+
*/
36+
readonly responseState?: 'FAILURE' | 'REPROMPT';
37+
/**
38+
* Optional field to store session attributes and their values.
39+
* @see {@link https://docs.aws.amazon.com/bedrock/latest/userguide/agents-session-state.html#session-state-attributes | Bedrock Agent Session State Attributes} for more details.
40+
*/
41+
readonly sessionAttributes: Record<string, string>;
42+
/**
43+
* Optional field to instruct the agent to prompt attributes and their values.
44+
* @see {@link https://docs.aws.amazon.com/bedrock/latest/userguide/agents-session-state.html#session-state-attributes | Bedrock Agent Session State Attributes} for more details.
45+
*/
46+
readonly promptSessionAttributes: Record<string, string>;
47+
48+
constructor({
49+
actionGroup,
50+
func,
51+
body,
52+
responseState = undefined,
53+
sessionAttributes = {},
54+
promptSessionAttributes = {},
55+
}: {
56+
actionGroup: string;
57+
func: string;
58+
body: string;
59+
responseState?: 'FAILURE' | 'REPROMPT';
60+
sessionAttributes?: Record<string, string>;
61+
promptSessionAttributes?: Record<string, string>;
62+
}) {
63+
this.actionGroup = actionGroup;
64+
this.func = func;
65+
this.body = body;
66+
this.responseState = responseState;
67+
this.sessionAttributes = sessionAttributes;
68+
this.promptSessionAttributes = promptSessionAttributes;
69+
}
70+
71+
/**
72+
* Builds the Bedrock function response object according to the Bedrock agent function {@link https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html#agents-lambda-response | response format}.
73+
*/
74+
build() {
75+
return {
76+
messageVersion: '1.0',
77+
response: {
78+
actionGroup: this.actionGroup,
79+
function: this.func,
80+
functionResponse: {
81+
...(this.responseState && { responseState: this.responseState }),
82+
responseBody: {
83+
TEXT: {
84+
body: this.body,
85+
},
86+
},
87+
},
88+
},
89+
...(this.sessionAttributes && {
90+
sessionAttributes: this.sessionAttributes,
91+
}),
92+
...(this.promptSessionAttributes && {
93+
promptSessionAttributes: this.promptSessionAttributes,
94+
}),
95+
};
96+
}
97+
}
98+
99+
export { BedrockFunctionResponse };

0 commit comments

Comments
 (0)