Merge branch 'feat/bedrock_functions' of https://github.com/aws-powertools/powertools-lambda-python into feat/bedrock_functions

Author: anafalcao

.cfnlintrc.yaml

@@ -0,0 +1,2 @@
+ignore_templates:
+  - examples/event_handler_appsync_events/sam/getting_started_with_appsync_events.yaml

CHANGELOG.md

@@ -4,8 +4,40 @@
 <a name="unreleased"></a>
 # Unreleased
 
-## Maintenance
+## Documentation
+
+* **appsync_events:** improve AppSync events documentation ([#6572](https://github.com/aws-powertools/powertools-lambda-python/issues/6572))
+* **i-made-this:** adding Michael's MCP server ([#6591](https://github.com/aws-powertools/powertools-lambda-python/issues/6591))
+
+## Features
 
+* **bedrock_agents:** add optional fields to response payload ([#6336](https://github.com/aws-powertools/powertools-lambda-python/issues/6336))
+
+## Maintenance
+
+* **ci:** new pre-release 3.11.1a3 ([#6582](https://github.com/aws-powertools/powertools-lambda-python/issues/6582))
+* **ci:** new pre-release 3.11.1a0 ([#6561](https://github.com/aws-powertools/powertools-lambda-python/issues/6561))
+* **ci:** new pre-release 3.11.1a4 ([#6589](https://github.com/aws-powertools/powertools-lambda-python/issues/6589))
+* **ci:** new pre-release 3.11.1a1 ([#6574](https://github.com/aws-powertools/powertools-lambda-python/issues/6574))
+* **ci:** new pre-release 3.11.1a2 ([#6578](https://github.com/aws-powertools/powertools-lambda-python/issues/6578))
+* **ci:** new pre-release 3.11.1a5 ([#6598](https://github.com/aws-powertools/powertools-lambda-python/issues/6598))
+* **deps:** bump pydantic from 2.11.3 to 2.11.4 ([#6585](https://github.com/aws-powertools/powertools-lambda-python/issues/6585))
+* **deps-dev:** bump boto3-stubs from 1.38.2 to 1.38.3 ([#6569](https://github.com/aws-powertools/powertools-lambda-python/issues/6569))
+* **deps-dev:** bump boto3-stubs from 1.38.5 to 1.38.6 ([#6587](https://github.com/aws-powertools/powertools-lambda-python/issues/6587))
+* **deps-dev:** bump boto3-stubs from 1.38.3 to 1.38.4 ([#6577](https://github.com/aws-powertools/powertools-lambda-python/issues/6577))
+* **deps-dev:** bump aws-cdk from 2.1012.0 to 2.1013.0 ([#6588](https://github.com/aws-powertools/powertools-lambda-python/issues/6588))
+* **deps-dev:** bump h11 from 0.14.0 to 0.16.0 ([#6575](https://github.com/aws-powertools/powertools-lambda-python/issues/6575))
+* **deps-dev:** bump aws-cdk-aws-lambda-python-alpha from 2.192.0a0 to 2.193.0a0 ([#6586](https://github.com/aws-powertools/powertools-lambda-python/issues/6586))
+* **deps-dev:** bump ruff from 0.11.7 to 0.11.8 ([#6595](https://github.com/aws-powertools/powertools-lambda-python/issues/6595))
+* **deps-dev:** bump boto3-stubs from 1.38.4 to 1.38.5 ([#6581](https://github.com/aws-powertools/powertools-lambda-python/issues/6581))
+* **deps-dev:** bump cfn-lint from 1.34.1 to 1.34.2 ([#6568](https://github.com/aws-powertools/powertools-lambda-python/issues/6568))
+* **deps-dev:** bump cdklabs-generative-ai-cdk-constructs from 0.1.305 to 0.1.306 ([#6567](https://github.com/aws-powertools/powertools-lambda-python/issues/6567))
+* **deps-dev:** bump aws-cdk-aws-lambda-python-alpha from 2.191.0a0 to 2.192.0a0 ([#6566](https://github.com/aws-powertools/powertools-lambda-python/issues/6566))
+* **deps-dev:** bump aws-cdk-lib from 2.191.0 to 2.192.0 ([#6565](https://github.com/aws-powertools/powertools-lambda-python/issues/6565))
+* **deps-dev:** bump boto3-stubs from 1.38.6 to 1.38.7 ([#6594](https://github.com/aws-powertools/powertools-lambda-python/issues/6594))
+* **deps-dev:** bump cdklabs-generative-ai-cdk-constructs from 0.1.307 to 0.1.308 ([#6597](https://github.com/aws-powertools/powertools-lambda-python/issues/6597))
+* **deps-dev:** bump cdklabs-generative-ai-cdk-constructs from 0.1.306 to 0.1.307 ([#6580](https://github.com/aws-powertools/powertools-lambda-python/issues/6580))
+* **docs:** fix youtube embed link in we made this ([#6593](https://github.com/aws-powertools/powertools-lambda-python/issues/6593))
 
 
 <a name="v3.11.0"></a>

aws_lambda_powertools/event_handler/__init__.py

@@ -11,11 +11,12 @@
     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.bedrock_agent_function import (
     BedrockAgentFunctionResolver,
     BedrockFunctionResponse,
 )
+
 from aws_lambda_powertools.event_handler.events_appsync.appsync_events import AppSyncEventsResolver
 from aws_lambda_powertools.event_handler.lambda_function_url import (
     LambdaFunctionUrlResolver,
@@ -31,10 +32,11 @@
     "ApiGatewayResolver",
     "BedrockAgentResolver",
     "BedrockAgentFunctionResolver",
+    "BedrockResponse",
+    "BedrockFunctionResponse",
     "CORSConfig",
     "LambdaFunctionUrlResolver",
     "Response",
-    "BedrockFunctionResponse",
     "VPCLatticeResolver",
     "VPCLatticeV2Resolver",
 ]

aws_lambda_powertools/event_handler/api_gateway.py

@@ -73,6 +73,7 @@
 _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION = "Successful Response"
 _ROUTE_REGEX = "^{}$"
 _JSON_DUMP_CALL = partial(json.dumps, separators=(",", ":"), cls=Encoder)
+_DEFAULT_CONTENT_TYPE = "application/json"
 
 ResponseEventT = TypeVar("ResponseEventT", bound=BaseProxyEvent)
 ResponseT = TypeVar("ResponseT")
@@ -255,6 +256,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 = _DEFAULT_CONTENT_TYPE,
+        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"""
 
@@ -300,7 +330,7 @@ def is_json(self) -> bool:
         content_type = self.headers.get("Content-Type", "")
         if isinstance(content_type, list):
             content_type = content_type[0]
-        return content_type.startswith("application/json")
+        return content_type.startswith(_DEFAULT_CONTENT_TYPE)
 
 
 class Route:
@@ -572,7 +602,7 @@ def _get_openapi_path(
         operation_responses: dict[int, OpenAPIResponse] = {
             422: {
                 "description": "Validation Error",
-                "content": {"application/json": {"schema": {"$ref": f"{COMPONENT_REF_PREFIX}HTTPValidationError"}}},
+                "content": {_DEFAULT_CONTENT_TYPE: {"schema": {"$ref": f"{COMPONENT_REF_PREFIX}HTTPValidationError"}}},
             },
         }
 
@@ -581,7 +611,9 @@ def _get_openapi_path(
             http_code = self.custom_response_validation_http_code.value
             operation_responses[http_code] = {
                 "description": "Response Validation Error",
-                "content": {"application/json": {"schema": {"$ref": f"{COMPONENT_REF_PREFIX}ResponseValidationError"}}},
+                "content": {
+                    _DEFAULT_CONTENT_TYPE: {"schema": {"$ref": f"{COMPONENT_REF_PREFIX}ResponseValidationError"}},
+                },
             }
             # Add model definition
             definitions["ResponseValidationError"] = response_validation_error_response_definition
@@ -594,7 +626,7 @@ def _get_openapi_path(
                 # Case 1: there is not 'content' key
                 if "content" not in response:
                     response["content"] = {
-                        "application/json": self._openapi_operation_return(
+                        _DEFAULT_CONTENT_TYPE: self._openapi_operation_return(
                             param=dependant.return_param,
                             model_name_map=model_name_map,
                             field_mapping=field_mapping,
@@ -645,7 +677,7 @@ def _get_openapi_path(
             # Add the response schema to the OpenAPI 200 response
             operation_responses[200] = {
                 "description": self.response_description or _DEFAULT_OPENAPI_RESPONSE_DESCRIPTION,
-                "content": {"application/json": response_schema},
+                "content": {_DEFAULT_CONTENT_TYPE: response_schema},
             }
 
         operation["responses"] = operation_responses
@@ -1474,7 +1506,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
@@ -1632,7 +1667,7 @@ def _add_resolver_response_validation_error_response_to_route(
             response_validation_error_response = {
                 "description": "Response Validation Error",
                 "content": {
-                    "application/json": {
+                    _DEFAULT_CONTENT_TYPE: {
                         "schema": {"$ref": f"{COMPONENT_REF_PREFIX}ResponseValidationError"},
                     },
                 },
@@ -2151,7 +2186,7 @@ def swagger_handler():
             if query_params.get("format") == "json":
                 return Response(
                     status_code=200,
-                    content_type="application/json",
+                    content_type=_DEFAULT_CONTENT_TYPE,
                     body=escaped_spec,
                 )
 
@@ -2538,7 +2573,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 +2662,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 +2673,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 +3006,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,
 )
@@ -32,14 +33,11 @@ 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)
-
         body = self.response.body
         if self.response.is_json() and not isinstance(self.response.body, str):
             body = self.serializer(self.response.body)
 
-        return {
+        response = {
             "messageVersion": "1.0",
             "response": {
                 "actionGroup": event.action_group,
@@ -54,6 +52,19 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]:
             },
         }
 
+        # Add Bedrock-specific attributes
+        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
+
 
 class BedrockAgentResolver(ApiGatewayResolver):
     """Bedrock Agent Resolver

aws_lambda_powertools/shared/version.py

@@ -1,3 +1,3 @@
 """Exposes version constant to avoid circular dependencies."""
 
-VERSION = "3.11.1a0"
+VERSION = "3.11.1a6"