From c3b9612020bc027f027a1cbbc637ef9fd97b2126 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 17 Apr 2025 16:29:05 +0100 Subject: [PATCH 1/4] Adding OpenAPI extensions to Bedrock --- .../event_handler/bedrock_agent.py | 8 ++++++-- .../_pydantic/test_bedrock_agent.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 31d1eb17213..31eec22535a 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,6 +192,7 @@ 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, @@ -233,6 +234,7 @@ 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, @@ -274,6 +276,7 @@ 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, @@ -325,6 +328,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. diff --git a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py index 6dcc55c2da5..4a0190071e3 100644 --- a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py +++ b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py @@ -200,3 +200,21 @@ 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 = app.get_openapi_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" From cf47f2e9308508d3c7441b583b1c191962cbfb49 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 17 Apr 2025 16:48:52 +0100 Subject: [PATCH 2/4] Adding OpenAPI extensions to Bedrock --- docs/core/event_handler/bedrock_agents.md | 10 +++++++ .../src/enabling_user_confirmation.py | 26 +++++++++++++++++++ 2 files changed, 36 insertions(+) create mode 100644 examples/event_handler_bedrock_agents/src/enabling_user_confirmation.py 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()) From 36a0456ba779d991d7394c74b7c5817e624c98fc Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 17 Apr 2025 16:53:07 +0100 Subject: [PATCH 3/4] Adding OpenAPI extensions to Bedrock --- aws_lambda_powertools/event_handler/bedrock_agent.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 31eec22535a..221163b1ae4 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -197,7 +197,6 @@ def put( # type: ignore[override] custom_response_validation_http_code: int | HTTPStatus | None = None, middlewares: list[Callable[..., Any]] | None = None, ): - openapi_extensions = None security = None return super().put( @@ -239,7 +238,6 @@ def patch( # type: ignore[override] custom_response_validation_http_code: int | HTTPStatus | None = None, middlewares: list[Callable] | None = None, ): - openapi_extensions = None security = None return super().patch( @@ -281,7 +279,6 @@ def delete( # type: ignore[override] custom_response_validation_http_code: int | HTTPStatus | None = None, middlewares: list[Callable[..., Any]] | None = None, ): - openapi_extensions = None security = None return super().delete( @@ -369,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, From e5791541e81e5177b59fb59d7b2e744fa650c17c Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 17 Apr 2025 16:54:54 +0100 Subject: [PATCH 4/4] Adding OpenAPI extensions to Bedrock --- tests/functional/event_handler/_pydantic/test_bedrock_agent.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py index 4a0190071e3..eaa666b3d36 100644 --- a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py +++ b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py @@ -212,8 +212,6 @@ def handler() -> Optional[Dict]: pass # WHEN we get the schema - schema = app.get_openapi_schema() - schema = json.loads(app.get_openapi_json_schema()) # THEN the OpenAPI schema must contain the "x-requireConfirmation" extension at the operation level