Repository: microsoft/promptflow
File: src/promptflow-tracing/promptflow/tracing/_trace.py (Python)

Code before:
from importlib.metadata import version
from threading import Lock
from typing import Callable, Dict, List, Optional

import opentelemetry.trace as otel_trace
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.trace import Link, Span
from opentelemetry.trace.span import NonRecordingSpan, format_trace_id
from opentelemetry.trace.status import StatusCode

from ._openai_utils import OpenAIMetricsCalculator, OpenAIResponseParser
from ._operation_context import OperationContext
from ._tracer import Tracer, _create_trace_from_function_call, get_node_name_from_context
from ._utils import get_input_names_for_prompt_template, get_prompt_param_name_from_func, serialize
from .contracts.generator_proxy import AsyncGeneratorProxy, GeneratorProxy
from .contracts.trace import Trace, TraceType

IS_LEGACY_OPENAI = version("openai").startswith("0.")


@contextlib.contextmanager
def _record_keyboard_interrupt_to_span(span: Span):
    try:
        yield
    except KeyboardInterrupt as ex:
        span.record_exception(ex)
        span.set_status(StatusCode.ERROR, "KeyboardInterrupt received, execution cancelled.")
        raise


class TokenCollector:
    _lock = Lock()

    def __init__(self):
        self._span_id_to_tokens = {}

    def collect_openai_tokens(self, span, output):
        span_id = span.get_span_context().span_id
        if not inspect.isgenerator(output) and hasattr(output, "usage") and output.usage is not None:
            tokens = output.usage.dict()
            if tokens:
                with self._lock:
                    self._span_id_to_tokens[span_id] = tokens

    def collect_openai_tokens_for_streaming(self, span, inputs, output, is_chat):
        span_id = span.get_span_context().span_id
        calculator = OpenAIMetricsCalculator()
        if is_chat:
            tokens = calculator.get_openai_metrics_for_chat_api(inputs, output)
        else:
            tokens = calculator.get_openai_metrics_for_completion_api(inputs, output)

Review comment (suggestion):
We might be able to add `if span.is_recording():` here to align with the default error handler of the otel sdk.

Below are picked-up code from the otel sdk:

```python
@contextmanager
def use_span(
    span: Span,
    end_on_exit: bool = False,
    record_exception: bool = True,
    set_status_on_exception: bool = True,
) -> Iterator[Span]:
    """Takes a non-active span and activates it in the current context.

    Args:
        span: The span that should be activated in the current context.
        end_on_exit: Whether to end the span automatically when leaving the
            context manager scope.
        record_exception: Whether to record any exceptions raised within the
            context as error event on the span.
        set_status_on_exception: Only relevant if the returned span is used
            in a with/context manager. Defines whether the span status will
            be automatically set to ERROR when an uncaught exception is
            raised in the span with block. The span status won't be set by
            this mechanism if it was previously set manually.
    """
    try:
        token = context_api.attach(context_api.set_value(_SPAN_KEY, span))
        print(f"==== attached, token={token}")
        try:
            yield span
        finally:
            print(f"==== detaching, token={token}")
            context_api.detach(token)

    except Exception as exc:  # pylint: disable=broad-except
        if isinstance(span, Span) and span.is_recording():
            # Record the exception as an event
            if record_exception:
                span.record_exception(exc)

            # Set status in case exception was raised
            if set_status_on_exception:
                span.set_status(
                    Status(
                        status_code=StatusCode.ERROR,
                        description=f"{type(exc).__name__}: {exc}",
                    )
                )

        # This causes parent spans to set their status to ERROR and to record
        # an exception as an event if a child span raises an exception even if
        # such child span was started with both record_exception and
        # set_status_on_exception attributes set to False.
        raise

    finally:
        if end_on_exit:
            span.end()
```

Code after:
from importlib.metadata import version
from threading import Lock
from typing import Callable, Dict, List, Optional

import opentelemetry.trace as otel_trace
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.trace import Link, Span
from opentelemetry.trace.span import NonRecordingSpan, format_trace_id
from opentelemetry.trace.status import StatusCode

from ._openai_utils import OpenAIMetricsCalculator, OpenAIResponseParser
from ._operation_context import OperationContext
from ._tracer import Tracer, _create_trace_from_function_call, get_node_name_from_context
from ._utils import get_input_names_for_prompt_template, get_prompt_param_name_from_func, serialize
from .contracts.generator_proxy import AsyncGeneratorProxy, GeneratorProxy
from .contracts.trace import Trace, TraceType

IS_LEGACY_OPENAI = version("openai").startswith("0.")


@contextlib.contextmanager
def _record_keyboard_interrupt_to_span(span: Span):
    try:
        yield
    except KeyboardInterrupt as ex:
        if span.is_recording():
            span.record_exception(ex)
            span.set_status(StatusCode.ERROR, "Execution cancelled.")
        raise


class TokenCollector:
    _lock = Lock()

    def __init__(self):
        self._span_id_to_tokens = {}

    def collect_openai_tokens(self, span, output):
        span_id = span.get_span_context().span_id
        if not inspect.isgenerator(output) and hasattr(output, "usage") and output.usage is not None:
            tokens = output.usage.dict()
            if tokens:
                with self._lock:
                    self._span_id_to_tokens[span_id] = tokens

    def collect_openai_tokens_for_streaming(self, span, inputs, output, is_chat):
        span_id = span.get_span_context().span_id
        calculator = OpenAIMetricsCalculator()
        if is_chat:
            tokens = calculator.get_openai_metrics_for_chat_api(inputs, output)
        else: