params injection

Author: anafalcao

aws_lambda_powertools/event_handler/bedrock_agent_function.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+import inspect
+import warnings
+from typing import TYPE_CHECKING, Any, Literal
+
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -19,7 +23,7 @@ class BedrockFunctionResponse:
         Session attributes to include in the response
     prompt_session_attributes : dict[str, str] | None
         Prompt session attributes to include in the response
-    response_state : str | None
+    response_state : Literal["FAILURE", "REPROMPT"] | None
         Response state ("FAILURE" or "REPROMPT")
 
     Examples
@@ -41,10 +45,10 @@ def __init__(
         session_attributes: dict[str, str] | None = None,
         prompt_session_attributes: dict[str, str] | None = None,
         knowledge_bases: list[dict[str, Any]] | None = None,
-        response_state: str | None = None,
+        response_state: Literal["FAILURE", "REPROMPT"] | 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'")
+            raise ValueError("responseState must be 'FAILURE' or 'REPROMPT'")
 
         self.body = body
         self.session_attributes = session_attributes
@@ -78,6 +82,8 @@ def build(self, event: BedrockAgentFunctionEvent) -> dict[str, Any]:
             knowledge_bases = None
             response_state = None
 
+        # Per AWS Bedrock documentation, currently only "TEXT" is supported as the responseBody content type
+        # https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html
         response: dict[str, Any] = {
             "messageVersion": "1.0",
             "response": {
@@ -147,12 +153,13 @@ def tool(
         """
 
         def decorator(func: Callable) -> Callable:
-            if not description:
-                raise ValueError("Tool description is required")
-
             function_name = name or func.__name__
             if function_name in self._tools:
-                raise ValueError(f"Tool '{function_name}' already registered")
+                warnings.warn(
+                    f"Tool '{function_name}' already registered. Overwriting with new definition.",
+                    PowertoolsUserWarning,
+                    stacklevel=2,
+                )
 
             self._tools[function_name] = {
                 "function": func,
@@ -178,7 +185,20 @@ def _resolve(self) -> dict[str, Any]:
         function_name = self.current_event.function
 
         try:
-            result = self._tools[function_name]["function"]()
+            parameters = {}
+            if hasattr(self.current_event, "parameters"):
+                for param in self.current_event.parameters:
+                    parameters[param.name] = param.value
+
+            func = self._tools[function_name]["function"]
+            sig = inspect.signature(func)
+
+            valid_params = {}
+            for name, value in parameters.items():
+                if name in sig.parameters:
+                    valid_params[name] = value
+
+            result = func(**valid_params)
             return BedrockFunctionsResponseBuilder(result).build(self.current_event)
         except Exception as e:
             return BedrockFunctionsResponseBuilder(

tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py

@@ -4,6 +4,7 @@
 
 from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver, BedrockFunctionResponse
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
 from tests.functional.utils import load_event
 
 
@@ -59,22 +60,25 @@ def test_bedrock_agent_function_registration():
     # GIVEN a Bedrock Agent Function resolver
     app = BedrockAgentFunctionResolver()
 
-    # WHEN registering without description or with duplicate name
-    with pytest.raises(ValueError, match="Tool description is required"):
-
-        @app.tool()
-        def test_function():
-            return "test"
-
+    # WHEN registering with duplicate name
     @app.tool(name="custom", description="First registration")
     def first_function():
-        return "test"
+        return "first test"
 
-    with pytest.raises(ValueError, match="Tool 'custom' already registered"):
+    # THEN a warning should be issued when registering a duplicate
+    with pytest.warns(PowertoolsUserWarning, match="Tool 'custom' already registered"):
 
         @app.tool(name="custom", description="Second registration")
         def second_function():
-            return "test"
+            return "second test"
+
+    # AND the most recent function should be registered
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    raw_event["function"] = "custom"
+    result = app.resolve(raw_event, {})
+
+    # The second function should be used
+    assert result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"] == "second test"
 
 
 def test_bedrock_agent_function_with_optional_fields():
@@ -156,7 +160,7 @@ def test_resolve_with_no_registered_function():
 
 def test_bedrock_function_response_state_validation():
     # GIVEN invalid and valid response states
-    valid_states = [None, "FAILURE", "REPROMPT"]
+    valid_states = ["FAILURE", "REPROMPT"]
     invalid_state = "INVALID"
 
     # WHEN creating responses with valid states
@@ -172,4 +176,32 @@ def test_bedrock_function_response_state_validation():
     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'"
+    assert str(exc_info.value) == "responseState must be 'FAILURE' or 'REPROMPT'"
+
+
+def test_bedrock_agent_function_with_parameters():
+    # GIVEN a Bedrock Agent Function resolver
+    app = BedrockAgentFunctionResolver()
+
+    # Track received parameters
+    received_params = {}
+
+    @app.tool(description="Function that accepts parameters")
+    def vacation_request(startDate, endDate):
+        # Store received parameters for assertion
+        received_params["startDate"] = startDate
+        received_params["endDate"] = endDate
+        return f"Vacation request from {startDate} to {endDate} submitted"
+
+    # WHEN calling the event handler with parameters
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    raw_event["function"] = "vacation_request"
+    result = app.resolve(raw_event, {})
+
+    # THEN parameters should be correctly passed to the function
+    assert received_params["startDate"] == "2024-03-15"
+    assert received_params["endDate"] == "2024-03-20"
+    assert (
+        "Vacation request from 2024-03-15 to 2024-03-20 submitted"
+        in result["response"]["functionResponse"]["responseBody"]["TEXT"]["body"]
+    )