add graph.execute and graph.execute_async functions

Author: awsarron

src/strands/multiagent/base.py

@@ -75,6 +75,12 @@ class MultiAgentBase(ABC):
 
     @abstractmethod
     # TODO: for task - multi-modal input (Message), list of messages
-    async def execute(self, task: str) -> MultiAgentResult:
-        """Execute task."""
+    async def execute_async(self, task: str) -> MultiAgentResult:
+        """Execute task asynchronously."""
+        raise NotImplementedError("execute_async not implemented")
+
+    @abstractmethod
+    # TODO: for task - multi-modal input (Message), list of messages
+    def execute(self, task: str) -> MultiAgentResult:
+        """Execute task synchronously."""
         raise NotImplementedError("execute not implemented")

src/strands/multiagent/graph.py

@@ -14,8 +14,10 @@
 - Supports nested graphs (Graph as a node in another Graph)
 """
 
+import asyncio
 import logging
 import time
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from typing import Any, Callable, Tuple
 
@@ -229,8 +231,18 @@ def __init__(self, nodes: dict[str, GraphNode], edges: set[GraphEdge], entry_poi
         self.entry_points = entry_points
         self.state = GraphState()
 
-    async def execute(self, task: str) -> GraphResult:
-        """Execute the graph."""
+    def execute(self, task: str) -> GraphResult:
+        """Execute task synchronously."""
+
+        def execute() -> GraphResult:
+            return asyncio.run(self.execute_async(task))
+
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(execute)
+            return future.result()
+
+    async def execute_async(self, task: str) -> GraphResult:
+        """Execute the graph asynchronously."""
         logger.debug("task=<%s> | starting graph execution", task)
 
         # Initialize state
@@ -319,7 +331,7 @@ async def _execute_node(self, node: "GraphNode") -> None:
 
             # Execute based on node type and create unified NodeResult
             if isinstance(node.executor, MultiAgentBase):
-                multi_agent_result = await node.executor.execute(node_input)
+                multi_agent_result = await node.executor.execute_async(node_input)
 
                 # Create NodeResult with MultiAgentResult directly
                 node_result = NodeResult(
@@ -332,7 +344,13 @@ async def _execute_node(self, node: "GraphNode") -> None:
                 )
 
             elif isinstance(node.executor, Agent):
-                agent_response = node.executor(node_input)
+                agent_response = None  # Initialize with None to handle case where no result is yielded
+                async for event in node.executor.stream_async(node_input):
+                    if "result" in event:
+                        agent_response = event["result"]
+
+                if not agent_response:
+                    raise ValueError(f"Node '{node.node_id}' did not return a result")
 
                 # Extract metrics from agent response
                 usage = Usage(inputTokens=0, outputTokens=0, totalTokens=0)

tests-integ/test_multiagent_graph.py

@@ -108,7 +108,7 @@ def proceed_to_second_summary(state):
         "Calculate 15 + 27 and 8 * 6, analyze both results, perform additional calculations, validate everything, "
         "and provide a comprehensive summary"
     )
-    result = await graph.execute(task)
+    result = await graph.execute_async(task)
 
     # Verify results
     assert result.status.value == "completed"

tests/strands/multiagent/test_base.py

@@ -138,7 +138,10 @@ class IncompleteMultiAgent(MultiAgentBase):
 
     # Test that complete implementations can be instantiated
     class CompleteMultiAgent(MultiAgentBase):
-        async def execute(self, task: str) -> MultiAgentResult:
+        async def execute_async(self, task: str) -> MultiAgentResult:
+            return MultiAgentResult(results={})
+
+        def execute(self, task: str) -> MultiAgentResult:
             return MultiAgentResult(results={})
 
     # Should not raise an exception

tests/strands/multiagent/test_graph.py

@@ -1,4 +1,4 @@
-from unittest.mock import AsyncMock, Mock
+from unittest.mock import AsyncMock, MagicMock, Mock
 
 import pytest
 
@@ -25,8 +25,15 @@ def create_mock_agent(name, response_text="Default response", metrics=None, agen
         state={},
         metrics=metrics,
     )
+
     agent.return_value = mock_result
     agent.__call__ = Mock(return_value=mock_result)
+
+    async def mock_stream_async(*args, **kwargs):
+        yield {"result": mock_result}
+
+    agent.stream_async = MagicMock(side_effect=mock_stream_async)
+
     return agent
 
 
@@ -51,7 +58,8 @@ def create_mock_multi_agent(name, response_text="Multi-agent response"):
         execution_count=1,
         execution_time=150,
     )
-    multi_agent.execute = AsyncMock(return_value=mock_result)
+    multi_agent.execute_async = AsyncMock(return_value=mock_result)
+    multi_agent.execute = Mock(return_value=mock_result)
     return multi_agent
 
 
@@ -158,7 +166,7 @@ async def test_graph_execution(mock_graph, mock_agents, string_content_agent):
     start_node = mock_graph.nodes["start_agent"]
     assert conditional_edge.should_traverse(GraphState(completed_nodes={start_node}))
 
-    result = await mock_graph.execute("Test comprehensive execution")
+    result = await mock_graph.execute_async("Test comprehensive execution")
 
     # Verify execution results
     assert result.status == Status.COMPLETED
@@ -169,14 +177,14 @@ async def test_graph_execution(mock_graph, mock_agents, string_content_agent):
     assert result.execution_order[0].node_id == "start_agent"
 
     # Verify agent calls
-    mock_agents["start_agent"].assert_called_once()
-    mock_agents["multi_agent"].execute.assert_called_once()
-    mock_agents["conditional_agent"].assert_called_once()
-    mock_agents["final_agent"].assert_called_once()
-    mock_agents["no_metrics_agent"].assert_called_once()
-    mock_agents["partial_metrics_agent"].assert_called_once()
-    string_content_agent.assert_called_once()
-    mock_agents["blocked_agent"].assert_not_called()
+    mock_agents["start_agent"].stream_async.assert_called_once()
+    mock_agents["multi_agent"].execute_async.assert_called_once()
+    mock_agents["conditional_agent"].stream_async.assert_called_once()
+    mock_agents["final_agent"].stream_async.assert_called_once()
+    mock_agents["no_metrics_agent"].stream_async.assert_called_once()
+    mock_agents["partial_metrics_agent"].stream_async.assert_called_once()
+    string_content_agent.stream_async.assert_called_once()
+    mock_agents["blocked_agent"].stream_async.assert_not_called()
 
     # Verify metrics aggregation
     assert result.accumulated_usage["totalTokens"] > 0
@@ -219,7 +227,7 @@ class UnsupportedExecutor:
     graph = builder.build()
 
     with pytest.raises(ValueError, match="Node 'unsupported_node' of type.*is not supported"):
-        await graph.execute("test task")
+        await graph.execute_async("test task")
 
 
 @pytest.mark.asyncio
@@ -228,9 +236,15 @@ async def test_graph_execution_with_failures():
     failing_agent = Mock(spec=Agent)
     failing_agent.name = "failing_agent"
     failing_agent.id = "fail_node"
-    failing_agent.side_effect = Exception("Simulated failure")
     failing_agent.__call__ = Mock(side_effect=Exception("Simulated failure"))
 
+    # Create a proper failing async generator for stream_async
+    async def mock_stream_failure(*args, **kwargs):
+        raise Exception("Simulated failure")
+        yield  # This will never be reached
+
+    failing_agent.stream_async = mock_stream_failure
+
     success_agent = create_mock_agent("success_agent", "Success")
 
     builder = GraphBuilder()
@@ -242,7 +256,7 @@ async def test_graph_execution_with_failures():
     graph = builder.build()
 
     with pytest.raises(Exception, match="Simulated failure"):
-        await graph.execute("Test error handling")
+        await graph.execute_async("Test error handling")
 
     assert graph.state.status == Status.FAILED
     assert any(node.node_id == "fail_node" for node in graph.state.failed_nodes)
@@ -259,10 +273,10 @@ async def test_graph_edge_cases():
     builder.add_node(entry_agent, "entry_only")
     graph = builder.build()
 
-    result = await graph.execute("Original task")
+    result = await graph.execute_async("Original task")
 
     # Verify entry node was called with original task
-    entry_agent.assert_called_once_with("Original task")
+    entry_agent.stream_async.assert_called_once_with("Original task")
     assert result.status == Status.COMPLETED
 
 
@@ -415,6 +429,37 @@ def test_condition(state):
     assert len(node.dependencies) == 0
 
 
+def test_graph_synchronous_execution(mock_agents):
+    """Test synchronous graph execution using execute method."""
+    builder = GraphBuilder()
+    builder.add_node(mock_agents["start_agent"], "start_agent")
+    builder.add_node(mock_agents["final_agent"], "final_agent")
+    builder.add_edge("start_agent", "final_agent")
+    builder.set_entry_point("start_agent")
+
+    graph = builder.build()
+
+    # Test synchronous execution
+    result = graph.execute("Test synchronous execution")
+
+    # Verify execution results
+    assert result.status == Status.COMPLETED
+    assert result.total_nodes == 2
+    assert result.completed_nodes == 2
+    assert result.failed_nodes == 0
+    assert len(result.execution_order) == 2
+    assert result.execution_order[0].node_id == "start_agent"
+    assert result.execution_order[1].node_id == "final_agent"
+
+    # Verify agent calls
+    mock_agents["start_agent"].stream_async.assert_called_once()
+    mock_agents["final_agent"].stream_async.assert_called_once()
+
+    # Verify return type is GraphResult
+    assert isinstance(result, GraphResult)
+    assert isinstance(result, MultiAgentResult)
+
+
 def test_graph_result_str_representation():
     """Test GraphResult string representation."""
     mock_agent = create_mock_agent("test_agent")