1from __future__ import annotations23import asyncio4from collections.abc import AsyncIterator5from typing import Any, Literal, cast67from langchain_core.callbacks import AsyncCallbackHandler8from langchain_core.outputs import LLMResult9from typing_extensions import override1011# TODO: If used by two LLM runs in parallel this won't work as expected121314class AsyncIteratorCallbackHandler(AsyncCallbackHandler):15    """Callback handler that returns an async iterator."""1617    queue: asyncio.Queue[str]1819    done: asyncio.Event2021    @property22    def always_verbose(self) -> bool:23        """Always verbose."""24        return True2526    def __init__(self) -> None:27        """Instantiate AsyncIteratorCallbackHandler."""28        self.queue = asyncio.Queue()29        self.done = asyncio.Event()3031    @override32    async def on_llm_start(33        self,34        serialized: dict[str, Any],35        prompts: list[str],36        **kwargs: Any,37    ) -> None:38        # If two calls are made in a row, this resets the state39        self.done.clear()4041    @override42    async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:43        if token is not None and token != "":44            self.queue.put_nowait(token)4546    @override47    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:48        self.done.set()4950    @override51    async def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:52        self.done.set()5354    # TODO: implement the other methods5556    async def aiter(self) -> AsyncIterator[str]:57        """Asynchronous iterator that yields tokens."""58        while not self.queue.empty() or not self.done.is_set():59            # Wait for the next token in the queue,60            # but stop waiting if the done event is set61            done, other = await asyncio.wait(62                [63                    # NOTE: If you add other tasks here, update the code below,64                    # which assumes each set has exactly one task each65                    asyncio.ensure_future(self.queue.get()),66                    asyncio.ensure_future(self.done.wait()),67                ],68                return_when=asyncio.FIRST_COMPLETED,69            )7071            # Cancel the other task72            if other:73                other.pop().cancel()7475            # Extract the value of the first completed task76            token_or_done = cast("str | Literal[True]", done.pop().result())7778            # If the extracted value is the boolean True, the done event was set79            if token_or_done is True:80                break8182            # Otherwise, the extracted value is a token, which we yield83            yield token_or_done