chore: tests and coverage

Author: dreamorosi

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

@@ -1,5 +1,8 @@
+import type {
+  BedrockAgentFunctionEvent,
+  ResponseState,
+} from '../types/bedrock-agent.js';
 import type { BedrockAgentFunctionResolver } from './BedrockAgentFunctionResolver.js';
-
 /**
  * Class representing a response from a Bedrock agent function.
  *
@@ -23,35 +26,35 @@ class BedrockFunctionResponse {
    * - `FAILURE`: The agent throws a `DependencyFailedException` for the current session.
    * - `REPROMPT`: The agent passes a response string to the model to reprompt it.
    */
-  readonly responseState?: 'FAILURE' | 'REPROMPT';
+  readonly responseState?: ResponseState;
   /**
    * Optional field to store session attributes and their values.
    * @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.
    */
-  readonly sessionAttributes: Record<string, string>;
+  readonly sessionAttributes: BedrockAgentFunctionEvent['sessionAttributes'];
   /**
    * Optional field to instruct the agent to prompt attributes and their values.
    * @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.
    */
-  readonly promptSessionAttributes: Record<string, string>;
+  readonly promptSessionAttributes: BedrockAgentFunctionEvent['promptSessionAttributes'];
   /**
    * Optional field to configure knowledge bases for the agent.
    * @see {@link https://docs.aws.amazon.com/bedrock/latest/userguide/agents-session-state.html#session-state-kb | Bedrock Agent Knowledge Bases} for more details.
    */
-  readonly knowledgeBasesConfiguration?: Record<string, unknown>;
+  readonly knowledgeBasesConfiguration?: BedrockAgentFunctionEvent['knowledgeBasesConfiguration'];
 
   constructor({
     body,
     responseState = undefined,
     sessionAttributes = {},
     promptSessionAttributes = {},
-    knowledgeBasesConfiguration = {},
+    knowledgeBasesConfiguration = undefined,
   }: {
     body: string;
-    responseState?: 'FAILURE' | 'REPROMPT';
-    sessionAttributes?: Record<string, string>;
-    promptSessionAttributes?: Record<string, string>;
-    knowledgeBasesConfiguration?: Record<string, unknown>;
+    responseState?: ResponseState;
+    sessionAttributes?: BedrockAgentFunctionEvent['sessionAttributes'];
+    promptSessionAttributes?: BedrockAgentFunctionEvent['promptSessionAttributes'];
+    knowledgeBasesConfiguration?: BedrockAgentFunctionEvent['knowledgeBasesConfiguration'];
   }) {
     this.body = body;
     this.responseState = responseState;

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

@@ -114,9 +114,9 @@ type BedrockAgentFunctionEvent = {
   parameters?: Array<Parameter>;
   inputText: string;
   sessionId: string;
-  sessionAttributes: Record<string, string>;
-  promptSessionAttributes: Record<string, string>;
-  knowledgeBasesConfiguration?: Record<string, unknown>;
+  sessionAttributes: Record<string, JSONValue>;
+  promptSessionAttributes: Record<string, JSONValue>;
+  knowledgeBasesConfiguration?: Record<string, JSONValue>;
 };
 
 /**
@@ -143,8 +143,9 @@ type BedrockAgentFunctionResponse = {
       };
     };
   };
-  sessionAttributes?: Record<string, string>;
-  promptSessionAttributes?: Record<string, string>;
+  sessionAttributes?: BedrockAgentFunctionEvent['sessionAttributes'];
+  promptSessionAttributes?: BedrockAgentFunctionEvent['promptSessionAttributes'];
+  knowledgeBasesConfiguration?: BedrockAgentFunctionEvent['knowledgeBasesConfiguration'];
 };
 
 /**
@@ -169,4 +170,5 @@ export type {
   BedrockAgentFunctionEvent,
   BedrockAgentFunctionResponse,
   ResolverOptions,
+  ResponseState,
 };

packages/event-handler/src/types/index.ts

@@ -14,6 +14,7 @@ export type {
   BedrockAgentFunctionResponse,
   ResolverOptions,
   Parameter,
+  ResponseState,
 } from './bedrock-agent.js';
 
 export type {

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

@@ -1,10 +1,8 @@
 import context from '@aws-lambda-powertools/testing-utils/context';
-import type { Context } from 'aws-lambda';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { BedrockFunctionResponse } from '../../../src/bedrock-agent/BedrockFunctionResponse.js';
 import { BedrockAgentFunctionResolver } from '../../../src/bedrock-agent/index.js';
 import type {
-  BedrockAgentFunctionEvent,
   Configuration,
   Parameter,
   ToolFunction,
@@ -575,14 +573,6 @@ describe('Class: BedrockAgentFunctionResolver', () => {
       }
     );
 
-    const customSessionAttrs = {
-      sessionAttr: '12345',
-    };
-
-    const customPromptAttrs = {
-      promptAttr: 'promptAttr',
-    };
-
     const customEvent = {
       ...createEvent('greeting', [
         {
@@ -592,8 +582,16 @@ describe('Class: BedrockAgentFunctionResolver', () => {
         },
       ]),
       actionGroup: 'actionGroup',
-      sessionAttributes: customSessionAttrs,
-      promptSessionAttributes: customPromptAttrs,
+      sessionAttributes: {
+        sessionAttr: '12345',
+      },
+      promptSessionAttributes: {
+        promptAttr: 'promptAttr',
+      },
+      knowledgeBasesConfiguration: {
+        knowledgeBase1: { enabled: true },
+        knowledgeBase2: { enabled: false },
+      },
     };
 
     // Act
@@ -613,8 +611,9 @@ describe('Class: BedrockAgentFunctionResolver', () => {
           },
         },
       },
-      sessionAttributes: customSessionAttrs,
-      promptSessionAttributes: customPromptAttrs,
+      sessionAttributes: customEvent.sessionAttributes,
+      promptSessionAttributes: customEvent.promptSessionAttributes,
+      knowledgeBasesConfiguration: customEvent.knowledgeBasesConfiguration,
     });
   });
 });