feat(bedrock_agent): create bedrock agents functions data class

Author: anafalcao

aws_lambda_powertools/utilities/data_classes/__init__.py

@@ -9,6 +9,7 @@
 from .appsync_resolver_events_event import AppSyncResolverEventsEvent
 from .aws_config_rule_event import AWSConfigRuleEvent
 from .bedrock_agent_event import BedrockAgentEvent
+from .bedrock_agent_function_event import BedrockAgentFunctionEvent
 from .cloud_watch_alarm_event import (
     CloudWatchAlarmConfiguration,
     CloudWatchAlarmData,
@@ -59,6 +60,7 @@
     "AppSyncResolverEventsEvent",
     "ALBEvent",
     "BedrockAgentEvent",
+    "BedrockAgentFunctionEvent",
     "CloudWatchAlarmData",
     "CloudWatchAlarmEvent",
     "CloudWatchAlarmMetric",

aws_lambda_powertools/utilities/data_classes/bedrock_agent_function_event.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from typing import Any
+
+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 BedrockAgentFunctionParameter(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 BedrockAgentFunctionEvent(DictWrapper):
+    """
+    Bedrock Agent Function input event
+
+    Documentation:
+    https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html
+    """
+
+    @classmethod
+    def validate_required_fields(cls, data: dict[str, Any]) -> None:
+        required_fields = {
+            "messageVersion": str,
+            "agent": dict,
+            "inputText": str,
+            "sessionId": str,
+            "actionGroup": str,
+            "function": str,
+        }
+
+        for field, field_type in required_fields.items():
+            if field not in data:
+                raise ValueError(f"Missing required field: {field}")
+            if not isinstance(data[field], field_type):
+                raise TypeError(f"Field {field} must be of type {field_type}")
+
+        # Validate agent structure
+        required_agent_fields = {"name", "id", "alias", "version"}
+        if not all(field in data["agent"] for field in required_agent_fields):
+            raise ValueError("Agent object missing required fields")
+
+    def __init__(self, data: dict[str, Any]) -> None:
+        super().__init__(data)
+        self.validate_required_fields(data)
+
+    @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 function(self) -> str:
+        return self["function"]
+
+    @property
+    def parameters(self) -> list[BedrockAgentFunctionParameter]:
+        parameters = self.get("parameters") or []
+        return [BedrockAgentFunctionParameter(x) for x in parameters]
+
+    @property
+    def agent(self) -> BedrockAgentInfo:
+        return BedrockAgentInfo(self["agent"])
+
+    @property
+    def session_attributes(self) -> dict[str, str]:
+        return self.get("sessionAttributes", {}) or {}
+
+    @property
+    def prompt_session_attributes(self) -> dict[str, str]:
+        return self.get("promptSessionAttributes", {}) or {}

tests/events/bedrockAgentFunctionEvent.json

@@ -0,0 +1,33 @@
+{
+    "messageVersion": "1.0",
+    "agent": {
+        "alias": "PROD",
+        "name": "hr-assistant-function-def",
+        "version": "1",
+        "id": "1234abcd"
+    },
+    "sessionId": "123456789123458",
+    "sessionAttributes": {
+        "employeeId": "EMP123",
+        "department": "Engineering"
+    },
+    "promptSessionAttributes": {
+        "lastInteraction": "2024-02-01T15:30:00Z",
+        "requestType": "vacation"
+    },
+    "inputText": "I want to request vacation from March 15 to March 20",
+    "actionGroup": "VacationsActionGroup",
+    "function": "submitVacationRequest",
+    "parameters": [
+        {
+            "name": "startDate",
+            "type": "string",
+            "value": "2024-03-15"
+        },
+        {
+            "name": "endDate",
+            "type": "string",
+            "value": "2024-03-20"
+        }
+    ]
+}

tests/unit/data_classes/required_dependencies/test_bedrock_agent_function_event.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import pytest
+
+from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
+from tests.functional.utils import load_event
+
+
+def test_bedrock_agent_function_event():
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    parsed_event = BedrockAgentFunctionEvent(raw_event)
+
+    # Test basic event properties
+    assert parsed_event.message_version == raw_event["messageVersion"]
+    assert parsed_event.session_id == raw_event["sessionId"]
+    assert parsed_event.input_text == raw_event["inputText"]
+    assert parsed_event.action_group == raw_event["actionGroup"]
+    assert parsed_event.function == raw_event["function"]
+
+    # Test agent information
+    agent = parsed_event.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"]
+
+    # Test session attributes
+    assert parsed_event.session_attributes == raw_event["sessionAttributes"]
+    assert parsed_event.prompt_session_attributes == raw_event["promptSessionAttributes"]
+
+    # Test parameters
+    parameters = parsed_event.parameters
+    raw_parameters = raw_event["parameters"]
+    assert len(parameters) == len(raw_parameters)
+
+    for param, raw_param in zip(parameters, raw_parameters):
+        assert param.name == raw_param["name"]
+        assert param.type == raw_param["type"]
+        assert param.value == raw_param["value"]
+
+
+def test_bedrock_agent_function_event_minimal():
+    """Test with minimal required fields"""
+    minimal_event = {
+        "messageVersion": "1.0",
+        "agent": {
+            "alias": "PROD",
+            "name": "hr-assistant-function-def",
+            "version": "1",
+            "id": "1234abcd-56ef-78gh-90ij-klmn12345678",
+        },
+        "sessionId": "87654321-abcd-efgh-ijkl-mnop12345678",
+        "inputText": "I want to request vacation",
+        "actionGroup": "VacationsActionGroup",
+        "function": "submitVacationRequest",
+    }
+
+    parsed_event = BedrockAgentFunctionEvent(minimal_event)
+
+    assert parsed_event.session_attributes == {}
+    assert parsed_event.prompt_session_attributes == {}
+    assert parsed_event.parameters == []
+
+
+def test_bedrock_agent_function_event_validation():
+    """Test validation of required fields"""
+    # Test missing required field
+    with pytest.raises(ValueError, match="Missing required field: messageVersion"):
+        BedrockAgentFunctionEvent({})
+
+    # Test invalid field type
+    invalid_event = {
+        "messageVersion": 1,  # should be string
+        "agent": {"alias": "PROD", "name": "hr-assistant", "version": "1", "id": "1234"},
+        "inputText": "",
+        "sessionId": "",
+        "actionGroup": "",
+        "function": "",
+    }
+    with pytest.raises(TypeError, match="Field messageVersion must be of type <class 'str'>"):
+        BedrockAgentFunctionEvent(invalid_event)
+
+    # Test missing agent fields
+    invalid_agent_event = {
+        "messageVersion": "1.0",
+        "agent": {"name": "test"},  # missing required agent fields
+        "inputText": "",
+        "sessionId": "",
+        "actionGroup": "",
+        "function": "",
+    }
+    with pytest.raises(ValueError, match="Agent object missing required fields"):
+        BedrockAgentFunctionEvent(invalid_agent_event)