add response

Author: anafalcao

aws_lambda_powertools/event_handler/bedrock_agent_function.py

@@ -2,19 +2,64 @@
 
 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
 
+from enum import Enum
+
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
 
 
+class ResponseState(Enum):
+    FAILURE = "FAILURE"
+    REPROMPT = "REPROMPT"
+
+
+class BedrockFunctionsResponseBuilder(ResponseBuilder):
+    """
+    Bedrock Functions Response Builder. This builds the response dict to be returned by Lambda
+    when using Bedrock Agent Functions.
+
+    Since the payload format is different from the standard API Gateway Proxy event,
+    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)
+
+        body = self.response.body
+        if self.response.is_json() and not isinstance(self.response.body, str):
+            body = self.serializer(body)
+
+        response: dict[str, Any] = {
+            "messageVersion": "1.0",
+            "response": {
+                "actionGroup": event.action_group,
+                "function": event.function,
+                "functionResponse": {"responseBody": {"TEXT": {"body": str(body)}}},
+            },
+        }
+
+        # Add responseState if it's an error
+        if self.response.status_code >= 400:
+            response["response"]["functionResponse"]["responseState"] = (
+                ResponseState.REPROMPT.value if self.response.status_code == 400 else ResponseState.FAILURE.value
+            )
+
+        return response
+
+
 class BedrockAgentFunctionResolver:
     """Bedrock Agent Function resolver that handles function definitions
 
     Examples
     --------
-    Simple example with a custom lambda handler
-
     ```python
     from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver
 
@@ -33,6 +78,7 @@ def lambda_handler(event, context):
     def __init__(self) -> None:
         self._tools: dict[str, dict[str, Any]] = {}
         self.current_event: BedrockAgentFunctionEvent | None = None
+        self._response_builder_class = BedrockFunctionsResponseBuilder
 
     def tool(self, description: str | None = None) -> Callable:
         """Decorator to register a tool function"""
@@ -67,32 +113,28 @@ def _resolve(self) -> dict[str, Any]:
             raise ValueError("No event to process")
 
         function_name = self.current_event.function
-        action_group = self.current_event.action_group
 
         if function_name not in self._tools:
-            return self._create_response(
-                action_group=action_group,
-                function_name=function_name,
-                result=f"Function not found: {function_name}",
-            )
+            return self._response_builder_class(
+                Response(
+                    status_code=400,  # Using 400 to trigger REPROMPT
+                    body=f"Function not found: {function_name}",
+                ),
+            ).build(self.current_event)
 
         try:
             result = self._tools[function_name]["function"]()
-            return self._create_response(action_group=action_group, function_name=function_name, result=result)
+            # 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)
         except Exception as e:
-            return self._create_response(
-                action_group=action_group,
-                function_name=function_name,
-                result=f"Error: {str(e)}",
-            )
-
-    def _create_response(self, action_group: str, function_name: str, result: Any) -> dict[str, Any]:
-        """Create response in Bedrock Agent format"""
-        return {
-            "messageVersion": "1.0",
-            "response": {
-                "actionGroup": action_group,
-                "function": function_name,
-                "functionResponse": {"responseBody": {"TEXT": {"body": str(result)}}},
-            },
-        }
+            return self._response_builder_class(
+                Response(
+                    status_code=500,  # Using 500 to trigger FAILURE
+                    body=f"Error: {str(e)}",
+                ),
+            ).build(self.current_event)

tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py

@@ -1,31 +1,34 @@
 from __future__ import annotations
 
+import json
+
 import pytest
 
-from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver
+from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver, Response, content_types
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
 from tests.functional.utils import load_event
 
 
-def test_bedrock_agent_function():
+def test_bedrock_agent_function_with_string_response():
     # GIVEN a Bedrock Agent Function resolver
     app = BedrockAgentFunctionResolver()
 
-    @app.tool(description="Gets the current time")
-    def get_current_time():
+    @app.tool(description="Returns a string")
+    def test_function():
         assert isinstance(app.current_event, BedrockAgentFunctionEvent)
-        return "2024-02-01T12:00:00Z"
+        return "Hello from string"
 
     # WHEN calling the event handler
     raw_event = load_event("bedrockAgentFunctionEvent.json")
-    raw_event["function"] = "get_current_time"  # ensure function name matches
+    raw_event["function"] = "test_function"
     result = app.resolve(raw_event, {})
 
-    # THEN process event correctly
+    # THEN process event correctly with string response
     assert result["messageVersion"] == "1.0"
     assert result["response"]["actionGroup"] == raw_event["actionGroup"]
-    assert result["response"]["function"] == "get_current_time"
-    assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "2024-02-01T12:00:00Z"
+    assert result["response"]["function"] == "test_function"
+    assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "Hello from string"
+    assert "responseState" not in result["response"]["functionResponse"]  # Success has no state
 
 
 def test_bedrock_agent_function_with_error():
@@ -46,29 +49,53 @@ def error_function():
     assert result["response"]["actionGroup"] == raw_event["actionGroup"]
     assert result["response"]["function"] == "error_function"
     assert "Error: Something went wrong" in result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"]
+    assert result["response"]["functionResponse"]["responseState"] == "FAILURE"
 
 
 def test_bedrock_agent_function_not_found():
     # GIVEN a Bedrock Agent Function resolver
     app = BedrockAgentFunctionResolver()
 
-    @app.tool(description="Test function")
-    def test_function():
-        return "test"
-
     # WHEN calling the event handler with a non-existent function
     raw_event = load_event("bedrockAgentFunctionEvent.json")
     raw_event["function"] = "nonexistent_function"
     result = app.resolve(raw_event, {})
 
-    # THEN return function not found response
+    # THEN return function not found response with REPROMPT state
     assert result["messageVersion"] == "1.0"
     assert result["response"]["actionGroup"] == raw_event["actionGroup"]
     assert result["response"]["function"] == "nonexistent_function"
     assert "Function not found" in result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"]
+    assert result["response"]["functionResponse"]["responseState"] == "REPROMPT"
 
 
-def test_bedrock_agent_function_missing_description():
+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()
 
@@ -80,32 +107,23 @@ def test_bedrock_agent_function_missing_description():
         def test_function():
             return "test"
 
-
-def test_bedrock_agent_function_duplicate_registration():
-    # GIVEN a Bedrock Agent Function resolver
-    app = BedrockAgentFunctionResolver()
-
     # WHEN registering the same function twice
+    # THEN raise ValueError
     @app.tool(description="First registration")
-    def test_function():
+    def duplicate_function():
         return "test"
 
-    # THEN raise ValueError on second registration
-    with pytest.raises(ValueError, match="Tool 'test_function' already registered"):
+    with pytest.raises(ValueError, match="Tool 'duplicate_function' already registered"):
 
         @app.tool(description="Second registration")
-        def test_function():  # noqa: F811
+        def duplicate_function():  # noqa: F811
             return "test"
 
 
 def test_bedrock_agent_function_invalid_event():
     # GIVEN a Bedrock Agent Function resolver
     app = BedrockAgentFunctionResolver()
 
-    @app.tool(description="Test function")
-    def test_function():
-        return "test"
-
     # WHEN calling with invalid event
     # THEN raise ValueError
     with pytest.raises(ValueError, match="Missing required field"):