tracing.py

# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Module for configuring objects used to create OpenTelemetry traces."""

import os
from opentelemetry import trace, context
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.propagate import set_global_textmap, get_global_textmap
from opentelemetry.propagators.composite import CompositePropagator
from tools.observability.llamaindex import opentelemetry_callback
import llama_index
from llama_index.callbacks.base import CallbackManager
from functools import wraps

# Configure tracer used by the Chain Server to create spans
resource = Resource.create({SERVICE_NAME: "chain-server"})
provider = TracerProvider(resource=resource)
if os.environ.get("ENABLE_TRACING") == "true":
    processor = SimpleSpanProcessor(OTLPSpanExporter())
    provider.add_span_processor(processor)
trace.set_tracer_provider(provider)
tracer = trace.get_tracer("chain-server")

# Configure Propagator used for processing trace context received by the Chain Server
if os.environ.get("ENABLE_TRACING") == "true":
    propagator = TraceContextTextMapPropagator()
    # Llamaindex global handler set to pass callbacks into the OpenTelemetry handler
    llama_index.global_handler = opentelemetry_callback.OpenTelemetryCallbackHandler(tracer)
else:
    propagator = CompositePropagator([]) # No-op propagator
set_global_textmap(propagator)

# Wrapper Function to perform instrumentation
def instrumentation_wrapper(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        request = kwargs.get("request")
        prompt = kwargs.get("prompt")
        ctx = get_global_textmap().extract(request.headers)
        if ctx is not None:
            context.attach(ctx)
        if prompt is not None and prompt.use_knowledge_base == False:
            # Hack to get the LLM event for no knowledge base queries to show up.
            # A trace is not generated by Llamaindex for these calls so we need to generate it instead.
            callback_manager = CallbackManager([])
            with callback_manager.as_trace("query"):
                result = func(*args, **kwargs)
        else:
            result = func(*args, **kwargs)
        return await result

    return wrapper