code: # 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
  
apis: llama_index.callbacks.base.CallbackManager
extract_api: [((1536, 1583), 'opentelemetry.sdk.resources.Resource.create', 'Resource.create', (["{SERVICE_NAME: 'chain-server'}"], {}), "({SERVICE_NAME: 'chain-server'})\n", (1551, 1583), False, 'from opentelemetry.sdk.resources import SERVICE_NAME, Resource\n'), ((1595, 1628), 'opentelemetry.sdk.trace.TracerProvider', 'TracerProvider', ([], {'resource': 'resource'}), '(resource=resource)\n', (1609, 1628), False, 'from opentelemetry.sdk.trace import TracerProvider\n'), ((1775, 1810), 'opentelemetry.trace.set_tracer_provider', 'trace.set_tracer_provider', (['provider'], {}), '(provider)\n', (1800, 1810), False, 'from opentelemetry import trace, context\n'), ((1820, 1852), 'opentelemetry.trace.get_tracer', 'trace.get_tracer', (['"""chain-server"""'], {}), "('chain-server')\n", (1836, 1852), False, 'from opentelemetry import trace, context\n'), ((2280, 2310), 'opentelemetry.propagate.set_global_textmap', 'set_global_textmap', (['propagator'], {}), '(propagator)\n', (2298, 2310), False, 'from opentelemetry.propagate import set_global_textmap, get_global_textmap\n'), ((1632, 1664), 'os.environ.get', 'os.environ.get', (['"""ENABLE_TRACING"""'], {}), "('ENABLE_TRACING')\n", (1646, 1664), False, 'import os\n'), ((1943, 1975), 'os.environ.get', 'os.environ.get', (['"""ENABLE_TRACING"""'], {}), "('ENABLE_TRACING')\n", (1957, 1975), False, 'import os\n'), ((2004, 2035), 'opentelemetry.trace.propagation.tracecontext.TraceContextTextMapPropagator', 'TraceContextTextMapPropagator', ([], {}), '()\n', (2033, 2035), False, 'from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator\n'), ((2154, 2213), 'tools.observability.llamaindex.opentelemetry_callback.OpenTelemetryCallbackHandler', 'opentelemetry_callback.OpenTelemetryCallbackHandler', (['tracer'], {}), '(tracer)\n', (2205, 2213), False, 'from tools.observability.llamaindex import opentelemetry_callback\n'), ((2237, 2260), 'opentelemetry.propagators.composite.CompositePropagator', 'CompositePropagator', (['[]'], {}), '([])\n', (2256, 2260), False, 'from opentelemetry.propagators.composite import CompositePropagator\n'), ((2398, 2409), 'functools.wraps', 'wraps', (['func'], {}), '(func)\n', (2403, 2409), False, 'from functools import wraps\n'), ((1712, 1730), 'opentelemetry.exporter.otlp.proto.grpc.trace_exporter.OTLPSpanExporter', 'OTLPSpanExporter', ([], {}), '()\n', (1728, 1730), False, 'from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter\n'), ((2628, 2647), 'opentelemetry.context.attach', 'context.attach', (['ctx'], {}), '(ctx)\n', (2642, 2647), False, 'from opentelemetry import trace, context\n'), ((2935, 2954), 'llama_index.callbacks.base.CallbackManager', 'CallbackManager', (['[]'], {}), '([])\n', (2950, 2954), False, 'from llama_index.callbacks.base import CallbackManager\n'), ((2542, 2562), 'opentelemetry.propagate.get_global_textmap', 'get_global_textmap', ([], {}), '()\n', (2560, 2562), False, 'from opentelemetry.propagate import set_global_textmap, get_global_textmap\n')]