1"""Asynchronous iterator utilities.23Adapted from4https://github.com/maxfischer2781/asyncstdlib/blob/master/asyncstdlib/itertools.py5MIT License.6"""78from collections import deque9from collections.abc import (10    AsyncGenerator,11    AsyncIterable,12    AsyncIterator,13    Awaitable,14    Callable,15    Iterator,16)17from contextlib import AbstractAsyncContextManager18from types import TracebackType19from typing import (20    Any,21    Generic,22    TypeVar,23    cast,24    overload,25)2627from typing_extensions import override2829from langchain_core._api.deprecation import deprecated3031T = TypeVar("T")3233_no_default = object()343536# https://github.com/python/cpython/blob/main/Lib/test/test_asyncgen.py#L5437@deprecated(since="1.1.2", removal="2.0.0")38def py_anext(39    iterator: AsyncIterator[T], default: T | Any = _no_default40) -> Awaitable[T | Any | None]:41    """Pure-Python implementation of `anext()` for testing purposes.4243    Closely matches the builtin `anext()` C implementation.4445    Can be used to compare the built-in implementation of the inner coroutines machinery46    to C-implementation of `__anext__()` and `send()` or `throw()` on the returned47    generator.4849    Args:50        iterator: The async iterator to advance.51        default: The value to return if the iterator is exhausted.5253            If not provided, a `StopAsyncIteration` exception is raised.5455    Returns:56        The next value from the iterator, or the default value if the iterator is57            exhausted.5859    Raises:60        TypeError: If the iterator is not an async iterator.61    """62    try:63        __anext__ = cast(64            "Callable[[AsyncIterator[T]], Awaitable[T]]", type(iterator).__anext__65        )66    except AttributeError as e:67        msg = f"{iterator!r} is not an async iterator"68        raise TypeError(msg) from e6970    if default is _no_default:71        return __anext__(iterator)7273    async def anext_impl() -> T | Any:74        try:75            # The C code is way more low-level than this, as it implements76            # all methods of the iterator protocol. In this implementation77            # we're relying on higher-level coroutine concepts, but that's78            # exactly what we want -- crosstest pure-Python high-level79            # implementation and low-level C anext() iterators.80            return await __anext__(iterator)81        except StopAsyncIteration:82            return default8384    return anext_impl()858687class NoLock:88    """Dummy lock that provides the proper interface but no protection."""8990    async def __aenter__(self) -> None:91        """Do nothing."""9293    async def __aexit__(94        self,95        exc_type: type[BaseException] | None,96        exc_val: BaseException | None,97        exc_tb: TracebackType | None,98    ) -> bool:99        """Return False, exception not suppressed."""100        return False101102103async def tee_peer(104    iterator: AsyncIterator[T],105    # the buffer specific to this peer106    buffer: deque[T],107    # the buffers of all peers, including our own108    peers: list[deque[T]],109    lock: AbstractAsyncContextManager[Any],110) -> AsyncGenerator[T, None]:111    """An individual iterator of a `tee`.112113    This function is a generator that yields items from the shared iterator114    `iterator`. It buffers items until the least advanced iterator has yielded them as115    well.116117    The buffer is shared with all other peers.118119    Args:120        iterator: The shared iterator.121        buffer: The buffer for this peer.122        peers: The buffers of all peers.123        lock: The lock to synchronise access to the shared buffers.124125    Yields:126        The next item from the shared iterator.127    """128    try:129        while True:130            if not buffer:131                async with lock:132                    # Another peer produced an item while we were waiting for the lock.133                    # Proceed with the next loop iteration to yield the item.134                    if buffer:135                        continue136                    try:137                        item = await anext(iterator)138                    except StopAsyncIteration:139                        break140                    else:141                        # Append to all buffers, including our own. We'll fetch our142                        # item from the buffer again, instead of yielding it directly.143                        # This ensures the proper item ordering if any of our peers144                        # are fetching items concurrently. They may have buffered their145                        # item already.146                        for peer_buffer in peers:147                            peer_buffer.append(item)148            yield buffer.popleft()149    finally:150        async with lock:151            # this peer is done - remove its buffer152            for idx, peer_buffer in enumerate(peers):  # pragma: no branch153                if peer_buffer is buffer:154                    peers.pop(idx)155                    break156            # if we are the last peer, try and close the iterator157            if not peers and hasattr(iterator, "aclose"):158                await iterator.aclose()159160161class Tee(Generic[T]):162    """Create `n` separate asynchronous iterators over `iterable`.163164    This splits a single `iterable` into multiple iterators, each providing165    the same items in the same order.166167    All child iterators may advance separately but share the same items from `iterable`168    -- when the most advanced iterator retrieves an item, it is buffered until the least169    advanced iterator has yielded it as well.170171    A `tee` works lazily and can handle an infinite `iterable`, provided172    that all iterators advance.173174    ```python175    async def derivative(sensor_data):176        previous, current = a.tee(sensor_data, n=2)177        await a.anext(previous)  # advance one iterator178        return a.map(operator.sub, previous, current)179    ```180181    Unlike `itertools.tee`, `.tee` returns a custom type instead of a `tuple`. Like a182    tuple, it can be indexed, iterated and unpacked to get the child iterators. In183    addition, its `.tee.aclose` method immediately closes all children, and it can be184    used in an `async with` context for the same effect.185186    If `iterable` is an iterator and read elsewhere, `tee` will *not* provide these187    items. Also, `tee` must internally buffer each item until the last iterator has188    yielded it; if the most and least advanced iterator differ by most data, using a189    `list` is more efficient (but not lazy).190191    If the underlying iterable is concurrency safe (`anext` may be awaited concurrently)192    the resulting iterators are concurrency safe as well. Otherwise, the iterators are193    safe if there is only ever one single "most advanced" iterator.194195    To enforce sequential use of `anext`, provide a `lock`196197    - e.g. an `asyncio.Lock` instance in an `asyncio` application - and access is198        automatically synchronised.199200    """201202    def __init__(203        self,204        iterable: AsyncIterator[T],205        n: int = 2,206        *,207        lock: AbstractAsyncContextManager[Any] | None = None,208    ):209        """Create a `tee`.210211        Args:212            iterable: The iterable to split.213            n: The number of iterators to create.214            lock: The lock to synchronise access to the shared buffers.215216        """217        self._iterator = iterable.__aiter__()  # before 3.10 aiter() doesn't exist218        self._buffers: list[deque[T]] = [deque() for _ in range(n)]219        self._children = tuple(220            tee_peer(221                iterator=self._iterator,222                buffer=buffer,223                peers=self._buffers,224                lock=lock if lock is not None else NoLock(),225            )226            for buffer in self._buffers227        )228229    def __len__(self) -> int:230        """Return the number of child iterators."""231        return len(self._children)232233    @overload234    def __getitem__(self, item: int) -> AsyncIterator[T]: ...235236    @overload237    def __getitem__(self, item: slice) -> tuple[AsyncIterator[T], ...]: ...238239    def __getitem__(240        self, item: int | slice241    ) -> AsyncIterator[T] | tuple[AsyncIterator[T], ...]:242        """Return the child iterator(s) for the given index or slice."""243        return self._children[item]244245    def __iter__(self) -> Iterator[AsyncIterator[T]]:246        """Iterate over the child iterators.247248        Yields:249            The child iterators.250        """251        yield from self._children252253    async def __aenter__(self) -> "Tee[T]":254        """Return the tee instance."""255        return self256257    async def __aexit__(258        self,259        exc_type: type[BaseException] | None,260        exc_val: BaseException | None,261        exc_tb: TracebackType | None,262    ) -> bool:263        """Close all child iterators.264265        Returns:266            `False`, exceptions not suppressed.267        """268        await self.aclose()269        return False270271    async def aclose(self) -> None:272        """Async close all child iterators."""273        for child in self._children:274            await child.aclose()275276277atee = Tee278279280class aclosing(AbstractAsyncContextManager):  # noqa: N801281    """Async context manager to wrap an `AsyncGenerator` that has a `aclose()` method.282283    Code like this:284285    ```python286    async with aclosing(<module>.fetch(<arguments>)) as agen:287        <block>288    ```289290    ...is equivalent to this:291292    ```python293    agen = <module>.fetch(<arguments>)294    try:295        <block>296    finally:297        await agen.aclose()298299    ```300    """301302    def __init__(self, thing: AsyncGenerator[Any, Any] | AsyncIterator[Any]) -> None:303        """Create the context manager.304305        Args:306            thing: The resource to wrap.307        """308        self.thing = thing309310    @override311    async def __aenter__(self) -> AsyncGenerator[Any, Any] | AsyncIterator[Any]:312        return self.thing313314    @override315    async def __aexit__(316        self,317        exc_type: type[BaseException] | None,318        exc_value: BaseException | None,319        traceback: TracebackType | None,320    ) -> None:321        if hasattr(self.thing, "aclose"):322            await self.thing.aclose()323324325async def abatch_iterate(326    size: int, iterable: AsyncIterable[T]327) -> AsyncIterator[list[T]]:328    """Utility batching function for async iterables.329330    Args:331        size: The size of the batch.332        iterable: The async iterable to batch.333334    Yields:335        The batches.336    """337    batch: list[T] = []338    async for element in iterable:339        if len(batch) < size:340            batch.append(element)341342        if len(batch) >= size:343            yield batch344            batch = []345346    if batch:347        yield batch