feat(event_handler): add Bedrock Agent event handler

Author: rubenfonseca

Diff for: aws_lambda_powertools/event_handler/__init__.py

@@ -11,6 +11,7 @@
     Response,
 )
 from aws_lambda_powertools.event_handler.appsync import AppSyncResolver
+from aws_lambda_powertools.event_handler.bedrock_agent import BedrockAgentResolver
 from aws_lambda_powertools.event_handler.lambda_function_url import (
     LambdaFunctionUrlResolver,
 )
@@ -22,6 +23,7 @@
     "APIGatewayHttpResolver",
     "ALBResolver",
     "ApiGatewayResolver",
+    "BedrockAgentResolver",
     "CORSConfig",
     "LambdaFunctionUrlResolver",
     "Response",

Diff for: aws_lambda_powertools/event_handler/api_gateway.py

@@ -42,6 +42,7 @@
     ALBEvent,
     APIGatewayProxyEvent,
     APIGatewayProxyEventV2,
+    BedrockAgentEvent,
     LambdaFunctionUrlEvent,
     VPCLatticeEvent,
     VPCLatticeEventV2,
@@ -83,6 +84,7 @@ class ProxyEventType(Enum):
     APIGatewayProxyEvent = "APIGatewayProxyEvent"
     APIGatewayProxyEventV2 = "APIGatewayProxyEventV2"
     ALBEvent = "ALBEvent"
+    BedrockAgentEvent = "BedrockAgentEvent"
     VPCLatticeEvent = "VPCLatticeEvent"
     VPCLatticeEventV2 = "VPCLatticeEventV2"
     LambdaFunctionUrlEvent = "LambdaFunctionUrlEvent"
@@ -1291,6 +1293,7 @@ def __init__(
         self._strip_prefixes = strip_prefixes
         self.context: Dict = {}  # early init as customers might add context before event resolution
         self.processed_stack_frames = []
+        self.response_builder_class = ResponseBuilder
 
         # Allow for a custom serializer or a concise json serialization
         self._serializer = serializer or partial(json.dumps, separators=(",", ":"), cls=Encoder)
@@ -1647,14 +1650,17 @@ def _compile_regex(rule: str, base_regex: str = _ROUTE_REGEX):
         rule_regex: str = re.sub(_DYNAMIC_ROUTE_PATTERN, _NAMED_GROUP_BOUNDARY_PATTERN, rule)
         return re.compile(base_regex.format(rule_regex))
 
-    def _to_proxy_event(self, event: Dict) -> BaseProxyEvent:
+    def _to_proxy_event(self, event: Dict) -> BaseProxyEvent:  # noqa: PLR0911
         """Convert the event dict to the corresponding data class"""
         if self._proxy_type == ProxyEventType.APIGatewayProxyEvent:
             logger.debug("Converting event to API Gateway REST API contract")
             return APIGatewayProxyEvent(event)
         if self._proxy_type == ProxyEventType.APIGatewayProxyEventV2:
             logger.debug("Converting event to API Gateway HTTP API contract")
             return APIGatewayProxyEventV2(event)
+        if self._proxy_type == ProxyEventType.BedrockAgentEvent:
+            logger.debug("Converting event to Bedrock Agent contract")
+            return BedrockAgentEvent(event)
         if self._proxy_type == ProxyEventType.LambdaFunctionUrlEvent:
             logger.debug("Converting event to Lambda Function URL contract")
             return LambdaFunctionUrlEvent(event)
@@ -1732,9 +1738,9 @@ def _not_found(self, method: str) -> ResponseBuilder:
 
         handler = self._lookup_exception_handler(NotFoundError)
         if handler:
-            return ResponseBuilder(handler(NotFoundError()))
+            return self.response_builder_class(handler(NotFoundError()))
 
-        return ResponseBuilder(
+        return self.response_builder_class(
             Response(
                 status_code=HTTPStatus.NOT_FOUND.value,
                 content_type=content_types.APPLICATION_JSON,
@@ -1749,7 +1755,7 @@ def _call_route(self, route: Route, route_arguments: Dict[str, str]) -> Response
             # Reset Processed stack for Middleware (for debugging purposes)
             self._reset_processed_stack()
 
-            return ResponseBuilder(
+            return self.response_builder_class(
                 self._to_response(
                     route(router_middlewares=self._router_middlewares, app=self, route_arguments=route_arguments),
                 ),
@@ -1766,7 +1772,7 @@ def _call_route(self, route: Route, route_arguments: Dict[str, str]) -> Response
                 # If the user has turned on debug mode,
                 # we'll let the original exception propagate, so
                 # they get more information about what went wrong.
-                return ResponseBuilder(
+                return self.response_builder_class(
                     Response(
                         status_code=500,
                         content_type=content_types.TEXT_PLAIN,
@@ -1805,12 +1811,12 @@ def _call_exception_handler(self, exp: Exception, route: Route) -> Optional[Resp
         handler = self._lookup_exception_handler(type(exp))
         if handler:
             try:
-                return ResponseBuilder(handler(exp), route)
+                return self.response_builder_class(handler(exp), route)
             except ServiceError as service_error:
                 exp = service_error
 
         if isinstance(exp, ServiceError):
-            return ResponseBuilder(
+            return self.response_builder_class(
                 Response(
                     status_code=exp.status_code,
                     content_type=content_types.APPLICATION_JSON,

Diff for: aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -0,0 +1,71 @@
+import logging
+from typing import Any, Dict, Optional, cast
+
+from aws_lambda_powertools.event_handler import ApiGatewayResolver
+from aws_lambda_powertools.event_handler.api_gateway import CORSConfig, ProxyEventType, ResponseBuilder
+from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
+from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent
+
+logger = logging.getLogger(__name__)
+
+
+class BedrockResponseBuilder(ResponseBuilder):
+    def build(self, event: BaseProxyEvent, cors: Optional[CORSConfig] = None) -> Dict[str, Any]:
+        """Build the full response dict to be returned by the lambda"""
+        self._route(event, cors)
+
+        bedrock_event = cast(BedrockAgentEvent, event)
+
+        return {
+            "messageVersion": "1.0",
+            "response": {
+                "actionGroup": bedrock_event.action_group,
+                "apiPath": bedrock_event.api_path,
+                "httpMethod": bedrock_event.http_method,
+                "httpStatusCode": self.response.status_code,
+                "responseBody": {
+                    "application/json": {
+                        "body": self.response.body,
+                    },
+                },
+            },
+        }
+
+
+class BedrockAgentResolver(ApiGatewayResolver):
+    """Bedrock Agent Resolver
+
+    See https://aws.amazon.com/bedrock/agents/ for more information.
+
+    Examples
+    --------
+    Simple example with a custom lambda handler using the Tracer capture_lambda_handler decorator
+
+    ```python
+    from aws_lambda_powertools import Tracer
+    from aws_lambda_powertools.event_handler import BedrockAgentResolver
+
+    tracer = Tracer()
+    app = BedrockAgentResolver()
+
+    @app.get("/claims")
+    def simple_get():
+        return "You have 3 claims"
+
+    @tracer.capture_lambda_handler
+    def lambda_handler(event, context):
+        return app.resolve(event, context)
+    """
+
+    current_event: BedrockAgentEvent
+
+    def __init__(self, debug: bool = False, enable_validation: bool = True):
+        super().__init__(
+            ProxyEventType.BedrockAgentEvent,
+            None,
+            debug,
+            None,
+            None,
+            enable_validation,
+        )
+        self.response_builder_class = BedrockResponseBuilder

Diff for: aws_lambda_powertools/shared/headers_serializer.py

@@ -123,3 +123,12 @@ def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Coo
                 payload["headers"][key] = values[-1]
 
         return payload
+
+
+class NoopSerializer(BaseHeadersSerializer):
+    """
+    Noop serializer that doesn't do anything. This is useful for resolvers that don't need to set headers or cookies.
+    """
+
+    def serialize(self, headers: Dict[str, Union[str, List[str]]], cookies: List[Cookie]) -> Dict[str, Any]:
+        return {}

Diff for: aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py

@@ -1,6 +1,7 @@
 from typing import Dict, List, Optional
 
-from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
+from aws_lambda_powertools.shared.headers_serializer import BaseHeadersSerializer, NoopSerializer
+from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent, DictWrapper
 
 
 class BedrockAgentInfo(DictWrapper):
@@ -47,7 +48,7 @@ def content(self) -> Dict[str, BedrockAgentRequestMedia]:
         return {k: BedrockAgentRequestMedia(v) for k, v in self["content"].items()}
 
 
-class BedrockAgentEvent(DictWrapper):
+class BedrockAgentEvent(BaseProxyEvent):
     """
     Bedrock Agent input event
 
@@ -97,3 +98,11 @@ def session_attributes(self) -> Dict[str, str]:
     @property
     def prompt_session_attributes(self) -> Dict[str, str]:
         return self["promptSessionAttributes"]
+
+    # For compatibility with BaseProxyEvent
+    @property
+    def path(self) -> str:
+        return self["apiPath"]
+
+    def header_serializer(self) -> BaseHeadersSerializer:
+        return NoopSerializer()

Diff for: tests/functional/event_handler/test_bedrock_agent.py

@@ -0,0 +1,79 @@
+import json
+from typing import Any, Dict
+
+from aws_lambda_powertools.event_handler import BedrockAgentResolver, Response, content_types
+from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
+from tests.functional.utils import load_event
+
+claims_response = "You have 3 claims"
+
+
+def test_bedrock_agent_event():
+    # GIVEN a Bedrock Agent event
+    app = BedrockAgentResolver()
+
+    @app.get("/claims")
+    def claims() -> Dict[str, Any]:
+        assert isinstance(app.current_event, BedrockAgentEvent)
+        assert app.lambda_context == {}
+        return {"output": claims_response}
+
+    # WHEN calling the event handler
+    result = app(load_event("bedrockAgentEvent.json"), {})
+
+    # THEN process event correctly
+    # AND set the current_event type as BedrockAgentEvent
+    assert result["messageVersion"] == "1.0"
+    assert result["response"]["apiPath"] == "/claims"
+    assert result["response"]["actionGroup"] == "ClaimManagementActionGroup"
+    assert result["response"]["httpMethod"] == "GET"
+    assert result["response"]["httpStatusCode"] == 200
+
+    body = result["response"]["responseBody"]["application/json"]["body"]
+    assert body == {"output": claims_response}
+
+
+def test_bedrock_agent_event_with_response():
+    # GIVEN a Bedrock Agent event
+    app = BedrockAgentResolver()
+    output = json.dumps({"output": claims_response})
+
+    @app.get("/claims")
+    def claims():
+        assert isinstance(app.current_event, BedrockAgentEvent)
+        assert app.lambda_context == {}
+        return Response(200, content_types.APPLICATION_JSON, output)
+
+    # WHEN calling the event handler
+    result = app(load_event("bedrockAgentEvent.json"), {})
+
+    # THEN process event correctly
+    # AND set the current_event type as BedrockAgentEvent
+    assert result["messageVersion"] == "1.0"
+    assert result["response"]["apiPath"] == "/claims"
+    assert result["response"]["actionGroup"] == "ClaimManagementActionGroup"
+    assert result["response"]["httpMethod"] == "GET"
+    assert result["response"]["httpStatusCode"] == 200
+
+    body = result["response"]["responseBody"]["application/json"]["body"]
+    assert body == output
+
+
+def test_bedrock_agent_event_with_no_matches():
+    # GIVEN a Bedrock Agent event
+    app = BedrockAgentResolver()
+
+    @app.get("/no_match")
+    def claims():
+        raise RuntimeError()
+
+    # WHEN calling the event handler
+    result = app(load_event("bedrockAgentEvent.json"), {})
+
+    # THEN process event correctly
+    # AND return 404 because the event doesn't match any known rule
+    assert result["messageVersion"] == "1.0"
+    assert result["response"]["apiPath"] == "/claims"
+    assert result["response"]["actionGroup"] == "ClaimManagementActionGroup"
+    assert result["response"]["httpMethod"] == "GET"
+    assert result["response"]["httpStatusCode"] == 404