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(43        self, token: str | list[str | dict[str, Any]], **kwargs: Any44    ) -> None:45        token_str = token if isinstance(token, str) else str(token)46        if token_str != "":47            self.queue.put_nowait(token_str)4849    @override50    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:51        self.done.set()5253    @override54    async def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:55        self.done.set()5657    # TODO: implement the other methods5859    async def aiter(self) -> AsyncIterator[str]:60        """Asynchronous iterator that yields tokens."""61        while not self.queue.empty() or not self.done.is_set():62            # Wait for the next token in the queue,63            # but stop waiting if the done event is set64            done, other = await asyncio.wait(65                [66                    # NOTE: If you add other tasks here, update the code below,67                    # which assumes each set has exactly one task each68                    asyncio.ensure_future(self.queue.get()),69                    asyncio.ensure_future(self.done.wait()),70                ],71                return_when=asyncio.FIRST_COMPLETED,72            )7374            # Cancel the other task75            if other:76                other.pop().cancel()7778            # Extract the value of the first completed task79            token_or_done = cast("str | Literal[True]", done.pop().result())8081            # If the extracted value is the boolean True, the done event was set82            if token_or_done is True:83                break8485            # Otherwise, the extracted value is a token, which we yield86            yield token_or_done