feat(parser): add BedrockEventModel parser and envelope (#3286)

Author: rubenfonseca

Diff for: aws_lambda_powertools/utilities/parser/envelopes/__init__.py

@@ -1,6 +1,7 @@
 from .apigw import ApiGatewayEnvelope
 from .apigwv2 import ApiGatewayV2Envelope
 from .base import BaseEnvelope
+from .bedrock_agent import BedrockAgentEnvelope
 from .cloudwatch import CloudWatchLogsEnvelope
 from .dynamodb import DynamoDBStreamEnvelope
 from .event_bridge import EventBridgeEnvelope
@@ -16,6 +17,7 @@
 __all__ = [
     "ApiGatewayEnvelope",
     "ApiGatewayV2Envelope",
+    "BedrockAgentEnvelope",
     "CloudWatchLogsEnvelope",
     "DynamoDBStreamEnvelope",
     "EventBridgeEnvelope",

Diff for: aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py

@@ -0,0 +1,32 @@
+import logging
+from typing import Any, Dict, Optional, Type, Union
+
+from ..models import BedrockAgentEventModel
+from ..types import Model
+from .base import BaseEnvelope
+
+logger = logging.getLogger(__name__)
+
+
+class BedrockAgentEnvelope(BaseEnvelope):
+    """Bedrock Agent envelope to extract data within input_text key"""
+
+    def parse(self, data: Optional[Union[Dict[str, Any], Any]], model: Type[Model]) -> Optional[Model]:
+        """Parses data found with model provided
+
+        Parameters
+        ----------
+        data : Dict
+            Lambda event to be parsed
+        model : Type[Model]
+            Data model provided to parse after extracting data using envelope
+
+        Returns
+        -------
+        Optional[Model]
+            Parsed detail payload with model provided
+        """
+        logger.debug(f"Parsing incoming data with Bedrock Agent model {BedrockAgentEventModel}")
+        parsed_envelope: BedrockAgentEventModel = BedrockAgentEventModel.parse_obj(data)
+        logger.debug(f"Parsing event payload in `input_text` with {model}")
+        return self._parse(data=parsed_envelope.input_text, model=model)

Diff for: aws_lambda_powertools/utilities/parser/models/__init__.py

@@ -18,6 +18,13 @@
     RequestContextV2AuthorizerJwt,
     RequestContextV2Http,
 )
+from .bedrock_agent import (
+    BedrockAgentEventModel,
+    BedrockAgentModel,
+    BedrockAgentPropertyModel,
+    BedrockAgentRequestBodyModel,
+    BedrockAgentRequestMediaModel,
+)
 from .cloudformation_custom_resource import (
     CloudFormationCustomResourceBaseModel,
     CloudFormationCustomResourceCreateModel,
@@ -165,4 +172,9 @@
     "CloudFormationCustomResourceBaseModel",
     "VpcLatticeModel",
     "VpcLatticeV2Model",
+    "BedrockAgentModel",
+    "BedrockAgentPropertyModel",
+    "BedrockAgentEventModel",
+    "BedrockAgentRequestBodyModel",
+    "BedrockAgentRequestMediaModel",
 ]

Diff for: aws_lambda_powertools/utilities/parser/models/bedrock_agent.py

@@ -0,0 +1,38 @@
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class BedrockAgentModel(BaseModel):
+    name: str
+    id_: str = Field(..., alias="id")
+    alias: str
+    version: str
+
+
+class BedrockAgentPropertyModel(BaseModel):
+    name: str
+    type_: str = Field(..., alias="type")
+    value: str
+
+
+class BedrockAgentRequestMediaModel(BaseModel):
+    properties: List[BedrockAgentPropertyModel]
+
+
+class BedrockAgentRequestBodyModel(BaseModel):
+    content: Dict[str, BedrockAgentRequestMediaModel]
+
+
+class BedrockAgentEventModel(BaseModel):
+    message_version: str = Field(..., alias="messageVersion")
+    input_text: str = Field(..., alias="inputText")
+    session_id: str = Field(..., alias="sessionId")
+    action_group: str = Field(..., alias="actionGroup")
+    api_path: str = Field(..., alias="apiPath")
+    http_method: str = Field(..., alias="httpMethod")
+    session_attributes: Dict[str, str] = Field({}, alias="sessionAttributes")
+    prompt_session_attributes: Dict[str, str] = Field({}, alias="promptSessionAttributes")
+    agent: BedrockAgentModel
+    parameters: Optional[List[BedrockAgentPropertyModel]] = None
+    request_body: Optional[BedrockAgentRequestBodyModel] = Field(None, alias="requestBody")

Diff for: docs/utilities/parser.md

@@ -178,6 +178,7 @@ Parser comes with the following built-in models:
 | **AlbModel**                                | Lambda Event Source payload for Amazon Application Load Balancer                      |
 | **APIGatewayProxyEventModel**               | Lambda Event Source payload for Amazon API Gateway                                    |
 | **APIGatewayProxyEventV2Model**             | Lambda Event Source payload for Amazon API Gateway v2 payload                         |
+| **BedrockAgentEventModel**                  | Lambda Event Source payload for Bedrock Agents                                        |
 | **CloudFormationCustomResourceCreateModel** | Lambda Event Source payload for AWS CloudFormation `CREATE` operation                 |
 | **CloudFormationCustomResourceUpdateModel** | Lambda Event Source payload for AWS CloudFormation `UPDATE` operation                 |
 | **CloudFormationCustomResourceDeleteModel** | Lambda Event Source payload for AWS CloudFormation `DELETE` operation                 |
@@ -356,6 +357,7 @@ Parser comes with the following built-in envelopes, where `Model` in the return
 | **LambdaFunctionUrlEnvelope** | 1. Parses data using `LambdaFunctionUrlModel`. <br/> 2. Parses `body` key using your model and returns it.                                                                                                  | `Model`                            |
 | **KafkaEnvelope**             | 1. Parses data using `KafkaRecordModel`. <br/> 2. Parses `value` key using your model and returns it.                                                                                                       | `Model`                            |
 | **VpcLatticeEnvelope**        | 1. Parses data using `VpcLatticeModel`. <br/> 2. Parses `value` key using your model and returns it.                                                                                                       | `Model`                            |
+| **BedrockAgentEnvelope**      | 1. Parses data using `BedrockAgentEventModel`. <br/> 2. Parses `inputText` key using your model and returns it.                                                                                                       | `Model`                            |
 
 #### Bringing your own envelope
 

Diff for: tests/unit/parser/schemas.py

@@ -104,3 +104,8 @@ class MyKinesisFirehoseBusiness(BaseModel):
 class MyVpcLatticeBusiness(BaseModel):
     username: str
     name: str
+
+
+class MyBedrockAgentBusiness(BaseModel):
+    username: str
+    name: str

Diff for: tests/unit/parser/test_bedrock_agent.py

@@ -0,0 +1,78 @@
+from aws_lambda_powertools.utilities.parser import envelopes, parse
+from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel
+from tests.functional.utils import load_event
+from tests.unit.parser.schemas import MyBedrockAgentBusiness
+
+
+def test_bedrock_agent_event_with_envelope():
+    raw_event = load_event("bedrockAgentEvent.json")
+    raw_event["inputText"] = '{"username": "Ruben", "name": "Fonseca"}'
+    parsed_event: MyBedrockAgentBusiness = parse(
+        event=raw_event,
+        model=MyBedrockAgentBusiness,
+        envelope=envelopes.BedrockAgentEnvelope,
+    )
+
+    assert parsed_event.username == "Ruben"
+    assert parsed_event.name == "Fonseca"
+
+
+def test_bedrock_agent_event():
+    raw_event = load_event("bedrockAgentEvent.json")
+    model = BedrockAgentEventModel(**raw_event)
+
+    assert model.message_version == raw_event["messageVersion"]
+    assert model.session_id == raw_event["sessionId"]
+    assert model.input_text == raw_event["inputText"]
+    assert model.message_version == raw_event["messageVersion"]
+    assert model.http_method == raw_event["httpMethod"]
+    assert model.api_path == raw_event["apiPath"]
+    assert model.session_attributes == {}
+    assert model.prompt_session_attributes == {}
+    assert model.action_group == raw_event["actionGroup"]
+
+    assert model.request_body is None
+
+    agent = model.agent
+    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")
+    model = BedrockAgentEventModel(**raw_event)
+
+    assert model.session_id == raw_event["sessionId"]
+    assert model.input_text == raw_event["inputText"]
+    assert model.message_version == raw_event["messageVersion"]
+    assert model.http_method == raw_event["httpMethod"]
+    assert model.api_path == raw_event["apiPath"]
+    assert model.session_attributes == {}
+    assert model.prompt_session_attributes == {}
+    assert model.action_group == raw_event["actionGroup"]
+
+    agent = model.agent
+    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 = model.request_body.content
+    assert "application/json" in request_body
+
+    json_request = request_body["application/json"]
+    properties = json_request.properties
+    assert len(properties) == 2
+
+    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 == raw_properties[1]["name"]
+    assert properties[1].type_ == raw_properties[1]["type"]
+    assert properties[1].value == raw_properties[1]["value"]