diff --git a/aws_lambda_powertools/utilities/parser/envelopes/__init__.py b/aws_lambda_powertools/utilities/parser/envelopes/__init__.py index affffd9817..d5754481ee 100644 --- a/aws_lambda_powertools/utilities/parser/envelopes/__init__.py +++ b/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 --git a/aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py b/aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py new file mode 100644 index 0000000000..3fd8a3beb8 --- /dev/null +++ b/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 --git a/aws_lambda_powertools/utilities/parser/models/__init__.py b/aws_lambda_powertools/utilities/parser/models/__init__.py index 3c707fda61..db3aa52437 100644 --- a/aws_lambda_powertools/utilities/parser/models/__init__.py +++ b/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 --git a/aws_lambda_powertools/utilities/parser/models/bedrock_agent.py b/aws_lambda_powertools/utilities/parser/models/bedrock_agent.py new file mode 100644 index 0000000000..6246516216 --- /dev/null +++ b/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 --git a/docs/utilities/parser.md b/docs/utilities/parser.md index 8f0a7bbd06..bfd64f8b7e 100644 --- a/docs/utilities/parser.md +++ b/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`.
2. Parses `body` key using your model and returns it. | `Model` | | **KafkaEnvelope** | 1. Parses data using `KafkaRecordModel`.
2. Parses `value` key using your model and returns it. | `Model` | | **VpcLatticeEnvelope** | 1. Parses data using `VpcLatticeModel`.
2. Parses `value` key using your model and returns it. | `Model` | +| **BedrockAgentEnvelope** | 1. Parses data using `BedrockAgentEventModel`.
2. Parses `inputText` key using your model and returns it. | `Model` | #### Bringing your own envelope diff --git a/tests/unit/parser/schemas.py b/tests/unit/parser/schemas.py index fd2f29697d..65499d319a 100644 --- a/tests/unit/parser/schemas.py +++ b/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 --git a/tests/unit/parser/test_bedrock_agent.py b/tests/unit/parser/test_bedrock_agent.py new file mode 100644 index 0000000000..f3c208469e --- /dev/null +++ b/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"]