create resolver

Author: anafalcao

aws_lambda_powertools/event_handler/__init__.py

@@ -12,6 +12,7 @@
 )
 from aws_lambda_powertools.event_handler.appsync import AppSyncResolver
 from aws_lambda_powertools.event_handler.bedrock_agent import BedrockAgentResolver
+from aws_lambda_powertools.event_handler.bedrock_agent_function import BedrockAgentFunctionResolver
 from aws_lambda_powertools.event_handler.events_appsync.appsync_events import AppSyncEventsResolver
 from aws_lambda_powertools.event_handler.lambda_function_url import (
     LambdaFunctionUrlResolver,
@@ -26,6 +27,7 @@
     "ALBResolver",
     "ApiGatewayResolver",
     "BedrockAgentResolver",
+    "BedrockAgentFunctionResolver",
     "CORSConfig",
     "LambdaFunctionUrlResolver",
     "Response",

aws_lambda_powertools/event_handler/bedrock_agent_function.py

@@ -0,0 +1,102 @@
+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 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
+
+    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
+
+    def tool(self, description: str | None = None) -> Callable:
+        """Decorator to register a tool function"""
+        def decorator(func: Callable) -> Callable:
+            if not description:
+                raise ValueError("Tool description is required")
+
+            function_name = 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"""
+        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}"
+            )
+
+        try:
+            result = self._tools[function_name]["function"]()
+            return self._create_response(
+                action_group=action_group,
+                function_name=function_name,
+                result=result
+            )
+        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)
+                        }
+                    }
+                }
+            }
+        }

tests/functional/event_handler/required_dependencies/test_bedrock_agent_functions.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import pytest
+from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver
+from aws_lambda_powertools.utilities.data_classes import BedrockAgentFunctionEvent
+from tests.functional.utils import load_event
+
+
+def test_bedrock_agent_function():
+    # GIVEN a Bedrock Agent Function resolver
+    app = BedrockAgentFunctionResolver()
+
+    @app.tool(description="Gets the current time")
+    def get_current_time():
+        assert isinstance(app.current_event, BedrockAgentFunctionEvent)
+        return "2024-02-01T12:00:00Z"
+
+    # WHEN calling the event handler
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    raw_event["function"] = "get_current_time"  # ensure function name matches
+    result = app.resolve(raw_event, {})
+
+    # THEN process event correctly
+    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"
+
+
+def test_bedrock_agent_function_with_error():
+    # GIVEN a Bedrock Agent Function resolver
+    app = BedrockAgentFunctionResolver()
+
+    @app.tool(description="Function that raises error")
+    def error_function():
+        raise ValueError("Something went wrong")
+
+    # WHEN calling the event handler with a function that raises an error
+    raw_event = load_event("bedrockAgentFunctionEvent.json")
+    raw_event["function"] = "error_function"
+    result = app.resolve(raw_event, {})
+
+    # THEN process the error correctly
+    assert result["messageVersion"] == "1.0"
+    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"]
+
+
+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
+    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"]
+
+
+def test_bedrock_agent_function_missing_description():
+    # GIVEN a Bedrock Agent Function resolver
+    app = BedrockAgentFunctionResolver()
+
+    # WHEN registering a tool without description
+    # THEN raise ValueError
+    with pytest.raises(ValueError, match="Tool description is required"):
+        @app.tool()
+        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
+    @app.tool(description="First registration")
+    def test_function():
+        return "test"
+
+    # THEN raise ValueError on second registration
+    with pytest.raises(ValueError, match="Tool 'test_function' already registered"):
+        @app.tool(description="Second registration")
+        def test_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"):
+        app.resolve({}, {})