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

Code before:
        trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION.
        name (str, optional): The name of the trace, will set to func name if not provided.


    Returns:
        Callable: The traced function.
    """

    def create_trace(func, args, kwargs):
        return _create_trace_from_function_call(
            func,
            args=args,
            kwargs=kwargs,
            args_to_ignore=args_to_ignore,
            trace_type=trace_type,
            name=name,
        )

    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        trace = create_trace(func, args, kwargs)
        # For node span we set the span name to node name, otherwise we use the function name.
        span_name = get_node_name_from_context(used_for_span_name=True) or trace.name
        # need to get everytime to ensure tracer is latest
        otel_tracer = otel_trace.get_tracer("promptflow")
        with otel_tracer.start_as_current_span(
            span_name, end_on_exit=False
        ) as span, _record_cancellation_exceptions_to_span(span):
            # Store otel trace id in context for correlation
            OperationContext.get_instance()["otel_trace_id"] = f"0x{format_trace_id(span.get_span_context().trace_id)}"
            enrich_span_with_trace(span, trace)
            enrich_span_with_prompt_info(span, func, kwargs)

            # Should not extract these codes to a separate function here.
            # We directly call func instead of calling Tracer.invoke,
            # because we want to avoid long stack trace when hitting an exception.
            try:
                Tracer.push(trace)
                enrich_span_with_input(span, trace.inputs)
                output = func(*args, **kwargs)
                if isinstance(output, Iterator):
                    output = Tracer.pop(output)
                    output = traced_generator(span, trace.inputs, output, trace_type)
                else:
                    enrich_span_with_trace_type(span, trace.inputs, output, trace_type)
                    span.set_status(StatusCode.OK)
                    output = Tracer.pop(output)
            except Exception as e:
                Tracer.pop(None, e)
                raise
        token_collector.collect_openai_tokens_for_parent_span(span)

Review comment (refactor):
Maybe we can have our own `start_as_current_span` to avoid duplicated if/end logic in multiple places?

Code after:
        trace_type (TraceType, optional): The type of the trace. Defaults to TraceType.FUNCTION.
        name (str, optional): The name of the trace, will set to func name if not provided.

    Returns:
        Callable: The traced function.
    """

    def create_trace(func, args, kwargs):
        return _create_trace_from_function_call(
            func,
            args=args,
            kwargs=kwargs,
            args_to_ignore=args_to_ignore,
            trace_type=trace_type,
            name=name,
        )

    @functools.wraps(func)
    async def wrapped(*args, **kwargs):
        trace = create_trace(func, args, kwargs)
        # For node span we set the span name to node name, otherwise we use the function name.
        span_name = get_node_name_from_context(used_for_span_name=True) or trace.name
        # need to get everytime to ensure tracer is latest
        otel_tracer = otel_trace.get_tracer("promptflow")
        with start_as_current_span(otel_tracer, span_name) as span:
            # Store otel trace id in context for correlation
            OperationContext.get_instance()["otel_trace_id"] = f"0x{format_trace_id(span.get_span_context().trace_id)}"
            enrich_span_with_trace(span, trace)
            enrich_span_with_prompt_info(span, func, kwargs)

            # Should not extract these codes to a separate function here.
            # We directly call func instead of calling Tracer.invoke,
            # because we want to avoid long stack trace when hitting an exception.
            try:
                Tracer.push(trace)
                enrich_span_with_input(span, trace.inputs)
                output = await func(*args, **kwargs)
                output = handle_output(span, trace.inputs, output, trace_type)
            except Exception as e:
                Tracer.pop(None, e)
                raise
        token_collector.collect_openai_tokens_for_parent_span(span)
        return output

    wrapped.__original_function = func

    return wrapped


def _traced_sync(
    func: Callable = None,