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

braintrust/audit.py

@@ -5,7 +5,7 @@ Utilities for working with audit headers.
 import base64
 import gzip
 import json
-from typing import List, TypedDict
+from typing import TypedDict
 
 
 class AuditResource(TypedDict):
@@ -14,7 +14,7 @@ class AuditResource(TypedDict):
     name: str
 
 
-def parse_audit_resources(hdr: str) -> List[AuditResource]:
+def parse_audit_resources(hdr: str) -> list[AuditResource]:
     j = json.loads(hdr)
     if j["v"] == 1:
         return json.loads(gzip.decompress(base64.b64decode(j["p"])))

braintrust/bt_json.py

@@ -1,6 +1,7 @@
 import dataclasses
 import json
-from typing import Any, cast
+import math
+from typing import Any, Callable, Mapping, NamedTuple, cast, overload
 
 # Try to import orjson for better performance
 # If not available, we'll use standard json
@@ -12,39 +13,184 @@ except ImportError:
     _HAS_ORJSON = False
 
 
-def _to_dict(obj: Any) -> Any:
-    """
-    Function-based default handler for non-JSON-serializable objects.
 
-    Handles:
-    - dataclasses
-    - Pydantic v2 BaseModel
-    - Pydantic v1 BaseModel
-    - Falls back to str() for unknown types
+def _to_bt_safe(v: Any) -> Any:
+    """
+    Converts the object to a Braintrust-safe representation (i.e. Attachment objects are safe (specially handled by background logger)).
     """
-    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
-        return dataclasses.asdict(obj)
+    # avoid circular imports
+    from braintrust.logger import BaseAttachment, Dataset, Experiment, Logger, ReadonlyAttachment, Span
+
+    if isinstance(v, Span):
+        return "<span>"
+
+    if isinstance(v, Experiment):
+        return "<experiment>"
+
+    if isinstance(v, Dataset):
+        return "<dataset>"
+
+    if isinstance(v, Logger):
+        return "<logger>"
+
+    if isinstance(v, BaseAttachment):
+        return v
+
+    if isinstance(v, ReadonlyAttachment):
+        return v.reference
+
+    if dataclasses.is_dataclass(v) and not isinstance(v, type):
+        # Use manual field iteration instead of dataclasses.asdict() because
+        # asdict() deep-copies values, which breaks objects like Attachment
+        # that contain non-copyable items (thread locks, file handles, etc.)
+        return {f.name: _to_bt_safe(getattr(v, f.name)) for f in dataclasses.fields(v)}
+
+    # Pydantic model classes (not instances) with model_json_schema
+    if isinstance(v, type) and hasattr(v, "model_json_schema") and callable(cast(Any, v).model_json_schema):
+        try:
+            return cast(Any, v).model_json_schema()
+        except Exception:
+            pass
 
     # Attempt to dump a Pydantic v2 `BaseModel`.
     try:
-        return cast(Any, obj).model_dump()
+        return cast(Any, v).model_dump(exclude_none=True)
     except (AttributeError, TypeError):
         pass
 
     # Attempt to dump a Pydantic v1 `BaseModel`.
     try:
-        return cast(Any, obj).dict()
+        return cast(Any, v).dict(exclude_none=True)
     except (AttributeError, TypeError):
         pass
 
-    # When everything fails, try to return the string representation of the object
+    if isinstance(v, float):
+        # Handle NaN and Infinity for JSON compatibility
+        if math.isnan(v):
+            return "NaN"
+
+        if math.isinf(v):
+            return "Infinity" if v > 0 else "-Infinity"
+
+        return v
+
+    if isinstance(v, (int, str, bool)) or v is None:
+        # Skip roundtrip for primitive types.
+        return v
+
+    # Note: we avoid using copy.deepcopy, because it's difficult to
+    # guarantee the independence of such copied types from their origin.
+    # E.g. the original type could have a `__del__` method that alters
+    # some shared internal state, and we need this deep copy to be
+    # fully-independent from the original.
+
+    # We pass `encoder=_str_encoder` since we've already tried converting rich objects to json safe objects.
+    return bt_loads(bt_dumps(v, encoder=_str_encoder))
+
+@overload
+def bt_safe_deep_copy(
+    obj: Mapping[str, Any],
+    max_depth: int = ...,
+) -> dict[str, Any]: ...
+
+@overload
+def bt_safe_deep_copy(
+    obj: list[Any],
+    max_depth: int = ...,
+) -> list[Any]: ...
+
+@overload
+def bt_safe_deep_copy(
+    obj: Any,
+    max_depth: int = ...,
+) -> Any: ...
+def bt_safe_deep_copy(obj: Any, max_depth: int=200):
+    """
+    Creates a deep copy of the given object and converts rich objects to Braintrust-safe representations. See `_to_bt_safe` for more details.
+
+    Args:
+        obj: Object to deep copy and sanitize.
+        to_json_safe: Function to ensure the object is json safe.
+        max_depth: Maximum depth to copy.
+
+    Returns:
+        Deep copy of the object.
+    """
+    # Track visited objects to detect circular references
+    visited: set[int] = set()
+
+    def _deep_copy_object(v: Any, depth: int = 0) -> Any:
+        # Check depth limit - use >= to stop before exceeding
+        if depth >= max_depth:
+            return "<max depth exceeded>"
+
+        # Check for circular references in mutable containers
+        # Use id() to track object identity
+        if isinstance(v, (Mapping, list, tuple, set)):
+            obj_id = id(v)
+            if obj_id in visited:
+                return "<circular reference>"
+            visited.add(obj_id)
+            try:
+                if isinstance(v, Mapping):
+                    # Prevent dict keys from holding references to user data. Note that
+                    # `bt_json` already coerces keys to string, a behavior that comes from
+                    # `json.dumps`. However, that runs at log upload time, while we want to
+                    # cut out all the references to user objects synchronously in this
+                    # function.
+                    result = {}
+                    for k in v:
+                        try:
+                            key_str = str(k)
+                        except Exception:
+                            # If str() fails on the key, use a fallback representation
+                            key_str = f"<non-stringifiable-key: {type(k).__name__}>"
+                        result[key_str] = _deep_copy_object(v[k], depth + 1)
+                    return result
+                elif isinstance(v, (list, tuple, set)):
+                    return [_deep_copy_object(x, depth + 1) for x in v]
+            finally:
+                # Remove from visited set after processing to allow the same object
+                # to appear in different branches of the tree
+                visited.discard(obj_id)
+
+        try:
+            return _to_bt_safe(v)
+        except Exception:
+            return f"<non-sanitizable: {type(v).__name__}>"
+
+    return _deep_copy_object(obj)
+
+def _safe_str(obj: Any) -> str:
     try:
         return str(obj)
     except Exception:
-        # If str() fails, return an error placeholder
         return f"<non-serializable: {type(obj).__name__}>"
 
 
+def _to_json_safe(obj: Any) -> Any:
+    """
+    Handler for non-JSON-serializable objects. Returns a string representation of the object.
+    """
+    # avoid circular imports
+    from braintrust.logger import BaseAttachment
+
+    try:
+        v = _to_bt_safe(obj)
+
+        # JSON-safe representation of Attachment objects are their reference.
+        # If we get this object at this point, we have to assume someone has already uploaded the attachment!
+        if isinstance(v, BaseAttachment):
+            v = v.reference
+
+        return v
+    except Exception:
+        pass
+
+    # When everything fails, try to return the string representation of the object
+    return _safe_str(obj)
+
+
 class BraintrustJSONEncoder(json.JSONEncoder):
     """
     Custom JSON encoder for standard json library.
@@ -53,10 +199,22 @@ class BraintrustJSONEncoder(json.JSONEncoder):
     """
 
     def default(self, o: Any):
-        return _to_dict(o)
+        return _to_json_safe(o)
+
+
+class BraintrustStrEncoder(json.JSONEncoder):
+    def default(self, o: Any):
+        return _safe_str(o)
+
+
+class Encoder(NamedTuple):
+    native: type[json.JSONEncoder]
+    orjson: Callable[[Any], Any]
 
+_json_encoder = Encoder(native=BraintrustJSONEncoder, orjson=_to_json_safe)
+_str_encoder = Encoder(native=BraintrustStrEncoder, orjson=_safe_str)
 
-def bt_dumps(obj, **kwargs) -> str:
+def bt_dumps(obj: Any, encoder: Encoder | None=_json_encoder, **kwargs: Any) -> str:
     """
     Serialize obj to a JSON-formatted string.
 
@@ -65,6 +223,7 @@ def bt_dumps(obj, **kwargs) -> str:
 
     Args:
         obj: Object to serialize
+        encoder: Encoder to use, defaults to `_default_encoder`
         **kwargs: Additional arguments (passed to json.dumps in fallback path)
 
     Returns:
@@ -76,7 +235,7 @@ def bt_dumps(obj, **kwargs) -> str:
             # pylint: disable=no-member  # orjson is a C extension, pylint can't introspect it
             return orjson.dumps(  # type: ignore[possibly-unbound]
                 obj,
-                default=_to_dict,
+                default=encoder.orjson if encoder else None,
                 # options match json.dumps behavior for bc
                 option=orjson.OPT_SORT_KEYS | orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS,  # type: ignore[possibly-unbound]
             ).decode("utf-8")
@@ -86,7 +245,7 @@ def bt_dumps(obj, **kwargs) -> str:
 
     # Use standard json (either orjson not available or it failed)
     # Use sort_keys=True for deterministic output (matches orjson OPT_SORT_KEYS)
-    return json.dumps(obj, cls=BraintrustJSONEncoder, allow_nan=False, sort_keys=True, **kwargs)
+    return json.dumps(obj, cls=encoder.native if encoder else None, allow_nan=False, sort_keys=True, **kwargs)
 
 
 def bt_loads(s: str, **kwargs) -> Any:

braintrust/cli/eval.py

@@ -6,7 +6,6 @@ import os
 import sys
 from dataclasses import dataclass, field
 from threading import Lock
-from typing import Dict, List, Optional, Union
 
 from .. import login
 from ..framework import (
@@ -70,7 +69,7 @@ class EvaluatorOpts:
     no_progress_bars: bool
     terminate_on_failure: bool
     watch: bool
-    filters: List[str]
+    filters: list[str]
     list: bool
     jsonl: bool
 
@@ -79,13 +78,13 @@ class EvaluatorOpts:
 class LoadedEvaluator:
     handle: FileHandle
     evaluator: Evaluator
-    reporter: Optional[Union[ReporterDef, str]] = None
+    reporter: ReporterDef | str | None = None
 
 
 @dataclass
 class EvaluatorState:
-    evaluators: List[LoadedEvaluator] = field(default_factory=list)
-    reporters: Dict[str, ReporterDef] = field(default_factory=dict)
+    evaluators: list[LoadedEvaluator] = field(default_factory=list)
+    reporters: dict[str, ReporterDef] = field(default_factory=dict)
 
 
 def update_evaluators(eval_state: EvaluatorState, handles, terminate_on_failure):
@@ -157,7 +156,7 @@ async def run_evaluator_task(evaluator, position, opts: EvaluatorOpts):
             experiment.flush()
 
 
-def resolve_reporter(reporter: Optional[Union[ReporterDef, str]], reporters: Dict[str, ReporterDef]) -> ReporterDef:
+def resolve_reporter(reporter: ReporterDef | str | None, reporters: dict[str, ReporterDef]) -> ReporterDef:
     if isinstance(reporter, str):
         if reporter not in reporters:
             raise ValueError(f"Reporter {reporter} not found")
@@ -179,7 +178,7 @@ def add_report(eval_reports, reporter, report):
     eval_reports[reporter.name]["results"].append(report)
 
 
-async def run_once(handles: List[FileHandle], evaluator_opts: EvaluatorOpts) -> bool:
+async def run_once(handles: list[FileHandle], evaluator_opts: EvaluatorOpts) -> bool:
     objects = EvaluatorState()
     update_evaluators(objects, handles, terminate_on_failure=evaluator_opts.terminate_on_failure)
 

braintrust/cli/push.py

@@ -13,7 +13,7 @@ import sys
 import tempfile
 import textwrap
 import zipfile
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import requests
 from braintrust.framework import _set_lazy_load
@@ -24,7 +24,7 @@ from ..generated_types import IfExists
 from ..util import add_azure_blob_headers
 
 
-def _pkg_install_arg(pkg) -> Optional[str]:
+def _pkg_install_arg(pkg) -> str | None:
     try:
         dist = importlib.metadata.distribution(pkg)
         direct_url = dist._path / "direct_url.json"  # type: ignore
@@ -73,11 +73,11 @@ class _ProjectRootImporter(importlib.abc.MetaPathFinder):
         self._project_root, self._path_rest = sys.path[0], sys.path[1:]
         self._sources = []
 
-    def _under_project_root(self, path: List[str]) -> bool:
+    def _under_project_root(self, path: list[str]) -> bool:
         """Returns true if all paths in `path` are under the project root."""
         return all(p.startswith(self._project_root) for p in path)
 
-    def _under_rest(self, path: List[str]) -> bool:
+    def _under_rest(self, path: list[str]) -> bool:
         """Returns true if any path in `path` is under one of the remaining paths in `sys.path`."""
         return any(p.startswith(pr) for p in path for pr in self._path_rest)
 
@@ -94,11 +94,11 @@ class _ProjectRootImporter(importlib.abc.MetaPathFinder):
             self._sources.append(spec.origin)
         return spec
 
-    def sources(self) -> List[str]:
+    def sources(self) -> list[str]:
         return self._sources
 
 
-def _import_module(name: str, path: str) -> List[str]:
+def _import_module(name: str, path: str) -> list[str]:
     """Imports the module and returns the list of source files
     of all modules imported in the process.
 
@@ -120,7 +120,7 @@ def _py_version() -> str:
     return f"{sys.version_info.major}.{sys.version_info.minor}"
 
 
-def _run_install(install_args: List[str], packages_dir: str):
+def _run_install(install_args: list[str], packages_dir: str):
     subprocess.run(
         [
             "uv",
@@ -138,7 +138,7 @@ def _run_install(install_args: List[str], packages_dir: str):
     )
 
 
-def _upload_bundle(entry_module_name: str, sources: List[str], requirements: Optional[str]) -> str:
+def _upload_bundle(entry_module_name: str, sources: list[str], requirements: str | None) -> str:
     _check_uv()
 
     resp = proxy_conn().post_json(
@@ -212,7 +212,7 @@ def _upload_bundle(entry_module_name: str, sources: List[str], requirements: Opt
 
 
 def _collect_function_function_defs(
-    project_ids: ProjectIdCache, functions: List[Dict[str, Any]], bundle_id: str, if_exists: IfExists
+    project_ids: ProjectIdCache, functions: list[dict[str, Any]], bundle_id: str, if_exists: IfExists
 ) -> None:
     for i, f in enumerate(global_.functions):
         source = inspect.getsource(f.handler)
@@ -262,7 +262,7 @@ def _collect_function_function_defs(
 
 
 def _collect_prompt_function_defs(
-    project_ids: ProjectIdCache, functions: List[Dict[str, Any]], if_exists: IfExists
+    project_ids: ProjectIdCache, functions: list[dict[str, Any]], if_exists: IfExists
 ) -> None:
     for p in global_.prompts:
         functions.append(p.to_function_definition(if_exists, project_ids))
@@ -300,7 +300,7 @@ def run(args):
         raise
 
     project_ids = ProjectIdCache()
-    functions: List[Dict[str, Any]] = []
+    functions: list[dict[str, Any]] = []
     if len(global_.functions) > 0:
         bundle_id = _upload_bundle(module_name, sources, args.requirements)
         _collect_function_function_defs(project_ids, functions, bundle_id, args.if_exists)

braintrust/context.py

@@ -5,12 +5,13 @@ import os
 from abc import ABC, abstractmethod
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Any, List, Optional
+from typing import Any
 
 
 @dataclass
 class SpanInfo:
     """Information about a span in the context."""
+
     trace_id: str
     span_id: str
     span_object: Any = None
@@ -19,7 +20,7 @@ class SpanInfo:
 @dataclass
 class ParentSpanIds:
     root_span_id: str
-    span_parents: List[str]
+    span_parents: list[str]
 
 
 class ContextManager(ABC):
@@ -30,7 +31,7 @@ class ContextManager(ABC):
     """
 
     @abstractmethod
-    def get_current_span_info(self) -> Optional[Any]:
+    def get_current_span_info(self) -> Any | None:
         """Get information about the currently active span.
 
         Returns:
@@ -40,7 +41,7 @@ class ContextManager(ABC):
         pass
 
     @abstractmethod
-    def get_parent_span_ids(self) -> Optional[ParentSpanIds]:
+    def get_parent_span_ids(self) -> ParentSpanIds | None:
         """Get parent span IDs for creating a new Braintrust span.
 
         Returns:
@@ -75,32 +76,25 @@ class BraintrustContextManager(ContextManager):
     """Braintrust-only context manager using contextvars when OTEL is not available."""
 
     def __init__(self):
-        self._current_span: ContextVar[Optional[Any]] = ContextVar('braintrust_current_span', default=None)
+        self._current_span: ContextVar[Any | None] = ContextVar("braintrust_current_span", default=None)
 
-    def get_current_span_info(self) -> Optional[SpanInfo]:
+    def get_current_span_info(self) -> SpanInfo | None:
         """Get information about the currently active span."""
         current_span = self._current_span.get()
         if not current_span:
             return None
 
         # Return SpanInfo for BT spans
-        return SpanInfo(
-            trace_id=current_span.root_span_id,
-            span_id=current_span.span_id,
-            span_object=current_span
-        )
+        return SpanInfo(trace_id=current_span.root_span_id, span_id=current_span.span_id, span_object=current_span)
 
-    def get_parent_span_ids(self) -> Optional[ParentSpanIds]:
+    def get_parent_span_ids(self) -> ParentSpanIds | None:
         """Get parent information for creating a new Braintrust span."""
         current_span = self._current_span.get()
         if not current_span:
             return None
 
         # If current span is a BT span, use it as parent
-        return ParentSpanIds(
-            root_span_id=current_span.root_span_id,
-            span_parents=[current_span.span_id]
-        )
+        return ParentSpanIds(root_span_id=current_span.root_span_id, span_parents=[current_span.span_id])
 
     def set_current_span(self, span_object: Any) -> Any:
         """Set the current active span."""
@@ -123,9 +117,10 @@ def get_context_manager() -> ContextManager:
     """
 
     # Check if OTEL should be explicitly enabled via environment variable
-    if os.environ.get('BRAINTRUST_OTEL_COMPAT', '').lower() in ('1', 'true', 'yes'):
+    if os.environ.get("BRAINTRUST_OTEL_COMPAT", "").lower() in ("1", "true", "yes"):
         try:
             from braintrust.otel.context import ContextManager as OtelContextManager
+
             return OtelContextManager()
         except ImportError:
             logging.warning("OTEL not available, falling back to Braintrust-only version")

braintrust/contrib/temporal/__init__.py

@@ -80,7 +80,8 @@ The integration will automatically:
 """
 
 import dataclasses
-from typing import Any, Dict, Mapping, Optional, Type
+from collections.abc import Mapping
+from typing import Any
 
 import braintrust
 import temporalio.activity
@@ -129,7 +130,7 @@ class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Int
     - Ensures replay safety (no duplicate spans during workflow replay)
     """
 
-    def __init__(self, logger: Optional[Any] = None) -> None:
+    def __init__(self, logger: Any | None = None) -> None:
         """Initialize interceptor.
 
         Args:
@@ -142,7 +143,7 @@ class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Int
             braintrust.logger._state._override_bg_logger.logger = logger
         self._logger = braintrust.current_logger()
 
-    def _get_logger(self) -> Optional[Any]:
+    def _get_logger(self) -> Any | None:
         """Get logger for creating spans.
 
         Sets thread-local override if background logger provided (for testing),
@@ -152,9 +153,7 @@ class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Int
             braintrust.logger._state._override_bg_logger.logger = self._bg_logger
         return self._logger
 
-    def intercept_client(
-        self, next: temporalio.client.OutboundInterceptor
-    ) -> temporalio.client.OutboundInterceptor:
+    def intercept_client(self, next: temporalio.client.OutboundInterceptor) -> temporalio.client.OutboundInterceptor:
         """Intercept client calls to propagate span context to workflows."""
         return _BraintrustClientOutboundInterceptor(next, self)
 
@@ -166,14 +165,14 @@ class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Int
 
     def workflow_interceptor_class(
         self, input: temporalio.worker.WorkflowInterceptorClassInput
-    ) -> Optional[Type["BraintrustWorkflowInboundInterceptor"]]:
+    ) -> type["BraintrustWorkflowInboundInterceptor"] | None:
         """Return workflow interceptor class to propagate context to activities."""
         input.unsafe_extern_functions["__braintrust_get_logger"] = self._get_logger
         return BraintrustWorkflowInboundInterceptor
 
     def _span_context_to_headers(
         self,
-        span_context: Dict[str, Any],
+        span_context: dict[str, Any],
         headers: Mapping[str, temporalio.api.common.v1.Payload],
     ) -> Mapping[str, temporalio.api.common.v1.Payload]:
         """Add span context to headers."""
@@ -188,7 +187,7 @@ class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Int
 
     def _span_context_from_headers(
         self, headers: Mapping[str, temporalio.api.common.v1.Payload]
-    ) -> Optional[Dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Extract span context from headers."""
         if _HEADER_KEY not in headers:
             return None
@@ -204,9 +203,7 @@ class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Int
 class _BraintrustClientOutboundInterceptor(temporalio.client.OutboundInterceptor):
     """Client interceptor that propagates span context to workflows."""
 
-    def __init__(
-        self, next: temporalio.client.OutboundInterceptor, root: BraintrustInterceptor
-    ) -> None:
+    def __init__(self, next: temporalio.client.OutboundInterceptor, root: BraintrustInterceptor) -> None:
         super().__init__(next)
         self.root = root
 
@@ -233,9 +230,7 @@ class _BraintrustActivityInboundInterceptor(temporalio.worker.ActivityInboundInt
         super().__init__(next)
         self.root = root
 
-    async def execute_activity(
-        self, input: temporalio.worker.ExecuteActivityInput
-    ) -> Any:
+    async def execute_activity(self, input: temporalio.worker.ExecuteActivityInput) -> Any:
         info = temporalio.activity.info()
 
         # Extract parent span context from headers
@@ -281,14 +276,12 @@ class BraintrustWorkflowInboundInterceptor(temporalio.worker.WorkflowInboundInte
     def __init__(self, next: temporalio.worker.WorkflowInboundInterceptor) -> None:
         super().__init__(next)
         self.payload_converter = temporalio.converter.PayloadConverter.default
-        self._parent_span_context: Optional[Dict[str, Any]] = None
+        self._parent_span_context: dict[str, Any] | None = None
 
     def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None:
         super().init(_BraintrustWorkflowOutboundInterceptor(outbound, self))
 
-    async def execute_workflow(
-        self, input: temporalio.worker.ExecuteWorkflowInput
-    ) -> Any:
+    async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) -> Any:
         # Extract parent span context from workflow headers (set by client)
         parent_span_context = None
         if _HEADER_KEY in input.headers:
@@ -342,9 +335,7 @@ class BraintrustWorkflowInboundInterceptor(temporalio.worker.WorkflowInboundInte
                     span.end()
 
 
-class _BraintrustWorkflowOutboundInterceptor(
-    temporalio.worker.WorkflowOutboundInterceptor
-):
+class _BraintrustWorkflowOutboundInterceptor(temporalio.worker.WorkflowOutboundInterceptor):
     """Outbound workflow interceptor that propagates span context to activities."""
 
     def __init__(
@@ -371,9 +362,7 @@ class _BraintrustWorkflowOutboundInterceptor(
                 return {**headers, _HEADER_KEY: payloads[0]}
         return headers
 
-    def start_activity(
-        self, input: temporalio.worker.StartActivityInput
-    ) -> temporalio.workflow.ActivityHandle:
+    def start_activity(self, input: temporalio.worker.StartActivityInput) -> temporalio.workflow.ActivityHandle:
         input.headers = self._add_span_context_to_headers(input.headers)
         return super().start_activity(input)
 
@@ -390,7 +379,7 @@ class _BraintrustWorkflowOutboundInterceptor(
         return super().start_child_workflow(input)
 
 
-def _modify_workflow_runner(existing: Optional[WorkflowRunner]) -> Optional[WorkflowRunner]:
+def _modify_workflow_runner(existing: WorkflowRunner | None) -> WorkflowRunner | None:
     """Add braintrust to sandbox passthrough modules."""
     if isinstance(existing, SandboxedWorkflowRunner):
         new_restrictions = existing.restrictions.with_passthrough_modules("braintrust")
@@ -420,7 +409,7 @@ class BraintrustPlugin(SimplePlugin):
     Requires temporalio >= 1.19.0.
     """
 
-    def __init__(self, logger: Optional[Any] = None) -> None:
+    def __init__(self, logger: Any | None = None) -> None:
         """Initialize the Braintrust plugin.
 
         Args:

braintrust/contrib/temporal/test_temporal.py

@@ -279,11 +279,15 @@ class TestBraintrustPluginIntegration:
 
             # Verify workflow span was created
             workflow_spans = [s for s in spans if "temporal.workflow" in s.get("span_attributes", {}).get("name", "")]
-            assert len(workflow_spans) > 0, f"Expected workflow span to be created. Span names: {[s.get('span_attributes', {}).get('name', 'unknown') for s in spans]}"
+            assert len(workflow_spans) > 0, (
+                f"Expected workflow span to be created. Span names: {[s.get('span_attributes', {}).get('name', 'unknown') for s in spans]}"
+            )
 
             # Verify activity span was created
             activity_spans = [s for s in spans if "temporal.activity" in s.get("span_attributes", {}).get("name", "")]
-            assert len(activity_spans) > 0, f"Expected activity span to be created. Span names: {[s.get('span_attributes', {}).get('name', 'unknown') for s in spans]}"
+            assert len(activity_spans) > 0, (
+                f"Expected activity span to be created. Span names: {[s.get('span_attributes', {}).get('name', 'unknown') for s in spans]}"
+            )
 
     @pytest.mark.asyncio
     async def test_plugin_context_propagation(self, temporal_env, memory_logger):
@@ -498,5 +502,6 @@ class TestBraintrustPluginIntegration:
         # Verify parent-child relationship: workflow should have client span as parent
         workflow_span = workflow_spans[0]
         client_span = client_spans[0]
-        assert workflow_span.get("root_span_id") == client_span.get("root_span_id"), \
+        assert workflow_span.get("root_span_id") == client_span.get("root_span_id"), (
             "Workflow span should be in same trace as client span"
+        )