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.
  
  from opentelemetry.trace import Tracer, get_tracer, set_span_in_context, Status, StatusCode
  from opentelemetry.trace.span import Span
  from opentelemetry.context import Context, get_current, attach, detach
  from typing import Any, Dict, List, Optional, Callable
  from llama_index.callbacks.base_handler import BaseCallbackHandler
  from llama_index.callbacks.base import CallbackManager
  from llama_index.callbacks.schema import CBEventType, EventPayload, BASE_TRACE_EVENT
  from llama_index.callbacks.token_counting import get_llm_token_counts, TokenCountingEvent
  from llama_index.utilities.token_counting import TokenCounter
  from llama_index.utils import get_tokenizer
  from dataclasses import dataclass
  from contextvars import ContextVar
  import threading
  
  global_root_trace = ContextVar("trace", default=None)
  
  @dataclass
  class SpanWithContext:
      """Object for tracking a span, its context, and its context token"""
      span: Span
      context: Context
      token: object
  
      def __init__(self, span: Span, context: Context, token: object, thread_identity):
          self.span = span
          self.context = context
          self.token = token
          self.thread_identity = thread_identity
  
  class OpenTelemetryCallbackHandler(BaseCallbackHandler):
      """Callback handler for creating OpenTelemetry traces from llamaindex traces and events."""
  
      def __init__(
          self,
          tracer: Optional[Tracer] = get_tracer(__name__),
          tokenizer: Optional[Callable[[str], List]] = None,
      ) -> None:
          """Initializes the OpenTelemetryCallbackHandler.
  
          Args:
              tracer: Optional[Tracer]: A OpenTelemetry tracer used to create OpenTelemetry spans
          """
          super().__init__(event_starts_to_ignore=[], event_ends_to_ignore=[])
          self._tracer = tracer
          self._event_map: Dict[str, SpanWithContext] = {}
          self.tokenizer = tokenizer or get_tokenizer()
          self._token_counter = TokenCounter(tokenizer=self.tokenizer)
          
      def start_trace(self, trace_id: Optional[str] = None) -> None:
          trace_name = "llamaindex.trace"
          if trace_id is not None:
              trace_name = "llamaindex.trace." + trace_id
          span = self._tracer.start_span(trace_name)
          ctx = set_span_in_context(span)
          token = attach(ctx)
          global_root_trace.set(SpanWithContext(span=span, context=ctx, token=token, thread_identity=threading.get_ident()))
  
      def end_trace(
          self,
          trace_id: Optional[str] = None,
          trace_map: Optional[Dict[str, List[str]]] = None,
      ) -> None:
          root_trace = global_root_trace.get()
          if root_trace is not None:
              if root_trace.thread_identity == threading.get_ident():
                  detach(root_trace.token)
              root_trace.span.end()
  
      def on_event_start(
          self,
          event_type: CBEventType,
          payload: Optional[Dict[str, Any]] = None,
          event_id: str = "",
          parent_id: str = "",
          **kwargs: Any,
      ) -> str:
          parent_ctx = None
          # Case where the parent of this event is another event
          if parent_id in self._event_map:
              parent_ctx = self._event_map[parent_id].context
          # Case where the parent of this event is the root trace, and the root trace exists
          elif parent_id is BASE_TRACE_EVENT and global_root_trace.get() is not None:
              parent_ctx = global_root_trace.get().context
          # Case where the parent of this event is the root trace, but the trace does not exist
          else:
              return
          
          span_prefix = "llamaindex.event."
          span = self._tracer.start_span(span_prefix + event_type.value, context=parent_ctx)
          ctx = set_span_in_context(span)
          token = attach(ctx)
          self._event_map[event_id] = SpanWithContext(span=span, context=ctx, token=token, thread_identity=threading.get_ident())
  
          span.set_attribute("event_id", event_id)
          if payload is not None:
              if event_type is CBEventType.QUERY:
                  span.set_attribute("query.text", payload[EventPayload.QUERY_STR])
              elif event_type is CBEventType.RETRIEVE:
                  pass
              elif event_type is CBEventType.EMBEDDING:
                  span.set_attribute("embedding.model", payload[EventPayload.SERIALIZED]['model_name'])
                  span.set_attribute("embedding.batch_size", payload[EventPayload.SERIALIZED]['embed_batch_size'])
                  span.set_attribute("embedding.class_name", payload[EventPayload.SERIALIZED]['class_name'])
              elif event_type is CBEventType.SYNTHESIZE:
                  span.set_attribute("synthesize.query_text", payload[EventPayload.QUERY_STR])
              elif event_type is CBEventType.CHUNKING:
                  for i, chunk in enumerate(payload[EventPayload.CHUNKS]):
                      span.set_attribute(f"chunk.{i}", chunk)
              elif event_type is CBEventType.TEMPLATING:
                  if payload[EventPayload.QUERY_WRAPPER_PROMPT]:
                      span.set_attribute("query_wrapper_prompt", payload[EventPayload.QUERY_WRAPPER_PROMPT])
                  if payload[EventPayload.SYSTEM_PROMPT]:
                      span.set_attribute("system_prompt", payload[EventPayload.SYSTEM_PROMPT])
                  if payload[EventPayload.TEMPLATE]:
                      span.set_attribute("template", payload[EventPayload.TEMPLATE])
                  if payload[EventPayload.TEMPLATE_VARS]:
                      for key, var in payload[EventPayload.TEMPLATE_VARS].items():
                          span.set_attribute(f"template_variables.{key}", var)
              elif event_type is CBEventType.LLM:
                  span.set_attribute("llm.class_name", payload[EventPayload.SERIALIZED]['class_name'])
                  span.set_attribute("llm.formatted_prompt", payload[EventPayload.PROMPT])
                  span.set_attribute("llm.additional_kwargs", str(payload[EventPayload.ADDITIONAL_KWARGS]))
              elif event_type is CBEventType.NODE_PARSING:
                  span.set_attribute("node_parsing.num_documents", len(payload[EventPayload.DOCUMENTS]))
              elif event_type is CBEventType.EXCEPTION:
                  span.set_status(Status(StatusCode.ERROR))
                  span.record_exception(payload[EventPayload.EXCEPTION])
          return event_id
  
      def on_event_end(
          self,
          event_type: CBEventType,
          payload: Optional[Dict[str, Any]] = None,
          event_id: str = "",
          **kwargs: Any,
      ) -> None:
          if event_id in self._event_map:
              span = self._event_map[event_id].span
              span.set_attribute("event_id", event_id)
              if payload is not None:
                  if event_type is CBEventType.QUERY:
                      pass
                  elif event_type is CBEventType.RETRIEVE:
                      for i, node_with_score in enumerate(payload[EventPayload.NODES]):
                          node = node_with_score.node
                          score = node_with_score.score
                          span.set_attribute(f"query.node.{i}.id", node.hash)
                          span.set_attribute(f"query.node.{i}.score", score)
                          span.set_attribute(f"query.node.{i}.text", node.text)
                  elif event_type is CBEventType.EMBEDDING:
                      texts = payload[EventPayload.CHUNKS]
                      vectors = payload[EventPayload.EMBEDDINGS]
                      total_chunk_tokens = 0
                      for text, vector in zip(texts, vectors) :
                          span.set_attribute(f"embedding_text_{texts.index(text)}", text) 
                          span.set_attribute(f"embedding_vector_{vectors.index(vector)}", vector) 
                          total_chunk_tokens +=self._token_counter.get_string_tokens(text)
                      span.set_attribute(f"embedding_token_usage", total_chunk_tokens) 
                  elif event_type is CBEventType.SYNTHESIZE:
                      pass
                  elif event_type is CBEventType.CHUNKING:
                      pass
                  elif event_type is CBEventType.TEMPLATING:
                      pass
                  elif event_type is CBEventType.LLM:
                      span.set_attribute("response.text", str(
                              payload.get(EventPayload.RESPONSE, "")
                          ) or str(payload.get(EventPayload.COMPLETION, ""))
                      )
                      token_counts = get_llm_token_counts(self._token_counter, payload, event_id)   
                      span.set_attribute("llm_prompt.token_usage", token_counts.prompt_token_count)
                      span.set_attribute("llm_completion.token_usage", token_counts.completion_token_count)
                      span.set_attribute("total_tokens_used", token_counts.total_token_count)
                  elif event_type is CBEventType.NODE_PARSING:
                      span.set_attribute("node_parsing.num_nodes", len(payload[EventPayload.NODES]))
                  elif event_type is CBEventType.EXCEPTION:
                      span.set_status(Status(StatusCode.ERROR))
                      span.record_exception(payload[EventPayload.EXCEPTION])
              if self._event_map[event_id].thread_identity == threading.get_ident():
                  detach(self._event_map[event_id].token)
              self._event_map.pop(event_id, None)
              span.end()
          
apis: llama_index.utils.get_tokenizer llama_index.callbacks.token_counting.get_llm_token_counts llama_index.utilities.token_counting.TokenCounter
extract_api: [((1450, 1483), 'contextvars.ContextVar', 'ContextVar', (['"""trace"""'], {'default': 'None'}), "('trace', default=None)\n", (1460, 1483), False, 'from contextvars import ContextVar\n'), ((2085, 2105), 'opentelemetry.trace.get_tracer', 'get_tracer', (['__name__'], {}), '(__name__)\n', (2095, 2105), False, 'from opentelemetry.trace import Tracer, get_tracer, set_span_in_context, Status, StatusCode\n'), ((2609, 2647), 'llama_index.utilities.token_counting.TokenCounter', 'TokenCounter', ([], {'tokenizer': 'self.tokenizer'}), '(tokenizer=self.tokenizer)\n', (2621, 2647), False, 'from llama_index.utilities.token_counting import TokenCounter\n'), ((2918, 2943), 'opentelemetry.trace.set_span_in_context', 'set_span_in_context', (['span'], {}), '(span)\n', (2937, 2943), False, 'from opentelemetry.trace import Tracer, get_tracer, set_span_in_context, Status, StatusCode\n'), ((2960, 2971), 'opentelemetry.context.attach', 'attach', (['ctx'], {}), '(ctx)\n', (2966, 2971), False, 'from opentelemetry.context import Context, get_current, attach, detach\n'), ((4386, 4411), 'opentelemetry.trace.set_span_in_context', 'set_span_in_context', (['span'], {}), '(span)\n', (4405, 4411), False, 'from opentelemetry.trace import Tracer, get_tracer, set_span_in_context, Status, StatusCode\n'), ((4428, 4439), 'opentelemetry.context.attach', 'attach', (['ctx'], {}), '(ctx)\n', (4434, 4439), False, 'from opentelemetry.context import Context, get_current, attach, detach\n'), ((2563, 2578), 'llama_index.utils.get_tokenizer', 'get_tokenizer', ([], {}), '()\n', (2576, 2578), False, 'from llama_index.utils import get_tokenizer\n'), ((3367, 3388), 'threading.get_ident', 'threading.get_ident', ([], {}), '()\n', (3386, 3388), False, 'import threading\n'), ((3406, 3430), 'opentelemetry.context.detach', 'detach', (['root_trace.token'], {}), '(root_trace.token)\n', (3412, 3430), False, 'from opentelemetry.context import Context, get_current, attach, detach\n'), ((4545, 4566), 'threading.get_ident', 'threading.get_ident', ([], {}), '()\n', (4564, 4566), False, 'import threading\n'), ((9898, 9919), 'threading.get_ident', 'threading.get_ident', ([], {}), '()\n', (9917, 9919), False, 'import threading\n'), ((9937, 9976), 'opentelemetry.context.detach', 'detach', (['self._event_map[event_id].token'], {}), '(self._event_map[event_id].token)\n', (9943, 9976), False, 'from opentelemetry.context import Context, get_current, attach, detach\n'), ((3071, 3092), 'threading.get_ident', 'threading.get_ident', ([], {}), '()\n', (3090, 3092), False, 'import threading\n'), ((9123, 9183), 'llama_index.callbacks.token_counting.get_llm_token_counts', 'get_llm_token_counts', (['self._token_counter', 'payload', 'event_id'], {}), '(self._token_counter, payload, event_id)\n', (9143, 9183), False, 'from llama_index.callbacks.token_counting import get_llm_token_counts, TokenCountingEvent\n'), ((6876, 6900), 'opentelemetry.trace.Status', 'Status', (['StatusCode.ERROR'], {}), '(StatusCode.ERROR)\n', (6882, 6900), False, 'from opentelemetry.trace import Tracer, get_tracer, set_span_in_context, Status, StatusCode\n'), ((9737, 9761), 'opentelemetry.trace.Status', 'Status', (['StatusCode.ERROR'], {}), '(StatusCode.ERROR)\n', (9743, 9761), False, 'from opentelemetry.trace import Tracer, get_tracer, set_span_in_context, Status, StatusCode\n')]