feat(bedrock_agents): add optional fields to response payload

Author: anafalcao

aws_lambda_powertools/event_handler/__init__.py

@@ -11,7 +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.bedrock_agent import BedrockAgentResolver, BedrockResponse
 from aws_lambda_powertools.event_handler.lambda_function_url import (
     LambdaFunctionUrlResolver,
 )
@@ -24,6 +24,7 @@
     "ALBResolver",
     "ApiGatewayResolver",
     "BedrockAgentResolver",
+    "BedrockResponse",
     "CORSConfig",
     "LambdaFunctionUrlResolver",
     "Response",

aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -21,6 +21,43 @@
     from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
 
 
+class BedrockResponse:
+    """
+    Response class for Bedrock Agents Lambda functions.
+
+    Parameters
+    ----------
+    status_code: int
+        HTTP status code for the response
+    body: Any
+        Response body content
+    content_type: str, optional
+        Content type of the response (default: application/json)
+    session_attributes: dict[str, Any], optional
+        Session attributes to maintain state
+    prompt_session_attributes: dict[str, Any], optional
+        Prompt-specific session attributes
+    knowledge_bases_configuration: dict[str, Any], optional
+        Knowledge base configuration settings
+    """
+
+    def __init__(
+        self,
+        status_code: int,
+        body: Any,
+        content_type: str = "application/json",
+        session_attributes: dict[str, Any] | None = None,
+        prompt_session_attributes: dict[str, Any] | None = None,
+        knowledge_bases_configuration: dict[str, Any] | None = None,
+    ) -> None:
+        self.status_code = status_code
+        self.body = body
+        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.
@@ -30,14 +67,12 @@ class BedrockResponseBuilder(ResponseBuilder):
 
     @override
     def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]:
-        """Build the full response dict to be returned by the lambda"""
+        """
+        Build the response dictionary to be returned by the Lambda function.
+        """
         self._route(event, None)
 
-        body = self.response.body
-        if self.response.is_json() and not isinstance(self.response.body, str):
-            body = self.serializer(self.response.body)
-
-        return {
+        base_response = {
             "messageVersion": "1.0",
             "response": {
                 "actionGroup": event.action_group,
@@ -46,12 +81,44 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]:
                 "httpStatusCode": self.response.status_code,
                 "responseBody": {
                     self.response.content_type: {
-                        "body": body,
+                        "body": self._get_formatted_body(),
                     },
                 },
             },
         }
 
+        if isinstance(self.response, BedrockResponse):
+            self._add_bedrock_specific_configs(base_response)
+
+        return base_response
+
+    def _get_formatted_body(self) -> Any:
+        """Format the response body based on content type"""
+        if not isinstance(self.response, BedrockResponse):
+            if self.response.is_json() and not isinstance(self.response.body, str):
+                return self.serializer(self.response.body)
+        return self.response.body
+
+    def _add_bedrock_specific_configs(self, response: dict[str, Any]) -> None:
+        """
+        Add Bedrock-specific configurations to the response if present.
+
+        Parameters
+        ----------
+        response: dict[str, Any]
+            The base response dictionary to be updated
+        """
+        if not isinstance(self.response, BedrockResponse):
+            return
+
+        optional_configs = {
+            "sessionAttributes": self.response.session_attributes,
+            "promptSessionAttributes": self.response.prompt_session_attributes,
+            "knowledgeBasesConfiguration": self.response.knowledge_bases_configuration,
+        }
+
+        response.update({k: v for k, v in optional_configs.items() if v is not None})
+
 
 class BedrockAgentResolver(ApiGatewayResolver):
     """Bedrock Agent Resolver

docs/core/event_handler/bedrock_agents.md

@@ -313,6 +313,19 @@ To implement these customizations, include extra parameters when defining your r
 --8<-- "examples/event_handler_bedrock_agents/src/customizing_bedrock_api_operations.py"
 ```
 
+### Fine grained responses
+
+`BedrockResponse` class that provides full control over Bedrock Agent responses.
+
+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).
+
+???+ 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"
+```
+
 ## Testing your code
 
 Test your routes by passing an [Agent for Amazon Bedrock proxy event](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html#agents-lambda-input) request:

examples/event_handler_bedrock_agents/src/bedrockresponse.py

@@ -0,0 +1,14 @@
+from http import HTTPStatus
+
+from aws_lambda_powertools.event_handler import BedrockResponse
+
+response = 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"}},
+    },
+)

tests/functional/event_handler/_pydantic/test_bedrock_agent.py

@@ -4,7 +4,7 @@
 import pytest
 from typing_extensions import Annotated
 
-from aws_lambda_powertools.event_handler import BedrockAgentResolver, Response, content_types
+from aws_lambda_powertools.event_handler import BedrockAgentResolver, BedrockResponse, Response, content_types
 from aws_lambda_powertools.event_handler.openapi.params import Body
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
 from tests.functional.utils import load_event
@@ -200,3 +200,43 @@ def handler() -> Optional[Dict]:
     # THEN the schema must be a valid 3.0.3 version
     assert openapi30_schema(schema)
     assert schema.get("openapi") == "3.0.3"
+
+
+def test_bedrock_agent_with_bedrock_response():
+    # GIVEN a Bedrock Agent resolver
+    app = BedrockAgentResolver()
+
+    @app.get("/claims", description="Gets claims")
+    def claims():
+        return BedrockResponse(
+            status_code=200,
+            body={"output": claims_response},
+            session_attributes={"last_request": "get_claims"},
+            prompt_session_attributes={"context": "claims_query"},
+            knowledge_bases_configuration={
+                "knowledgeBaseId": "kb-123",
+                "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 3}},
+            },
+        )
+
+    # WHEN calling the event handler
+    result = app(load_event("bedrockAgentEvent.json"), {})
+
+    # THEN process event correctly
+    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
+
+    # AND return the correct body
+    body = result["response"]["responseBody"]["application/json"]["body"]
+    assert json.loads(body) == {"output": claims_response}
+
+    # AND include the optional configurations
+    assert result["sessionAttributes"] == {"last_request": "get_claims"}
+    assert result["promptSessionAttributes"] == {"context": "claims_query"}
+    assert result["knowledgeBasesConfiguration"] == {
+        "knowledgeBaseId": "kb-123",
+        "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 3}},
+    }