def wrap_function(self, func):
        """Create a wrapper of the given function to perform cost tracking."""

        if python_utils.safe_hasattr(func, INSTRUMENT):
            # Store the types of callback classes that will handle calls to the
            # wrapped function in the INSTRUMENT attribute. This will be used to
            # invoke appropriate callbacks when the wrapped function gets
            # called.

            # If INSTRUMENT is set, we don't need to instrument the method again
            # but we may need to add the additional callback class to expected
            # handlers stored at the attribute.

            registered_callback_classes = getattr(func, INSTRUMENT)

            if self.callback_class in registered_callback_classes:
                # If our callback class is already in the list, dont bother
                # adding it again.

                logger.debug(
                    "%s already instrumented for callbacks of type %s",
                    func.__name__,
                    self.callback_class.__name__,
                )

                return func

            else:
                # Otherwise add our callback class but don't instrument again.
                registered_callback_classes += [self.callback_class]
                setattr(func, INSTRUMENT, registered_callback_classes)

                return func

        # If INSTRUMENT is not set, create a wrapper method and return it.
        @functools.wraps(func)
        def tru_wrapper(*args, **kwargs):
            logger.debug(
                "Calling instrumented method %s of type %s, "
                "iscoroutinefunction=%s, "
                "isasyncgeneratorfunction=%s",
                func,
                type(func),
                python_utils.is_really_coroutinefunction(func),
                inspect.isasyncgenfunction(func),
            )

            if not Endpoint._have_context():
                return func(*args, **kwargs)

            endpoints = Endpoint._context_endpoints.get()

            # Context vars as they are now before we reset them. They are
            # captured in the closure of the below.
            context_vars = {
                Endpoint._context_endpoints: Endpoint._context_endpoints.get()
            }

            # Get the result of the wrapped function:
            # response = context_vars.run(func, *args, **kwargs)
            response = func(*args, **kwargs)

            # if len(Endpoint._context_endpoints.get()) == 0:
            #    raise ValueError("No endpoints.")

            bindings = inspect.signature(func).bind(*args, **kwargs)

            # Get all of the callback classes suitable for handling this
            # call. Note that we stored this in the INSTRUMENT attribute of
            # the wrapper method.
            registered_callback_classes = getattr(tru_wrapper, INSTRUMENT)

            # If wrapped method was not called from within _track_costs, we
            # will get None here and do nothing but return wrapped
            # function's response.
            if len(endpoints) == 0:
                logger.debug("No endpoints found.")
                return response

            def update_response(response):
                if python_utils.is_lazy(response):
                    return python_utils.wrap_lazy(
                        response,
                        wrap=update_response,
                        context_vars=context_vars,
                    )

                for callback_class in registered_callback_classes:
                    if callback_class not in endpoints:
                        logger.warning(
                            "Callback class %s is registered for handling %s"
                            " but there are no endpoints waiting to receive the result.",
                            callback_class.__name__,
                            func.__name__,
                        )
                        continue

                    for endpoint, callback in endpoints[callback_class]:
                        logger.debug("Handling endpoint %s.", endpoint.name)

                        with python_utils.with_context(context_vars):
                            response_ = endpoint.handle_wrapped_call(
                                func=func,
                                bindings=bindings,
                                response=response,
                                callback=callback,
                            )
                            if response_ is not None:
                                # Handler is allowed to override the response in
                                # case it wants to wrap some generators or similar
                                # lazy structures.
                                response = response_

                return response

            # Problem here: OpenAI returns generators inside its own special
            # classes. These are thus handled in
            # OpenAIEndpoint.wrapped_call .

            return update_response(response)

        # Set our tracking attribute to tell whether something is already
        # instrumented onto both the sync and async version since either one
        # could be returned from this method.
        setattr(tru_wrapper, INSTRUMENT, [self.callback_class])

        logger.debug("Instrumenting %s for %s.", func.__name__, self.name)

        return tru_wrapper
Create a wrapper of the given function to perform cost tracking.
src/core/trulens/core/feedback/endpoint.py