feat: add span customizers (SDK-316)

add customizers and re-implement masking fn with a customizer

Author: realark

py/src/braintrust/__init__.py

@@ -86,5 +86,8 @@ def is_equal(expected, output):
 from .sandbox import RegisterSandboxResult as RegisterSandboxResult
 from .sandbox import SandboxConfig as SandboxConfig
 from .sandbox import register_sandbox as register_sandbox
+from .span_customizer import SpanCustomizer as SpanCustomizer
+from .span_customizer import SpanExportData as SpanExportData
+from .span_customizer import set_span_customizers as set_span_customizers
 from .util import BT_IS_ASYNC_ATTRIBUTE as BT_IS_ASYNC_ATTRIBUTE
 from .util import MarkAsyncWrapper as MarkAsyncWrapper

py/src/braintrust/auto.py

@@ -5,6 +5,7 @@
 """
 
 import logging
+from collections.abc import Sequence
 from contextlib import contextmanager
 
 from braintrust.integrations import (
@@ -40,6 +41,7 @@
     TypeSafeIntegration,
 )
 from braintrust.integrations.base import BaseIntegration
+from braintrust.span_customizer import SpanCustomizer, set_span_customizers
 
 
 __all__ = ["auto_instrument"]
@@ -90,6 +92,7 @@ def auto_instrument(
     livekit_agents: bool = True,
     pipecat: bool = True,
     typesafe: bool = True,
+    span_customizers: Sequence[SpanCustomizer] | None = None,
 ) -> dict[str, bool]:
     """
     Auto-instrument supported AI/ML libraries for Braintrust tracing.
@@ -131,6 +134,9 @@ def auto_instrument(
         livekit_agents: Enable LiveKit Agents instrumentation (default: True)
         pipecat: Enable Pipecat AI instrumentation (default: True)
         typesafe: Enable TypeSafe instrumentation (default: True)
+        span_customizers: Ordered synchronous export customizers for instrumentation
+            spans. Copies and replaces the global list when provided; None leaves
+            existing configuration unchanged. Pass [] to disable.
 
     Returns:
         Dict mapping integration name to whether it was successfully instrumented.
@@ -176,6 +182,9 @@ def auto_instrument(
         client.models.generate_content(model="gemini-2.0-flash", contents="Hello!")
         ```
     """
+    if span_customizers is not None:
+        set_span_customizers(span_customizers)
+
     results: dict[str, bool] = {}
 
     if openai:

py/src/braintrust/integrations/auto_test_scripts/test_auto_google_discoveryengine.py

@@ -18,7 +18,11 @@
     print("SUCCESS")
     sys.exit(0)
 
-options = {name: False for name in inspect.signature(auto_instrument).parameters}
+options = {
+    name: False
+    for name, parameter in inspect.signature(auto_instrument).parameters.items()
+    if isinstance(parameter.default, bool)
+}
 assert auto_instrument(**options) == {}
 RankServiceClient = None
 if sys.argv[1] == "before":

py/src/braintrust/integrations/pipecat/test_pipecat.py

@@ -8,7 +8,7 @@
 from pathlib import Path
 
 import pytest
-from braintrust import logger
+from braintrust import SpanCustomizer, logger, set_span_customizers
 from braintrust.integrations.pipecat import (
     BraintrustPipecatObserver,
     PipecatIntegration,
@@ -58,6 +58,45 @@ def _single_span(logs, name):
     return matches[0]
 
 
+@pytest.mark.asyncio
+async def test_span_customizer_redacts_incremental_tts_input(memory_logger):
+    TTSStartedFrame = _import("pipecat.frames.frames.TTSStartedFrame")
+    TTSTextFrame = _import("pipecat.frames.frames.TTSTextFrame")
+    TTSStoppedFrame = _import("pipecat.frames.frames.TTSStoppedFrame")
+
+    class Redact(SpanCustomizer):
+        def on_span_export(self, data):
+            if "input" in data:
+                data["input"] = "[redacted]"
+            return data
+
+    observer = BraintrustPipecatObserver()
+    set_span_customizers([Redact()])
+    try:
+        await observer.on_pipeline_started()
+        await observer._handle_frame(TTSStartedFrame(context_id="ctx"))
+        first_frame = TTSTextFrame("private first", aggregated_by="sentence", context_id="ctx")
+        await observer._handle_frame(first_frame)
+        initial = memory_logger.pop()
+        pipeline = _single_span(initial, "pipecat_pipeline")
+        tts = _single_span(initial, "tts_response")
+        assert tts["input"] == "[redacted]"
+        assert tts["span_parents"] == [pipeline["span_id"]]
+        assert first_frame.text == "private first"
+
+        second_frame = TTSTextFrame("private second", aggregated_by="sentence", context_id="ctx")
+        await observer._handle_frame(second_frame)
+        await observer._handle_frame(TTSStoppedFrame(context_id="ctx"))
+        await observer.cleanup()
+        updates = memory_logger.pop()
+        updated_tts = next(row for row in updates if row["id"] == tts["id"])
+        assert updated_tts["input"] == "[redacted]"
+        assert updated_tts["span_id"] == tts["span_id"]
+        assert second_frame.text == "private second"
+    finally:
+        set_span_customizers(None)
+
+
 def _pipeline_worker_kwargs(**overrides):
     PipelineWorker = _import("pipecat.pipeline.worker.PipelineWorker")
     signature = inspect.signature(PipelineWorker)

py/src/braintrust/integrations/pipecat/tracing.py

@@ -9,7 +9,18 @@
     _pcm_to_wav,
     _resolve_audio_attachment_options,
 )
-from braintrust.logger import NOOP_SPAN, Attachment, SpanTypeAttribute, current_span, start_span
+from braintrust.logger import NOOP_SPAN, Attachment, SpanTypeAttribute, current_span
+from braintrust.logger import start_span as _bt_start_span
+
+
+_INSTRUMENTATION = "pipecat-auto"
+
+
+def start_span(*args, **kwargs):
+    internal = dict(kwargs.get("internal") or {})
+    internal.setdefault("instrumentation", _INSTRUMENTATION)
+    kwargs["internal"] = internal
+    return _bt_start_span(*args, **kwargs)
 
 
 try:

py/src/braintrust/logger.py

@@ -101,6 +101,7 @@
 )
 from .queue import DEFAULT_QUEUE_SIZE, LogQueue
 from .serializable_data_class import SerializableDataClass
+from .span_customizer import SpanCustomizer, _customize_span_export, _get_span_customizers, _MaskingCustomizer
 from .span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3
 from .span_identifier_v4 import SpanComponentsV4
 from .span_origin import SpanOriginEnvironment, detect_environment, merge_span_origin_context
@@ -124,10 +125,6 @@
 from .xact_ids import prettify_xact
 
 
-# Fields that should be passed to the masking function
-# Note: "tags" field is intentionally excluded, but can be added if needed
-REDACTION_FIELDS = ["input", "output", "expected", "metadata", "context", "scores", "metrics"]
-
 DATA_API_VERSION = 2
 LOGS3_OVERFLOW_REFERENCE_TYPE = "logs3_overflow"
 # 6 MB for the AWS lambda gateway (from our own testing).
@@ -1002,40 +999,6 @@ def utf8_byte_length(value: str) -> int:
     return len(value.encode("utf-8"))
 
 
-class _MaskingError:
-    """Internal class to signal masking errors that need special handling."""
-
-    def __init__(self, field_name: str, error_type: str):
-        self.field_name = field_name
-        self.error_type = error_type
-        self.error_msg = f"ERROR: Failed to mask field '{field_name}' - {error_type}"
-
-
-def _apply_masking_to_field(masking_function: Callable[[Any], Any], data: Any, field_name: str) -> Any:
-    """Apply masking function to data and handle errors gracefully.
-
-    If the masking function raises an exception, returns an error message.
-    Returns _MaskingError for scores/metrics fields to signal they should be dropped.
-    """
-    try:
-        return masking_function(data)
-    except Exception as mask_error:
-        # Return a generic error message without the stack trace to avoid leaking PII
-        error_type = type(mask_error).__name__
-
-        # For scores and metrics fields, return a special error object
-        # to signal the field should be dropped and error logged
-        if field_name in ["scores", "metrics"]:
-            return _MaskingError(field_name, error_type)
-
-        # For metadata field that expects dict type, return a dict with error key
-        if field_name == "metadata":
-            return {"error": f"ERROR: Failed to mask field '{field_name}' - {error_type}"}
-
-        # For other fields, return the error message as a string
-        return f"ERROR: Failed to mask field '{field_name}' - {error_type}"
-
-
 class _BackgroundLogger(ABC):
     @abstractmethod
     def log(self, *args: LazyValue[dict[str, Any]]) -> None:
@@ -1050,7 +1013,7 @@ class _MemoryBackgroundLogger(_BackgroundLogger):
     def __init__(self):
         self.lock = threading.Lock()
         self.logs = []
-        self.masking_function: Callable[[Any], Any] | None = None
+        self._export_customizers: tuple[SpanCustomizer, ...] = ()
         self.upload_attempts: list[BaseAttachment] = []  # Track upload attempts
 
     def enforce_queue_size_limit(self, enforce: bool) -> None:
@@ -1062,7 +1025,7 @@ def log(self, *args: LazyValue[dict[str, Any]]) -> None:
 
     def set_masking_function(self, masking_function: Callable[[Any], Any] | None) -> None:
         """Set the masking function for the memory logger."""
-        self.masking_function = masking_function
+        self._export_customizers = (_MaskingCustomizer(masking_function),) if masking_function is not None else ()
 
     def flush(self, batch_size: int | None = None):
         """Flush the memory logger, extracting attachments and tracking upload attempts."""
@@ -1093,28 +1056,8 @@ def pop(self):
             # here
             batch = merge_row_batch(logs)
 
-            # Apply masking after merge, similar to HTTPBackgroundLogger
-            if self.masking_function:
-                for i in range(len(batch)):
-                    item = batch[i]
-                    masked_item = item.copy()
-
-                    # Only mask specific fields if they exist
-                    for field in REDACTION_FIELDS:
-                        if field in item:
-                            masked_value = _apply_masking_to_field(self.masking_function, item[field], field)
-                            if isinstance(masked_value, _MaskingError):
-                                # Drop the field and add error message
-                                if field in masked_item:
-                                    del masked_item[field]
-                                if "error" in masked_item:
-                                    masked_item["error"] = f"{masked_item['error']}; {masked_value.error_msg}"
-                                else:
-                                    masked_item["error"] = masked_value.error_msg
-                            else:
-                                masked_item[field] = masked_value
-
-                    batch[i] = masked_item
+            if self._export_customizers:
+                batch = [_customize_span_export(item, self._export_customizers) for item in batch]
 
             return batch
 
@@ -1129,7 +1072,7 @@ def pop(self):
 class _HTTPBackgroundLogger:
     def __init__(self, api_conn: LazyValue[HTTPConnection]):
         self.api_conn = api_conn
-        self.masking_function: Callable[[Any], Any] | None = None
+        self._export_customizers: tuple[SpanCustomizer, ...] = ()
         self.outfile = sys.stderr
         self.flush_lock = threading.RLock()
         self._max_request_size_override: int | None = None
@@ -1318,28 +1261,9 @@ def _unwrap_lazy_values(
                 unwrapped_items = [item.get() for item in wrapped_items]
                 merged_items = merge_row_batch(unwrapped_items)
 
-                # Apply masking after merging but before sending to backend
-                if self.masking_function:
-                    for item_idx in range(len(merged_items)):
-                        item = merged_items[item_idx]
-                        masked_item = item.copy()
-
-                        # Only mask specific fields if they exist
-                        for field in REDACTION_FIELDS:
-                            if field in item:
-                                masked_value = _apply_masking_to_field(self.masking_function, item[field], field)
-                                if isinstance(masked_value, _MaskingError):
-                                    # Drop the field and add error message
-                                    if field in masked_item:
-                                        del masked_item[field]
-                                    if "error" in masked_item:
-                                        masked_item["error"] = f"{masked_item['error']}; {masked_value.error_msg}"
-                                    else:
-                                        masked_item["error"] = masked_value.error_msg
-                                else:
-                                    masked_item[field] = masked_value
-
-                        merged_items[item_idx] = masked_item
+                # Logger-local hooks run after instrumentation hooks and merging.
+                if self._export_customizers:
+                    merged_items = [_customize_span_export(item, self._export_customizers) for item in merged_items]
 
                 attachments: list["BaseAttachment"] = []
                 for item in merged_items:
@@ -1553,7 +1477,7 @@ def internal_replace_api_conn(self, api_conn: HTTPConnection):
 
     def set_masking_function(self, masking_function: Callable[[Any], Any] | None):
         """Set or update the masking function."""
-        self.masking_function = masking_function
+        self._export_customizers = (_MaskingCustomizer(masking_function),) if masking_function is not None else ()
 
 
 def _internal_reset_global_state() -> None:
@@ -2565,6 +2489,8 @@ def set_masking_function(masking_function: Callable[[Any], Any] | None) -> None:
     """
     Set a global masking function that will be applied to all logged data before sending to Braintrust.
     The masking function will be applied after records are merged but before they are sent to the backend.
+    Internally, masking is a logger-local export customizer that runs after instrumentation
+    customizers and also covers manually logged records.
 
     :param masking_function: A function that takes a JSON-serializable object and returns a masked version.
                            Set to None to disable masking.
@@ -4975,36 +4901,68 @@ def log_internal(self, event: dict[str, Any] | None = None, internal_data: dict[
         if serializable_partial_record.get("metrics", {}).get("end") is not None:
             self._logged_end_time = serializable_partial_record["metrics"]["end"]
 
-        # Write to local span cache for scorer access
-        # Only cache experiment spans - regular logs don't need caching
-        if self.parent_object_type == SpanObjectTypeV3.EXPERIMENT:
+        # Snapshot at log time so the span cache and the export agree on whether
+        # (and how) this record is customized.
+        customizers = _get_span_customizers() if self._instrumentation != "braintrust-python-logger" else ()
+        pending_cache_key = (
+            object()
+            if customizers
+            and self.parent_object_type == SpanObjectTypeV3.EXPERIMENT
+            and not self.state.span_cache.disabled
+            else None
+        )
+
+        def write_span_cache(record: dict[str, Any]) -> None:
+            # Write to local span cache for scorer access
+            # Only cache experiment spans - regular logs don't need caching
+            if self.parent_object_type != SpanObjectTypeV3.EXPERIMENT:
+                return
             from braintrust.span_cache import CachedSpan
 
             cached_span = CachedSpan(
                 span_id=self.span_id,
-                input=serializable_partial_record.get("input"),
-                output=serializable_partial_record.get("output"),
-                metadata=serializable_partial_record.get("metadata"),
+                input=record.get("input"),
+                output=record.get("output"),
+                metadata=record.get("metadata"),
                 span_parents=self.span_parents,
-                span_attributes=serializable_partial_record.get("span_attributes"),
-                error=serializable_partial_record.get("error"),
-                metrics=serializable_partial_record.get("metrics"),
-                tags=serializable_partial_record.get("tags"),
+                span_attributes=record.get("span_attributes"),
+                error=record.get("error"),
+                metrics=record.get("metrics"),
+                tags=record.get("tags"),
             )
             self.state.span_cache.queue_write(self.root_span_id, self.span_id, cached_span)
 
+        # Customized records are cached after export customization instead, so
+        # local scorers never see content that a customizer redacted.
+        if not customizers:
+            write_span_cache(serializable_partial_record)
+
         def compute_record() -> dict[str, Any]:
             exporter = _get_exporter()
-            return dict(
+            record = dict(
                 **serializable_partial_record,
                 **{k: v.get() for k, v in lazy_partial_record.items()},
                 **exporter(
                     object_type=self.parent_object_type,
                     object_id=self.parent_object_id.get(),
                 ).object_id_fields(),
             )
-
-        self.state.global_bg_logger().log(LazyValue(compute_record, use_mutex=False))
+            # Resolve and customize inside the cached LazyValue: every incremental
+            # instrumentation record is transformed once, before the background
+            # logger merges, masks, extracts attachments, or retries delivery.
+            if customizers:
+                record = _customize_span_export(record, customizers)
+                write_span_cache(record)
+                if pending_cache_key is not None:
+                    self.state.span_cache._forget_pending_record(self.root_span_id, pending_cache_key)
+            return record
+
+        # Cache readers and the publisher share resolution, including the cache
+        # write, so concurrent reads cannot customize a record twice.
+        lazy_record = LazyValue(compute_record, use_mutex=pending_cache_key is not None)
+        if pending_cache_key is not None:
+            self.state.span_cache._track_pending_record(self.root_span_id, pending_cache_key, lazy_record)
+        self.state.global_bg_logger().log(lazy_record)
 
     def log_feedback(self, **event: Any) -> None:
         return _log_feedback_impl(