braintrust 0.3.15__py3-none-any.whl → 0.4.1__py3-none-any.whl

braintrust/merge_row_batch.py

@@ -1,10 +1,11 @@
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
+from collections.abc import Mapping, Sequence
+from typing import Any, Optional
 
 from .db_fields import IS_MERGE_FIELD, PARENT_ID_FIELD
 from .graph_util import UndirectedGraph, topological_sort, undirected_connected_components
 from .util import merge_dicts
 
-_MergedRowKey = Tuple[Optional[Any], ...]
+_MergedRowKey = tuple[Optional[Any], ...]
 
 
 def _generate_merged_row_key(row: Mapping[str, Any], use_parent_id_for_id: bool = False) -> _MergedRowKey:
@@ -33,7 +34,7 @@ MERGE_ROW_SKIP_FIELDS = [
 ]
 
 
-def _pop_merge_row_skip_fields(row: Dict[str, Any]) -> Dict[str, Any]:
+def _pop_merge_row_skip_fields(row: dict[str, Any]) -> dict[str, Any]:
     popped = {}
     for field in MERGE_ROW_SKIP_FIELDS:
         if field in row:
@@ -41,14 +42,14 @@ def _pop_merge_row_skip_fields(row: Dict[str, Any]) -> Dict[str, Any]:
     return popped
 
 
-def _restore_merge_row_skip_fields(row: Dict[str, Any], skip_fields: Dict[str, Any]):
+def _restore_merge_row_skip_fields(row: dict[str, Any], skip_fields: dict[str, Any]):
     for field in MERGE_ROW_SKIP_FIELDS:
         row.pop(field, None)
         if field in skip_fields:
             row[field] = skip_fields[field]
 
 
-def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]]:
+def merge_row_batch(rows: Sequence[dict[str, Any]]) -> list[list[dict[str, Any]]]:
     """Given a batch of rows, merges conflicting rows together to end up with a
     set of rows to insert. Returns a set of de-conflicted rows, as a list of
     lists, where separate lists contain "independent" rows which can be
@@ -98,7 +99,7 @@ def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]
                 "Logged row is missing an id. This is an internal braintrust error. Please contact us at info@braintrust.dev for help"
             )
 
-    row_groups: Dict[_MergedRowKey, Dict[str, Any]] = {}
+    row_groups: dict[_MergedRowKey, dict[str, Any]] = {}
     for row in rows:
         key = _generate_merged_row_key(row)
         existing_row = row_groups.get(key)
@@ -138,7 +139,7 @@ def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]
     # all groups of rows which each row in a group has a PARENT_ID_FIELD
     # relationship with at least one other row in the group.
     connected_components = undirected_connected_components(
-        UndirectedGraph(vertices=set(graph.keys()), edges=set((k, v) for k, vs in graph.items() for v in vs))
+        UndirectedGraph(vertices=set(graph.keys()), edges={(k, v) for k, vs in graph.items() for v in vs})
     )
 
     # For each connected row group, run topological sort over that subgraph to
@@ -148,8 +149,8 @@ def merge_row_batch(rows: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]
 
 
 def batch_items(
-    items: List[List[str]], batch_max_num_items: Optional[int] = None, batch_max_num_bytes: Optional[int] = None
-) -> List[List[List[str]]]:
+    items: list[list[str]], batch_max_num_items: int | None = None, batch_max_num_bytes: int | None = None
+) -> list[list[list[str]]]:
     """Repartition the given list of items into sets of batches which can be
     published in parallel or in sequence.
 

braintrust/oai.py

@@ -2,7 +2,8 @@ import abc
 import base64
 import re
 import time
-from typing import Any, Callable, Dict, List, Optional, Union
+from collections.abc import Callable
+from typing import Any
 
 from .logger import Attachment, Span, start_span
 from .span_types import SpanTypeAttribute
@@ -70,7 +71,7 @@ def log_headers(response: Any, span: Span):
         )
 
 
-def _convert_data_url_to_attachment(data_url: str, filename: Optional[str] = None) -> Union[Attachment, str]:
+def _convert_data_url_to_attachment(data_url: str, filename: str | None = None) -> Attachment | str:
     """Helper function to convert data URL to an Attachment."""
     data_url_match = re.match(r"^data:([^;]+);base64,(.+)$", data_url)
     if not data_url_match:
@@ -140,7 +141,7 @@ def _process_attachments_in_input(input_data: Any) -> Any:
 
 
 class ChatCompletionWrapper:
-    def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]]):
+    def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None):
         self.create_fn = create_fn
         self.acreate_fn = acreate_fn
 
@@ -254,7 +255,7 @@ class ChatCompletionWrapper:
                 span.end()
 
     @classmethod
-    def _parse_params(cls, params: Dict[str, Any]) -> Dict[str, Any]:
+    def _parse_params(cls, params: dict[str, Any]) -> dict[str, Any]:
         # First, destructively remove span_info
         ret = params.pop("span_info", {})
 
@@ -274,12 +275,12 @@ class ChatCompletionWrapper:
         )
 
     @classmethod
-    def _postprocess_streaming_results(cls, all_results: List[Dict[str, Any]]) -> Dict[str, Any]:
+    def _postprocess_streaming_results(cls, all_results: list[dict[str, Any]]) -> dict[str, Any]:
         role = None
         content = None
-        tool_calls: Optional[List[Any]] = None
+        tool_calls: list[Any] | None = None
         finish_reason = None
-        metrics: Dict[str, float] = {}
+        metrics: dict[str, float] = {}
         for result in all_results:
             usage = result.get("usage")
             if usage:
@@ -338,7 +339,7 @@ class ChatCompletionWrapper:
 
 
 class ResponseWrapper:
-    def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]], name: str = "openai.responses.create"):
+    def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None, name: str = "openai.responses.create"):
         self.create_fn = create_fn
         self.acreate_fn = acreate_fn
         self.name = name
@@ -447,7 +448,7 @@ class ResponseWrapper:
                 span.end()
 
     @classmethod
-    def _parse_params(cls, params: Dict[str, Any]) -> Dict[str, Any]:
+    def _parse_params(cls, params: dict[str, Any]) -> dict[str, Any]:
         # First, destructively remove span_info
         ret = params.pop("span_info", {})
 
@@ -467,7 +468,7 @@ class ResponseWrapper:
         )
 
     @classmethod
-    def _parse_event_from_result(cls, result: Dict[str, Any]) -> Dict[str, Any]:
+    def _parse_event_from_result(cls, result: dict[str, Any]) -> dict[str, Any]:
         """Parse event from response result"""
         data = {"metrics": {}}
 
@@ -487,7 +488,7 @@ class ResponseWrapper:
         return data
 
     @classmethod
-    def _postprocess_streaming_results(cls, all_results: List[Any]) -> Dict[str, Any]:
+    def _postprocess_streaming_results(cls, all_results: list[Any]) -> dict[str, Any]:
         """Process streaming results - minimal version focused on metrics extraction."""
         metrics = {}
         output = []
@@ -570,13 +571,13 @@ class ResponseWrapper:
 
 
 class BaseWrapper(abc.ABC):
-    def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]], name: str):
+    def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None, name: str):
         self._create_fn = create_fn
         self._acreate_fn = acreate_fn
         self._name = name
 
     @abc.abstractmethod
-    def process_output(self, response: Dict[str, Any], span: Span):
+    def process_output(self, response: dict[str, Any], span: Span):
         """Process the API response and log relevant information to the span."""
         pass
 
@@ -614,7 +615,7 @@ class BaseWrapper(abc.ABC):
             return raw_response
 
     @classmethod
-    def _parse_params(cls, params: Dict[str, Any]) -> Dict[str, Any]:
+    def _parse_params(cls, params: dict[str, Any]) -> dict[str, Any]:
         # First, destructively remove span_info
         ret = params.pop("span_info", {})
 
@@ -634,10 +635,10 @@ class BaseWrapper(abc.ABC):
 
 
 class EmbeddingWrapper(BaseWrapper):
-    def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]]):
+    def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None):
         super().__init__(create_fn, acreate_fn, "Embedding")
 
-    def process_output(self, response: Dict[str, Any], span: Span):
+    def process_output(self, response: dict[str, Any], span: Span):
         usage = response.get("usage")
         metrics = _parse_metrics_from_usage(usage)
         span.log(
@@ -649,7 +650,7 @@ class EmbeddingWrapper(BaseWrapper):
 
 
 class ModerationWrapper(BaseWrapper):
-    def __init__(self, create_fn: Optional[Callable[..., Any]], acreate_fn: Optional[Callable[..., Any]]):
+    def __init__(self, create_fn: Callable[..., Any] | None, acreate_fn: Callable[..., Any] | None):
         super().__init__(create_fn, acreate_fn, "Moderation")
 
     def process_output(self, response: Any, span: Span):
@@ -896,7 +897,7 @@ TOKEN_PREFIX_MAP = {
 }
 
 
-def _parse_metrics_from_usage(usage: Any) -> Dict[str, Any]:
+def _parse_metrics_from_usage(usage: Any) -> dict[str, Any]:
     # For simplicity, this function handles all the different APIs
     metrics = {}
 
@@ -930,7 +931,7 @@ def _is_numeric(v):
     return isinstance(v, (int, float, complex)) and not isinstance(v, bool)
 
 
-def prettify_params(params: Dict[str, Any]) -> Dict[str, Any]:
+def prettify_params(params: dict[str, Any]) -> dict[str, Any]:
     # Filter out NOT_GIVEN parameters
     # https://linear.app/braintrustdata/issue/BRA-2467
     ret = {k: v for k, v in params.items() if not _is_not_given(v)}
@@ -940,7 +941,7 @@ def prettify_params(params: Dict[str, Any]) -> Dict[str, Any]:
     return ret
 
 
-def _try_to_dict(obj: Any) -> Dict[str, Any]:
+def _try_to_dict(obj: Any) -> dict[str, Any]:
     if isinstance(obj, dict):
         return obj
     # convert a pydantic object to a dict

braintrust/otel/__init__.py

@@ -1,7 +1,6 @@
 import logging
 import os
 import warnings
-from typing import Dict, Optional
 from urllib.parse import urljoin
 
 INSTALL_ERR_MSG = (
@@ -138,10 +137,10 @@ class OtelExporter(OTLPSpanExporter):
 
     def __init__(
         self,
-        url: Optional[str] = None,
-        api_key: Optional[str] = None,
-        parent: Optional[str] = None,
-        headers: Optional[Dict[str, str]] = None,
+        url: str | None = None,
+        api_key: str | None = None,
+        parent: str | None = None,
+        headers: dict[str, str] | None = None,
         **kwargs,
     ):
         """
@@ -189,13 +188,14 @@ class OtelExporter(OTLPSpanExporter):
         super().__init__(endpoint=endpoint, headers=exporter_headers, **kwargs)
 
 
-def add_braintrust_span_processor(tracer_provider,
-    api_key: Optional[str] = None,
-    parent: Optional[str] = None,
-    api_url: Optional[str] = None,
+def add_braintrust_span_processor(
+    tracer_provider,
+    api_key: str | None = None,
+    parent: str | None = None,
+    api_url: str | None = None,
     filter_ai_spans: bool = False,
     custom_filter=None,
-    headers: Optional[Dict[str, str]] = None,
+    headers: dict[str, str] | None = None,
 ):
     processor = BraintrustSpanProcessor(
         api_key=api_key,
@@ -225,13 +225,13 @@ class BraintrustSpanProcessor:
 
     def __init__(
         self,
-        api_key: Optional[str] = None,
-        parent: Optional[str] = None,
-        api_url: Optional[str] = None,
+        api_key: str | None = None,
+        parent: str | None = None,
+        api_url: str | None = None,
         filter_ai_spans: bool = False,
-        custom_filter = None,
-        headers: Optional[Dict[str, str]] = None,
-        SpanProcessor: Optional[type] = None,
+        custom_filter=None,
+        headers: dict[str, str] | None = None,
+        SpanProcessor: type | None = None,
     ):
         """
         Initialize the BraintrustSpanProcessor.
@@ -279,16 +279,17 @@ class BraintrustSpanProcessor:
 
             # Priority 1: Check if braintrust.parent is in current OTEL context
             from opentelemetry import baggage, context
+
             current_context = context.get_current()
-            parent_value = context.get_value('braintrust.parent', current_context)
+            parent_value = context.get_value("braintrust.parent", current_context)
 
             # Priority 2: Check OTEL baggage (propagates automatically across contexts)
             if not parent_value:
-                parent_value = baggage.get_baggage('braintrust.parent', context=current_context)
+                parent_value = baggage.get_baggage("braintrust.parent", context=current_context)
 
             # Priority 3: Check if parent_context has braintrust.parent (backup)
             if not parent_value and parent_context:
-                parent_value = context.get_value('braintrust.parent', parent_context)
+                parent_value = context.get_value("braintrust.parent", parent_context)
 
             # Priority 4: Check if parent OTEL span has braintrust.parent attribute
             if not parent_value and parent_context:
@@ -304,7 +305,6 @@ class BraintrustSpanProcessor:
 
         self._processor.on_start(span, parent_context)
 
-
     def _get_parent_otel_braintrust_parent(self, parent_context):
         """Get braintrust.parent attribute from parent OTEL span if it exists."""
         try:
@@ -313,7 +313,7 @@ class BraintrustSpanProcessor:
             # Get the current span from the parent context
             current_span = trace.get_current_span(parent_context)
 
-            if current_span and hasattr(current_span, 'attributes') and current_span.attributes:
+            if current_span and hasattr(current_span, "attributes") and current_span.attributes:
                 # Check if parent span has braintrust.parent attribute
                 attributes = dict(current_span.attributes)
                 return attributes.get("braintrust.parent")
@@ -346,11 +346,7 @@ class BraintrustSpanProcessor:
         return self._processor
 
 
-def _get_braintrust_parent(
-    object_type,
-    object_id: Optional[str] = None,
-    compute_args: Optional[Dict] = None
-) -> Optional[str]:
+def _get_braintrust_parent(object_type, object_id: str | None = None, compute_args: dict | None = None) -> str | None:
     """
     Construct a braintrust.parent identifier string from span components.
 
@@ -405,11 +401,10 @@ def context_from_span_export(export_str: str):
     if not OTEL_AVAILABLE:
         raise ImportError(INSTALL_ERR_MSG)
 
+    from braintrust.span_identifier_v4 import SpanComponentsV4
     from opentelemetry import baggage, trace
     from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
 
-    from braintrust.span_identifier_v4 import SpanComponentsV4
-
     # Parse the export string (handles V3/V4 automatically)
     components = SpanComponentsV4.from_str(export_str)
 
@@ -417,7 +412,7 @@ def context_from_span_export(export_str: str):
     braintrust_parent = _get_braintrust_parent(
         object_type=components.object_type,
         object_id=components.object_id,
-        compute_args=components.compute_object_metadata_args
+        compute_args=components.compute_object_metadata_args,
     )
 
     # Convert hex strings to OTEL integers
@@ -429,7 +424,7 @@ def context_from_span_export(export_str: str):
         trace_id=trace_id_int,
         span_id=span_id_int,
         is_remote=True,  # Critical: mark as remote for distributed tracing
-        trace_flags=TraceFlags(TraceFlags.SAMPLED)
+        trace_flags=TraceFlags(TraceFlags.SAMPLED),
     )
 
     # Create NonRecordingSpan and set in context
@@ -438,7 +433,7 @@ def context_from_span_export(export_str: str):
 
     # Set braintrust.parent in OTEL baggage so it propagates automatically
     if braintrust_parent:
-        ctx = baggage.set_baggage('braintrust.parent', braintrust_parent, context=ctx)
+        ctx = baggage.set_baggage("braintrust.parent", braintrust_parent, context=ctx)
 
     return ctx
 
@@ -475,7 +470,7 @@ def add_parent_to_baggage(parent: str, ctx=None):
     from opentelemetry import baggage, context
 
     # Set in baggage so it propagates via inject()
-    new_ctx = baggage.set_baggage('braintrust.parent', parent, context=ctx)
+    new_ctx = baggage.set_baggage("braintrust.parent", parent, context=ctx)
     token = context.attach(new_ctx)
     return token
 
@@ -511,11 +506,11 @@ def add_span_parent_to_baggage(span, ctx=None):
         raise ImportError(INSTALL_ERR_MSG)
 
     # Get braintrust.parent from span attributes
-    if not span or not hasattr(span, 'attributes') or not span.attributes:
+    if not span or not hasattr(span, "attributes") or not span.attributes:
         logging.warning("add_span_parent_to_baggage: span has no attributes")
         return None
 
-    parent_value = span.attributes.get('braintrust.parent')
+    parent_value = span.attributes.get("braintrust.parent")
     if not parent_value:
         logging.warning(
             "add_span_parent_to_baggage: braintrust.parent attribute not found. "
@@ -527,7 +522,7 @@ def add_span_parent_to_baggage(span, ctx=None):
     return add_parent_to_baggage(parent_value, ctx=ctx)
 
 
-def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
+def parent_from_headers(headers: dict[str, str]) -> str | None:
     """
     Extract a Braintrust-compatible parent string from W3C Trace Context headers.
 
@@ -554,17 +549,16 @@ def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
     if not OTEL_AVAILABLE:
         raise ImportError(INSTALL_ERR_MSG)
 
+    from braintrust.span_identifier_v4 import SpanComponentsV4
     from opentelemetry import baggage, trace
     from opentelemetry.propagate import extract
 
-    from braintrust.span_identifier_v4 import SpanComponentsV4
-
     # Extract context from headers using W3C Trace Context propagator
     ctx = extract(headers)
 
     # Get span from context
     span = trace.get_current_span(ctx)
-    if not span or not hasattr(span, 'get_span_context'):
+    if not span or not hasattr(span, "get_span_context"):
         logging.error("parent_from_headers: No valid span found in headers")
         return None
 
@@ -574,19 +568,19 @@ def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
         return None
 
     # Convert OTEL IDs to hex strings
-    trace_id_hex = format(span_context.trace_id, '032x')
-    span_id_hex = format(span_context.span_id, '016x')
+    trace_id_hex = format(span_context.trace_id, "032x")
+    span_id_hex = format(span_context.span_id, "016x")
 
     # Validate trace_id and span_id are not all zeros
-    if trace_id_hex == '00000000000000000000000000000000':
+    if trace_id_hex == "00000000000000000000000000000000":
         logging.error("parent_from_headers: Invalid trace_id (all zeros)")
         return None
-    if span_id_hex == '0000000000000000':
+    if span_id_hex == "0000000000000000":
         logging.error("parent_from_headers: Invalid span_id (all zeros)")
         return None
 
     # Get braintrust.parent from baggage if present
-    braintrust_parent = baggage.get_baggage('braintrust.parent', context=ctx)
+    braintrust_parent = baggage.get_baggage("braintrust.parent", context=ctx)
 
     # Parse braintrust.parent to extract object_type and object_id
     object_type = None
@@ -607,22 +601,28 @@ def parent_from_headers(headers: Dict[str, str]) -> Optional[str]:
         # Parse braintrust.parent format: "project_id:abc", "project_name:xyz", or "experiment_id:123"
         if braintrust_parent.startswith("project_id:"):
             object_type = SpanObjectTypeV3.PROJECT_LOGS
-            object_id = braintrust_parent[len("project_id:"):]
+            object_id = braintrust_parent[len("project_id:") :]
             if not object_id:
-                logging.error(f"parent_from_headers: Invalid braintrust.parent format (empty project_id): {braintrust_parent}")
+                logging.error(
+                    f"parent_from_headers: Invalid braintrust.parent format (empty project_id): {braintrust_parent}"
+                )
                 return None
         elif braintrust_parent.startswith("project_name:"):
             object_type = SpanObjectTypeV3.PROJECT_LOGS
-            project_name = braintrust_parent[len("project_name:"):]
+            project_name = braintrust_parent[len("project_name:") :]
             if not project_name:
-                logging.error(f"parent_from_headers: Invalid braintrust.parent format (empty project_name): {braintrust_parent}")
+                logging.error(
+                    f"parent_from_headers: Invalid braintrust.parent format (empty project_name): {braintrust_parent}"
+                )
                 return None
             compute_args = {"project_name": project_name}
         elif braintrust_parent.startswith("experiment_id:"):
             object_type = SpanObjectTypeV3.EXPERIMENT
-            object_id = braintrust_parent[len("experiment_id:"):]
+            object_id = braintrust_parent[len("experiment_id:") :]
             if not object_id:
-                logging.error(f"parent_from_headers: Invalid braintrust.parent format (empty experiment_id): {braintrust_parent}")
+                logging.error(
+                    f"parent_from_headers: Invalid braintrust.parent format (empty experiment_id): {braintrust_parent}"
+                )
                 return None
         else:
             logging.error(

braintrust/otel/context.py

@@ -3,23 +3,21 @@
 import logging
 from typing import Any, Optional
 
-from opentelemetry import context, trace
-from opentelemetry.trace import SpanContext, TraceFlags
-
 from braintrust.context import ParentSpanIds, SpanInfo
 from braintrust.logger import Span
+from opentelemetry import context, trace
+from opentelemetry.trace import SpanContext, TraceFlags
 
 log = logging.getLogger(__name__)
 
 
-
 class ContextManager:
     """Context manager that uses OTEL's built-in context as single storage."""
 
     def __init__(self):
         pass
 
-    def get_current_span_info(self) -> Optional['SpanInfo']:
+    def get_current_span_info(self) -> Optional["SpanInfo"]:
         """Get information about the currently active span from OTEL context."""
 
         # Get the current span from OTEL context
@@ -35,25 +33,17 @@ class ContextManager:
         if span_context and span_context.span_id != 0:
             # Always prioritize the actual current OTEL span over stored BT span
             # Only use stored BT span if the current OTEL span IS the BT span wrapper
-            bt_span = context.get_value('braintrust_span')
+            bt_span = context.get_value("braintrust_span")
 
             # If there's a BT span stored AND the current OTEL span is a NonRecordingSpan
             # (which means it's our BT->OTEL wrapper), then return BT span info
-            if (bt_span and isinstance(current_span, trace.NonRecordingSpan)):
-                return SpanInfo(
-                    trace_id=bt_span.root_span_id,
-                    span_id=bt_span.span_id,
-                    span_object=bt_span
-                )
+            if bt_span and isinstance(current_span, trace.NonRecordingSpan):
+                return SpanInfo(trace_id=bt_span.root_span_id, span_id=bt_span.span_id, span_object=bt_span)
             else:
                 # Return OTEL span info - this is a real OTEL span, not our wrapper
-                otel_trace_id = format(span_context.trace_id, '032x')
-                otel_span_id = format(span_context.span_id, '016x')
-                return SpanInfo(
-                    trace_id=otel_trace_id,
-                    span_id=otel_span_id,
-                    span_object=current_span
-                )
+                otel_trace_id = format(span_context.trace_id, "032x")
+                otel_span_id = format(span_context.span_id, "016x")
+                return SpanInfo(trace_id=otel_trace_id, span_id=otel_span_id, span_object=current_span)
 
         return None
 
@@ -61,11 +51,10 @@ class ContextManager:
         """Set the current active span in OTEL context."""
         from opentelemetry import context, trace
 
-        if hasattr(span, 'get_span_context'):
+        if hasattr(span, "get_span_context"):
             # This is an OTEL span - it will manage its own context
             return None
         else:
-
             try:
                 trace_id_int = int(span.root_span_id, 16)
             except ValueError:
@@ -80,15 +69,12 @@ class ContextManager:
 
             # This is a BT span - store it in OTEL context AND set as current OTEL span
             # First store the BT span
-            ctx = context.set_value('braintrust_span', span)
+            ctx = context.set_value("braintrust_span", span)
             parent_value = span._get_otel_parent()
-            ctx = context.set_value('braintrust.parent', parent_value, ctx)
+            ctx = context.set_value("braintrust.parent", parent_value, ctx)
 
             otel_span_context = SpanContext(
-                trace_id=trace_id_int,
-                span_id=span_id_int,
-                is_remote=False,
-                trace_flags=TraceFlags(TraceFlags.SAMPLED)
+                trace_id=trace_id_int, span_id=span_id_int, is_remote=False, trace_flags=TraceFlags(TraceFlags.SAMPLED)
             )
 
             # Create a non-recording span to represent the BT span in OTEL context
@@ -110,9 +96,9 @@ class ContextManager:
         else:
             # No token means we need to explicitly clear the span
             # This shouldn't normally happen, but handle it gracefully
-            context.attach(context.set_value('braintrust_span', None))
+            context.attach(context.set_value("braintrust_span", None))
 
-    def get_parent_span_ids(self) -> Optional[ParentSpanIds]:
+    def get_parent_span_ids(self) -> ParentSpanIds | None:
         """Get parent information for creating a new BT span."""
         span_info = self.get_current_span_info()
         if not span_info:
@@ -125,4 +111,4 @@ class ContextManager:
 
 def _is_otel_span(span: Any) -> bool:
     """Check if the span object is an OTEL span."""
-    return hasattr(span, 'get_span_context')
+    return hasattr(span, "get_span_context")

braintrust/otel/test_distributed_tracing.py

@@ -8,7 +8,6 @@ is exported from one service and imported in another service.
 import os
 
 import pytest
-
 from braintrust.logger import _internal_with_memory_background_logger
 from braintrust.otel import BraintrustSpanProcessor, context_from_span_export
 from braintrust.test_helpers import init_test_logger, preserve_env_vars
@@ -19,6 +18,7 @@ try:
     from opentelemetry.sdk.trace.export import SimpleSpanProcessor
     from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 except ImportError:
+
     class InMemorySpanExporter:
         def __init__(self):
             pass
@@ -47,10 +47,10 @@ def otel_fixture():
     if not OTEL_AVAILABLE:
         pytest.skip("OpenTelemetry not installed")
 
-    with preserve_env_vars('BRAINTRUST_OTEL_COMPAT', 'BRAINTRUST_API_KEY'):
+    with preserve_env_vars("BRAINTRUST_OTEL_COMPAT", "BRAINTRUST_API_KEY"):
         # Enable OTEL compatibility mode
-        os.environ['BRAINTRUST_OTEL_COMPAT'] = 'true'
-        os.environ['BRAINTRUST_API_KEY'] = 'test-api-key-for-fixture'
+        os.environ["BRAINTRUST_OTEL_COMPAT"] = "true"
+        os.environ["BRAINTRUST_API_KEY"] = "test-api-key-for-fixture"
 
         # Set up memory logger for BT spans
         with _internal_with_memory_background_logger() as memory_logger:
@@ -103,6 +103,7 @@ def test_bt_to_otel_simple_distributed_trace(otel_fixture):
     # ===== Service B: Import context and create OTEL child span =====
     # Simulate receiving exported_context over network (e.g., in HTTP header)
     from opentelemetry import context as otel_context
+
     ctx = context_from_span_export(exported_context)
 
     # Attach the context to make it current, then create the span
@@ -125,22 +126,24 @@ def test_bt_to_otel_simple_distributed_trace(otel_fixture):
     service_b_exported = otel_spans[0]
 
     # Convert OTEL IDs to hex for comparison
-    service_b_trace_id = format(service_b_exported.context.trace_id, '032x')
-    service_b_parent_span_id = format(service_b_exported.parent.span_id, '016x') if service_b_exported.parent else None
+    service_b_trace_id = format(service_b_exported.context.trace_id, "032x")
+    service_b_parent_span_id = format(service_b_exported.parent.span_id, "016x") if service_b_exported.parent else None
 
     # Assert unified trace ID
-    assert service_a_trace_id == service_b_trace_id, \
+    assert service_a_trace_id == service_b_trace_id, (
         f"Trace IDs should match: {service_a_trace_id} != {service_b_trace_id}"
+    )
 
     # Assert Service B span has Service A span as parent
-    assert service_b_parent_span_id == service_a_span_id, \
+    assert service_b_parent_span_id == service_a_span_id, (
         f"Service B parent should be Service A span: {service_b_parent_span_id} != {service_a_span_id}"
+    )
 
     # Assert braintrust.parent attribute is set on OTEL span
-    assert "braintrust.parent" in service_b_exported.attributes, \
-        "OTEL span should have braintrust.parent attribute"
-    assert service_b_exported.attributes["braintrust.parent"] == f"project_name:{project_name}", \
+    assert "braintrust.parent" in service_b_exported.attributes, "OTEL span should have braintrust.parent attribute"
+    assert service_b_exported.attributes["braintrust.parent"] == f"project_name:{project_name}", (
         f"braintrust.parent should be 'project_name:{project_name}'"
+    )
 
 
 if __name__ == "__main__":