Fix middleware validation

Author: leandrodamascena

aws_lambda_powertools/event_handler/api_gateway.py

@@ -255,6 +255,35 @@ def build_allow_methods(methods: set[str]) -> str:
         return ",".join(sorted(methods))
 
 
+class BedrockResponse(Generic[ResponseT]):
+    """
+    Contains the response body, status code, content type, and optional attributes
+    for session management and knowledge base configuration.
+    """
+
+    def __init__(
+        self,
+        body: Any = None,
+        status_code: int = 200,
+        content_type: str = "application/json",
+        session_attributes: dict[str, Any] | None = None,
+        prompt_session_attributes: dict[str, Any] | None = None,
+        knowledge_bases_configuration: list[dict[str, Any]] | None = None,
+    ) -> None:
+        self.body = body
+        self.status_code = status_code
+        self.content_type = content_type
+        self.session_attributes = session_attributes
+        self.prompt_session_attributes = prompt_session_attributes
+        self.knowledge_bases_configuration = knowledge_bases_configuration
+
+    def is_json(self) -> bool:
+        """
+        Returns True if the response is JSON, based on the Content-Type.
+        """
+        return True
+
+
 class Response(Generic[ResponseT]):
     """Response data class that provides greater control over what is returned from the proxy event"""
 
@@ -1474,7 +1503,10 @@ def __call__(self, app: ApiGatewayResolver) -> dict | tuple | Response:
         return self.current_middleware(app, self.next_middleware)
 
 
-def _registered_api_adapter(app: ApiGatewayResolver, next_middleware: Callable[..., Any]) -> dict | tuple | Response:
+def _registered_api_adapter(
+    app: ApiGatewayResolver,
+    next_middleware: Callable[..., Any],
+) -> dict | tuple | Response | BedrockResponse:
     """
     Calls the registered API using the "_route_args" from the Resolver context to ensure the last call
     in the chain will match the API route function signature and ensure that Powertools passes the API
@@ -2538,7 +2570,7 @@ def _call_route(self, route: Route, route_arguments: dict[str, str]) -> Response
             self._reset_processed_stack()
 
             return self._response_builder_class(
-                response=self._to_response(
+                response=self._to_response(  # type: ignore[arg-type]
                     route(router_middlewares=self._router_middlewares, app=self, route_arguments=route_arguments),
                 ),
                 serializer=self._serializer,
@@ -2627,7 +2659,7 @@ def _call_exception_handler(self, exp: Exception, route: Route) -> ResponseBuild
 
         return None
 
-    def _to_response(self, result: dict | tuple | Response) -> Response:
+    def _to_response(self, result: dict | tuple | Response | BedrockResponse) -> Response | BedrockResponse:
         """Convert the route's result to a Response
 
          3 main result types are supported:
@@ -2638,7 +2670,7 @@ def _to_response(self, result: dict | tuple | Response) -> Response:
         - Response: returned as is, and allows for more flexibility
         """
         status_code = HTTPStatus.OK
-        if isinstance(result, Response):
+        if isinstance(result, (Response, BedrockResponse)):
             return result
         elif isinstance(result, tuple) and len(result) == 2:
             # Unpack result dict and status code from tuple
@@ -2971,8 +3003,9 @@ def _get_base_path(self) -> str:
         # ALB doesn't have a stage variable, so we just return an empty string
         return ""
 
+    # BedrockResponse is not used here but adding the same signature to keep strong typing
     @override
-    def _to_response(self, result: dict | tuple | Response) -> Response:
+    def _to_response(self, result: dict | tuple | Response | BedrockResponse) -> Response | BedrockResponse:
         """Convert the route's result to a Response
 
         ALB requires a non-null body otherwise it converts as HTTP 5xx

aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -8,6 +8,7 @@
 from aws_lambda_powertools.event_handler import ApiGatewayResolver
 from aws_lambda_powertools.event_handler.api_gateway import (
     _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
+    BedrockResponse,
     ProxyEventType,
     ResponseBuilder,
 )
@@ -23,29 +24,6 @@
     from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
 
 
-class BedrockResponse:
-    """
-    Contains the response body, status code, content type, and optional attributes
-    for session management and knowledge base configuration.
-    """
-
-    def __init__(
-        self,
-        body: Any = None,
-        status_code: int = 200,
-        content_type: str = "application/json",
-        session_attributes: dict[str, Any] | None = None,
-        prompt_session_attributes: dict[str, Any] | None = None,
-        knowledge_bases_configuration: list[dict[str, Any]] | None = None,
-    ) -> None:
-        self.body = body
-        self.status_code = status_code
-        self.content_type = content_type
-        self.session_attributes = session_attributes
-        self.prompt_session_attributes = prompt_session_attributes
-        self.knowledge_bases_configuration = knowledge_bases_configuration
-
-
 class BedrockResponseBuilder(ResponseBuilder):
     """
     Bedrock Response Builder. This builds the response dict to be returned by Lambda when using Bedrock Agents.
@@ -55,18 +33,9 @@ class BedrockResponseBuilder(ResponseBuilder):
 
     @override
     def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]:
-        """Build the full response dict to be returned by the lambda"""
-        self._route(event, None)
-
-        bedrock_response = None
-        if isinstance(self.response.body, dict) and "body" in self.response.body:
-            bedrock_response = BedrockResponse(**self.response.body)
-            body = bedrock_response.body
-        else:
-            body = self.response.body
-
-        if self.response.is_json() and not isinstance(body, str):
-            body = self.serializer(body)
+        body = self.response.body
+        if self.response.is_json() and not isinstance(self.response.body, str):
+            body = self.serializer(self.response.body)
 
         response = {
             "messageVersion": "1.0",
@@ -84,13 +53,15 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]:
         }
 
         # Add Bedrock-specific attributes
-        if bedrock_response:
-            if bedrock_response.session_attributes:
-                response["sessionAttributes"] = bedrock_response.session_attributes
-            if bedrock_response.prompt_session_attributes:
-                response["promptSessionAttributes"] = bedrock_response.prompt_session_attributes
-            if bedrock_response.knowledge_bases_configuration:
-                response["knowledgeBasesConfiguration"] = bedrock_response.knowledge_bases_configuration  # type: ignore
+        if isinstance(self.response, BedrockResponse):
+            if self.response.session_attributes:
+                response["sessionAttributes"] = self.response.session_attributes
+
+            if self.response.prompt_session_attributes:
+                response["promptSessionAttributes"] = self.response.prompt_session_attributes
+
+            if self.response.knowledge_bases_configuration:
+                response["knowledgeBasesConfiguration"] = self.response.knowledge_bases_configuration
 
         return response
 

docs/core/event_handler/bedrock_agents.md

@@ -321,20 +321,18 @@ You can enable user confirmation with Bedrock Agents to have your application as
 --8<-- "examples/event_handler_bedrock_agents/src/enabling_user_confirmation.py"
 ```
 
-### Fine grained responses
-
-`BedrockResponse` class that provides full control over Bedrock Agent responses.
+1. Add an openapi extension
 
-You can use this class to add additional fields as needed, such as [session attributes, prompt session attributes, and knowledge base configurations](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html#agents-lambda-response).
+### Fine grained responses
 
 ???+ info "Note"
 	The default response only includes the essential fields to keep the payload size minimal, as AWS Lambda has a maximum response size of 25 KB.
 
-```python title="bedrockresponse.py" title="Customzing your Bedrock Response"
---8<-- "examples/event_handler_bedrock_agents/src/bedrockresponse.py"
-```
+You can use `BedrockResponse` class to add additional fields as needed, such as [session attributes, prompt session attributes, and knowledge base configurations](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html#agents-lambda-response){target="_blank"}.
 
-1. Add an openapi extension
+```python title="working_with_bedrockresponse.py" title="Customzing your Bedrock Response" hl_lines="5 16"
+--8<-- "examples/event_handler_bedrock_agents/src/working_with_bedrockresponse.py"
+```
 
 ## Testing your code
 

examples/event_handler_bedrock_agents/src/bedrockresponse.py

@@ -0,0 +1,35 @@
+from http import HTTPStatus
+
+from aws_lambda_powertools import Logger, Tracer
+from aws_lambda_powertools.event_handler import BedrockAgentResolver
+from aws_lambda_powertools.event_handler.api_gateway import BedrockResponse
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+tracer = Tracer()
+logger = Logger()
+app = BedrockAgentResolver()
+
+
+@app.get("/return_with_session", description="Returns a hello world with session attributes")
+@tracer.capture_method
+def hello_world():
+    return BedrockResponse(
+        status_code=HTTPStatus.OK.value,
+        body={"message": "Hello from Bedrock!"},
+        session_attributes={"user_id": "123"},
+        prompt_session_attributes={"context": "testing"},
+        knowledge_bases_configuration=[
+            {
+                "knowledgeBaseId": "kb-123",
+                "retrievalConfiguration": {
+                    "vectorSearchConfiguration": {"numberOfResults": 3, "overrideSearchType": "HYBRID"},
+                },
+            },
+        ],
+    )
+
+
+@logger.inject_lambda_context
+@tracer.capture_lambda_handler
+def lambda_handler(event: dict, context: LambdaContext):
+    return app.resolve(event, context)

examples/event_handler_bedrock_agents/src/working_with_bedrockresponse.py

@@ -266,7 +266,7 @@ def test_bedrock_agent_with_partial_bedrock_response():
     app = BedrockAgentResolver()
 
     @app.get("/claims", description="Gets claims")
-    def claims():
+    def claims() -> Dict[str, Any]:
         return BedrockResponse(
             body={"message": "test"},
             session_attributes={"user_id": "123"},
@@ -289,7 +289,7 @@ def test_bedrock_agent_with_different_attributes_combination():
     app = BedrockAgentResolver()
 
     @app.get("/claims", description="Gets claims")
-    def claims():
+    def claims() -> Dict[str, Any]:
         return BedrockResponse(
             body={"message": "test"},
             prompt_session_attributes={"context": "testing"},