Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 6 commits into from
Nov 2, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -16,6 +17,7 @@
__all__ = [
"ApiGatewayEnvelope",
"ApiGatewayV2Envelope",
"BedrockAgentEnvelope",
"CloudWatchLogsEnvelope",
"DynamoDBStreamEnvelope",
"EventBridgeEnvelope",
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
12 changes: 12 additions & 0 deletions aws_lambda_powertools/utilities/parser/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,13 @@
RequestContextV2AuthorizerJwt,
RequestContextV2Http,
)
from .bedrock_agent import (
BedrockAgentEventModel,
BedrockAgentModel,
BedrockAgentPropertyModel,
BedrockAgentRequestBodyModel,
BedrockAgentRequestMediaModel,
)
from .cloudformation_custom_resource import (
CloudFormationCustomResourceBaseModel,
CloudFormationCustomResourceCreateModel,
Expand Down Expand Up @@ -165,4 +172,9 @@
"CloudFormationCustomResourceBaseModel",
"VpcLatticeModel",
"VpcLatticeV2Model",
"BedrockAgentModel",
"BedrockAgentPropertyModel",
"BedrockAgentEventModel",
"BedrockAgentRequestBodyModel",
"BedrockAgentRequestMediaModel",
]
38 changes: 38 additions & 0 deletions aws_lambda_powertools/utilities/parser/models/bedrock_agent.py
Original file line number Diff line number Diff line change
@@ -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")
2 changes: 2 additions & 0 deletions docs/utilities/parser.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
Expand Down Expand Up @@ -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

Expand Down
5 changes: 5 additions & 0 deletions tests/unit/parser/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,3 +104,8 @@ class MyKinesisFirehoseBusiness(BaseModel):
class MyVpcLatticeBusiness(BaseModel):
username: str
name: str


class MyBedrockAgentBusiness(BaseModel):
username: str
name: str
78 changes: 78 additions & 0 deletions tests/unit/parser/test_bedrock_agent.py
Original file line number Diff line number Diff line change
@@ -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"]