From 44544cf08b9d4cfc8c14d33467ece9750e4579fc Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Thu, 20 Mar 2025 19:48:18 -0300 Subject: [PATCH 01/10] feat(bedrock_agents): add optional fields to response payload --- .../event_handler/__init__.py | 3 +- .../event_handler/bedrock_agent.py | 81 +++++++++++++++++-- docs/core/event_handler/bedrock_agents.md | 13 +++ .../src/bedrockresponse.py | 14 ++++ .../_pydantic/test_bedrock_agent.py | 42 +++++++++- 5 files changed, 144 insertions(+), 9 deletions(-) create mode 100644 examples/event_handler_bedrock_agents/src/bedrockresponse.py diff --git a/aws_lambda_powertools/event_handler/__init__.py b/aws_lambda_powertools/event_handler/__init__.py index ffbb2abe4ae..c61cf4df1e1 100644 --- a/aws_lambda_powertools/event_handler/__init__.py +++ b/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", diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 215199e0022..492faa8808f 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/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 diff --git a/docs/core/event_handler/bedrock_agents.md b/docs/core/event_handler/bedrock_agents.md index d5c1d1c9534..0d673a560ea 100644 --- a/docs/core/event_handler/bedrock_agents.md +++ b/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: diff --git a/examples/event_handler_bedrock_agents/src/bedrockresponse.py b/examples/event_handler_bedrock_agents/src/bedrockresponse.py new file mode 100644 index 00000000000..683221d4d21 --- /dev/null +++ b/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"}}, + }, +) diff --git a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py index 6dcc55c2da5..85151f35c25 100644 --- a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py +++ b/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}}, + } From 1f728746ffa13f8c6893dd85de02592b6f6399e5 Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Mon, 24 Mar 2025 15:45:55 -0300 Subject: [PATCH 02/10] reformat bedrock response --- .../event_handler/bedrock_agent.py | 79 ++++++++++--------- .../_pydantic/test_bedrock_agent.py | 35 ++------ 2 files changed, 49 insertions(+), 65 deletions(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 492faa8808f..0ea384e43b0 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -22,41 +22,55 @@ 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, + status_code: int = 200, 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, + knowledge_bases_configuration: list[dict[str, Any]] | None = None, ) -> None: - self.status_code = status_code 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 + + if knowledge_bases_configuration is not None: + if not isinstance(knowledge_bases_configuration, list) or not all( + isinstance(item, dict) for item in knowledge_bases_configuration + ): + raise ValueError("knowledge_bases_configuration must be a list of dictionaries") + self.knowledge_bases_configuration = knowledge_bases_configuration + def to_dict(self, event) -> dict[str, Any]: + result = { + "messageVersion": "1.0", + "response": { + "apiPath": event.api_path, + "actionGroup": event.action_group, + "httpMethod": event.http_method, + "httpStatusCode": self.status_code, + "responseBody": { + self.content_type: {"body": json.dumps(self.body) if isinstance(self.body, dict) else self.body}, + }, + }, + } + + # Add optional attributes if they exist + if self.session_attributes is not None: + result["sessionAttributes"] = self.session_attributes + + if self.prompt_session_attributes is not None: + result["promptSessionAttributes"] = self.prompt_session_attributes + + if self.knowledge_bases_configuration is not None: + result["knowledgeBasesConfiguration"] = self.knowledge_bases_configuration + + return result + class BedrockResponseBuilder(ResponseBuilder): """ @@ -72,6 +86,10 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: """ 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) + base_response = { "messageVersion": "1.0", "response": { @@ -81,7 +99,7 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: "httpStatusCode": self.response.status_code, "responseBody": { self.response.content_type: { - "body": self._get_formatted_body(), + "body": body, }, }, }, @@ -92,22 +110,7 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: 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 diff --git a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py index 85151f35c25..a0d6606601a 100644 --- a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py +++ b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py @@ -203,40 +203,21 @@ def handler() -> Optional[Dict]: def test_bedrock_agent_with_bedrock_response(): - # GIVEN a Bedrock Agent resolver + # GIVEN a Bedrock Agent event app = BedrockAgentResolver() + # WHEN using BedrockResponse @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}}, - }, - ) + def claims() -> Dict[str, Any]: + assert isinstance(app.current_event, BedrockAgentEvent) + assert app.lambda_context == {} + return BedrockResponse(body={"message": "success"}, session_attributes={"last_request": "get_claims"}) - # WHEN calling the event handler result = app(load_event("bedrockAgentEvent.json"), {}) + print(result) - # THEN process event correctly + # To be implemented: check if session_attributes 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}}, - } From 3b1f76739b2d58c6223ee3282b8b8305ac4d52f8 Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Mon, 24 Mar 2025 15:53:04 -0300 Subject: [PATCH 03/10] fix type knowledge_base_config --- aws_lambda_powertools/event_handler/bedrock_agent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 0ea384e43b0..a792c3244c1 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -45,8 +45,8 @@ def __init__( self.knowledge_bases_configuration = knowledge_bases_configuration - def to_dict(self, event) -> dict[str, Any]: - result = { + def to_dict(self, event: BedrockAgentEvent) -> dict[str, Any]: + result: dict[str, Any] = { "messageVersion": "1.0", "response": { "apiPath": event.api_path, From 41468d6ca9d5b12b3095f822bfe62ca44218bdfc Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Mon, 24 Mar 2025 15:57:12 -0300 Subject: [PATCH 04/10] fix type knowledge_base_config --- .../src/bedrockresponse.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/examples/event_handler_bedrock_agents/src/bedrockresponse.py b/examples/event_handler_bedrock_agents/src/bedrockresponse.py index 683221d4d21..51ead4ee05e 100644 --- a/examples/event_handler_bedrock_agents/src/bedrockresponse.py +++ b/examples/event_handler_bedrock_agents/src/bedrockresponse.py @@ -7,8 +7,12 @@ 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"}}, - }, + knowledge_bases_configuration=[ + { + "knowledgeBaseId": "kb-123", + "retrievalConfiguration": { + "vectorSearchConfiguration": {"numberOfResults": 3, "overrideSearchType": "HYBRID"}, + }, + }, + ], ) From 4ca4e2d9d63f0c5f2bcff5beb38338a9b950c4c9 Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Tue, 25 Mar 2025 19:02:40 -0300 Subject: [PATCH 05/10] fix bedrock response --- .../event_handler/bedrock_agent.py | 91 +++++++------------ .../_pydantic/test_bedrock_agent.py | 27 +++++- 2 files changed, 54 insertions(+), 64 deletions(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index a792c3244c1..cc2ce630033 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -24,7 +24,7 @@ class BedrockResponse: def __init__( self, - body: Any, + body: Any = None, status_code: int = 200, content_type: str = "application/json", session_attributes: dict[str, Any] | None = None, @@ -36,61 +36,38 @@ def __init__( self.content_type = content_type self.session_attributes = session_attributes self.prompt_session_attributes = prompt_session_attributes - - if knowledge_bases_configuration is not None: - if not isinstance(knowledge_bases_configuration, list) or not all( - isinstance(item, dict) for item in knowledge_bases_configuration - ): - raise ValueError("knowledge_bases_configuration must be a list of dictionaries") - self.knowledge_bases_configuration = knowledge_bases_configuration - def to_dict(self, event: BedrockAgentEvent) -> dict[str, Any]: - result: dict[str, Any] = { - "messageVersion": "1.0", - "response": { - "apiPath": event.api_path, - "actionGroup": event.action_group, - "httpMethod": event.http_method, - "httpStatusCode": self.status_code, - "responseBody": { - self.content_type: {"body": json.dumps(self.body) if isinstance(self.body, dict) else self.body}, - }, - }, + def is_json(self) -> bool: + return self.content_type == "application/json" + + def to_dict(self) -> dict: + return { + "body": self.body, + "status_code": self.status_code, + "content_type": self.content_type, + "session_attributes": self.session_attributes, + "prompt_session_attributes": self.prompt_session_attributes, + "knowledge_bases_configuration": self.knowledge_bases_configuration, } - # Add optional attributes if they exist - if self.session_attributes is not None: - result["sessionAttributes"] = self.session_attributes - - if self.prompt_session_attributes is not None: - result["promptSessionAttributes"] = self.prompt_session_attributes - - if self.knowledge_bases_configuration is not None: - result["knowledgeBasesConfiguration"] = self.knowledge_bases_configuration - - return result - class BedrockResponseBuilder(ResponseBuilder): - """ - Bedrock Response Builder. This builds the response dict to be returned by Lambda when using Bedrock Agents. - - Since the payload format is different from the standard API Gateway Proxy event, we override the build method. - """ - @override def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: - """ - 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) + bedrock_response = None + if isinstance(self.response.body, dict) and "body" in self.response.body: + bedrock_response = BedrockResponse(**self.response.body) + body = bedrock_response.body + else: + body = self.response.body + + if self.response.is_json() and not isinstance(body, str): + body = self.serializer(body) - base_response = { + response = { "messageVersion": "1.0", "response": { "actionGroup": event.action_group, @@ -105,22 +82,16 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: }, } - if isinstance(self.response, BedrockResponse): - self._add_bedrock_specific_configs(base_response) - - return base_response - - def _add_bedrock_specific_configs(self, response: dict[str, Any]) -> None: - 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, - } + # Add Bedrock-specific attributes + if bedrock_response: + if bedrock_response.session_attributes: + response["sessionAttributes"] = bedrock_response.session_attributes + if bedrock_response.prompt_session_attributes: + response["promptSessionAttributes"] = bedrock_response.prompt_session_attributes + if bedrock_response.knowledge_bases_configuration: + response["knowledgeBasesConfiguration"] = bedrock_response.knowledge_bases_configuration - response.update({k: v for k, v in optional_configs.items() if v is not None}) + return response class BedrockAgentResolver(ApiGatewayResolver): diff --git a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py index a0d6606601a..2d096486fb0 100644 --- a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py +++ b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py @@ -208,16 +208,35 @@ def test_bedrock_agent_with_bedrock_response(): # WHEN using BedrockResponse @app.get("/claims", description="Gets claims") - def claims() -> Dict[str, Any]: + def claims(): assert isinstance(app.current_event, BedrockAgentEvent) assert app.lambda_context == {} - return BedrockResponse(body={"message": "success"}, session_attributes={"last_request": "get_claims"}) + return BedrockResponse( + session_attributes={"user_id": "123"}, + prompt_session_attributes={"context": "testing"}, + knowledge_bases_configuration=[ + { + "knowledgeBaseId": "kb-123", + "retrievalConfiguration": { + "vectorSearchConfiguration": {"numberOfResults": 3, "overrideSearchType": "HYBRID"}, + }, + }, + ], + ) result = app(load_event("bedrockAgentEvent.json"), {}) - print(result) - # To be implemented: check if session_attributes assert result["messageVersion"] == "1.0" assert result["response"]["apiPath"] == "/claims" assert result["response"]["actionGroup"] == "ClaimManagementActionGroup" assert result["response"]["httpMethod"] == "GET" + assert result["sessionAttributes"] == {"user_id": "123"} + assert result["promptSessionAttributes"] == {"context": "testing"} + assert result["knowledgeBasesConfiguration"] == [ + { + "knowledgeBaseId": "kb-123", + "retrievalConfiguration": { + "vectorSearchConfiguration": {"numberOfResults": 3, "overrideSearchType": "HYBRID"}, + }, + }, + ] From 24099d3f3e67a3607cd0b867634ada8bfb8c26c4 Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Tue, 25 Mar 2025 23:00:44 -0300 Subject: [PATCH 06/10] mypy --- aws_lambda_powertools/event_handler/bedrock_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index cc2ce630033..70788e6b53d 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -41,7 +41,7 @@ def __init__( def is_json(self) -> bool: return self.content_type == "application/json" - def to_dict(self) -> dict: + def to_dict(self) -> dict[str, Any]: return { "body": self.body, "status_code": self.status_code, From 2f9ce0bbc5310c1dd9044e9dbc5b7710d9878e4d Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Wed, 26 Mar 2025 08:33:59 -0300 Subject: [PATCH 07/10] mypy --- aws_lambda_powertools/event_handler/bedrock_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 70788e6b53d..bdf8de2fa29 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -89,7 +89,7 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: if bedrock_response.prompt_session_attributes: response["promptSessionAttributes"] = bedrock_response.prompt_session_attributes if bedrock_response.knowledge_bases_configuration: - response["knowledgeBasesConfiguration"] = bedrock_response.knowledge_bases_configuration + response["knowledgeBasesConfiguration"] = bedrock_response.knowledge_bases_configuration # type: ignore return response From d5bd8dc269b4f8169b7ba3de1207471cae6b9e56 Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Wed, 26 Mar 2025 10:10:46 -0300 Subject: [PATCH 08/10] remove unnecessary attributes and add docstrings --- .../event_handler/bedrock_agent.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index bdf8de2fa29..18daa6c0e31 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -22,6 +22,10 @@ class BedrockResponse: + """ + Contains the response body, status code, content type, and optional attributes + for session management and knowledge base configuration. + """ def __init__( self, body: Any = None, @@ -38,9 +42,6 @@ def __init__( self.prompt_session_attributes = prompt_session_attributes self.knowledge_bases_configuration = knowledge_bases_configuration - def is_json(self) -> bool: - return self.content_type == "application/json" - def to_dict(self) -> dict[str, Any]: return { "body": self.body, @@ -53,6 +54,11 @@ def to_dict(self) -> dict[str, Any]: class BedrockResponseBuilder(ResponseBuilder): + """ + Bedrock Response Builder. This builds the response dict to be returned by Lambda when using Bedrock Agents. + + Since the payload format is different from the standard API Gateway Proxy event, we override the build method. + """ @override def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: self._route(event, None) @@ -89,7 +95,7 @@ def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: if bedrock_response.prompt_session_attributes: response["promptSessionAttributes"] = bedrock_response.prompt_session_attributes if bedrock_response.knowledge_bases_configuration: - response["knowledgeBasesConfiguration"] = bedrock_response.knowledge_bases_configuration # type: ignore + response["knowledgeBasesConfiguration"] = bedrock_response.knowledge_bases_configuration # type: ignore return response From 5cba20afae498fd5016656eb029b493855f44dda Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Wed, 26 Mar 2025 10:14:04 -0300 Subject: [PATCH 09/10] remove unnecessary attributes and add docstrings --- .../event_handler/bedrock_agent.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 18daa6c0e31..6baff500567 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -26,6 +26,7 @@ class BedrockResponse: Contains the response body, status code, content type, and optional attributes for session management and knowledge base configuration. """ + def __init__( self, body: Any = None, @@ -42,23 +43,13 @@ def __init__( self.prompt_session_attributes = prompt_session_attributes self.knowledge_bases_configuration = knowledge_bases_configuration - def to_dict(self) -> dict[str, Any]: - return { - "body": self.body, - "status_code": self.status_code, - "content_type": self.content_type, - "session_attributes": self.session_attributes, - "prompt_session_attributes": self.prompt_session_attributes, - "knowledge_bases_configuration": self.knowledge_bases_configuration, - } - - class BedrockResponseBuilder(ResponseBuilder): """ Bedrock Response Builder. This builds the response dict to be returned by Lambda when using Bedrock Agents. Since the payload format is different from the standard API Gateway Proxy event, we override the build method. """ + @override def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]: self._route(event, None) From a743685f80cf99f13af7f44a5f31479d601245f5 Mon Sep 17 00:00:00 2001 From: Ana Falcao Date: Thu, 27 Mar 2025 09:47:22 -0300 Subject: [PATCH 10/10] add more tests --- .../event_handler/bedrock_agent.py | 2 + .../_pydantic/test_bedrock_agent.py | 71 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/aws_lambda_powertools/event_handler/bedrock_agent.py b/aws_lambda_powertools/event_handler/bedrock_agent.py index 6baff500567..f4e3face605 100644 --- a/aws_lambda_powertools/event_handler/bedrock_agent.py +++ b/aws_lambda_powertools/event_handler/bedrock_agent.py @@ -43,6 +43,7 @@ def __init__( 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. @@ -52,6 +53,7 @@ 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) bedrock_response = None diff --git a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py index 2d096486fb0..24430b8c441 100644 --- a/tests/functional/event_handler/_pydantic/test_bedrock_agent.py +++ b/tests/functional/event_handler/_pydantic/test_bedrock_agent.py @@ -240,3 +240,74 @@ def claims(): }, }, ] + + +def test_bedrock_agent_with_empty_bedrock_response(): + # GIVEN a Bedrock Agent event + app = BedrockAgentResolver() + + @app.get("/claims", description="Gets claims") + def claims(): + return BedrockResponse(body={"message": "test"}) + + # WHEN calling the event handler + result = app(load_event("bedrockAgentEvent.json"), {}) + + # THEN process event correctly without optional attributes + assert result["messageVersion"] == "1.0" + assert result["response"]["httpStatusCode"] == 200 + assert "sessionAttributes" not in result + assert "promptSessionAttributes" not in result + assert "knowledgeBasesConfiguration" not in result + + +def test_bedrock_agent_with_partial_bedrock_response(): + # GIVEN a Bedrock Agent event + app = BedrockAgentResolver() + + @app.get("/claims", description="Gets claims") + def claims(): + return BedrockResponse( + body={"message": "test"}, + session_attributes={"user_id": "123"}, + # Only include session_attributes to test partial response + ) + + # WHEN calling the event handler + result = app(load_event("bedrockAgentEvent.json"), {}) + + # THEN process event correctly with only session_attributes + assert result["messageVersion"] == "1.0" + assert result["response"]["httpStatusCode"] == 200 + assert result["sessionAttributes"] == {"user_id": "123"} + assert "promptSessionAttributes" not in result + assert "knowledgeBasesConfiguration" not in result + + +def test_bedrock_agent_with_different_attributes_combination(): + # GIVEN a Bedrock Agent event + app = BedrockAgentResolver() + + @app.get("/claims", description="Gets claims") + def claims(): + return BedrockResponse( + body={"message": "test"}, + prompt_session_attributes={"context": "testing"}, + knowledge_bases_configuration=[ + { + "knowledgeBaseId": "kb-123", + "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 3}}, + }, + ], + # Omit session_attributes to test different combination + ) + + # WHEN calling the event handler + result = app(load_event("bedrockAgentEvent.json"), {}) + + # THEN process event correctly with specific attributes + assert result["messageVersion"] == "1.0" + assert result["response"]["httpStatusCode"] == 200 + assert "sessionAttributes" not in result + assert result["promptSessionAttributes"] == {"context": "testing"} + assert result["knowledgeBasesConfiguration"][0]["knowledgeBaseId"] == "kb-123"