diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 31d1eb17213..221163b1ae4 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/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, diff --git a/docs/core/event_handler/bedrock_agents.md b/docs/core/event_handler/bedrock_agents.md index 6c514012f72..9665628ff30 100644 --- a/docs/core/event_handler/bedrock_agents.md +++ b/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: diff --git a/examples/event_handler_bedrock_agents/src/enabling_user_confirmation.py b/examples/event_handler_bedrock_agents/src/enabling_user_confirmation.py new file mode 100644 index 00000000000..21a10666014 --- /dev/null +++ b/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()) diff --git a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py index 6dcc55c2da5..eaa666b3d36 100644 --- a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py +++ b/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"