comments

Author: dbschmigelski

src/strands/multiagent/function_node.py

@@ -7,7 +7,7 @@
 
 import logging
 import time
-from typing import Any, Callable, Union
+from typing import Any, Protocol, Union
 
 from opentelemetry import trace as trace_api
 
@@ -21,7 +21,14 @@
 logger = logging.getLogger(__name__)
 
 
-FunctionNodeCallable = Callable[[Union[str, list[ContentBlock]], dict[str, Any] | None], str]
+class FunctionNodeCallable(Protocol):
+    """Protocol for functions that can be executed within FunctionNode."""
+
+    def __call__(
+        self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any
+    ) -> str | list[ContentBlock] | Message:
+        """Execute deterministic logic within the multiagent system."""
+        ...
 
 
 class FunctionNode(MultiAgentBase):
@@ -80,8 +87,16 @@ async def invoke_async(
                     execution_time,
                 )
 
-                # Convert function result to Message
-                message = Message(role="assistant", content=[ContentBlock(text=str(function_result))])
+                # Convert function result to Message based on type
+                if isinstance(function_result, dict) and "role" in function_result and "content" in function_result:
+                    # Already a Message
+                    message = function_result
+                elif isinstance(function_result, list):
+                    # List of ContentBlocks
+                    message = Message(role="assistant", content=function_result)
+                else:
+                    # String or other type - convert to string
+                    message = Message(role="assistant", content=[ContentBlock(text=str(function_result))])
                 agent_result = AgentResult(
                     stop_reason="end_turn",  # "Normal completion of the response" - function executed successfully
                     message=message,

tests/strands/multiagent/test_function_node.py

@@ -6,7 +6,7 @@
 
 from strands.multiagent.base import Status
 from strands.multiagent.function_node import FunctionNode
-from strands.types.content import ContentBlock
+from strands.types.content import ContentBlock, Message
 
 
 @pytest.fixture
@@ -92,3 +92,54 @@ def failing_function(task, invocation_state=None, **kwargs):
         assert node_result.status == Status.FAILED
         assert isinstance(node_result.result, ValueError)
         assert str(node_result.result) == "Test exception"
+
+
+@pytest.mark.asyncio
+async def test_function_returns_string(mock_tracer):
+    """Test function returning string."""
+
+    def string_function(task, invocation_state=None, **kwargs):
+        return "Hello World"
+
+    node = FunctionNode(string_function, "string_node")
+
+    with patch.object(node, "tracer", mock_tracer):
+        result = await node.invoke_async("test")
+
+        agent_result = result.results["string_node"].result
+        assert agent_result.message["content"][0]["text"] == "Hello World"
+
+
+@pytest.mark.asyncio
+async def test_function_returns_content_blocks(mock_tracer):
+    """Test function returning list of ContentBlocks."""
+
+    def content_block_function(task, invocation_state=None, **kwargs):
+        return [ContentBlock(text="Block 1"), ContentBlock(text="Block 2")]
+
+    node = FunctionNode(content_block_function, "content_node")
+
+    with patch.object(node, "tracer", mock_tracer):
+        result = await node.invoke_async("test")
+
+        agent_result = result.results["content_node"].result
+        assert len(agent_result.message["content"]) == 2
+        assert agent_result.message["content"][0]["text"] == "Block 1"
+        assert agent_result.message["content"][1]["text"] == "Block 2"
+
+
+@pytest.mark.asyncio
+async def test_function_returns_message(mock_tracer):
+    """Test function returning Message."""
+
+    def message_function(task, invocation_state=None, **kwargs):
+        return Message(role="user", content=[ContentBlock(text="Custom message")])
+
+    node = FunctionNode(message_function, "message_node")
+
+    with patch.object(node, "tracer", mock_tracer):
+        result = await node.invoke_async("test")
+
+        agent_result = result.results["message_node"].result
+        assert agent_result.message["role"] == "user"
+        assert agent_result.message["content"][0]["text"] == "Custom message"