|
23 | 23 | import asyncio |
24 | 24 | import json |
25 | 25 | import os |
| 26 | +import queue as _queue |
| 27 | +import threading |
26 | 28 | import time |
27 | 29 | import traceback |
28 | 30 | 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 |
30 | 32 |
|
31 | 33 | from loguru import logger |
32 | 34 | from pydantic import BaseModel, Field |
@@ -1894,3 +1896,165 @@ async def arun( |
1894 | 1896 | *args, |
1895 | 1897 | **kwargs, |
1896 | 1898 | ) |
| 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