feat(bedrock_agent): add new Amazon Bedrock Agents Functions Resolver

aws_lambda_powertools/event_handler/__init__.py

@@ -12,6 +12,10 @@
 )
 from aws_lambda_powertools.event_handler.appsync import AppSyncResolver
 from aws_lambda_powertools.event_handler.bedrock_agent import BedrockAgentResolver, BedrockResponse
+from aws_lambda_powertools.event_handler.bedrock_agent_function import (
+    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,
@@ -26,7 +30,9 @@
     "ALBResolver",
     "ApiGatewayResolver",
     "BedrockAgentResolver",
+    "BedrockAgentFunctionResolver",
     "BedrockResponse",
+    "BedrockFunctionResponse",
     "CORSConfig",
     "LambdaFunctionUrlResolver",
     "Response",

aws_lambda_powertools/event_handler/bedrock_agent_function.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import inspect
+import warnings
+from collections.abc import Callable
+from typing import Any, Literal, TypeVar
+
+from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
+
+# Define a generic type for the function
+T = TypeVar("T", bound=Callable[..., Any])
+
+
+class BedrockFunctionResponse:
+    """Response class for Bedrock Agent Functions.
+
+    Parameters
+    ----------
+    body : Any, optional
+        Response body to be returned to the caller.
+    session_attributes : dict[str, str] or None, optional
+        Session attributes to include in the response for maintaining state.
+    prompt_session_attributes : dict[str, str] or None, optional
+        Prompt session attributes to include in the response.
+    knowledge_bases : list[dict[str, Any]] or None, optional
+        Knowledge bases to include in the response.
+    response_state : {"FAILURE", "REPROMPT"} or None, optional
+        Response state indicating if the function failed or needs reprompting.
+
+    Examples
+    --------
+    >>> @app.tool(description="Function that uses session attributes")
+    >>> def test_function():
+    ...     return BedrockFunctionResponse(
+    ...         body="Hello",
+    ...         session_attributes={"userId": "123"},
+    ...         prompt_session_attributes={"lastAction": "login"}
+    ...     )
+
+    Notes
+    -----
+    The `response_state` parameter can only be set to "FAILURE" or "REPROMPT".
+    """
+
+    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,
+        response_state: Literal["FAILURE", "REPROMPT"] | None = None,
+    ) -> None:
+        if response_state and response_state not in ["FAILURE", "REPROMPT"]:
+            raise ValueError("responseState must be 'FAILURE' or 'REPROMPT'")
+
+        self.body = body
+        self.session_attributes = session_attributes
+        self.prompt_session_attributes = prompt_session_attributes
+        self.knowledge_bases = knowledge_bases
+        self.response_state = response_state
+
+
+class BedrockFunctionsResponseBuilder:
+    """
+    Bedrock Functions Response Builder. This builds the response dict to be returned by Lambda
+    when using Bedrock Agent Functions.
+    """
+
+    def __init__(self, result: BedrockFunctionResponse | Any) -> None:
+        self.result = result
+
+    def build(self, event: BedrockAgentFunctionEvent) -> dict[str, Any]:
+        result_obj = self.result
+
+        # Extract attributes from BedrockFunctionResponse or use defaults
+        body = getattr(result_obj, "body", result_obj)
+        session_attributes = getattr(result_obj, "session_attributes", None)
+        prompt_session_attributes = getattr(result_obj, "prompt_session_attributes", None)
+        knowledge_bases = getattr(result_obj, "knowledge_bases", None)
+        response_state = getattr(result_obj, "response_state", None)
+
+        # Build base response structure
+        # 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": {
+                "actionGroup": event.action_group,
+                "function": event.function,
+                "functionResponse": {
+                    "responseBody": {"TEXT": {"body": str(body if body is not None else "")}},
+                },
+            },
+            "sessionAttributes": session_attributes or event.session_attributes or {},
+            "promptSessionAttributes": prompt_session_attributes or event.prompt_session_attributes or {},
+        }
+
+        # Add optional fields when present
+        if response_state:
+            response["response"]["functionResponse"]["responseState"] = response_state
+
+        if knowledge_bases:
+            response["knowledgeBasesConfiguration"] = knowledge_bases
+
+        return response
+
+
+class BedrockAgentFunctionResolver:
+    """Bedrock Agent Function resolver that handles function definitions
+
+    Examples
+    --------
+    ```python
+    from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver
+
+    app = BedrockAgentFunctionResolver()
+
+    @app.tool(description="Gets the current UTC time")
+    def get_current_time():
+        from datetime import datetime
+        return datetime.utcnow().isoformat()
+
+    def lambda_handler(event, context):
+        return app.resolve(event, context)
+    ```
+    """
+
+    context: dict
+
+    def __init__(self) -> None:
+        self._tools: dict[str, dict[str, Any]] = {}
+        self.current_event: BedrockAgentFunctionEvent | None = None
+        self.context = {}
+        self._response_builder_class = BedrockFunctionsResponseBuilder
+
+    def tool(
+        self,
+        name: str | None = None,
+        description: str | None = None,
+    ) -> Callable[[T], T]:
+        """Decorator to register a tool function
+
+        Parameters
+        ----------
+        name : str | None
+            Custom name for the tool. If not provided, uses the function name
+        description : str | None
+            Description of what the tool does
+
+        Returns
+        -------
+        Callable
+            Decorator function that registers and returns the original function
+        """
+
+        def decorator(func: T) -> T:
+            function_name = name or func.__name__
+            if function_name in self._tools:
+                warnings.warn(
+                    f"Tool '{function_name}' already registered. Overwriting with new definition.",
+                    PowertoolsUserWarning,
+                    stacklevel=2,
+                )
+
+            self._tools[function_name] = {
+                "function": func,
+                "description": description,
+            }
+            return func
+
+        return decorator
+
+    def resolve(self, event: dict[str, Any], context: Any) -> dict[str, Any]:
+        """Resolves the function call from Bedrock Agent event"""
+        try:
+            self.current_event = BedrockAgentFunctionEvent(event)
+            return self._resolve()
+        except KeyError as e:
+            raise ValueError(f"Missing required field: {str(e)}") from e
+
+    def _resolve(self) -> dict[str, Any]:
+        """Internal resolution logic"""
+        if self.current_event is None:
+            raise ValueError("No event to process")
+
+        function_name = self.current_event.function
+
+        try:
+            # Extract parameters from the event
+            parameters = {param.name: param.value for param in getattr(self.current_event, "parameters", [])}
+
+            func = self._tools[function_name]["function"]
+            # Filter parameters to only include those expected by the function
+            sig = inspect.signature(func)
+            valid_params = {name: value for name, value in parameters.items() if name in sig.parameters}
+
+            # Call the function with the filtered parameters
+            result = func(**valid_params)
+
+            self.clear_context()
+
+            # Build and return the response
+            return BedrockFunctionsResponseBuilder(result).build(self.current_event)
+        except Exception as error:
+            # Return a formatted error response
+            error_response = BedrockFunctionResponse(body=f"Error: {str(error)}", response_state="FAILURE")
+            return BedrockFunctionsResponseBuilder(error_response).build(self.current_event)
+
+    def append_context(self, **additional_context):
+        """Append key=value data as routing context"""
+        self.context.update(**additional_context)
+
+    def clear_context(self):
+        """Resets routing context"""
+        self.context.clear()

aws_lambda_powertools/utilities/data_classes/__init__.py

@@ -9,6 +9,7 @@
 from .appsync_resolver_events_event import AppSyncResolverEventsEvent
 from .aws_config_rule_event import AWSConfigRuleEvent
 from .bedrock_agent_event import BedrockAgentEvent
+from .bedrock_agent_function_event import BedrockAgentFunctionEvent
 from .cloud_watch_alarm_event import (
     CloudWatchAlarmConfiguration,
     CloudWatchAlarmData,
@@ -59,6 +60,7 @@
     "AppSyncResolverEventsEvent",
     "ALBEvent",
     "BedrockAgentEvent",
+    "BedrockAgentFunctionEvent",
     "CloudWatchAlarmData",
     "CloudWatchAlarmEvent",
     "CloudWatchAlarmMetric",

aws_lambda_powertools/utilities/data_classes/bedrock_agent_function_event.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
+
+
+class BedrockAgentInfo(DictWrapper):
+    @property
+    def name(self) -> str:
+        return self["name"]
+
+    @property
+    def id(self) -> str:  # noqa: A003
+        return self["id"]
+
+    @property
+    def alias(self) -> str:
+        return self["alias"]
+
+    @property
+    def version(self) -> str:
+        return self["version"]
+
+
+class BedrockAgentFunctionParameter(DictWrapper):
+    @property
+    def name(self) -> str:
+        return self["name"]
+
+    @property
+    def type(self) -> str:  # noqa: A003
+        return self["type"]
+
+    @property
+    def value(self) -> str:
+        return self["value"]
+
+
+class BedrockAgentFunctionEvent(DictWrapper):
+    """
+    Bedrock Agent Function input event
+
+    Documentation:
+    https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html
+    """
+
+    @property
+    def message_version(self) -> str:
+        return self["messageVersion"]
+
+    @property
+    def input_text(self) -> str:
+        return self["inputText"]
+
+    @property
+    def session_id(self) -> str:
+        return self["sessionId"]
+
+    @property
+    def action_group(self) -> str:
+        return self["actionGroup"]
+
+    @property
+    def function(self) -> str:
+        return self["function"]
+
+    @property
+    def parameters(self) -> list[BedrockAgentFunctionParameter]:
+        parameters = self.get("parameters") or []
+        return [BedrockAgentFunctionParameter(x) for x in parameters]
+
+    @property
+    def agent(self) -> BedrockAgentInfo:
+        return BedrockAgentInfo(self["agent"])
+
+    @property
+    def session_attributes(self) -> dict[str, str]:
+        return self.get("sessionAttributes", {}) or {}
+
+    @property
+    def prompt_session_attributes(self) -> dict[str, str]:
+        return self.get("promptSessionAttributes", {}) or {}

aws_lambda_powertools/utilities/parser/envelopes/__init__.py

@@ -2,7 +2,7 @@
 from .apigw_websocket import ApiGatewayWebSocketEnvelope
 from .apigwv2 import ApiGatewayV2Envelope
 from .base import BaseEnvelope
-from .bedrock_agent import BedrockAgentEnvelope
+from .bedrock_agent import BedrockAgentEnvelope, BedrockAgentFunctionEnvelope
 from .cloudwatch import CloudWatchLogsEnvelope
 from .dynamodb import DynamoDBStreamEnvelope
 from .event_bridge import EventBridgeEnvelope
@@ -20,6 +20,7 @@
     "ApiGatewayV2Envelope",
     "ApiGatewayWebSocketEnvelope",
     "BedrockAgentEnvelope",
+    "BedrockAgentFunctionEnvelope",
     "CloudWatchLogsEnvelope",
     "DynamoDBStreamEnvelope",
     "EventBridgeEnvelope",

aws_lambda_powertools/utilities/parser/envelopes/bedrock_agent.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Any
 
 from aws_lambda_powertools.utilities.parser.envelopes.base import BaseEnvelope
-from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel
+from aws_lambda_powertools.utilities.parser.models import BedrockAgentEventModel, BedrockAgentFunctionEventModel
 
 if TYPE_CHECKING:
     from aws_lambda_powertools.utilities.parser.types import Model
@@ -34,3 +34,27 @@ def parse(self, data: dict[str, Any] | Any | None, model: type[Model]) -> Model
         parsed_envelope: BedrockAgentEventModel = BedrockAgentEventModel.model_validate(data)
         logger.debug(f"Parsing event payload in `input_text` with {model}")
         return self._parse(data=parsed_envelope.input_text, model=model)
+
+
+class BedrockAgentFunctionEnvelope(BaseEnvelope):
+    """Bedrock Agent Function envelope to extract data within input_text key"""
+
+    def parse(self, data: dict[str, Any] | Any | None, model: type[Model]) -> Model | None:
+        """Parses data found with model provided
+
+        Parameters
+        ----------
+        data : dict
+            Lambda event to be parsed
+        model : type[Model]
+            Data model provided to parse after extracting data using envelope
+
+        Returns
+        -------
+        Model | None
+            Parsed detail payload with model provided
+        """
+        logger.debug(f"Parsing incoming data with Bedrock Agent Function model {BedrockAgentFunctionEventModel}")
+        parsed_envelope: BedrockAgentFunctionEventModel = BedrockAgentFunctionEventModel.model_validate(data)
+        logger.debug(f"Parsing event payload in `input_text` with {model}")
+        return self._parse(data=parsed_envelope.input_text, model=model)