def wrap_awaitable(
    awaitable: Awaitable[A],
    callback_class: Type[AwaitableCallbacks[A]] = AwaitableCallbacks,
    **kwargs: Dict[str, Any],
) -> Awaitable[A]:
    """Wrap an awaitable in another awaitable that will call callbacks before
    and after the given awaitable finishes.

    Note that the resulting awaitable needs to be awaited for all but the
    `on_now` callback to trigger.

    Args:
        awaitable: The awaitable to wrap.

        callback_class: The class that provides callbacks.

        **kwargs: All other arguments are passed to the callback class
            constructor and to `on_awaitable_wrapped`.
    """

    init_args: Dict[str, Any] = {"awaitable": awaitable, **kwargs}

    async def wrapper(awaitable: Awaitable[A]) -> A:
        cb: AwaitableCallbacks[A] = callback_class(**init_args)

        try:
            result = await awaitable
            result = cb.on_awaitable_result(result=result)
            cb.on_awaitable_end()
            return result

        except Exception as e:
            e_wrapped = cb.on_awaitable_exception(error=e)
            cb.on_awaitable_end()

            if e == e_wrapped:
                raise e

            raise e_wrapped from e

    w: Awaitable[A] = wrapper(awaitable)

    init_args["wrapper"] = w
    callback_class.on_awaitable_wrapped(**init_args)

    return w
Wrap an awaitable in another awaitable that will call callbacks before
    and after the given awaitable finishes.

    Note that the resulting awaitable needs to be awaited for all but the
    `on_now` callback to trigger.

    Args:
        awaitable: The awaitable to wrap.

        callback_class: The class that provides callbacks.

        **kwargs: All other arguments are passed to the callback class
            constructor and to `on_awaitable_wrapped`.
    
src/core/trulens/experimental/otel_tracing/core/_utils/wrap.py