fix: changed to the new payload schema

Author: rubenfonseca

aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py

@@ -47,28 +47,6 @@ def content(self) -> Dict[str, BedrockAgentRequestMedia]:
         return {k: BedrockAgentRequestMedia(v) for k, v in self["content"].items()}
 
 
-class BedrockAgentActionGroup(DictWrapper):
-    @property
-    def action_group(self) -> str:
-        return self["actionGroup"]
-
-    @property
-    def api_path(self) -> str:
-        return self["apiPath"]
-
-    @property
-    def http_method(self) -> str:
-        return self["httpMethod"]
-
-    @property
-    def parameters(self) -> List[BedrockAgentProperty]:
-        return [BedrockAgentProperty(x) for x in self["parameters"]]
-
-    @property
-    def request_body(self) -> BedrockAgentRequestBody:
-        return BedrockAgentRequestBody(self["requestBody"])
-
-
 class BedrockAgentEvent(DictWrapper):
     """
     Bedrock Agent input event
@@ -89,8 +67,24 @@ def session_id(self) -> str:
         return self["sessionId"]
 
     @property
-    def action_groups(self) -> List[BedrockAgentActionGroup]:
-        return [BedrockAgentActionGroup(x) for x in self["actionGroups"]]
+    def action_group(self) -> str:
+        return self["actionGroup"]
+
+    @property
+    def api_path(self) -> str:
+        return self["apiPath"]
+
+    @property
+    def http_method(self) -> str:
+        return self["httpMethod"]
+
+    @property
+    def parameters(self) -> List[BedrockAgentProperty]:
+        return [BedrockAgentProperty(x) for x in self["parameters"]]
+
+    @property
+    def request_body(self) -> BedrockAgentRequestBody:
+        return BedrockAgentRequestBody(self["requestBody"])
 
     @property
     def agent(self) -> BedrockAgentInfo:
@@ -100,6 +94,10 @@ def agent(self) -> BedrockAgentInfo:
     def session_attributes(self) -> Dict[str, str]:
         return self["sessionAttributes"]
 
+    @property
+    def prompt_session_attributes(self) -> Dict[str, str]:
+        return self["promptSessionAttributes"]
+
 
 class BedrockAgentResponseMedia(DictWrapper):
     @property
@@ -141,5 +139,5 @@ def message_version(self) -> str:
         return self["messageVersion"]
 
     @property
-    def responses(self) -> List[BedrockAgentResponsePayload]:
-        return [BedrockAgentResponsePayload(x) for x in self["response"]]
+    def response(self) -> BedrockAgentResponsePayload:
+        return BedrockAgentResponsePayload(self["response"])

tests/events/bedrockAgentEvent.json

@@ -1,26 +1,16 @@
 {
+  "actionGroup": "ClaimManagementActionGroup",
+  "messageVersion": "1.0",
   "sessionId": "12345678912345",
+  "sessionAttributes": {},
+  "promptSessionAttributes": {},
+  "inputText": "I want to claim my insurance",
   "agent": {
     "alias": "TSTALIASID",
     "name": "test",
     "version": "DRAFT",
     "id": "8ZXY0W8P1H"
   },
-  "inputText": "ABC12345",
-  "sessionAttributes": {},
-  "messageVersion": "1.0",
-  "actionGroups": [
-    {
-      "httpMethod": "GET",
-      "apiPath": "/claims/{claimId}/identify-missing-documents",
-      "actionGroup": "ClaimManagementActionGroup",
-      "parameters": [
-        {
-          "name": "claimId",
-          "type": "string",
-          "value": "ABC12345"
-        }
-      ]
-    }
-  ]
+  "httpMethod": "GET",
+  "apiPath": "/claims"
 }

tests/events/bedrockAgentResponseEvent.json

@@ -1,16 +1,14 @@
 {
   "messageVersion": "1.0",
-  "response": [
-    {
-      "actionGroup": "ClaimManagementActionGroup",
-      "apiPath": "/claims/{claimId}/identify-missing-documents",
-      "httpMethod": "GET",
-      "httpStatusCode": 200,
-      "responseBody": {
-        "application/json": {
-          "body": "This is the response"
-        }
+  "response": {
+    "actionGroup": "ClaimManagementActionGroup",
+    "apiPath": "/claims/{claimId}/identify-missing-documents",
+    "httpMethod": "GET",
+    "httpStatusCode": 200,
+    "responseBody": {
+      "application/json": {
+        "body": "This is the response"
       }
     }
-  ]
+  }
 }

tests/unit/data_classes/test_bedrock_agent_event.py

@@ -7,49 +7,33 @@ def test_bedrock_agent_event():
     parsed_event = BedrockAgentEvent(raw_event)
 
     assert parsed_event.session_id == "12345678912345"
-    assert parsed_event.input_text == "ABC12345"
-    assert parsed_event.session_attributes == {}
+    assert parsed_event.input_text == "I want to claim my insurance"
     assert parsed_event.message_version == "1.0"
+    assert parsed_event.http_method == "GET"
+    assert parsed_event.api_path == "/claims"
+    assert parsed_event.session_attributes == {}
+    assert parsed_event.prompt_session_attributes == {}
+    assert parsed_event.action_group == "ClaimManagementActionGroup"
 
     agent = parsed_event.agent
     assert agent.alias == "TSTALIASID"
     assert agent.name == "test"
     assert agent.version == "DRAFT"
     assert agent.id == "8ZXY0W8P1H"
 
-    action_groups = parsed_event.action_groups
-    assert len(action_groups) == 1
-
-    action_group = action_groups[0]
-    assert action_group.http_method == "GET"
-    assert action_group.api_path == "/claims/{claimId}/identify-missing-documents"
-    assert action_group.action_group == "ClaimManagementActionGroup"
-
-    parameters = action_group.parameters
-    assert len(parameters) == 1
-
-    parameter = parameters[0]
-    assert parameter.name == "claimId"
-    assert parameter.type == "string"
-    assert parameter.value == "ABC12345"
-
 
 def test_bedrock_agent_response_event():
     raw_event = load_event("bedrockAgentResponseEvent.json")
     parsed_event = BedrockAgentResponseEvent(raw_event)
 
     assert parsed_event.message_version == "1.0"
 
-    responses = parsed_event.responses
-    assert len(responses) == 1
-
-    response = responses[0]
-    assert response.action_group == "ClaimManagementActionGroup"
-    assert response.api_path == "/claims/{claimId}/identify-missing-documents"
-    assert response.http_method == "GET"
-    assert response.http_status_code == 200
+    assert parsed_event.response.action_group == "ClaimManagementActionGroup"
+    assert parsed_event.response.api_path == "/claims/{claimId}/identify-missing-documents"
+    assert parsed_event.response.http_method == "GET"
+    assert parsed_event.response.http_status_code == 200
 
-    response_body = response.response_body
+    response_body = parsed_event.response.response_body
     assert "application/json" in response_body
 
     json_response = response_body["application/json"]