Skip to content

Commit 1998e86

Browse files
svozzasvozza
committed
make tool method generic for type Record<string, ParameterValue>>
1 parent bb483f3 commit 1998e86

File tree

3 files changed

+242
-24
lines changed

3 files changed

+242
-24
lines changed

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

Lines changed: 41 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type {
44
BedrockAgentFunctionEvent,
55
BedrockAgentFunctionResponse,
66
Configuration,
7+
ParameterValue,
78
ResolverOptions,
89
ResponseOptions,
910
Tool,
@@ -13,7 +14,8 @@ import type { GenericLogger } from '../types/common.js';
1314
import { isPrimitive } from './utils.js';
1415

1516
export class BedrockAgentFunctionResolver {
16-
readonly #tools: Map<string, Tool> = new Map<string, Tool>();
17+
readonly #tools: Map<string, Tool<Record<string, ParameterValue>>> =
18+
new Map();
1719
readonly #envService: EnvironmentVariablesService;
1820
readonly #logger: Pick<GenericLogger, 'debug' | 'warn' | 'error'>;
1921

@@ -77,10 +79,15 @@ export class BedrockAgentFunctionResolver {
7779
* @param fn - The tool function
7880
* @param config - The configuration object for the tool
7981
*/
80-
public tool(fn: ToolFunction, config: Configuration): undefined;
81-
public tool(config: Configuration): MethodDecorator;
82-
public tool(
83-
fnOrConfig: ToolFunction | Configuration,
82+
public tool<TParams extends Record<string, ParameterValue>>(
83+
fn: ToolFunction<TParams>,
84+
config: Configuration
85+
): undefined;
86+
public tool<TParams extends Record<string, ParameterValue>>(
87+
config: Configuration
88+
): MethodDecorator;
89+
public tool<TParams extends Record<string, ParameterValue>>(
90+
fnOrConfig: ToolFunction<TParams> | Configuration,
8491
config?: Configuration
8592
): MethodDecorator | undefined {
8693
// When used as a method (not a decorator)
@@ -97,7 +104,10 @@ export class BedrockAgentFunctionResolver {
97104
};
98105
}
99106

100-
#registerTool(fn: ToolFunction, config: Configuration): void {
107+
#registerTool<TParams extends Record<string, ParameterValue>>(
108+
handler: ToolFunction<TParams>,
109+
config: Configuration
110+
): void {
101111
const { name } = config;
102112

103113
if (this.#tools.size >= 5) {
@@ -113,7 +123,10 @@ export class BedrockAgentFunctionResolver {
113123
);
114124
}
115125

116-
this.#tools.set(name, { function: fn, config });
126+
this.#tools.set(name, {
127+
handler: handler as ToolFunction<Record<string, ParameterValue>>,
128+
config,
129+
});
117130
this.#logger.debug(`Tool ${name} has been registered.`);
118131
}
119132

@@ -169,12 +182,29 @@ export class BedrockAgentFunctionResolver {
169182
});
170183
}
171184

172-
const parameterObject: Record<string, string> = Object.fromEntries(
173-
parameters.map((param) => [param.name, param.value])
174-
);
185+
const toolParams: Record<string, ParameterValue> = {};
186+
for (const param of parameters) {
187+
switch (param.type) {
188+
case 'boolean': {
189+
toolParams[param.name] = param.value === 'true';
190+
break;
191+
}
192+
case 'number':
193+
case 'integer': {
194+
toolParams[param.name] = Number(param.value);
195+
break;
196+
}
197+
// this default will also catch array types but we leave them as strings
198+
// because we cannot reliably parse them
199+
default: {
200+
toolParams[param.name] = param.value;
201+
break;
202+
}
203+
}
204+
}
175205

176206
try {
177-
const res = (await tool.function(parameterObject)) ?? '';
207+
const res = (await tool.handler(toolParams, event, context)) ?? '';
178208
const body = isPrimitive(res) ? String(res) : JSON.stringify(res);
179209
return this.#buildResponse({
180210
actionGroup,

‎packages/event-handler/src/types/bedrock-agent-function.ts

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import type { JSONValue } from '@aws-lambda-powertools/commons/types';
2+
import type { Context } from 'aws-lambda';
23
import type { GenericLogger } from '../types/common.js';
34

45
type Configuration = {
@@ -16,14 +17,18 @@ type ParameterPrimitives = string | number | boolean;
1617

1718
type ParameterValue = ParameterPrimitives | Array<ParameterValue>;
1819

19-
type Tool<TParams = Record<string, ParameterValue>> = {
20+
type ToolFunction<TParams = Record<string, ParameterValue>> = (
21+
params: TParams,
22+
event: BedrockAgentFunctionEvent,
23+
context: Context
2024
// biome-ignore lint/suspicious/noConfusingVoidType: we need to support async functions that don't have an explicit return value
21-
function: (params: TParams) => Promise<JSONValue | void>;
25+
) => Promise<JSONValue | void>;
26+
27+
type Tool<TParams = Record<string, ParameterValue>> = {
28+
handler: ToolFunction<TParams>;
2229
config: Configuration;
2330
};
2431

25-
type ToolFunction = Tool['function'];
26-
2732
type Attributes = Record<string, string>;
2833

2934
type FunctionIdentifier = {
@@ -96,6 +101,7 @@ export type {
96101
ToolFunction,
97102
Parameter,
98103
Attributes,
104+
ParameterValue,
99105
FunctionIdentifier,
100106
FunctionInvocation,
101107
BedrockAgentFunctionEvent,

‎packages/event-handler/tests/unit/bedrock-agent/BedrockAgentFunctionResolver.test.ts

Lines changed: 191 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,13 @@ import context from '@aws-lambda-powertools/testing-utils/context';
22
import type { Context } from 'aws-lambda';
33
import { beforeEach, describe, expect, it, vi } from 'vitest';
44
import { BedrockAgentFunctionResolver } from '../../../src/bedrock-agent-function/index.js';
5-
import type { BedrockAgentFunctionEvent, Parameter } from '../../../src/types';
5+
import type {
6+
BedrockAgentFunctionEvent,
7+
Configuration,
8+
Parameter,
9+
ParameterValue,
10+
ToolFunction,
11+
} from '../../../src/types/bedrock-agent-function.js';
612

713
function createEvent(functionName: string, parameters?: Parameter[]) {
814
return {
@@ -93,8 +99,8 @@ describe('Class: BedrockAgentFunctionResolver', () => {
9399
}
94100

95101
app.tool(
96-
async (params) => {
97-
return Number(params.a) + Number(params);
102+
async (params: { a: number; b: number }) => {
103+
return params.a + params.b;
98104
},
99105
{
100106
name: 'mult',
@@ -146,8 +152,8 @@ describe('Class: BedrockAgentFunctionResolver', () => {
146152
]);
147153

148154
app.tool(
149-
async (params) => {
150-
return Number(params.a) + Number(params.b);
155+
async (params: { a: number; b: number }) => {
156+
return params.a + params.b;
151157
},
152158
{
153159
name: 'math',
@@ -162,8 +168,8 @@ describe('Class: BedrockAgentFunctionResolver', () => {
162168
);
163169

164170
app.tool(
165-
async (params) => {
166-
return Number(params.a) * Number(params.b);
171+
async (params: { a: number; b: number }) => {
172+
return params.a * params.b;
167173
},
168174
{
169175
name: 'math',
@@ -178,6 +184,32 @@ describe('Class: BedrockAgentFunctionResolver', () => {
178184
).toEqual('20');
179185
});
180186

187+
it('tool function has access to the event variable', async () => {
188+
// Prepare
189+
const app = new BedrockAgentFunctionResolver();
190+
191+
app.tool(
192+
async (_params, event) => {
193+
return event;
194+
},
195+
{
196+
name: 'event-accessor',
197+
description: 'Accesses the event object',
198+
}
199+
);
200+
201+
const event = createEvent('event-accessor');
202+
203+
// Act
204+
const result = await app.resolve(event, context);
205+
206+
// Assess
207+
expect(result.response.function).toEqual('event-accessor');
208+
expect(result.response.functionResponse.responseBody.TEXT.body).toEqual(
209+
JSON.stringify(event)
210+
);
211+
});
212+
181213
it('can be invoked using the decorator pattern', async () => {
182214
// Prepare
183215
const app = new BedrockAgentFunctionResolver();
@@ -189,9 +221,9 @@ describe('Class: BedrockAgentFunctionResolver', () => {
189221
}
190222

191223
@app.tool({ name: 'add', description: 'Adds two numbers' })
192-
async add(params: { a: string; b: string }) {
224+
async add(params: { a: number; b: number }) {
193225
const { a, b } = params;
194-
return Number.parseInt(a) + Number.parseInt(b);
226+
return a + b;
195227
}
196228

197229
public async handler(event: BedrockAgentFunctionEvent, context: Context) {
@@ -327,6 +359,156 @@ describe('Class: BedrockAgentFunctionResolver', () => {
327359
}
328360
);
329361

362+
it('correctly parses boolean parameters', async () => {
363+
// Prepare
364+
const toolFunction: ToolFunction<{ arg: boolean }> = async (params) =>
365+
params.arg;
366+
367+
const toolParams: Configuration = {
368+
name: 'boolean',
369+
description: 'Handles boolean parameters',
370+
};
371+
372+
const parameters: Parameter[] = [
373+
{ name: 'arg', type: 'boolean', value: 'true' },
374+
];
375+
376+
const app = new BedrockAgentFunctionResolver();
377+
app.tool(toolFunction, toolParams);
378+
379+
//Act
380+
const actual = await app.resolve(
381+
createEvent(toolParams.name, parameters),
382+
context
383+
);
384+
385+
// Assess
386+
expect(actual.response.function).toEqual(toolParams.name);
387+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
388+
'true'
389+
);
390+
});
391+
392+
it('correctly parses number parameters', async () => {
393+
// Prepare
394+
const toolFunction: ToolFunction<{ arg: number }> = async (params) =>
395+
params.arg + 10;
396+
397+
const toolParams: Configuration = {
398+
name: 'number',
399+
description: 'Handles number parameters',
400+
};
401+
402+
const parameters: Parameter[] = [
403+
{ name: 'arg', type: 'number', value: '42' },
404+
];
405+
406+
const app = new BedrockAgentFunctionResolver();
407+
app.tool(toolFunction, toolParams);
408+
409+
// Act
410+
const actual = await app.resolve(
411+
createEvent(toolParams.name, parameters),
412+
context
413+
);
414+
415+
// Assess
416+
expect(actual.response.function).toEqual(toolParams.name);
417+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
418+
'52'
419+
);
420+
});
421+
422+
it('correctly parses integer parameters', async () => {
423+
// Prepare
424+
const toolFunction: ToolFunction<{ arg: number }> = async (params) =>
425+
params.arg + 10;
426+
427+
const toolParams: Configuration = {
428+
name: 'integer',
429+
description: 'Handles integer parameters',
430+
};
431+
432+
const parameters: Parameter[] = [
433+
{ name: 'arg', type: 'integer', value: '37' },
434+
];
435+
436+
const app = new BedrockAgentFunctionResolver();
437+
app.tool(toolFunction, toolParams);
438+
439+
// Act
440+
const actual = await app.resolve(
441+
createEvent(toolParams.name, parameters),
442+
context
443+
);
444+
445+
// Assess
446+
expect(actual.response.function).toEqual(toolParams.name);
447+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
448+
'47'
449+
);
450+
});
451+
452+
it('correctly parses string parameters', async () => {
453+
// Prepare
454+
const toolFunction: ToolFunction<{ arg: string }> = async (params) =>
455+
`String: ${params.arg}`;
456+
457+
const toolParams: Configuration = {
458+
name: 'string',
459+
description: 'Handles string parameters',
460+
};
461+
462+
const parameters: Parameter[] = [
463+
{ name: 'arg', type: 'string', value: 'hello world' },
464+
];
465+
466+
const app = new BedrockAgentFunctionResolver();
467+
app.tool(toolFunction, toolParams);
468+
469+
// Act
470+
const actual = await app.resolve(
471+
createEvent(toolParams.name, parameters),
472+
context
473+
);
474+
475+
// Assess
476+
expect(actual.response.function).toEqual(toolParams.name);
477+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
478+
'String: hello world'
479+
);
480+
});
481+
482+
it('correctly parses array parameters', async () => {
483+
// Prepare
484+
const toolFunction: ToolFunction<{ arg: string }> = async (params) =>
485+
`Array as string: ${params.arg}`;
486+
487+
const toolParams: Configuration = {
488+
name: 'array',
489+
description: 'Handles array parameters (as string)',
490+
};
491+
492+
const parameters: Parameter[] = [
493+
{ name: 'arg', type: 'array', value: '[1,2,3]' },
494+
];
495+
496+
const app = new BedrockAgentFunctionResolver();
497+
app.tool(toolFunction, toolParams);
498+
499+
// Act
500+
const actual = await app.resolve(
501+
createEvent(toolParams.name, parameters),
502+
context
503+
);
504+
505+
// Assess
506+
expect(actual.response.function).toEqual(toolParams.name);
507+
expect(actual.response.functionResponse.responseBody.TEXT.body).toEqual(
508+
'Array as string: [1,2,3]'
509+
);
510+
});
511+
330512
it('handles functions that throw errors', async () => {
331513
// Prepare
332514
const app = new BedrockAgentFunctionResolver();

0 commit comments

Comments
 (0)