add validation response state

Author: anafalcao

aws_lambda_powertools/event_handler/__init__.py

@@ -16,7 +16,6 @@
     BedrockAgentFunctionResolver,
     BedrockFunctionResponse,
 )
-
 from aws_lambda_powertools.event_handler.events_appsync.appsync_events import AppSyncEventsResolver
 from aws_lambda_powertools.event_handler.lambda_function_url import (
     LambdaFunctionUrlResolver,

aws_lambda_powertools/event_handler/bedrock_agent_function.py

@@ -43,6 +43,9 @@ def __init__(
         knowledge_bases: list[dict[str, Any]] | None = None,
         response_state: str | None = None,
     ) -> None:
+        if response_state is not None and response_state not in ["FAILURE", "REPROMPT"]:
+            raise ValueError("responseState must be None, 'FAILURE' or 'REPROMPT'")
+
         self.body = body
         self.session_attributes = session_attributes
         self.prompt_session_attributes = prompt_session_attributes

tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py

@@ -152,3 +152,24 @@ def test_resolve_with_no_registered_function():
 
     # THEN the response should contain an error message
     assert "Error: 'non_existent_function'" in result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"]
+
+
+def test_bedrock_function_response_state_validation():
+    # GIVEN invalid and valid response states
+    valid_states = [None, "FAILURE", "REPROMPT"]
+    invalid_state = "INVALID"
+
+    # WHEN creating responses with valid states
+    # THEN no error should be raised
+    for state in valid_states:
+        try:
+            BedrockFunctionResponse(body="test", response_state=state)
+        except ValueError:
+            pytest.fail(f"Unexpected ValueError for response_state={state}")
+
+    # WHEN creating a response with invalid state
+    # THEN ValueError should be raised with correct message
+    with pytest.raises(ValueError) as exc_info:
+        BedrockFunctionResponse(body="test", response_state=invalid_state)
+
+    assert str(exc_info.value) == "responseState must be None, 'FAILURE' or 'REPROMPT'"