From 030aed6f78b2e82e9c52772c3f9c6cf559eeca0e Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Fri, 27 Oct 2023 17:48:12 +0200 Subject: [PATCH 01/10] feat(data_classes): add support for Bedrock Agents events --- .../utilities/data_classes/__init__.py | 3 + .../data_classes/bedrock_agent_event.py | 119 ++++++++++++++++++ 2 files changed, 122 insertions(+) create mode 100644 aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py diff --git a/aws_lambda_powertools/utilities/data_classes/__init__.py b/aws_lambda_powertools/utilities/data_classes/__init__.py index d245bc35f0d..e8057ac8277 100644 --- a/aws_lambda_powertools/utilities/data_classes/__init__.py +++ b/aws_lambda_powertools/utilities/data_classes/__init__.py @@ -6,6 +6,7 @@ from .api_gateway_proxy_event import APIGatewayProxyEvent, APIGatewayProxyEventV2 from .appsync_resolver_event import AppSyncResolverEvent from .aws_config_rule_event import AWSConfigRuleEvent +from .bedrock_agent_event import BedrockAgentEvent, BedrockAgentResponseEvent from .cloud_watch_custom_widget_event import CloudWatchDashboardCustomWidgetEvent from .cloud_watch_logs_event import CloudWatchLogsEvent from .code_pipeline_job_event import CodePipelineJobEvent @@ -35,6 +36,8 @@ "SecretsManagerEvent", "AppSyncResolverEvent", "ALBEvent", + "BedrockAgentEvent", + "BedrockAgentResponseEvent", "CloudWatchDashboardCustomWidgetEvent", "CloudWatchLogsEvent", "CodePipelineJobEvent", diff --git a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py new file mode 100644 index 00000000000..d6ba1c338db --- /dev/null +++ b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py @@ -0,0 +1,119 @@ +from typing import Dict, List + +from aws_lambda_powertools.utilities.data_classes.common import DictWrapper + + +class BedrockAgentInfo(DictWrapper): + @property + def name(self) -> str: + return self["name"] + + @property + def id(self) -> str: # noqa: A003 + return self["id"] + + @property + def alias(self) -> str: + return self["alias"] + + @property + def version(self) -> str: + return self["version"] + + +class BedrockAgentProperty(DictWrapper): + @property + def name(self) -> str: + return self["name"] + + @property + def type(self) -> str: # noqa: A003 + return self["type"] + + @property + def value(self) -> str: + return self["value"] + + +class BedrockAgentRequestMedia(DictWrapper): + @property + def properties(self) -> List[BedrockAgentProperty]: + return [BedrockAgentProperty(x) for x in self["properties"]] + + +class BedrockAgentRequestBody(DictWrapper): + @property + def content(self) -> Dict[str, BedrockAgentRequestMedia]: + return {k: BedrockAgentRequestMedia(v) for k, v in self["content"].items()} + + +class BedrockAgentEvent(DictWrapper): + @property + def message_version(self) -> str: + return self["messageVersion"] + + @property + def input_text(self) -> str: + return self["inputText"] + + @property + def session_id(self) -> str: + return self["sessionId"] + + @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 session_attributes(self) -> Dict[str, str]: + return self["sessionAttributes"] + + +class BedrockAgentResponseMedia(DictWrapper): + @property + def body(self) -> str: + return self["body"] + + +class BedrockAgentResponsePayload(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 http_status_code(self) -> int: + return self["httpStatusCode"] + + @property + def response_body(self) -> Dict[str, BedrockAgentResponseMedia]: + return {k: BedrockAgentResponseMedia(v) for k, v in self["responseBody"].items()} + + +class BedrockAgentResponseEvent(DictWrapper): + @property + def message_version(self) -> str: + return self["messageVersion"] + + @property + def response(self) -> BedrockAgentResponsePayload: + return BedrockAgentResponsePayload(self["response"]) From 2b4cf350708b38caef751edcbd1c7438ab90899d Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Mon, 30 Oct 2023 14:01:48 +0100 Subject: [PATCH 02/10] fix: add tests --- .../data_classes/bedrock_agent_event.py | 54 +++++++++++++----- tests/events/bedrockAgentEvent.json | 26 +++++++++ tests/events/bedrockAgentResponseEvent.json | 16 ++++++ .../data_classes/test_bedrock_agent_event.py | 56 +++++++++++++++++++ 4 files changed, 138 insertions(+), 14 deletions(-) create mode 100644 tests/events/bedrockAgentEvent.json create mode 100644 tests/events/bedrockAgentResponseEvent.json create mode 100644 tests/unit/data_classes/test_bedrock_agent_event.py diff --git a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py index d6ba1c338db..48ef21e640f 100644 --- a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py +++ b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py @@ -47,7 +47,35 @@ 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 + + See https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.html + """ + @property def message_version(self) -> str: return self["messageVersion"] @@ -61,20 +89,12 @@ def session_id(self) -> str: return self["sessionId"] @property - def action_group(self) -> str: - return self["actionGroup"] + def action_groups(self) -> List[BedrockAgentActionGroup]: + return [BedrockAgentActionGroup(x) for x in self["actionGroups"]] @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"]] + def agent(self) -> BedrockAgentInfo: + return BedrockAgentInfo(self["agent"]) @property def session_attributes(self) -> Dict[str, str]: @@ -110,10 +130,16 @@ def response_body(self) -> Dict[str, BedrockAgentResponseMedia]: class BedrockAgentResponseEvent(DictWrapper): + """ + Bedrock Agent output event + + See: https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.html + """ + @property def message_version(self) -> str: return self["messageVersion"] @property - def response(self) -> BedrockAgentResponsePayload: - return BedrockAgentResponsePayload(self["response"]) + def responses(self) -> List[BedrockAgentResponsePayload]: + return [BedrockAgentResponsePayload(x) for x in self["response"]] diff --git a/tests/events/bedrockAgentEvent.json b/tests/events/bedrockAgentEvent.json new file mode 100644 index 00000000000..08473ec583e --- /dev/null +++ b/tests/events/bedrockAgentEvent.json @@ -0,0 +1,26 @@ +{ + "sessionId": "12345678912345", + "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" + } + ] + } + ] +} diff --git a/tests/events/bedrockAgentResponseEvent.json b/tests/events/bedrockAgentResponseEvent.json new file mode 100644 index 00000000000..49b42e2f9eb --- /dev/null +++ b/tests/events/bedrockAgentResponseEvent.json @@ -0,0 +1,16 @@ +{ + "messageVersion": "1.0", + "response": [ + { + "actionGroup": "ClaimManagementActionGroup", + "apiPath": "/claims/{claimId}/identify-missing-documents", + "httpMethod": "GET", + "httpStatusCode": 200, + "responseBody": { + "application/json": { + "body": "This is the response" + } + } + } + ] +} diff --git a/tests/unit/data_classes/test_bedrock_agent_event.py b/tests/unit/data_classes/test_bedrock_agent_event.py new file mode 100644 index 00000000000..178ef4bb13d --- /dev/null +++ b/tests/unit/data_classes/test_bedrock_agent_event.py @@ -0,0 +1,56 @@ +from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent, BedrockAgentResponseEvent +from tests.functional.utils import load_event + + +def test_bedrock_agent_event(): + raw_event = load_event("bedrockAgentEvent.json") + 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.message_version == "1.0" + + 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 + + response_body = response.response_body + assert "application/json" in response_body + + json_response = response_body["application/json"] + assert json_response.body == "This is the response" From 2671337f7e150b056b64f9bf9795a9f28a79c0a5 Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Mon, 30 Oct 2023 14:36:41 +0100 Subject: [PATCH 03/10] fix: add docs --- docs/utilities/data_classes.md | 9 +++++++ .../event_sources/src/bedrock_agent_event.py | 26 +++++++++++++++++++ 2 files changed, 35 insertions(+) create mode 100644 examples/event_sources/src/bedrock_agent_event.py diff --git a/docs/utilities/data_classes.md b/docs/utilities/data_classes.md index 7cc966313fb..37d3725967c 100644 --- a/docs/utilities/data_classes.md +++ b/docs/utilities/data_classes.md @@ -85,6 +85,7 @@ Log Data Event for Troubleshooting | [AppSync Authorizer](#appsync-authorizer) | `AppSyncAuthorizerEvent` | | [AppSync Resolver](#appsync-resolver) | `AppSyncResolverEvent` | | [AWS Config Rule](#aws-config-rule) | `AWSConfigRuleEvent` | +| [Bedrock Agent](#bedrock-agent) | `BedrockAgent` | | [CloudWatch Dashboard Custom Widget](#cloudwatch-dashboard-custom-widget) | `CloudWatchDashboardCustomWidgetEvent` | | [CloudWatch Logs](#cloudwatch-logs) | `CloudWatchLogsEvent` | | [CodePipeline Job Event](#codepipeline-job) | `CodePipelineJobEvent` | @@ -484,6 +485,14 @@ In this example, we also use the new Logger `correlation_id` and built-in `corre --8<-- "examples/event_sources/src/aws_config_rule_scheduled.json" ``` +### Bedrock Agent + +=== "app.py" + + ```python hl_lines="2 8 10" + --8<-- "examples/event_sources/src/bedrock_agent_event.py" + ``` + ### CloudWatch Dashboard Custom Widget === "app.py" diff --git a/examples/event_sources/src/bedrock_agent_event.py b/examples/event_sources/src/bedrock_agent_event.py new file mode 100644 index 00000000000..025e4448722 --- /dev/null +++ b/examples/event_sources/src/bedrock_agent_event.py @@ -0,0 +1,26 @@ +from aws_lambda_powertools import Logger +from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent, event_source +from aws_lambda_powertools.utilities.typing import LambdaContext + +logger = Logger() + + +@event_source(data_class=BedrockAgentEvent) +def lambda_handler(event: BedrockAgentEvent, context: LambdaContext) -> dict: + input_text = event.input_text + action = event.action_groups[0] + + logger.info(f"Bedrock Agent {action.action_group} invoked with input", input_text=input_text) + + return { + "message_version": "1.0", + "responses": [ + { + "action_group": action.action_group, + "api_path": action.api_path, + "http_method": action.http_method, + "http_status_code": 200, + "response_body": {"application/json": {"body": "This is the response"}}, + }, + ], + } From a7eb9e2a4c0f0a94b5f7814b2452811f92c10cf0 Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Wed, 1 Nov 2023 15:12:42 +0100 Subject: [PATCH 04/10] fix: changed to the new payload schema --- .../data_classes/bedrock_agent_event.py | 50 +++++++++---------- tests/events/bedrockAgentEvent.json | 24 +++------ tests/events/bedrockAgentResponseEvent.json | 20 ++++---- .../data_classes/test_bedrock_agent_event.py | 38 ++++---------- 4 files changed, 51 insertions(+), 81 deletions(-) diff --git a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py index 48ef21e640f..79afa20f920 100644 --- a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py +++ b/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"]) diff --git a/tests/events/bedrockAgentEvent.json b/tests/events/bedrockAgentEvent.json index 08473ec583e..b7ad75b3c43 100644 --- a/tests/events/bedrockAgentEvent.json +++ b/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" } diff --git a/tests/events/bedrockAgentResponseEvent.json b/tests/events/bedrockAgentResponseEvent.json index 49b42e2f9eb..f01bb976ac1 100644 --- a/tests/events/bedrockAgentResponseEvent.json +++ b/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" } } - ] + } } diff --git a/tests/unit/data_classes/test_bedrock_agent_event.py b/tests/unit/data_classes/test_bedrock_agent_event.py index 178ef4bb13d..57ae93561f5 100644 --- a/tests/unit/data_classes/test_bedrock_agent_event.py +++ b/tests/unit/data_classes/test_bedrock_agent_event.py @@ -7,9 +7,13 @@ 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" @@ -17,22 +21,6 @@ def test_bedrock_agent_event(): 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") @@ -40,16 +28,12 @@ def test_bedrock_agent_response_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"] From 5d5ce0aeb25c3a2a65a7df959c72a43c39c182d7 Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Wed, 1 Nov 2023 15:24:13 +0100 Subject: [PATCH 05/10] fix: example --- examples/event_sources/src/bedrock_agent_event.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/examples/event_sources/src/bedrock_agent_event.py b/examples/event_sources/src/bedrock_agent_event.py index 025e4448722..b16d3c86bad 100644 --- a/examples/event_sources/src/bedrock_agent_event.py +++ b/examples/event_sources/src/bedrock_agent_event.py @@ -8,17 +8,16 @@ @event_source(data_class=BedrockAgentEvent) def lambda_handler(event: BedrockAgentEvent, context: LambdaContext) -> dict: input_text = event.input_text - action = event.action_groups[0] - logger.info(f"Bedrock Agent {action.action_group} invoked with input", input_text=input_text) + logger.info(f"Bedrock Agent {event.action_group} invoked with input", input_text=input_text) return { "message_version": "1.0", "responses": [ { - "action_group": action.action_group, - "api_path": action.api_path, - "http_method": action.http_method, + "action_group": event.action_group, + "api_path": event.api_path, + "http_method": event.http_method, "http_status_code": 200, "response_body": {"application/json": {"body": "This is the response"}}, }, From db71067f83025617545971e47ad26a293237b4d2 Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Thu, 2 Nov 2023 12:50:30 +0100 Subject: [PATCH 06/10] fix: add more tests --- .../utilities/data_classes/__init__.py | 3 +- tests/events/bedrockAgentPostEvent.json | 35 ++++++++++++++++ .../data_classes/test_bedrock_agent_event.py | 41 +++++++++++++------ 3 files changed, 65 insertions(+), 14 deletions(-) create mode 100644 tests/events/bedrockAgentPostEvent.json diff --git a/aws_lambda_powertools/utilities/data_classes/__init__.py b/aws_lambda_powertools/utilities/data_classes/__init__.py index e8057ac8277..fd9294bc8bb 100644 --- a/aws_lambda_powertools/utilities/data_classes/__init__.py +++ b/aws_lambda_powertools/utilities/data_classes/__init__.py @@ -6,7 +6,7 @@ from .api_gateway_proxy_event import APIGatewayProxyEvent, APIGatewayProxyEventV2 from .appsync_resolver_event import AppSyncResolverEvent from .aws_config_rule_event import AWSConfigRuleEvent -from .bedrock_agent_event import BedrockAgentEvent, BedrockAgentResponseEvent +from .bedrock_agent_event import BedrockAgentEvent from .cloud_watch_custom_widget_event import CloudWatchDashboardCustomWidgetEvent from .cloud_watch_logs_event import CloudWatchLogsEvent from .code_pipeline_job_event import CodePipelineJobEvent @@ -37,7 +37,6 @@ "AppSyncResolverEvent", "ALBEvent", "BedrockAgentEvent", - "BedrockAgentResponseEvent", "CloudWatchDashboardCustomWidgetEvent", "CloudWatchLogsEvent", "CodePipelineJobEvent", diff --git a/tests/events/bedrockAgentPostEvent.json b/tests/events/bedrockAgentPostEvent.json new file mode 100644 index 00000000000..f223bfcd516 --- /dev/null +++ b/tests/events/bedrockAgentPostEvent.json @@ -0,0 +1,35 @@ +{ + "actionGroup": "ClaimManagementActionGroup", + "messageVersion": "1.0", + "sessionId": "12345678912345", + "sessionAttributes": {}, + "promptSessionAttributes": {}, + "inputText": "Send reminders to all pending documents", + "agent": { + "alias": "TSTALIASID", + "name": "test", + "version": "DRAFT", + "id": "8ZXY0W8P1H" + }, + "httpMethod": "POST", + "apiPath": "/send-reminders", + "requestBody": { + "content": { + "application/json": { + "properties": [ + { + "name": "claimId", + "type": "string", + "value": "20" + }, + { + "name": "pendingDocuments", + "type": "string", + "value": "social number and vat" + } + ] + } + } + }, + "parameters": [] +} diff --git a/tests/unit/data_classes/test_bedrock_agent_event.py b/tests/unit/data_classes/test_bedrock_agent_event.py index 57ae93561f5..66c18a4f8a2 100644 --- a/tests/unit/data_classes/test_bedrock_agent_event.py +++ b/tests/unit/data_classes/test_bedrock_agent_event.py @@ -1,4 +1,4 @@ -from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent, BedrockAgentResponseEvent +from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent from tests.functional.utils import load_event @@ -22,19 +22,36 @@ def test_bedrock_agent_event(): assert agent.id == "8ZXY0W8P1H" -def test_bedrock_agent_response_event(): - raw_event = load_event("bedrockAgentResponseEvent.json") - parsed_event = BedrockAgentResponseEvent(raw_event) +def test_bedrock_agent_event_with_post(): + raw_event = load_event("bedrockAgentPostEvent.json") + parsed_event = BedrockAgentEvent(raw_event) + assert parsed_event.session_id == "12345678912345" + assert parsed_event.input_text == "Send reminders to all pending documents" assert parsed_event.message_version == "1.0" + assert parsed_event.http_method == "POST" + assert parsed_event.api_path == "/send-reminders" + 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" + + request_body = parsed_event.request_body.content + assert "application/json" in request_body - 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 + json_request = request_body["application/json"] + properties = json_request.properties + assert len(properties) == 2 - response_body = parsed_event.response.response_body - assert "application/json" in response_body + assert properties[0].name == "claimId" + assert properties[0].type == "string" + assert properties[0].value == "20" - json_response = response_body["application/json"] - assert json_response.body == "This is the response" + assert properties[1].name == "pendingDocuments" + assert properties[1].type == "string" + assert properties[1].value == "social number and vat" From 4e9110d8e15113bb6b81df05c992fdd0c6e2e4dd Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Thu, 2 Nov 2023 13:49:36 +0100 Subject: [PATCH 07/10] fix: remove response --- .../data_classes/bedrock_agent_event.py | 39 ------------------- 1 file changed, 39 deletions(-) diff --git a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py index 79afa20f920..f75b5805c4b 100644 --- a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py +++ b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py @@ -98,46 +98,7 @@ def session_attributes(self) -> Dict[str, str]: def prompt_session_attributes(self) -> Dict[str, str]: return self["promptSessionAttributes"] - -class BedrockAgentResponseMedia(DictWrapper): - @property - def body(self) -> str: - return self["body"] - - -class BedrockAgentResponsePayload(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 http_status_code(self) -> int: - return self["httpStatusCode"] - - @property - def response_body(self) -> Dict[str, BedrockAgentResponseMedia]: - return {k: BedrockAgentResponseMedia(v) for k, v in self["responseBody"].items()} - - -class BedrockAgentResponseEvent(DictWrapper): - """ - Bedrock Agent output event - - See: https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.html - """ - - @property - def message_version(self) -> str: - return self["messageVersion"] - - @property - def response(self) -> BedrockAgentResponsePayload: - return BedrockAgentResponsePayload(self["response"]) From fd524b405d6db647d4d08d7411ad71153fc10205 Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Thu, 2 Nov 2023 13:51:17 +0100 Subject: [PATCH 08/10] fix: part 2 --- .../utilities/data_classes/bedrock_agent_event.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py index f75b5805c4b..101f8993b56 100644 --- a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py +++ b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py @@ -97,8 +97,3 @@ def session_attributes(self) -> Dict[str, str]: @property def prompt_session_attributes(self) -> Dict[str, str]: return self["promptSessionAttributes"] - - @property - def api_path(self) -> str: - return self["apiPath"] - From ead0e9a3286cc4c2137b9a8d6a2f41f8bc7cffaf Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Thu, 2 Nov 2023 13:59:40 +0100 Subject: [PATCH 09/10] fix: comments --- .../data_classes/bedrock_agent_event.py | 10 ++-- .../data_classes/test_bedrock_agent_event.py | 57 ++++++++++--------- 2 files changed, 36 insertions(+), 31 deletions(-) diff --git a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py index 101f8993b56..b482b5b2b3e 100644 --- a/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py +++ b/aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py @@ -1,4 +1,4 @@ -from typing import Dict, List +from typing import Dict, List, Optional from aws_lambda_powertools.utilities.data_classes.common import DictWrapper @@ -79,12 +79,12 @@ def http_method(self) -> str: return self["httpMethod"] @property - def parameters(self) -> List[BedrockAgentProperty]: - return [BedrockAgentProperty(x) for x in self["parameters"]] + def parameters(self) -> Optional[List[BedrockAgentProperty]]: + return [BedrockAgentProperty(x) for x in self["parameters"]] if self.get("parameters") else None @property - def request_body(self) -> BedrockAgentRequestBody: - return BedrockAgentRequestBody(self["requestBody"]) + def request_body(self) -> Optional[BedrockAgentRequestBody]: + return BedrockAgentRequestBody(self["requestBody"]) if self.get("requestBody") else None @property def agent(self) -> BedrockAgentInfo: diff --git a/tests/unit/data_classes/test_bedrock_agent_event.py b/tests/unit/data_classes/test_bedrock_agent_event.py index 66c18a4f8a2..c4b56695774 100644 --- a/tests/unit/data_classes/test_bedrock_agent_event.py +++ b/tests/unit/data_classes/test_bedrock_agent_event.py @@ -6,40 +6,44 @@ def test_bedrock_agent_event(): raw_event = load_event("bedrockAgentEvent.json") parsed_event = BedrockAgentEvent(raw_event) - assert parsed_event.session_id == "12345678912345" - 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_id == raw_event["sessionId"] + assert parsed_event.input_text == raw_event["inputText"] + assert parsed_event.message_version == raw_event["messageVersion"] + assert parsed_event.http_method == raw_event["httpMethod"] + assert parsed_event.api_path == raw_event["apiPath"] assert parsed_event.session_attributes == {} assert parsed_event.prompt_session_attributes == {} - assert parsed_event.action_group == "ClaimManagementActionGroup" + assert parsed_event.action_group == raw_event["actionGroup"] + + assert parsed_event.request_body is None agent = parsed_event.agent - assert agent.alias == "TSTALIASID" - assert agent.name == "test" - assert agent.version == "DRAFT" - assert agent.id == "8ZXY0W8P1H" + raw_agent = raw_event["agent"] + assert agent.alias == raw_agent["alias"] + assert agent.name == raw_agent["name"] + assert agent.version == raw_agent["version"] + assert agent.id == raw_agent["id"] def test_bedrock_agent_event_with_post(): raw_event = load_event("bedrockAgentPostEvent.json") parsed_event = BedrockAgentEvent(raw_event) - assert parsed_event.session_id == "12345678912345" - assert parsed_event.input_text == "Send reminders to all pending documents" - assert parsed_event.message_version == "1.0" - assert parsed_event.http_method == "POST" - assert parsed_event.api_path == "/send-reminders" + assert parsed_event.session_id == raw_event["sessionId"] + assert parsed_event.input_text == raw_event["inputText"] + assert parsed_event.message_version == raw_event["messageVersion"] + assert parsed_event.http_method == raw_event["httpMethod"] + assert parsed_event.api_path == raw_event["apiPath"] assert parsed_event.session_attributes == {} assert parsed_event.prompt_session_attributes == {} - assert parsed_event.action_group == "ClaimManagementActionGroup" + assert parsed_event.action_group == raw_event["actionGroup"] agent = parsed_event.agent - assert agent.alias == "TSTALIASID" - assert agent.name == "test" - assert agent.version == "DRAFT" - assert agent.id == "8ZXY0W8P1H" + raw_agent = raw_event["agent"] + assert agent.alias == raw_agent["alias"] + assert agent.name == raw_agent["name"] + assert agent.version == raw_agent["version"] + assert agent.id == raw_agent["id"] request_body = parsed_event.request_body.content assert "application/json" in request_body @@ -48,10 +52,11 @@ def test_bedrock_agent_event_with_post(): properties = json_request.properties assert len(properties) == 2 - assert properties[0].name == "claimId" - assert properties[0].type == "string" - assert properties[0].value == "20" + raw_properties = raw_event["requestBody"]["content"]["application/json"]["properties"] + assert properties[0].name == raw_properties[0]["name"] + assert properties[0].type == raw_properties[0]["type"] + assert properties[0].value == raw_properties[0]["value"] - assert properties[1].name == "pendingDocuments" - assert properties[1].type == "string" - assert properties[1].value == "social number and vat" + assert properties[1].name == raw_properties[1]["name"] + assert properties[1].type == raw_properties[1]["type"] + assert properties[1].value == raw_properties[1]["value"] From 89700ecb177567a49080a4f5dbd2e9c4a1ba4981 Mon Sep 17 00:00:00 2001 From: Ruben Fonseca Date: Thu, 2 Nov 2023 14:02:30 +0100 Subject: [PATCH 10/10] fix: remove response example --- tests/events/bedrockAgentResponseEvent.json | 14 -------------- 1 file changed, 14 deletions(-) delete mode 100644 tests/events/bedrockAgentResponseEvent.json diff --git a/tests/events/bedrockAgentResponseEvent.json b/tests/events/bedrockAgentResponseEvent.json deleted file mode 100644 index f01bb976ac1..00000000000 --- a/tests/events/bedrockAgentResponseEvent.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "messageVersion": "1.0", - "response": { - "actionGroup": "ClaimManagementActionGroup", - "apiPath": "/claims/{claimId}/identify-missing-documents", - "httpMethod": "GET", - "httpStatusCode": 200, - "responseBody": { - "application/json": { - "body": "This is the response" - } - } - } -}