Skip to content

Commit b1021d1

Browse files
Steve-Dustyclaude
andcommitted
feat: add arun_stream and run_stream to HierarchicalSwarm
Bridge the existing callback-based streaming in run() to async/sync generators so HierarchicalSwarm can be consumed the same way as AgentRearrange and SequentialWorkflow. arun_stream(with_events=True) emits agent_start / token / agent_end structured events. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent d0b6016 commit b1021d1

1 file changed

Lines changed: 165 additions & 1 deletion

File tree

swarms/structs/hiearchical_swarm.py

Lines changed: 165 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,12 @@
2323
import asyncio
2424
import json
2525
import os
26+
import queue as _queue
27+
import threading
2628
import time
2729
import traceback
2830
from concurrent.futures import ThreadPoolExecutor, as_completed
29-
from typing import Any, Callable, List, Optional, Union
31+
from typing import Any, Callable, Dict, List, Optional, Union
3032

3133
from loguru import logger
3234
from pydantic import BaseModel, Field
@@ -1894,3 +1896,165 @@ async def arun(
18941896
*args,
18951897
**kwargs,
18961898
)
1899+
1900+
async def arun_stream(
1901+
self,
1902+
task: Optional[str] = None,
1903+
img: Optional[str] = None,
1904+
with_events: bool = False,
1905+
**kwargs,
1906+
):
1907+
"""Async generator that streams tokens from worker agents.
1908+
1909+
Bridges the existing callback-based streaming in ``run()`` to an
1910+
async generator using an ``asyncio.Queue``. The synchronous
1911+
``run()`` executes in a background thread; each token callback
1912+
pushes an item to the queue which this generator yields.
1913+
1914+
Args:
1915+
task: The task to be processed by the swarm.
1916+
img: Optional image input for the agents.
1917+
with_events: When False (default), yield ``(agent_name, token)``
1918+
tuples. When True, yield structured event dicts:
1919+
``{"type": "agent_start", "agent": ...}``,
1920+
``{"type": "token", "agent": ..., "token": ...}``,
1921+
``{"type": "agent_end", "agent": ...}``.
1922+
1923+
Yields:
1924+
tuple | dict: Per-token streaming items.
1925+
"""
1926+
q: asyncio.Queue = asyncio.Queue()
1927+
DONE = object()
1928+
ERROR = object()
1929+
loop = asyncio.get_running_loop()
1930+
1931+
# Track which agents have started and accumulate their tokens
1932+
agent_chunks: Dict[str, List[str]] = {}
1933+
1934+
def _streaming_callback(
1935+
agent_name: str, chunk: str, is_final: bool
1936+
):
1937+
loop.call_soon_threadsafe(
1938+
q.put_nowait, (agent_name, chunk, is_final)
1939+
)
1940+
1941+
async def _run_in_thread():
1942+
try:
1943+
await asyncio.to_thread(
1944+
self.run,
1945+
task=task,
1946+
img=img,
1947+
streaming_callback=_streaming_callback,
1948+
**kwargs,
1949+
)
1950+
except Exception as e:
1951+
loop.call_soon_threadsafe(
1952+
q.put_nowait, (ERROR, e, False)
1953+
)
1954+
finally:
1955+
loop.call_soon_threadsafe(
1956+
q.put_nowait, (DONE, None, False)
1957+
)
1958+
1959+
runner = asyncio.create_task(_run_in_thread())
1960+
1961+
try:
1962+
while True:
1963+
item = await q.get()
1964+
kind, payload, is_final = item
1965+
1966+
if kind is DONE:
1967+
break
1968+
if kind is ERROR:
1969+
raise payload
1970+
1971+
agent_name: str = kind
1972+
chunk: str = payload
1973+
1974+
if agent_name not in agent_chunks:
1975+
agent_chunks[agent_name] = []
1976+
if with_events:
1977+
yield {
1978+
"type": "agent_start",
1979+
"agent": agent_name,
1980+
}
1981+
1982+
if is_final:
1983+
agent_output = "".join(agent_chunks.pop(agent_name, []))
1984+
if with_events:
1985+
yield {
1986+
"type": "agent_end",
1987+
"agent": agent_name,
1988+
"output": agent_output,
1989+
}
1990+
else:
1991+
if chunk:
1992+
agent_chunks[agent_name].append(chunk)
1993+
if with_events:
1994+
yield {
1995+
"type": "token",
1996+
"agent": agent_name,
1997+
"token": chunk,
1998+
}
1999+
else:
2000+
yield (agent_name, chunk)
2001+
finally:
2002+
await runner
2003+
2004+
def run_stream(
2005+
self,
2006+
task: Optional[str] = None,
2007+
img: Optional[str] = None,
2008+
with_events: bool = False,
2009+
**kwargs,
2010+
):
2011+
"""Sync generator version of ``arun_stream``.
2012+
2013+
Bridges the async generator to a sync iterator using a thread
2014+
and a ``queue.Queue``. Use ``arun_stream`` directly when you
2015+
already have a running event loop (e.g. inside a FastAPI handler).
2016+
2017+
Args:
2018+
task: The task to be processed by the swarm.
2019+
img: Optional image input for the agents.
2020+
with_events: When False (default), yield ``(agent_name, token)``
2021+
tuples. When True, yield structured event dicts.
2022+
2023+
Yields:
2024+
tuple | dict: Per-token streaming items.
2025+
"""
2026+
sync_q: _queue.Queue = _queue.Queue()
2027+
DONE = object()
2028+
exc_holder: List[Optional[Exception]] = [None]
2029+
2030+
def _runner():
2031+
async def _consume():
2032+
async for evt in self.arun_stream(
2033+
task=task,
2034+
img=img,
2035+
with_events=with_events,
2036+
**kwargs,
2037+
):
2038+
sync_q.put(evt)
2039+
2040+
try:
2041+
asyncio.run(_consume())
2042+
except Exception as e:
2043+
exc_holder[0] = e
2044+
finally:
2045+
sync_q.put(DONE)
2046+
2047+
t = threading.Thread(target=_runner, daemon=True)
2048+
t.start()
2049+
2050+
try:
2051+
while True:
2052+
item = sync_q.get()
2053+
if item is DONE:
2054+
break
2055+
yield item
2056+
finally:
2057+
t.join(timeout=5)
2058+
2059+
if exc_holder[0] is not None:
2060+
raise exc_holder[0]

0 commit comments

Comments
 (0)