def wrap_lazy_values(
        self,
        rets: Any,
        wrap: Callable[[T], T],
        on_done: Optional[Callable[[T], T]],
        context_vars: Optional[python_utils.ContextVarsOrValues] = None,
    ) -> Any:
        """Wrap any llamaindex specific lazy values with wrappers that have callback wrap."""

        # NOTE(piotrm): This is all very frail. We need to make sure we call
        # on_done on things which are not lazy and wrap things which are lazy.
        # It is not easy to tell which is which sometimes in llamaindex.

        was_lazy = False

        members = {k: v for k, v in getmembers_static(rets)}

        if hasattr(rets, "is_done") and rets.is_done:
            return on_done(rets)

        if isinstance(rets, (Response)):
            return on_done(rets)

        if isinstance(rets, (AgentChatResponse)):
            return on_done(rets)

        if "async_response_gen" in members and isinstance(
            rets.async_response_gen, AsyncGenerator
        ):
            rets.async_response_gen = python_utils.wrap_async_generator(
                rets.async_response_gen,
                wrap=wrap,
                on_done=on_done,
                context_vars=context_vars,
            )
            was_lazy = True

        if "achat_stream" in members and isinstance(
            rets.achat_stream, AsyncGenerator
        ):
            rets.achat_stream = python_utils.wrap_async_generator(
                rets.achat_stream,
                wrap=wrap,
                on_done=on_done,
                context_vars=context_vars,
            )
            was_lazy = True

        if "chat_stream" in members and isinstance(rets.chat_stream, Generator):
            rets.chat_stream = python_utils.wrap_generator(
                rets.chat_stream,
                wrap=wrap,
                on_done=on_done,
                context_vars=context_vars,
            )
            was_lazy = True

        if "response_gen" in members and isinstance(
            rets.response_gen, Generator
        ):
            wrapped_response_gen = python_utils.wrap_generator(
                rets.response_gen,
                wrap=wrap,
                on_done=on_done,
                context_vars=context_vars,
            )
            if isinstance(members["response_gen"], property):
                # NOTE(piotrm): problem here as this is a property in
                # StramingAgentChatResponse so we cannot set it. Instead we override the
                # class which has an overridden property.

                # TODO(piotrm): Figure out if there is an easier way to override
                # an attribute which is a property.

                class Wrappable(rets.__class__):
                    @property
                    def response_gen(self):
                        return wrapped_response_gen

                Wrappable.__name__ = rets.__class__.__name__
                rets.__class__ = Wrappable

            else:
                rets.response_gen = wrapped_response_gen

            was_lazy = True

        if was_lazy:
            return rets

        else:
            return on_done(rets)
Wrap any llamaindex specific lazy values with wrappers that have callback wrap.
src/apps/llamaindex/trulens/apps/llamaindex/tru_llama.py