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,188 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
+
+
+class BedrockFunctionResponse:
+    """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
+    response_state : str | None
+        Response state ("FAILURE" or "REPROMPT")
+
+    Examples
+    --------
+    ```python
+    @app.tool(description="Function that uses session attributes")
+    def test_function():
+        return BedrockFunctionResponse(
+            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,
+        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
+        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]:
+        """Build the full response dict to be returned by the lambda"""
+        if isinstance(self.result, BedrockFunctionResponse):
+            body = self.result.body
+            session_attributes = self.result.session_attributes
+            prompt_session_attributes = self.result.prompt_session_attributes
+            knowledge_bases = self.result.knowledge_bases
+            response_state = self.result.response_state
+
+        else:
+            body = self.result
+            session_attributes = None
+            prompt_session_attributes = None
+            knowledge_bases = None
+            response_state = None
+
+        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 "")}}},
+            },
+        }
+
+        # Add responseState if provided
+        if response_state:
+            response["response"]["functionResponse"]["responseState"] = response_state
+
+        # 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
+
+
+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)
+    ```
+    """
+
+    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,
+        name: str | None = None,
+    ) -> Callable:
+        """Decorator to register a tool function
+
+        Parameters
+        ----------
+        description : str | None
+            Description of what the tool does
+        name : str | None
+            Custom name for the tool. If not provided, uses the function name
+        """
+
+        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")
+
+            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)}")
+
+    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:
+            result = self._tools[function_name]["function"]()
+            return BedrockFunctionsResponseBuilder(result).build(self.current_event)
+        except Exception as e:
+            return BedrockFunctionsResponseBuilder(
+                BedrockFunctionResponse(
+                    body=f"Error: {str(e)}",
+                ),
+            ).build(self.current_event)

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)

aws_lambda_powertools/utilities/parser/models/__init__.py

@@ -32,6 +32,7 @@
 )
 from .bedrock_agent import (
     BedrockAgentEventModel,
+    BedrockAgentFunctionEventModel,
     BedrockAgentModel,
     BedrockAgentPropertyModel,
     BedrockAgentRequestBodyModel,
@@ -208,6 +209,7 @@
     "BedrockAgentEventModel",
     "BedrockAgentRequestBodyModel",
     "BedrockAgentRequestMediaModel",
+    "BedrockAgentFunctionEventModel",
     "S3BatchOperationJobModel",
     "S3BatchOperationModel",
     "S3BatchOperationTaskModel",

aws_lambda_powertools/utilities/parser/models/bedrock_agent.py

@@ -36,3 +36,21 @@ class BedrockAgentEventModel(BaseModel):
     agent: BedrockAgentModel
     parameters: Optional[List[BedrockAgentPropertyModel]] = None
     request_body: Optional[BedrockAgentRequestBodyModel] = Field(None, alias="requestBody")
+
+
+class BedrockAgentFunctionEventModel(BaseModel):
+    """Bedrock Agent Function event model
+
+    Documentation:
+    https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html
+    """
+
+    message_version: str = Field(..., alias="messageVersion")
+    agent: BedrockAgentModel
+    input_text: str = Field(..., alias="inputText")
+    session_id: str = Field(..., alias="sessionId")
+    action_group: str = Field(..., alias="actionGroup")
+    function: str
+    parameters: Optional[List[BedrockAgentPropertyModel]] = None
+    session_attributes: Dict[str, str] = Field({}, alias="sessionAttributes")
+    prompt_session_attributes: Dict[str, str] = Field({}, alias="promptSessionAttributes")