feat(bedrock): add openapi_extensions in BedrockAgentResolver (#6510)

* Adding OpenAPI extensions to Bedrock

* Adding OpenAPI extensions to Bedrock

* Adding OpenAPI extensions to Bedrock

* Adding OpenAPI extensions to Bedrock

Author: leandrodamascena

aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -110,11 +110,11 @@ def get(  # type: ignore[override]
         tags: list[str] | None = None,
         operation_id: str | None = None,
         include_in_schema: bool = True,
+        openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
-        openapi_extensions = None
         security = None
 
         return super().get(
@@ -151,11 +151,11 @@ def post(  # type: ignore[override]
         tags: list[str] | None = None,
         operation_id: str | None = None,
         include_in_schema: bool = True,
+        openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ):
-        openapi_extensions = None
         security = None
 
         return super().post(
@@ -192,11 +192,11 @@ def put(  # type: ignore[override]
         tags: list[str] | None = None,
         operation_id: str | None = None,
         include_in_schema: bool = True,
+        openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ):
-        openapi_extensions = None
         security = None
 
         return super().put(
@@ -233,11 +233,11 @@ def patch(  # type: ignore[override]
         tags: list[str] | None = None,
         operation_id: str | None = None,
         include_in_schema: bool = True,
+        openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable] | None = None,
     ):
-        openapi_extensions = None
         security = None
 
         return super().patch(
@@ -274,11 +274,11 @@ def delete(  # type: ignore[override]
         tags: list[str] | None = None,
         operation_id: str | None = None,
         include_in_schema: bool = True,
+        openapi_extensions: dict[str, Any] | None = None,
         deprecated: bool = False,
         custom_response_validation_http_code: int | HTTPStatus | None = None,
         middlewares: list[Callable[..., Any]] | None = None,
     ):
-        openapi_extensions = None
         security = None
 
         return super().delete(
@@ -325,6 +325,7 @@ def get_openapi_json_schema(  # type: ignore[override]
         license_info: License | None = None,
         security_schemes: dict[str, SecurityScheme] | None = None,
         security: list[dict[str, list[str]]] | None = None,
+        openapi_extensions: dict[str, Any] | None = None,
     ) -> str:
         """
         Returns the OpenAPI schema as a JSON serializable dict.
@@ -365,8 +366,6 @@ def get_openapi_json_schema(  # type: ignore[override]
         """
         from aws_lambda_powertools.event_handler.openapi.compat import model_json
 
-        openapi_extensions = None
-
         schema = super().get_openapi_schema(
             title=title,
             version=version,

docs/core/event_handler/bedrock_agents.md

@@ -313,6 +313,16 @@ To implement these customizations, include extra parameters when defining your r
 --8<-- "examples/event_handler_bedrock_agents/src/customizing_bedrock_api_operations.py"
 ```
 
+#### Enabling user confirmation
+
+You can enable user confirmation with Bedrock Agents to have your application ask for explicit user approval before invoking an action.
+
+```python hl_lines="14" title="enabling_user_confirmation.py" title="Enabling user confirmation"
+--8<-- "examples/event_handler_bedrock_agents/src/enabling_user_confirmation.py"
+```
+
+1. Add an openapi extension
+
 ## 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/enabling_user_confirmation.py

@@ -0,0 +1,26 @@
+from time import time
+
+from aws_lambda_powertools import Logger
+from aws_lambda_powertools.event_handler import BedrockAgentResolver
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+logger = Logger()
+app = BedrockAgentResolver()
+
+
+@app.get(
+    "/current_time",
+    description="Gets the current time in seconds",
+    openapi_extensions={"x-requireConfirmation": "ENABLED"},  # (1)!
+)
+def current_time() -> int:
+    return int(time())
+
+
+@logger.inject_lambda_context
+def lambda_handler(event: dict, context: LambdaContext):
+    return app.resolve(event, context)
+
+
+if __name__ == "__main__":
+    print(app.get_openapi_json_schema())

tests/functional/event_handler/_pydantic/test_bedrock_agent.py

@@ -200,3 +200,19 @@ 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_resolver_with_openapi_extensions():
+    # GIVEN BedrockAgentResolver is initialized with enable_validation=True
+    app = BedrockAgentResolver(enable_validation=True)
+
+    # WHEN we have a simple handler with openapi extension
+    @app.get("/", description="Testing", openapi_extensions={"x-requireConfirmation": "ENABLED"})
+    def handler() -> Optional[Dict]:
+        pass
+
+    # WHEN we get the schema
+    schema = json.loads(app.get_openapi_json_schema())
+
+    # THEN the OpenAPI schema must contain the "x-requireConfirmation" extension at the operation level
+    assert schema["paths"]["/"]["get"]["x-requireConfirmation"] == "ENABLED"