arize-phoenix 2.11.1__py3-none-any.whl → 3.0.1__py3-none-any.whl

phoenix/trace/llama_index/callback.py

@@ -1,240 +1,27 @@
-"""
-Callback handler for emitting trace data in OpenInference tracing format.
-OpenInference tracing is an open standard for capturing and storing
-LLM Application execution logs.
-
-It enables production LLMapp servers to seamlessly integrate with LLM
-observability solutions such as Arize and Phoenix.
-
-For more information on the specification, see
-https://github.com/Arize-ai/openinference
-"""
-import json
 import logging
-from collections import defaultdict
-from dataclasses import dataclass, field
-from datetime import datetime, timezone
+from importlib.metadata import PackageNotFoundError
+from importlib.util import find_spec
 from typing import (
     Any,
-    Callable,
-    Dict,
-    Iterable,
-    Iterator,
-    List,
-    Mapping,
-    Optional,
-    Tuple,
-    Union,
-    cast,
-)
-from uuid import uuid4
-
-import llama_index
-from llama_index.callbacks.base_handler import BaseCallbackHandler
-from llama_index.callbacks.schema import (
-    TIMESTAMP_FORMAT,
-    CBEvent,
-    CBEventType,
-    EventPayload,
 )
-from llama_index.llms.types import ChatMessage, ChatResponse
-from llama_index.response.schema import Response, StreamingResponse
-from llama_index.tools import ToolMetadata
-from typing_extensions import TypeGuard
 
-from phoenix.trace.exporter import HttpExporter
-from phoenix.trace.llama_index.streaming import (
-    instrument_streaming_response as _instrument_streaming_response,
-)
-from phoenix.trace.schemas import (
-    MimeType,
-    Span,
-    SpanEvent,
-    SpanException,
-    SpanID,
-    SpanKind,
-    SpanStatusCode,
-    TraceID,
+from openinference.instrumentation.llama_index._callback import (
+    OpenInferenceTraceCallbackHandler as _OpenInferenceTraceCallbackHandler,
 )
-from phoenix.trace.semantic_conventions import (
-    DOCUMENT_CONTENT,
-    DOCUMENT_ID,
-    DOCUMENT_METADATA,
-    DOCUMENT_SCORE,
-    EMBEDDING_EMBEDDINGS,
-    EMBEDDING_MODEL_NAME,
-    EMBEDDING_TEXT,
-    EMBEDDING_VECTOR,
-    INPUT_MIME_TYPE,
-    INPUT_VALUE,
-    LLM_INPUT_MESSAGES,
-    LLM_INVOCATION_PARAMETERS,
-    LLM_MODEL_NAME,
-    LLM_OUTPUT_MESSAGES,
-    LLM_PROMPT_TEMPLATE,
-    LLM_PROMPT_TEMPLATE_VARIABLES,
-    LLM_PROMPTS,
-    LLM_TOKEN_COUNT_COMPLETION,
-    LLM_TOKEN_COUNT_PROMPT,
-    LLM_TOKEN_COUNT_TOTAL,
-    MESSAGE_CONTENT,
-    MESSAGE_NAME,
-    MESSAGE_ROLE,
-    MESSAGE_TOOL_CALLS,
-    OUTPUT_MIME_TYPE,
-    OUTPUT_VALUE,
-    RERANKER_INPUT_DOCUMENTS,
-    RERANKER_MODEL_NAME,
-    RERANKER_OUTPUT_DOCUMENTS,
-    RERANKER_QUERY,
-    RERANKER_TOP_K,
-    RETRIEVAL_DOCUMENTS,
-    TOOL_CALL_FUNCTION_ARGUMENTS_JSON,
-    TOOL_CALL_FUNCTION_NAME,
-    TOOL_DESCRIPTION,
-    TOOL_NAME,
-    TOOL_PARAMETERS,
+from openinference.instrumentation.llama_index.version import (
+    __version__,
 )
-from phoenix.trace.tracer import SpanExporter, Tracer
-from phoenix.trace.utils import extract_version_triplet, get_stacktrace
-from phoenix.utilities.error_handling import graceful_fallback
-
-LLAMA_INDEX_MINIMUM_VERSION_TRIPLET = (0, 9, 8)
-logger = logging.getLogger(__name__)
-logger.addHandler(logging.NullHandler())
-
-CBEventID = str
-_LOCAL_TZINFO = datetime.now().astimezone().tzinfo
-
-
-@dataclass
-class CBEventData:
-    name: Optional[str] = field(default=None)
-    event_type: Optional[CBEventType] = field(default=None)
-    start_event: Optional[CBEvent] = field(default=None)
-    end_event: Optional[CBEvent] = field(default=None)
-    attributes: Dict[str, Any] = field(default_factory=dict)
-    span_id: Optional[CBEventID] = field(default=None)
-    parent_id: Optional[CBEventID] = field(default=None)
-    trace_id: Optional[TraceID] = field(default=None)
-    streaming_event: bool = field(default=False)
-
-    def set_if_unset(self, key: str, value: Any) -> None:
-        if not getattr(self, key):
-            setattr(self, key, value)
-
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 
-ChildEventIds = Dict[CBEventID, List[CBEventID]]
-EventData = Dict[CBEventID, CBEventData]
+from phoenix.trace.exporter import _OpenInferenceExporter
+from phoenix.trace.tracer import _show_deprecation_warnings
 
-
-def payload_to_semantic_attributes(
-    event_type: CBEventType,
-    payload: Dict[str, Any],
-    is_event_end: bool = False,
-) -> Dict[str, Any]:
-    """
-    Converts a LLMapp payload to a dictionary of semantic conventions compliant attributes.
-    """
-    attributes: Dict[str, Any] = {}
-    if event_type in (CBEventType.NODE_PARSING, CBEventType.CHUNKING):
-        # TODO(maybe): handle these events
-        return attributes
-    if EventPayload.CHUNKS in payload and EventPayload.EMBEDDINGS in payload:
-        attributes[EMBEDDING_EMBEDDINGS] = [
-            {EMBEDDING_TEXT: text, EMBEDDING_VECTOR: vector}
-            for text, vector in zip(payload[EventPayload.CHUNKS], payload[EventPayload.EMBEDDINGS])
-        ]
-    if event_type is not CBEventType.RERANKING and EventPayload.QUERY_STR in payload:
-        attributes[INPUT_VALUE] = payload[EventPayload.QUERY_STR]
-        attributes[INPUT_MIME_TYPE] = MimeType.TEXT
-    if event_type is not CBEventType.RERANKING and EventPayload.NODES in payload:
-        attributes[RETRIEVAL_DOCUMENTS] = [
-            {
-                DOCUMENT_ID: node_with_score.node.node_id,
-                DOCUMENT_SCORE: node_with_score.score,
-                DOCUMENT_CONTENT: node_with_score.node.text,
-                DOCUMENT_METADATA: node_with_score.node.metadata,
-            }
-            for node_with_score in payload[EventPayload.NODES]
-        ]
-    if EventPayload.PROMPT in payload:
-        attributes[LLM_PROMPTS] = [payload[EventPayload.PROMPT]]
-    if EventPayload.MESSAGES in payload:
-        messages = payload[EventPayload.MESSAGES]
-        # Messages is only relevant to the LLM invocation
-        if event_type is CBEventType.LLM:
-            attributes[LLM_INPUT_MESSAGES] = [
-                _message_payload_to_attributes(message_data) for message_data in messages
-            ]
-        elif event_type is CBEventType.AGENT_STEP and len(messages):
-            # the agent step contains a message that is actually the input
-            # akin to the query_str
-            attributes[INPUT_VALUE] = _message_payload_to_str(messages[0])
-    if response := (payload.get(EventPayload.RESPONSE) or payload.get(EventPayload.COMPLETION)):
-        attributes.update(_get_response_output(response))
-        if raw := getattr(response, "raw", None):
-            assert hasattr(raw, "get"), f"raw must be Mapping, found {type(raw)}"
-            attributes.update(_get_output_messages(raw))
-            if usage := raw.get("usage"):
-                # OpenAI token counts are available on raw.usage but can also be
-                # found in additional_kwargs. Thus the duplicate handling.
-                attributes.update(_get_token_counts(usage))
-        # Look for token counts in additional_kwargs of the completion payload
-        # This is needed for non-OpenAI models
-        if (additional_kwargs := getattr(response, "additional_kwargs", None)) is not None:
-            attributes.update(_get_token_counts(additional_kwargs))
-    if event_type is CBEventType.RERANKING:
-        if EventPayload.TOP_K in payload:
-            attributes[RERANKER_TOP_K] = payload[EventPayload.TOP_K]
-        if EventPayload.MODEL_NAME in payload:
-            attributes[RERANKER_MODEL_NAME] = payload[EventPayload.MODEL_NAME]
-        if EventPayload.QUERY_STR in payload:
-            attributes[RERANKER_QUERY] = payload[EventPayload.QUERY_STR]
-        if nodes := payload.get(EventPayload.NODES):
-            attributes[RERANKER_OUTPUT_DOCUMENTS if is_event_end else RERANKER_INPUT_DOCUMENTS] = [
-                {
-                    DOCUMENT_ID: node_with_score.node.node_id,
-                    DOCUMENT_SCORE: node_with_score.score,
-                    DOCUMENT_CONTENT: node_with_score.node.text,
-                    DOCUMENT_METADATA: node_with_score.node.metadata,
-                }
-                for node_with_score in nodes
-            ]
-    if EventPayload.TOOL in payload:
-        tool_metadata = cast(ToolMetadata, payload.get(EventPayload.TOOL))
-        attributes[TOOL_NAME] = tool_metadata.name
-        attributes[TOOL_DESCRIPTION] = tool_metadata.description
-        attributes[TOOL_PARAMETERS] = tool_metadata.to_openai_tool()["function"]["parameters"]
-    if EventPayload.SERIALIZED in payload:
-        serialized = payload[EventPayload.SERIALIZED]
-        if event_type is CBEventType.EMBEDDING:
-            if model_name := serialized.get("model_name"):
-                attributes[EMBEDDING_MODEL_NAME] = model_name
-        if event_type is CBEventType.LLM:
-            if model_name := serialized.get("model"):
-                attributes[LLM_MODEL_NAME] = model_name
-                invocation_parameters = _extract_invocation_parameters(serialized)
-                invocation_parameters["model"] = model_name
-                attributes[LLM_INVOCATION_PARAMETERS] = json.dumps(invocation_parameters)
-    return attributes
-
-
-def _extract_invocation_parameters(serialized: Mapping[str, Any]) -> Dict[str, Any]:
-    # FIXME: this is only based on openai. Other models have different parameters.
-    if not hasattr(serialized, "get"):
-        return {}
-    invocation_parameters: Dict[str, Any] = {}
-    additional_kwargs = serialized.get("additional_kwargs")
-    if additional_kwargs and isinstance(additional_kwargs, Mapping):
-        invocation_parameters.update(additional_kwargs)
-    for key in ("temperature", "max_tokens"):
-        if (value := serialized.get(key)) is not None:
-            invocation_parameters[key] = value
-    return invocation_parameters
+logger = logging.getLogger(__name__)
 
 
-class OpenInferenceTraceCallbackHandler(BaseCallbackHandler):
+class OpenInferenceTraceCallbackHandler(_OpenInferenceTraceCallbackHandler):
     """Callback handler for storing LLM application trace data in OpenInference format.
     OpenInference is an open standard for capturing and storing AI model
     inferences. It enables production LLMapp servers to seamlessly integrate
@@ -244,436 +31,12 @@ class OpenInferenceTraceCallbackHandler(BaseCallbackHandler):
     https://github.com/Arize-ai/openinference
     """
 
-    def __init__(
-        self,
-        callback: Optional[Callable[[List[Span]], None]] = None,
-        exporter: Optional[SpanExporter] = None,
-    ) -> None:
-        if (
-            hasattr(llama_index, "__version__")
-            and (version_triplet := extract_version_triplet(llama_index.__version__))
-            and version_triplet < LLAMA_INDEX_MINIMUM_VERSION_TRIPLET
-        ):
-            raise NotImplementedError(
-                f"minimum supported version of llama-index is "
-                f"{'.'.join(map(str, LLAMA_INDEX_MINIMUM_VERSION_TRIPLET))}, "
-                f"but yours is {llama_index.__version__}. "
-                f"you can update to the latest using `pip install llama-index --upgrade`, "
-                f"or install the minimum required for Phoenix using "
-                f'`pip install "arize-phoenix[llama-index]"`',
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        _show_deprecation_warnings(self, *args, **kwargs)
+        if find_spec("llama_index") is None:
+            raise PackageNotFoundError(
+                "Missing `llama-index`. Install with `pip install llama-index`."
             )
-        super().__init__(event_starts_to_ignore=[], event_ends_to_ignore=[])
-        self._tracer = Tracer(on_append=callback, exporter=exporter or HttpExporter())
-        self._event_id_to_event_data: EventData = defaultdict(lambda: CBEventData())
-
-    def _null_fallback(self, *args: Any, **kwargs: Any) -> None:
-        return
-
-    def _on_event_fallback(
-        self,
-        event_type: CBEventType,
-        payload: Optional[Dict[str, Any]] = None,
-        event_id: CBEventID = "",
-        **kwargs: Any,
-    ) -> CBEventID:
-        return event_id or str(uuid4())
-
-    @graceful_fallback(_on_event_fallback)
-    def on_event_start(
-        self,
-        event_type: CBEventType,
-        payload: Optional[Dict[str, Any]] = None,
-        event_id: CBEventID = "",
-        parent_id: CBEventID = "",
-        **kwargs: Any,
-    ) -> CBEventID:
-        event_id = event_id or str(uuid4())
-        if parent_data := self._event_id_to_event_data.get(parent_id):
-            trace_id = parent_data.trace_id
-        else:
-            trace_id = TraceID(uuid4())
-        event_data = self._event_id_to_event_data[event_id]
-        event_data.name = event_type.value
-        event_data.event_type = event_type
-        event_data.parent_id = None if parent_id == "root" else parent_id
-        event_data.span_id = event_id
-        event_data.trace_id = trace_id
-        event_data.start_event = CBEvent(
-            event_type=event_type,
-            payload=payload,
-            id_=event_id,
-        )
-        # Parse the payload to extract the parameters
-        if payload is not None:
-            event_data.attributes.update(
-                payload_to_semantic_attributes(event_type, payload),
-            )
-
-        return event_id
-
-    @graceful_fallback(_null_fallback)
-    def on_event_end(
-        self,
-        event_type: CBEventType,
-        payload: Optional[Dict[str, Any]] = None,
-        event_id: CBEventID = "",
-        **kwargs: Any,
-    ) -> None:
-        event_data = self._event_id_to_event_data[event_id]
-        event_data.set_if_unset("name", event_type.value)
-        event_data.set_if_unset("event_type", event_type)
-        event_data.end_event = CBEvent(
-            event_type=event_type,
-            payload=payload,
-            id_=event_id,
-        )
-
-        # Parse the payload to extract the parameters
-        if payload is not None:
-            event_data.attributes.update(
-                payload_to_semantic_attributes(event_type, payload, is_event_end=True),
-            )
-            response = payload.get(EventPayload.RESPONSE)
-            if _is_streaming_response(response):
-                event_data.streaming_event = True
-                response = _instrument_streaming_response(response, self._tracer, event_data)
-
-    @graceful_fallback(_null_fallback)
-    def start_trace(self, trace_id: Optional[str] = None) -> None:
-        self._event_id_to_event_data = defaultdict(lambda: CBEventData())
-
-    @graceful_fallback(_null_fallback)
-    def end_trace(
-        self,
-        trace_id: Optional[str] = None,
-        trace_map: Optional[ChildEventIds] = None,
-    ) -> None:
-        if not trace_map:
-            return  # TODO: investigate when empty or None trace_map is passed
-        _add_spans_to_tracer(
-            event_id_to_event_data=self._event_id_to_event_data,
-            trace_map=trace_map,
-            tracer=self._tracer,
-        )
-        self._event_id_to_event_data = defaultdict(lambda: CBEventData())
-
-    def get_spans(self) -> Iterator[Span]:
-        """
-        Returns the spans stored in the tracer. This is useful if you are running
-        LlamaIndex in a notebook environment and you want to inspect the spans.
-        """
-        return self._tracer.get_spans()
-
-
-def _add_spans_to_tracer(
-    event_id_to_event_data: EventData,
-    trace_map: ChildEventIds,
-    tracer: Tracer,
-) -> None:
-    """
-    Adds event data to the tracer, where it is converted to a span and stored in a buffer.
-
-    Args:
-        event_id_to_event_data (EventData): A map of event IDs to event data.
-
-        trace_map (ChildEventIds): A map of parent event IDs to child event IDs. The root event IDs
-        are stored under the key "root".
-
-        tracer (Tracer): The tracer that stores spans.
-    """
-
-    parent_child_id_stack: List[Tuple[Optional[SpanID], CBEventID]] = [
-        (None, root_event_id) for root_event_id in trace_map["root"]
-    ]
-    while parent_child_id_stack:
-        parent_span_id, event_id = parent_child_id_stack.pop()
-        event_data = event_id_to_event_data[event_id]
-        event_type = event_data.event_type
-        attributes = event_data.attributes
-        if not (start_event := event_data.start_event):
-            # if the callback system has broken its contract by calling
-            # on_event_end without on_event_start, do not create a span
-            continue
-
-        if event_type is CBEventType.LLM:
-            while parent_child_id_stack:
-                preceding_event_parent_span_id, preceding_event_id = parent_child_id_stack[-1]
-                if preceding_event_parent_span_id != parent_span_id:
-                    break
-                preceding_event_data = event_id_to_event_data[preceding_event_id]
-                if preceding_event_data.event_type is not CBEventType.TEMPLATING:
-                    break
-                parent_child_id_stack.pop()
-                if preceding_event_start := preceding_event_data.start_event:
-                    if preceding_payload := preceding_event_start.payload:
-                        # Add template attributes to the LLM span to which they belong.
-                        attributes.update(_template_attributes(preceding_payload))
-
-        start_time = _timestamp_to_tz_aware_datetime(start_event.time)
-        span_exceptions = _get_span_exceptions(event_data, start_time)
-        if event_data.streaming_event:
-            # Do not set the end time for streaming events so we can update the event later
-            end_time = None
-        else:
-            end_time = _get_end_time(event_data, span_exceptions)
-        start_time = start_time or end_time or datetime.now(timezone.utc)
-
-        name = event_name if (event_name := event_data.name) is not None else "unknown"
-        span_kind = _get_span_kind(event_type)
-        span = tracer.create_span(
-            name=name,
-            span_kind=span_kind,
-            trace_id=event_data.trace_id,
-            start_time=start_time,
-            end_time=end_time,
-            status_code=SpanStatusCode.ERROR if span_exceptions else SpanStatusCode.OK,
-            status_message="",
-            parent_id=parent_span_id,
-            attributes=attributes,
-            events=sorted(span_exceptions, key=lambda event: event.timestamp) or None,
-            conversation=None,
-            span_id=SpanID(event_data.span_id),
-        )
-        new_parent_span_id = span.context.span_id
-        for new_child_event_id in trace_map.get(event_id, []):
-            parent_child_id_stack.append((new_parent_span_id, new_child_event_id))
-
-
-def _get_span_kind(event_type: Optional[CBEventType]) -> SpanKind:
-    """Maps a CBEventType to a SpanKind.
-
-    Args:
-        event_type (CBEventType): LlamaIndex callback event type.
-
-    Returns:
-        SpanKind: The corresponding span kind.
-    """
-    if event_type is None:
-        return SpanKind.UNKNOWN
-    return {
-        CBEventType.EMBEDDING: SpanKind.EMBEDDING,
-        CBEventType.LLM: SpanKind.LLM,
-        CBEventType.RETRIEVE: SpanKind.RETRIEVER,
-        CBEventType.FUNCTION_CALL: SpanKind.TOOL,
-        CBEventType.AGENT_STEP: SpanKind.AGENT,
-        CBEventType.RERANKING: SpanKind.RERANKER,
-    }.get(event_type, SpanKind.CHAIN)
-
-
-def _message_payload_to_attributes(message: Any) -> Dict[str, Optional[str]]:
-    if isinstance(message, ChatMessage):
-        message_attributes = {
-            MESSAGE_ROLE: message.role.value,
-            MESSAGE_CONTENT: message.content,
-        }
-        # Parse the kwargs to extract the function name and parameters for function calling
-        # NB: these additional kwargs exist both for 'agent' and 'function' roles
-        if "name" in message.additional_kwargs:
-            message_attributes[MESSAGE_NAME] = message.additional_kwargs["name"]
-        if tool_calls := message.additional_kwargs.get("tool_calls"):
-            assert isinstance(
-                tool_calls, Iterable
-            ), f"tool_calls must be Iterable, found {type(tool_calls)}"
-            message_tool_calls = []
-            for tool_call in tool_calls:
-                if message_tool_call := dict(_get_tool_call(tool_call)):
-                    message_tool_calls.append(message_tool_call)
-            if message_tool_calls:
-                message_attributes[MESSAGE_TOOL_CALLS] = message_tool_calls
-        return message_attributes
-
-    return {
-        MESSAGE_ROLE: "user",  # assume user if not ChatMessage
-        MESSAGE_CONTENT: str(message),
-    }
-
-
-def _message_payload_to_str(message: Any) -> Optional[str]:
-    """Converts a message payload to a string, if possible"""
-    if isinstance(message, ChatMessage):
-        return message.content
-
-    return str(message)
-
-
-class _CustomJSONEncoder(json.JSONEncoder):
-    def default(self, obj: object) -> Any:
-        try:
-            return super().default(obj)
-        except TypeError:
-            if callable(as_dict := getattr(obj, "dict", None)):
-                return as_dict()
-            raise
-
-
-def _get_response_output(response: Any) -> Iterator[Tuple[str, Any]]:
-    """
-    Gets output from response objects. This is needed since the string representation of some
-    response objects includes extra information in addition to the content itself. In the
-    case of an agent's ChatResponse the output may be a `function_call` object specifying
-    the name of the function to call and the arguments to call it with.
-    """
-    if isinstance(response, ChatResponse):
-        message = response.message
-        if content := message.content:
-            yield OUTPUT_VALUE, content
-            yield OUTPUT_MIME_TYPE, MimeType.TEXT
-        else:
-            yield OUTPUT_VALUE, json.dumps(message.additional_kwargs, cls=_CustomJSONEncoder)
-            yield OUTPUT_MIME_TYPE, MimeType.JSON
-    elif isinstance(response, Response):
-        yield OUTPUT_VALUE, response.response or ""
-        yield OUTPUT_MIME_TYPE, MimeType.TEXT
-    elif isinstance(response, StreamingResponse):
-        # We cannot get the output from a streaming response without exhausting
-        # the stream, so we initially return an empty string. Additional work is
-        # needed to instrument the returned response object to update the span
-        # with the actual response once the stream has been exhausted:
-        # https://github.com/Arize-ai/phoenix/issues/1867
-        yield OUTPUT_VALUE, ""
-        yield OUTPUT_MIME_TYPE, MimeType.TEXT
-    else:  # if the response has unknown type, make a best-effort attempt to get the output
-        yield OUTPUT_VALUE, str(response)
-        yield OUTPUT_MIME_TYPE, MimeType.TEXT
-
-
-def _get_end_time(event_data: CBEventData, span_events: Iterable[SpanEvent]) -> Optional[datetime]:
-    """
-    A best-effort attempt to get the end time of an event.
-
-    LlamaIndex's callback system does not guarantee that the on_event_end hook is always called, for
-    example, when an error occurs mid-event.
-    """
-    if end_event := event_data.end_event:
-        tz_naive_end_time = _timestamp_to_tz_naive_datetime(end_event.time)
-    elif span_events:
-        last_span_event = sorted(span_events, key=lambda event: event.timestamp)[-1]
-        tz_naive_end_time = last_span_event.timestamp
-    else:
-        return None
-    return _tz_naive_to_tz_aware_datetime(tz_naive_end_time)
-
-
-def _get_span_exceptions(event_data: CBEventData, start_time: datetime) -> List[SpanException]:
-    """Collects exceptions from the start and end events, if present."""
-    span_exceptions = []
-    for event in [event_data.start_event, event_data.end_event]:
-        if event and (payload := event.payload) and (error := payload.get(EventPayload.EXCEPTION)):
-            span_exceptions.append(
-                SpanException(
-                    message=str(error),
-                    timestamp=start_time,
-                    exception_type=type(error).__name__,
-                    exception_stacktrace=get_stacktrace(error),
-                )
-            )
-    return span_exceptions
-
-
-def _timestamp_to_tz_aware_datetime(timestamp: str) -> datetime:
-    """Converts a timestamp string to a timezone-aware datetime."""
-    return _tz_naive_to_tz_aware_datetime(_timestamp_to_tz_naive_datetime(timestamp))
-
-
-def _timestamp_to_tz_naive_datetime(timestamp: str) -> datetime:
-    """Converts a timestamp string to a timezone-naive datetime."""
-    return datetime.strptime(timestamp, TIMESTAMP_FORMAT)
-
-
-def _tz_naive_to_tz_aware_datetime(timestamp: datetime) -> datetime:
-    """Converts a timezone-naive datetime to a timezone-aware datetime."""
-    return timestamp.replace(tzinfo=_LOCAL_TZINFO)
-
-
-def _get_message(message: object) -> Iterator[Tuple[str, Any]]:
-    if role := getattr(message, "role", None):
-        assert isinstance(role, str), f"content must be str, found {type(role)}"
-        yield MESSAGE_ROLE, role
-    if content := getattr(message, "content", None):
-        assert isinstance(content, str), f"content must be str, found {type(content)}"
-        yield MESSAGE_CONTENT, content
-    if tool_calls := getattr(message, "tool_calls", None):
-        assert isinstance(
-            tool_calls, Iterable
-        ), f"tool_calls must be Iterable, found {type(tool_calls)}"
-        message_tool_calls = []
-        for tool_call in tool_calls:
-            if message_tool_call := dict(_get_tool_call(tool_call)):
-                message_tool_calls.append(message_tool_call)
-        if message_tool_calls:
-            yield MESSAGE_TOOL_CALLS, message_tool_calls
-
-
-def _get_output_messages(raw: Mapping[str, Any]) -> Iterator[Tuple[str, Any]]:
-    assert hasattr(raw, "get"), f"raw must be Mapping, found {type(raw)}"
-    if not (choices := raw.get("choices")):
-        return
-    assert isinstance(choices, Iterable), f"choices must be Iterable, found {type(choices)}"
-    messages = [
-        dict(_get_message(message))
-        for choice in choices
-        if (message := getattr(choice, "message", None)) is not None
-    ]
-    yield LLM_OUTPUT_MESSAGES, messages
-
-
-def _get_token_counts(usage: Union[object, Mapping[str, Any]]) -> Iterator[Tuple[str, Any]]:
-    """
-    Yields token count attributes from a object or mapping
-    """
-    # Call the appropriate function based on the type of usage
-    if isinstance(usage, Mapping):
-        yield from _get_token_counts_from_mapping(usage)
-    elif isinstance(usage, object):
-        yield from _get_token_counts_from_object(usage)
-
-
-def _get_token_counts_from_object(usage: object) -> Iterator[Tuple[str, Any]]:
-    """
-    Yields token count attributes from response.raw.usage
-    """
-    if (prompt_tokens := getattr(usage, "prompt_tokens", None)) is not None:
-        yield LLM_TOKEN_COUNT_PROMPT, prompt_tokens
-    if (completion_tokens := getattr(usage, "completion_tokens", None)) is not None:
-        yield LLM_TOKEN_COUNT_COMPLETION, completion_tokens
-    if (total_tokens := getattr(usage, "total_tokens", None)) is not None:
-        yield LLM_TOKEN_COUNT_TOTAL, total_tokens
-
-
-def _get_token_counts_from_mapping(
-    usage_mapping: Mapping[str, Any],
-) -> Iterator[Tuple[str, Any]]:
-    """
-    Yields token count attributes from a mapping (e.x. completion kwargs payload)
-    """
-    if (prompt_tokens := usage_mapping.get("prompt_tokens")) is not None:
-        yield LLM_TOKEN_COUNT_PROMPT, prompt_tokens
-    if (completion_tokens := usage_mapping.get("completion_tokens")) is not None:
-        yield LLM_TOKEN_COUNT_COMPLETION, completion_tokens
-    if (total_tokens := usage_mapping.get("total_tokens")) is not None:
-        yield LLM_TOKEN_COUNT_TOTAL, total_tokens
-
-
-def _template_attributes(payload: Dict[str, Any]) -> Iterator[Tuple[str, Any]]:
-    """Yields template attributes if present"""
-    if template := payload.get(EventPayload.TEMPLATE):
-        yield LLM_PROMPT_TEMPLATE, template
-    if template_vars := payload.get(EventPayload.TEMPLATE_VARS):
-        yield LLM_PROMPT_TEMPLATE_VARIABLES, template_vars
-        # TODO(maybe): other keys in the same payload
-        # EventPayload.SYSTEM_PROMPT
-        # EventPayload.QUERY_WRAPPER_PROMPT
-
-
-def _get_tool_call(tool_call: object) -> Iterator[Tuple[str, Any]]:
-    if function := getattr(tool_call, "function", None):
-        if name := getattr(function, "name", None):
-            assert isinstance(name, str), f"name must be str, found {type(name)}"
-            yield TOOL_CALL_FUNCTION_NAME, name
-        if arguments := getattr(function, "arguments", None):
-            assert isinstance(arguments, str), f"arguments must be str, found {type(arguments)}"
-            yield TOOL_CALL_FUNCTION_ARGUMENTS_JSON, arguments
-
-
-def _is_streaming_response(response: Any) -> TypeGuard[StreamingResponse]:
-    return isinstance(response, StreamingResponse)
+        tracer_provider = trace_sdk.TracerProvider()
+        tracer_provider.add_span_processor(SimpleSpanProcessor(_OpenInferenceExporter()))
+        super().__init__(trace_api.get_tracer(__name__, __version__, tracer_provider))