add parser

Author: anafalcao

aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Any
 
 from aws_lambda_powertools.utilities.parser.envelopes.base import BaseEnvelope
-from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel
+from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel, BedrockAgentFunctionEventModel
 
 if TYPE_CHECKING:
     from aws_lambda_powertools.utilities.parser.types import Model
@@ -34,3 +34,27 @@ def parse(self, data: dict[str, Any] | Any | None, model: type[Model]) -> Model
         parsed_envelope: BedrockAgentEventModel = BedrockAgentEventModel.model_validate(data)
         logger.debug(f"Parsing event payload in `input_text` with {model}")
         return self._parse(data=parsed_envelope.input_text, model=model)
+
+
+class BedrockAgentFunctionEnvelope(BaseEnvelope):
+    """Bedrock Agent Function envelope to extract data within input_text key"""
+
+    def parse(self, data: dict[str, Any] | Any | None, model: type[Model]) -> Model | None:
+        """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
+        -------
+        Model | None
+            Parsed detail payload with model provided
+        """
+        logger.debug(f"Parsing incoming data with Bedrock Agent Function model {BedrockAgentFunctionEventModel}")
+        parsed_envelope: BedrockAgentFunctionEventModel = BedrockAgentFunctionEventModel.model_validate(data)
+        logger.debug(f"Parsing event payload in `input_text` with {model}")
+        return self._parse(data=parsed_envelope.input_text, model=model)

aws_lambda_powertools/utilities/parser/models/__init__.py

@@ -32,6 +32,7 @@
 )
 from .bedrock_agent import (
     BedrockAgentEventModel,
+    BedrockAgentFunctionEventModel,
     BedrockAgentModel,
     BedrockAgentPropertyModel,
     BedrockAgentRequestBodyModel,
@@ -208,6 +209,7 @@
     "BedrockAgentEventModel",
     "BedrockAgentRequestBodyModel",
     "BedrockAgentRequestMediaModel",
+    "BedrockAgentFunctionEventModel",
     "S3BatchOperationJobModel",
     "S3BatchOperationModel",
     "S3BatchOperationTaskModel",

aws_lambda_powertools/utilities/parser/models/bedrock_agent.py

@@ -36,3 +36,21 @@ class BedrockAgentEventModel(BaseModel):
     agent: BedrockAgentModel
     parameters: Optional[List[BedrockAgentPropertyModel]] = None
     request_body: Optional[BedrockAgentRequestBodyModel] = Field(None, alias="requestBody")
+
+
+class BedrockAgentFunctionEventModel(BaseModel):
+    """Bedrock Agent Function event model
+
+    Documentation:
+    https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html
+    """
+
+    message_version: str = Field(..., alias="messageVersion")
+    agent: BedrockAgentModel
+    input_text: str = Field(..., alias="inputText")
+    session_id: str = Field(..., alias="sessionId")
+    action_group: str = Field(..., alias="actionGroup")
+    function: str
+    parameters: Optional[List[BedrockAgentPropertyModel]] = None
+    session_attributes: Dict[str, str] = Field({}, alias="sessionAttributes")
+    prompt_session_attributes: Dict[str, str] = Field({}, alias="promptSessionAttributes")

tests/unit/parser/_pydantic/test_bedrock_agent.py

@@ -1,5 +1,5 @@
 from aws_lambda_powertools.utilities.parser import envelopes, parse
-from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel
+from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel, BedrockAgentFunctionEventModel
 from tests.functional.utils import load_event
 from tests.unit.parser._pydantic.schemas import MyBedrockAgentBusiness
 
@@ -76,3 +76,35 @@ def test_bedrock_agent_event_with_post():
     assert properties[1].name == raw_properties[1]["name"]
     assert properties[1].type_ == raw_properties[1]["type"]
     assert properties[1].value == raw_properties[1]["value"]
+
+
+def test_bedrock_agent_function_event():
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    model = BedrockAgentFunctionEventModel(**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.action_group == raw_event["actionGroup"]
+    assert model.function == raw_event["function"]
+    assert model.session_attributes == {"employeeId": "EMP123"}
+    assert model.prompt_session_attributes == {"lastInteraction": "2024-02-01T15:30:00Z", "requestType": "vacation"}
+
+    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"]
+
+    parameters = model.parameters
+    assert parameters is not None
+    assert len(parameters) == 2
+
+    assert parameters[0].name == "startDate"
+    assert parameters[0].type_ == "string"
+    assert parameters[0].value == "2024-03-15"
+
+    assert parameters[1].name == "endDate"
+    assert parameters[1].type_ == "string"
+    assert parameters[1].value == "2024-03-20"