Tee Class — langchain Architecture

Source Code

libs/core/langchain_core/utils/aiter.py lines 161–274

class Tee(Generic[T]):
    """Create `n` separate asynchronous iterators over `iterable`.

    This splits a single `iterable` into multiple iterators, each providing
    the same items in the same order.

    All child iterators may advance separately but share the same items from `iterable`
    -- when the most advanced iterator retrieves an item, it is buffered until the least
    advanced iterator has yielded it as well.

    A `tee` works lazily and can handle an infinite `iterable`, provided
    that all iterators advance.

    ```python
    async def derivative(sensor_data):
        previous, current = a.tee(sensor_data, n=2)
        await a.anext(previous)  # advance one iterator
        return a.map(operator.sub, previous, current)
    ```

    Unlike `itertools.tee`, `.tee` returns a custom type instead of a `tuple`. Like a
    tuple, it can be indexed, iterated and unpacked to get the child iterators. In
    addition, its `.tee.aclose` method immediately closes all children, and it can be
    used in an `async with` context for the same effect.

    If `iterable` is an iterator and read elsewhere, `tee` will *not* provide these
    items. Also, `tee` must internally buffer each item until the last iterator has
    yielded it; if the most and least advanced iterator differ by most data, using a
    `list` is more efficient (but not lazy).

    If the underlying iterable is concurrency safe (`anext` may be awaited concurrently)
    the resulting iterators are concurrency safe as well. Otherwise, the iterators are
    safe if there is only ever one single "most advanced" iterator.

    To enforce sequential use of `anext`, provide a `lock`

    - e.g. an `asyncio.Lock` instance in an `asyncio` application - and access is
        automatically synchronised.

    """

    def __init__(
        self,
        iterable: AsyncIterator[T],
        n: int = 2,
        *,
        lock: AbstractAsyncContextManager[Any] | None = None,
    ):
        """Create a `tee`.

        Args:
            iterable: The iterable to split.
            n: The number of iterators to create.
            lock: The lock to synchronise access to the shared buffers.

        """
        self._iterator = iterable.__aiter__()  # before 3.10 aiter() doesn't exist
        self._buffers: list[deque[T]] = [deque() for _ in range(n)]
        self._children = tuple(
            tee_peer(
                iterator=self._iterator,
                buffer=buffer,
                peers=self._buffers,
                lock=lock if lock is not None else NoLock(),
            )
            for buffer in self._buffers
        )

    def __len__(self) -> int:
        """Return the number of child iterators."""
        return len(self._children)

    @overload
    def __getitem__(self, item: int) -> AsyncIterator[T]: ...

    @overload
    def __getitem__(self, item: slice) -> tuple[AsyncIterator[T], ...]: ...

    def __getitem__(
        self, item: int | slice
    ) -> AsyncIterator[T] | tuple[AsyncIterator[T], ...]: