add response optional fields

Author: anafalcao

aws_lambda_powertools/event_handler/__init__.py

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

aws_lambda_powertools/event_handler/bedrock_agent_function.py

@@ -2,10 +2,6 @@
 
 from typing import TYPE_CHECKING, Any
 
-from typing_extensions import override
-
-from aws_lambda_powertools.event_handler.api_gateway import Response, ResponseBuilder
-
 if TYPE_CHECKING:
     from collections.abc import Callable
 
@@ -19,7 +15,49 @@ class ResponseState(Enum):
     REPROMPT = "REPROMPT"
 
 
-class BedrockFunctionsResponseBuilder(ResponseBuilder):
+class BedrockResponse:
+    """Response class for Bedrock Agent Functions
+
+    Parameters
+    ----------
+    body : Any, optional
+        Response body
+    session_attributes : dict[str, str] | None
+        Session attributes to include in the response
+    prompt_session_attributes : dict[str, str] | None
+        Prompt session attributes to include in the response
+    status_code : int
+        Status code to determine responseState (400 for REPROMPT, >=500 for FAILURE)
+
+    Examples
+    --------
+    ```python
+    @app.tool(description="Function that uses session attributes")
+    def test_function():
+        return BedrockResponse(
+            body="Hello",
+            session_attributes={"userId": "123"},
+            prompt_session_attributes={"lastAction": "login"}
+        )
+    ```
+    """
+
+    def __init__(
+        self,
+        body: Any = None,
+        session_attributes: dict[str, str] | None = None,
+        prompt_session_attributes: dict[str, str] | None = None,
+        knowledge_bases: list[dict[str, Any]] | None = None,
+        status_code: int = 200,
+    ) -> None:
+        self.body = body
+        self.session_attributes = session_attributes
+        self.prompt_session_attributes = prompt_session_attributes
+        self.knowledge_bases = knowledge_bases
+        self.status_code = status_code
+
+
+class BedrockFunctionsResponseBuilder:
     """
     Bedrock Functions Response Builder. This builds the response dict to be returned by Lambda
     when using Bedrock Agent Functions.
@@ -28,30 +66,50 @@ class BedrockFunctionsResponseBuilder(ResponseBuilder):
     we override the build method.
     """
 
-    @override
-    def build(self, event: BedrockAgentFunctionEvent, *args) -> dict[str, Any]:
-        """Build the full response dict to be returned by the lambda"""
-        self._route(event, None)
+    def __init__(self, result: BedrockResponse | Any, status_code: int = 200) -> None:
+        self.result = result
+        self.status_code = status_code if not isinstance(result, BedrockResponse) else result.status_code
 
-        body = self.response.body
-        if self.response.is_json() and not isinstance(self.response.body, str):
-            body = self.serializer(body)
+    def build(self, event: BedrockAgentFunctionEvent) -> dict[str, Any]:
+        """Build the full response dict to be returned by the lambda"""
+        if isinstance(self.result, BedrockResponse):
+            body = self.result.body
+            session_attributes = self.result.session_attributes
+            prompt_session_attributes = self.result.prompt_session_attributes
+            knowledge_bases = self.result.knowledge_bases
+        else:
+            body = self.result
+            session_attributes = None
+            prompt_session_attributes = None
+            knowledge_bases = None
 
         response: dict[str, Any] = {
             "messageVersion": "1.0",
             "response": {
                 "actionGroup": event.action_group,
                 "function": event.function,
-                "functionResponse": {"responseBody": {"TEXT": {"body": str(body)}}},
+                "functionResponse": {"responseBody": {"TEXT": {"body": str(body if body is not None else "")}}},
             },
         }
 
         # Add responseState if it's an error
-        if self.response.status_code >= 400:
+        if self.status_code >= 400:
             response["response"]["functionResponse"]["responseState"] = (
-                ResponseState.REPROMPT.value if self.response.status_code == 400 else ResponseState.FAILURE.value
+                ResponseState.REPROMPT.value if self.status_code == 400 else ResponseState.FAILURE.value
             )
 
+        # Add session attributes if provided in response or maintain from input
+        response.update(
+            {
+                "sessionAttributes": session_attributes or event.session_attributes or {},
+                "promptSessionAttributes": prompt_session_attributes or event.prompt_session_attributes or {},
+            },
+        )
+
+        # Add knowledge bases configuration if provided
+        if knowledge_bases:
+            response["knowledgeBasesConfiguration"] = knowledge_bases
+
         return response
 
 
@@ -127,26 +185,20 @@ def _resolve(self) -> dict[str, Any]:
         function_name = self.current_event.function
 
         if function_name not in self._tools:
-            return self._response_builder_class(
-                Response(
-                    status_code=400,  # Using 400 to trigger REPROMPT
+            return BedrockFunctionsResponseBuilder(
+                BedrockResponse(
                     body=f"Function not found: {function_name}",
+                    status_code=400,  # Using 400 to trigger REPROMPT
                 ),
             ).build(self.current_event)
 
         try:
             result = self._tools[function_name]["function"]()
-            # Always wrap the result in a Response object
-            if not isinstance(result, Response):
-                result = Response(
-                    status_code=200,  # Success
-                    body=result,
-                )
-            return self._response_builder_class(result).build(self.current_event)
+            return BedrockFunctionsResponseBuilder(result).build(self.current_event)
         except Exception as e:
-            return self._response_builder_class(
-                Response(
-                    status_code=500,  # Using 500 to trigger FAILURE
+            return BedrockFunctionsResponseBuilder(
+                BedrockResponse(
                     body=f"Error: {str(e)}",
+                    status_code=500,  # Using 500 to trigger FAILURE
                 ),
             ).build(self.current_event)

tests/events/bedrockAgentFunctionEvent.json

@@ -8,8 +8,7 @@
     },
     "sessionId": "123456789123458",
     "sessionAttributes": {
-        "employeeId": "EMP123",
-        "department": "Engineering"
+        "employeeId": "EMP123"
     },
     "promptSessionAttributes": {
         "lastInteraction": "2024-02-01T15:30:00Z",

tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-import json
+from typing import Any
 
 import pytest
 
-from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver, Response, content_types
+from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver, BedrockResponse
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
 from tests.functional.utils import load_event
 
@@ -69,32 +69,6 @@ def test_bedrock_agent_function_not_found():
     assert result["response"]["functionResponse"]["responseState"] == "REPROMPT"
 
 
-def test_bedrock_agent_function_with_response_object():
-    # GIVEN a Bedrock Agent Function resolver
-    app = BedrockAgentFunctionResolver()
-
-    @app.tool(description="Returns a Response object")
-    def test_function():
-        return Response(
-            status_code=200,
-            content_type=content_types.APPLICATION_JSON,
-            body={"message": "Hello from Response"},
-        )
-
-    # WHEN calling the event handler
-    raw_event = load_event("bedrockAgentFunctionEvent.json")
-    raw_event["function"] = "test_function"
-    result = app.resolve(raw_event, {})
-
-    # THEN process event correctly with Response object
-    assert result["messageVersion"] == "1.0"
-    assert result["response"]["actionGroup"] == raw_event["actionGroup"]
-    assert result["response"]["function"] == "test_function"
-    response_body = result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"]
-    assert json.loads(response_body) == {"message": "Hello from Response"}
-    assert "responseState" not in result["response"]["functionResponse"]  # Success has no state
-
-
 def test_bedrock_agent_function_registration():
     # GIVEN a Bedrock Agent Function resolver
     app = BedrockAgentFunctionResolver()
@@ -148,3 +122,77 @@ def test_function():
     assert result["response"]["actionGroup"] == raw_event["actionGroup"]
     assert result["response"]["function"] == "customName"
     assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "Hello from custom named function"
+
+
+def test_bedrock_agent_function_with_session_attributes():
+    # GIVEN a Bedrock Agent Function resolver
+    app = BedrockAgentFunctionResolver()
+
+    @app.tool(description="Function that uses session attributes")
+    def test_function() -> dict[str, Any]:
+        return BedrockResponse(
+            body="Hello",
+            session_attributes={"userId": "123"},
+            prompt_session_attributes={"lastAction": "login"},
+        )
+
+    # WHEN calling the event handler
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    raw_event["function"] = "test_function"
+    raw_event["parameters"] = []
+    result = app.resolve(raw_event, {})
+
+    # THEN include session attributes in response
+    assert result["messageVersion"] == "1.0"
+    assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "Hello"
+    assert result["sessionAttributes"] == {"userId": "123"}
+    assert result["promptSessionAttributes"] == {"lastAction": "login"}
+
+
+def test_bedrock_agent_function_with_error_response():
+    # GIVEN a Bedrock Agent Function resolver
+    app = BedrockAgentFunctionResolver()
+
+    @app.tool(description="Function that returns error")
+    def test_function() -> dict[str, Any]:
+        return BedrockResponse(
+            body="Invalid input",
+            status_code=400,  # This will trigger REPROMPT
+            session_attributes={"error": "true"},
+        )
+
+    # WHEN calling the event handler
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    raw_event["function"] = "test_function"
+    raw_event["parameters"] = []
+    result = app.resolve(raw_event, {})
+
+    # THEN include error state and session attributes
+    assert result["response"]["functionResponse"]["responseState"] == "REPROMPT"
+    assert result["sessionAttributes"] == {"error": "true"}
+
+
+def test_bedrock_agent_function_with_knowledge_bases():
+    # GIVEN a Bedrock Agent Function resolver
+    app = BedrockAgentFunctionResolver()
+
+    @app.tool(description="Returns response with knowledge bases config")
+    def test_function() -> dict[Any]:
+        return BedrockResponse(
+            knowledge_bases=[
+                {
+                    "knowledgeBaseId": "kb1",
+                    "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 5}},
+                },
+            ],
+        )
+
+    # WHEN calling the event handler
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    raw_event["function"] = "test_function"
+    result = app.resolve(raw_event, {})
+
+    # THEN include knowledge bases in response
+    assert "knowledgeBasesConfiguration" in result
+    assert len(result["knowledgeBasesConfiguration"]) == 1
+    assert result["knowledgeBasesConfiguration"][0]["knowledgeBaseId"] == "kb1"