def wrap_async_generator(
    gen: AsyncGenerator[T, None],
    on_iter: Optional[Callable[[], Any]] = None,
    wrap: Optional[Callable[[T], T]] = None,
    on_done: Optional[Callable[[List[T]], Any]] = None,
    context_vars: Optional[ContextVarsOrValues] = None,
) -> AsyncGenerator[T, None]:
    """Wrap a generator in another generator that will call callbacks at various
    points in the generation process.

    Args:
        gen: The generator to wrap.

        on_iter: The callback to call when the wrapper generator is created but
            before a first iteration is produced.

        wrap: The callback to call with the result of each iteration of the
            wrapped generator.

        on_done: The callback to call when the wrapped generator is exhausted.

        context_vars: The context variables to copy over to the wrapped
            generator. If None, all context variables are taken with their
            present values. See
            [with_context][trulens.core.utils.python.with_context].
    """

    async def wrapper(gen):
        with with_context(context_vars):
            if on_iter is not None:
                on_iter()

            vals = []

            async for val in gen:
                if wrap is not None:
                    val = wrap(val)  # allow handlers to rewrap the value

                vals.append(val)

                yield val

            if on_done is not None:
                on_done(vals)

    wrapper.__name__ = f"tru_wrapped_{gen.__class__.__name__}"

    return wrapper(gen)
Wrap a generator in another generator that will call callbacks at various
    points in the generation process.

    Args:
        gen: The generator to wrap.

        on_iter: The callback to call when the wrapper generator is created but
            before a first iteration is produced.

        wrap: The callback to call with the result of each iteration of the
            wrapped generator.

        on_done: The callback to call when the wrapped generator is exhausted.

        context_vars: The context variables to copy over to the wrapped
            generator. If None, all context variables are taken with their
            present values. See
            [with_context][trulens.core.utils.python.with_context].
    
src/core/trulens/core/utils/python.py