diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py b/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py index c798f36f736..76a2616b14d 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py @@ -172,6 +172,10 @@ async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowCon else: ret = await asyncio.to_thread(self._callback, results) + # Async callable objects and wrappers returning a coroutine are not coroutine functions + if inspect.isawaitable(ret): + ret = await ret + # If the callback returned a value, finalize the workflow with it if ret is not None: await ctx.yield_output(ret) diff --git a/python/packages/orchestrations/tests/test_concurrent.py b/python/packages/orchestrations/tests/test_concurrent.py index 11ea6bec0fa..7fc0f3ec3d7 100644 --- a/python/packages/orchestrations/tests/test_concurrent.py +++ b/python/packages/orchestrations/tests/test_concurrent.py @@ -145,6 +145,34 @@ def summarize_sync(results: list[AgentExecutorResponse], _ctx: WorkflowContext[A assert output == "One | Two" +class _AsyncCallableAggregator: + """Aggregator object whose __call__ is async, for example one that holds a client.""" + + async def __call__(self, results: list[AgentExecutorResponse]) -> str: + return " | ".join(sorted(r.agent_response.text for r in results)) + + +async def _join_texts(results: list[AgentExecutorResponse]) -> str: + return " | ".join(sorted(r.agent_response.text for r in results)) + + +@pytest.mark.parametrize( + "aggregator", + [ + pytest.param(_AsyncCallableAggregator(), id="async_callable_object"), + pytest.param(lambda results: _join_texts(results), id="callable_returning_coroutine"), + ], +) +async def test_concurrent_custom_aggregator_awaits_awaitable_result(aggregator: Any) -> None: + e1 = _FakeAgentExec("agentA", "One") + e2 = _FakeAgentExec("agentB", "Two") + + wf = ConcurrentBuilder(participants=[e1, e2]).with_aggregator(aggregator).build() + result = await wf.run("prompt: awaitable aggregator") + + assert result.get_outputs() == ["One | Two"] + + def test_concurrent_custom_aggregator_uses_callback_name_for_id() -> None: e1 = _FakeAgentExec("agentA", "One") e2 = _FakeAgentExec("agentB", "Two")