feat: add span customizers (SDK-316)

add customizers and re-implement masking fn with a customizer

Author: realark

py/README.md

@@ -41,6 +41,70 @@ Then run:
 BRAINTRUST_API_KEY=<YOUR_API_KEY> braintrust eval tutorial_eval.py
 ```
 
+## Customizing instrumentation exports
+
+Use synchronous span customizers to redact captured values or add metadata without
+changing provider responses:
+
+```python
+import braintrust
+
+
+class Redact(braintrust.SpanCustomizer):
+    def on_span_export(self, data: braintrust.SpanExportData) -> braintrust.SpanExportData:
+        for field in ("input", "output"):
+            if field in data:
+                data[field] = "[redacted]"
+        data.pop("error", None)
+        return data
+
+
+braintrust.auto_instrument(span_customizers=[Redact()])
+braintrust.init_logger(project="my-project")
+```
+
+For explicit wrappers or integration setup, call
+`braintrust.set_span_customizers([Redact()])` before instrumented calls. Both APIs
+replace the process-wide ordered list with a snapshot of the supplied sequence;
+the customizer objects themselves are not copied. `set_span_customizers(None)` or
+an empty sequence disables customization. Omitting `span_customizers` (or passing
+`None`) to `auto_instrument()` leaves existing configuration unchanged. There is no
+environment-variable registration.
+
+`SpanCustomizer` is an extensible base class with a default no-op
+`on_span_export(data)` method. Hooks run synchronously in registration order, each
+receiving its predecessor's result. Return a plain `dict` containing the data to
+export: either the input record or a replacement. Replacement is **not** a merge;
+retain the content fields you want to export. Returning `None` cannot drop a span.
+
+Hooks apply only to **native instrumentation-created span records**, not manually
+created spans (even children of instrumented spans), dataset rows, feedback, or
+spans exported through a separate OpenTelemetry exporter. They run after lazy
+values resolve and before attachment processing, merging, masking, and transport
+serialization. Existing attachment objects retain their identity; hooks can
+remove them before upload without copying their contents.
+
+A span can produce several incremental records, including records flushed before
+it ends. Input, output, metrics, and errors need not be present together. Removing
+a field does not retract an earlier export, so redact every record containing
+sensitive data. Each record uses the configuration at its first export resolution.
+Retries reuse the transformed record without calling hooks again.
+
+Identity, parentage, routing, and merge protocol fields are restored after every
+hook, including nested protocol arrays. Protected fields are `id`, `span_id`,
+`root_span_id`, `span_parents`, `org_id`, `project_id`, `experiment_id`, `dataset_id`,
+`prompt_session_id`, `log_id`, `function_data`, `_is_merge`, `_merge_paths`,
+`_parent_id`, `_object_delete`, `_array_delete`, and `_xact_id`. Hooks cannot add an
+absent protected field or change the destination.
+
+**Failures are fail-open.** Exceptions, `None`, awaitables, and other invalid return
+types are ignored; later hooks and export continue with the current record.
+In-place content mutations made before failure are retained, not rolled back.
+An asynchronous hook is not awaited. A failing redaction hook can therefore leak
+unredacted data: never throw to block export. Returned content must remain valid
+for SDK serialization. Keep hooks fast, avoid blocking I/O, and make them safe for
+background export threads rather than assuming the caller's thread or context.
+
 ## Optional Extras
 
 Install extras as needed for specific workflows:

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: