Skip to content

Commit 59db37b

Browse files
committed
wip(event-handler): add Amazon Bedrock Agents Functions Resolver
1 parent 0a4f3e6 commit 59db37b

File tree

9 files changed

+846
-16
lines changed

9 files changed

+846
-16
lines changed

packages/event-handler/package.json

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,16 @@
3939
"default": "./lib/esm/appsync-events/index.js"
4040
}
4141
},
42+
"./bedrock-agent-function": {
43+
"require": {
44+
"types": "./lib/cjs/bedrock-agent-function/index.d.ts",
45+
"default": "./lib/cjs/bedrock-agent-function/index.js"
46+
},
47+
"import": {
48+
"types": "./lib/esm/bedrock-agent-function/index.d.ts",
49+
"default": "./lib/esm/bedrock-agent-function/index.js"
50+
}
51+
},
4252
"./types": {
4353
"require": {
4454
"types": "./lib/cjs/types/index.d.ts",
@@ -56,6 +66,10 @@
5666
"./lib/cjs/appsync-events/index.d.ts",
5767
"./lib/esm/appsync-events/index.d.ts"
5868
],
69+
"bedrock-agent-function": [
70+
"./lib/cjs/bedrock-agent-function/index.d.ts",
71+
"./lib/esm/bedrock-agent-function/index.d.ts"
72+
],
5973
"types": [
6074
"./lib/cjs/types/index.d.ts",
6175
"./lib/esm/types/index.d.ts"
Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
import { EnvironmentVariablesService } from '@aws-lambda-powertools/commons';
2+
import type { Context } from 'aws-lambda';
3+
import type {
4+
BedrockAgentFunctionEvent,
5+
BedrockAgentFunctionResponse,
6+
Configuration,
7+
GenericLogger,
8+
ResolverOptions,
9+
ResponseOptions,
10+
Tool,
11+
ToolFunction,
12+
} from '../types/index.js';
13+
import { isPrimitive } from './utils.js';
14+
15+
export class BedrockAgentFunctionResolver {
16+
#tools: Map<string, Tool> = new Map<string, Tool>();
17+
#envService: EnvironmentVariablesService;
18+
#logger: Pick<GenericLogger, 'debug' | 'warn' | 'error'>;
19+
20+
constructor(options?: ResolverOptions) {
21+
this.#envService = new EnvironmentVariablesService();
22+
const alcLogLevel = this.#envService.get('AWS_LAMBDA_LOG_LEVEL');
23+
this.#logger = options?.logger ?? {
24+
debug: alcLogLevel === 'DEBUG' ? console.debug : () => {},
25+
error: console.error,
26+
warn: console.warn,
27+
};
28+
}
29+
30+
/**
31+
* Register a tool function for the Bedrock Agent.
32+
*
33+
* This method registers a function that can be invoked by a Bedrock Agent.
34+
*
35+
* @example
36+
* ```ts
37+
* import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/bedrock-agent-function';
38+
*
39+
* const app = new BedrockAgentFunctionResolver();
40+
*
41+
* app.tool(async (params) => {
42+
* const { name } = params;
43+
* return `Hello, ${name}!`;
44+
* }, {
45+
* name: 'greeting',
46+
* definition: 'Greets a person by name',
47+
* });
48+
*
49+
* export const handler = async (event, context) =>
50+
* app.resolve(event, context);
51+
* ```
52+
*
53+
* The method also works as a class method decorator:
54+
*
55+
* @example
56+
* ```ts
57+
* import { BedrockAgentFunctionResolver } from '@aws-lambda-powertools/event-handler/bedrock-agent-function';
58+
*
59+
* const app = new BedrockAgentFunctionResolver();
60+
*
61+
* class Lambda {
62+
* @app.tool({ name: 'greeting', definition: 'Greets a person by name' })
63+
* async greeting(params) {
64+
* const { name } = params;
65+
* return `Hello, ${name}!`;
66+
* }
67+
*
68+
* async handler(event, context) {
69+
* return app.resolve(event, context);
70+
* }
71+
* }
72+
*
73+
* const lambda = new Lambda();
74+
* export const handler = lambda.handler.bind(lambda);
75+
* ```
76+
*
77+
* @param fn - The tool function
78+
* @param config - The configuration object for the tool
79+
*/
80+
public tool(fn: ToolFunction, config: Configuration): void;
81+
public tool(config: Configuration): MethodDecorator;
82+
public tool(
83+
fnOrConfig: ToolFunction | Configuration,
84+
config?: Configuration
85+
): MethodDecorator | void {
86+
// When used as a method (not a decorator)
87+
if (typeof fnOrConfig === 'function') {
88+
return this.#registerTool(fnOrConfig, config as Configuration);
89+
}
90+
91+
// When used as a decorator
92+
return (_target, _propertyKey, descriptor: PropertyDescriptor) => {
93+
const toolFn = descriptor.value as ToolFunction;
94+
this.#registerTool(toolFn, fnOrConfig as Configuration);
95+
return descriptor;
96+
};
97+
}
98+
99+
#registerTool(fn: ToolFunction, config: Configuration): void {
100+
const { name } = config;
101+
102+
if (this.#tools.size >= 5) {
103+
this.#logger.warn(
104+
`The maximum number of tools that can be registered is 5. Tool ${name} will not be registered.`
105+
);
106+
return;
107+
}
108+
109+
if (this.#tools.has(name)) {
110+
this.#logger.warn(
111+
`Tool ${name} already registered. Overwriting with new definition.`
112+
);
113+
}
114+
115+
this.#tools.set(name, { function: fn, config });
116+
this.#logger.debug(`Tool ${name} has been registered.`);
117+
}
118+
119+
#buildResponse(options: ResponseOptions): BedrockAgentFunctionResponse {
120+
const {
121+
actionGroup,
122+
function: func,
123+
body,
124+
errorType,
125+
sessionAttributes,
126+
promptSessionAttributes,
127+
} = options;
128+
129+
return {
130+
messageVersion: '1.0',
131+
response: {
132+
actionGroup,
133+
function: func,
134+
functionResponse: {
135+
responseState: errorType,
136+
responseBody: {
137+
TEXT: {
138+
body,
139+
},
140+
},
141+
},
142+
},
143+
sessionAttributes,
144+
promptSessionAttributes,
145+
};
146+
}
147+
148+
async resolve(
149+
event: BedrockAgentFunctionEvent,
150+
context: Context
151+
): Promise<BedrockAgentFunctionResponse> {
152+
const {
153+
function: toolName,
154+
parameters = [],
155+
actionGroup,
156+
sessionAttributes,
157+
promptSessionAttributes,
158+
} = event;
159+
160+
const tool = this.#tools.get(toolName);
161+
162+
if (tool == null) {
163+
this.#logger.error(`Tool ${toolName} has not been registered.`);
164+
return this.#buildResponse({
165+
actionGroup,
166+
function: toolName,
167+
body: 'Error: tool has not been registered in handler.',
168+
});
169+
}
170+
171+
const parameterObject: Record<string, string> = Object.fromEntries(
172+
parameters.map((param) => [param.name, param.value])
173+
);
174+
175+
try {
176+
const res = (await tool.function(parameterObject)) ?? '';
177+
const body = isPrimitive(res) ? String(res) : JSON.stringify(res);
178+
return this.#buildResponse({
179+
actionGroup,
180+
function: toolName,
181+
body,
182+
sessionAttributes,
183+
promptSessionAttributes,
184+
});
185+
} catch (error) {
186+
this.#logger.error(`An error occurred in tool ${toolName}.`, error);
187+
return this.#buildResponse({
188+
actionGroup,
189+
function: toolName,
190+
body: `Error when invoking tool: ${error}`,
191+
sessionAttributes,
192+
promptSessionAttributes,
193+
});
194+
}
195+
}
196+
}
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export { BedrockAgentFunctionResolver } from './BedrockAgentFunctionResolver.js';
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import type {
2+
JSONPrimitive,
3+
JSONValue,
4+
} from '@aws-lambda-powertools/commons/types';
5+
6+
export function isPrimitive(value: JSONValue): value is JSONPrimitive {
7+
return (
8+
value === null ||
9+
value === undefined ||
10+
typeof value === 'string' ||
11+
typeof value === 'number' ||
12+
typeof value === 'boolean'
13+
);
14+
}

packages/event-handler/src/types/appsync-events.ts

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,7 @@
11
import type { Context } from 'aws-lambda';
22
import type { RouteHandlerRegistry } from '../appsync-events/RouteHandlerRegistry.js';
33
import type { Router } from '../appsync-events/Router.js';
4-
5-
// #region Shared
6-
7-
// 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.
8-
type Anything = any;
9-
10-
/**
11-
* Interface for a generic logger object.
12-
*/
13-
type GenericLogger = {
14-
trace?: (...content: Anything[]) => void;
15-
debug: (...content: Anything[]) => void;
16-
info?: (...content: Anything[]) => void;
17-
warn: (...content: Anything[]) => void;
18-
error: (...content: Anything[]) => void;
19-
};
4+
import type { Anything, GenericLogger } from './common.js';
205

216
// #region OnPublish fn
227

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import type { JSONValue } from '@aws-lambda-powertools/commons/types';
2+
import type { GenericLogger } from '../types/common.js';
3+
4+
type Configuration = {
5+
name: string;
6+
definition: string;
7+
};
8+
9+
// biome-ignore lint/suspicious/noExplicitAny: this is a generic type that is intentionally open
10+
type Tool<TParams = Record<string, any>> = {
11+
// biome-ignore lint/suspicious/noConfusingVoidType: we need to support async functions that don't have an explicit return value
12+
function: (params: TParams) => Promise<JSONValue | void>;
13+
config: Configuration;
14+
};
15+
16+
type ToolFunction = Tool['function'];
17+
18+
type Parameter = {
19+
name: string;
20+
type: string;
21+
value: string;
22+
};
23+
24+
type Attributes = Record<string, string>;
25+
26+
type FunctionIdentifier = {
27+
actionGroup: string;
28+
function: string;
29+
};
30+
31+
type FunctionInvocation = FunctionIdentifier & {
32+
parameters?: Array<Parameter>;
33+
};
34+
35+
type BedrockAgentFunctionEvent = FunctionInvocation & {
36+
messageVersion: string;
37+
agent: {
38+
name: string;
39+
id: string;
40+
alias: string;
41+
version: string;
42+
};
43+
inputText: string;
44+
sessionId: string;
45+
sessionAttributes: Attributes;
46+
promptSessionAttributes: Attributes;
47+
};
48+
49+
type ResponseState = 'ERROR' | 'REPROMPT';
50+
51+
type TextResponseBody = {
52+
TEXT: {
53+
body: string;
54+
};
55+
};
56+
57+
type SessionData = {
58+
sessionAttributes?: Attributes;
59+
promptSessionAttributes?: Attributes;
60+
};
61+
62+
type BedrockAgentFunctionResponse = SessionData & {
63+
messageVersion: string;
64+
response: FunctionIdentifier & {
65+
functionResponse: {
66+
responseState?: ResponseState;
67+
responseBody: TextResponseBody;
68+
};
69+
};
70+
};
71+
72+
type ResponseOptions = FunctionIdentifier &
73+
SessionData & {
74+
body: string;
75+
errorType?: ResponseState;
76+
};
77+
78+
/**
79+
* Options for the {@link BedrockAgentFunctionResolver} class
80+
*/
81+
type ResolverOptions = {
82+
/**
83+
* A logger instance to be used for logging debug, warning, and error messages.
84+
*
85+
* When no logger is provided, we'll only log warnings and errors using the global `console` object.
86+
*/
87+
logger?: GenericLogger;
88+
};
89+
90+
export type {
91+
Configuration,
92+
Tool,
93+
ToolFunction,
94+
Parameter,
95+
Attributes,
96+
FunctionIdentifier,
97+
FunctionInvocation,
98+
BedrockAgentFunctionEvent,
99+
BedrockAgentFunctionResponse,
100+
ResponseOptions,
101+
ResolverOptions,
102+
};
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
// 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.
2+
type Anything = any;
3+
4+
/**
5+
* Interface for a generic logger object.
6+
*/
7+
type GenericLogger = {
8+
trace?: (...content: Anything[]) => void;
9+
debug: (...content: Anything[]) => void;
10+
info?: (...content: Anything[]) => void;
11+
warn: (...content: Anything[]) => void;
12+
error: (...content: Anything[]) => void;
13+
};
14+
15+
export type { Anything, GenericLogger };

0 commit comments

Comments
 (0)