Can you help propose a PR with the following code (this is in braintrust/wrappers/openai.py? `` """ Exports BraintrustTracingProcessor, a tracing.TracingProcessor` that logs traces to Braintrust. """ import datetime from typing import Any, Callab...

[Copilot]

Plan to add span_name_fn parameter to BraintrustTracingProcessor:

• Add Callable import to typing imports
• Rename global _span_name function to _default_span_name
• Add span_name_fn parameter to __init__ method signature
• Add span_name_fn to instance variables in __init__
• Add instance method _span_name that uses custom function if provided, otherwise falls back to default
• Update on_span_start to use instance method self._span_name instead of global function
• Update class docstring to document the new span_name_fn parameter
• Run tests to ensure changes don't break existing functionality

Original prompt

Can you help propose a PR with the following code (this is in braintrust/wrappers/openai.py?

"""
Exports `BraintrustTracingProcessor`, a `tracing.TracingProcessor` that logs traces to Braintrust.
"""

import datetime
from typing import Any, Callable, Dict, Optional, Union

from agents import tracing

import braintrust
from braintrust.logger import NOOP_SPAN


def _span_type(span: tracing.Span[Any]) -> braintrust.SpanTypeAttribute:
    if span.span_data.type in ["agent", "handoff", "custom"]:
        return braintrust.SpanTypeAttribute.TASK
    elif span.span_data.type in ["function", "guardrail"]:
        return braintrust.SpanTypeAttribute.TOOL
    elif span.span_data.type in ["generation", "response"]:
        return braintrust.SpanTypeAttribute.LLM
    else:
        return braintrust.SpanTypeAttribute.TASK


def _default_span_name(span: tracing.Span[Any]) -> str:
    # TODO(sachin): span name should also come from the span_data.
    if (
        isinstance(span.span_data, tracing.AgentSpanData)
        or isinstance(span.span_data, tracing.FunctionSpanData)
        or isinstance(span.span_data, tracing.GuardrailSpanData)
        or isinstance(span.span_data, tracing.CustomSpanData)
    ):
        return span.span_data.name
    elif isinstance(span.span_data, tracing.GenerationSpanData):
        return "Generation"
    elif isinstance(span.span_data, tracing.ResponseSpanData):
        return "Response"
    elif isinstance(span.span_data, tracing.HandoffSpanData):
        return "Handoff"
    else:
        return "Unknown"


def _timestamp_from_maybe_iso(timestamp: Optional[str]) -> Optional[float]:
    if timestamp is None:
        return None
    return datetime.datetime.fromisoformat(timestamp).timestamp()


def _maybe_timestamp_elapsed(end: Optional[str], start: Optional[str]) -> Optional[float]:
    if start is None or end is None:
        return None
    return (datetime.datetime.fromisoformat(end) - datetime.datetime.fromisoformat(start)).total_seconds()


class BraintrustTracingProcessor(tracing.TracingProcessor):
    """
    `BraintrustTracingProcessor` is a `tracing.TracingProcessor` that logs traces to Braintrust.

    Args:
        logger: A `braintrust.Span` or `braintrust.Experiment` or `braintrust.Logger` to use for logging.
            If `None`, the current span, experiment, or logger will be selected exactly as in `braintrust.start_span`.
        span_name_fn: An optional callable that takes a `tracing.Span` and returns a custom span name as a string.
            If the callable returns `None`, the default span naming logic will be used as a fallback.
            This allows you to customize span names for specific span types (e.g., LLM calls) while preserving
            default behavior for others.
    """

    def __init__(
        self,
        logger: Optional[Union[braintrust.Span, braintrust.Experiment, braintrust.Logger]] = None,
        span_name_fn: Optional[Callable[[tracing.Span[Any]], Optional[str]]] = None,
    ):  
        self._logger = logger
        self._span_name_fn = span_name_fn
        self._spans: Dict[str, braintrust.Span] = {}
        self._first_input: Dict[str, Any] = {}
        self._last_output: Dict[str, Any] = {}

    def _span_name(self, span: tracing.Span[Any]) -> str:
        # Try custom naming function first if provided
        if self._span_name_fn is not None:
            custom_name = self._span_name_fn(span)
            if custom_name is not None:
                return custom_name
        
        # Fall back to default naming logic
        return _default_span_name(span)

    def on_trace_start(self, trace: tracing.Trace) -> None:
        trace_meta = trace.export() or {}
        metadata = {
            "group_id": trace_meta.get("group_id"),
            **(trace_meta.get("metadata") or {}),
        }

        current_context = braintrust.current_span()
        if current_context != NOOP_SPAN:
            self._spans[trace.trace_id] = current_context.start_span(
                name=trace.name,
                span_attributes={"type": "task", "name": trace.name},
                metadata=metadata,
            )
        elif self._logger is not None:
            self._spans[trace.trace_id] = self._logger.start_span(
                span_attributes={"type": "task", "name": trace.name},
                span_id=trace.trace_id,
                root_span_id=trace.trace_id,
                metadata=metadata,
                # TODO(sachin): Add start time when SDK provides it.
                # start_time=_timestamp_from_maybe_iso(trace.started_at),
            )
        else:
            self._spans[trace.trace_id] = braintrust.start_span(
                id=trace.trace_id,
                span_attributes={"type": "task", "name": trace.name},
                metadata=metadata,
                # TODO(sachin): Add start time when SDK provides it.
                # start_time=_timestamp_from_maybe_iso(trace.started_at),
            )

    def on_t...

</details>


<!-- START COPILOT CODING AGENT TIPS -->
---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more [Copilot coding agent tips](https://gh.io/copilot-coding-agent-tips) in the docs.