From 0bae7f4c27b948cc9f560ee5fe2e009a6a090926 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Mon, 21 Sep 2026 02:13:13 -0600 Subject: [PATCH] feat: add span customizers (SDK-316) add customizers and re-implement masking fn with a customizer --- py/README.md | 64 +++ py/src/braintrust/__init__.py | 3 + py/src/braintrust/auto.py | 9 + .../test_auto_google_discoveryengine.py | 6 +- .../integrations/pipecat/test_pipecat.py | 41 +- .../integrations/pipecat/tracing.py | 13 +- py/src/braintrust/logger.py | 138 ++---- py/src/braintrust/span_customizer.py | 171 +++++++ py/src/braintrust/test_logger.py | 166 +------ py/src/braintrust/test_span_customizer.py | 447 ++++++++++++++++++ py/src/braintrust/wrappers/langchain.py | 3 + 11 files changed, 819 insertions(+), 242 deletions(-) create mode 100644 py/src/braintrust/span_customizer.py create mode 100644 py/src/braintrust/test_span_customizer.py diff --git a/py/README.md b/py/README.md index b0b3242dc..441e71bff 100644 --- a/py/README.md +++ b/py/README.md @@ -41,6 +41,70 @@ Then run: BRAINTRUST_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: diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 2513a0e75..7af51764e 100644 --- a/py/src/braintrust/__init__.py +++ b/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 diff --git a/py/src/braintrust/auto.py b/py/src/braintrust/auto.py index 77d64ab25..d0721813d 100644 --- a/py/src/braintrust/auto.py +++ b/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: diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_google_discoveryengine.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_google_discoveryengine.py index dd7fa19fe..54ad287e2 100644 --- a/py/src/braintrust/integrations/auto_test_scripts/test_auto_google_discoveryengine.py +++ b/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": diff --git a/py/src/braintrust/integrations/pipecat/test_pipecat.py b/py/src/braintrust/integrations/pipecat/test_pipecat.py index b42764041..d60034a0c 100644 --- a/py/src/braintrust/integrations/pipecat/test_pipecat.py +++ b/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) diff --git a/py/src/braintrust/integrations/pipecat/tracing.py b/py/src/braintrust/integrations/pipecat/tracing.py index 1a56a712e..a60c04d2f 100644 --- a/py/src/braintrust/integrations/pipecat/tracing.py +++ b/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: diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 663e29688..042685ef3 100644 --- a/py/src/braintrust/logger.py +++ b/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,27 +4901,38 @@ 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 () + + 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( @@ -5003,6 +4940,13 @@ def compute_record() -> dict[str, Any]: object_id=self.parent_object_id.get(), ).object_id_fields(), ) + # 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) + return record self.state.global_bg_logger().log(LazyValue(compute_record, use_mutex=False)) diff --git a/py/src/braintrust/span_customizer.py b/py/src/braintrust/span_customizer.py new file mode 100644 index 000000000..ff5a16f99 --- /dev/null +++ b/py/src/braintrust/span_customizer.py @@ -0,0 +1,171 @@ +"""Synchronous transformations of native export records.""" + +import inspect +from collections.abc import Callable, Sequence +from typing import Any + +from .db_fields import ( + ARRAY_DELETE_FIELD, + ID_FIELD, + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + OBJECT_DELETE_FIELD, + OBJECT_ID_KEYS, + PARENT_ID_FIELD, + TRANSACTION_ID_FIELD, +) + + +__all__ = ["SpanCustomizer", "SpanExportData", "set_span_customizers"] + +SpanExportData = dict[str, Any] +"""A native, possibly incremental span export record, not a live span.""" + + +class SpanCustomizer: + """Extensible hooks for instrumentation-created spans. Omitted hooks are no-ops.""" + + def on_span_export(self, data: SpanExportData) -> SpanExportData: + """Return the original or a replacement export record, synchronously. + + Hooks run in registration order after lazy values resolve, before merging, + attachment processing, masking, and serialization. Identity, routing, and + merge protocol fields are restored after each hook. Exceptions and invalid + returns fail open: in-place edits survive and later hooks still run. + """ + return data + + +_span_customizers: tuple[SpanCustomizer, ...] = () +_PROTECTED_FIELDS = frozenset( + ( + ID_FIELD, + "span_id", + "root_span_id", + "span_parents", + "org_id", + *OBJECT_ID_KEYS, + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + PARENT_ID_FIELD, + OBJECT_DELETE_FIELD, + ARRAY_DELETE_FIELD, + TRANSACTION_ID_FIELD, + ) +) + +# Tags and record-level errors are intentionally outside the masking contract. +_MASKING_FIELDS = ("input", "output", "expected", "metadata", "context", "scores", "metrics") + + +class _MaskingCustomizer(SpanCustomizer): + """Adapt field-level masking to a logger-local hook on all merged records.""" + + def __init__(self, masking_function: Callable[[Any], Any]): + self._masking_function = masking_function + + def on_span_export(self, data: SpanExportData) -> SpanExportData: + for field in _MASKING_FIELDS: + if field not in data: + continue + try: + data[field] = self._masking_function(data[field]) + except Exception as error: + # Fail closed per field, without leaking exception messages or stacks. + message = f"ERROR: Failed to mask field '{field}' - {type(error).__name__}" + if field in ("scores", "metrics"): + del data[field] + data["error"] = f"{data['error']}; {message}" if "error" in data else message + else: + data[field] = {"error": message} if field == "metadata" else message + return data + + +def set_span_customizers(customizers: Sequence[SpanCustomizer] | None) -> None: + """Replace the process-wide ordered customizer list with a snapshot. + + Configure before making instrumented calls; pass None or an empty sequence to + disable. The sequence is copied, but customizer instances are not. Each record + uses the configuration active when it is logged, and retries reuse its + transformed data. There is no environment-variable registration. + + Raises TypeError for classes, objects without a callable ``on_span_export``, + and async hooks, since those would otherwise silently export unredacted data. + """ + global _span_customizers + snapshot = tuple(customizers) if customizers is not None else () + for customizer in snapshot: + _validate_customizer(customizer) + _span_customizers = snapshot + + +def _get_span_customizers() -> tuple[SpanCustomizer, ...]: + return _span_customizers + + +def _validate_customizer(customizer: Any) -> None: + # Reject misconfiguration eagerly: at export time these would silently fail + # open and ship unredacted data. + if isinstance(customizer, type): + raise TypeError(f"Span customizers must be instances, not classes; did you mean {customizer.__name__}()?") + hook = getattr(customizer, "on_span_export", None) + if not callable(hook): + raise TypeError( + f"Span customizer {type(customizer).__name__} must define a callable on_span_export(data) method" + ) + if inspect.iscoroutinefunction(hook): + raise TypeError(f"{type(customizer).__name__}.on_span_export must be synchronous") + + +def _copy_protocol_value(value: Any) -> Any: + # Copy only protocol containers. Payloads (especially Attachment objects) must + # retain their existing serialization behavior and must not be deep-copied. + if isinstance(value, list): + return [_copy_protocol_value(item) for item in value] + if isinstance(value, dict): + return {key: _copy_protocol_value(item) for key, item in value.items()} + return value + + +def _restore_protocol_fields(data: SpanExportData, protected: SpanExportData) -> SpanExportData: + restored = {key: value for key, value in data.items() if key not in _PROTECTED_FIELDS} + # Never expose the snapshot itself: a later hook may mutate nested arrays. + restored.update({key: _copy_protocol_value(value) for key, value in protected.items()}) + return restored + + +def _customize_span_export( + data: SpanExportData, customizers: Sequence[SpanCustomizer] | None = None +) -> SpanExportData: + if customizers is None: + customizers = _span_customizers + if not customizers: + return data + + protected = None + for customizer in customizers: + candidate = data + try: + hook = getattr(customizer, "on_span_export", None) + if hook is None: + continue + if protected is None: + protected = {key: _copy_protocol_value(data[key]) for key in _PROTECTED_FIELDS if key in data} + # Protocol containers may alias span state; don't expose those to hooks. + data = _restore_protocol_fields(data, protected) + candidate = data + result = hook(data) + if inspect.isawaitable(result): + # Do not execute asynchronous hooks or emit unawaited coroutine warnings. + if inspect.iscoroutine(result): + result.close() + elif type(result) is dict: + candidate = result + except Exception: + # Deliberately fail open, retaining in-place changes. Never log the + # exception, since its message may itself contain sensitive payloads. + pass + + if protected is not None: + data = _restore_protocol_fields(candidate, protected) + return data diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index c642d8a68..9745c17be 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -3211,154 +3211,36 @@ def masking_function(data): def test_masking_function_with_error(with_memory_logger, with_simulate_login): - """Test that masking errors are handled gracefully and stack traces are captured.""" - def broken_masking_function(data): - """A masking function that throws errors for certain data types.""" - if isinstance(data, dict): - # This will throw an error when trying to iterate - for key in data: - if key == "password": - # Simulate a complex error - raise ValueError(f"Cannot mask sensitive field '{key}' - internal masking error") - elif key == "accuracy": - # Trigger error for scores field - raise TypeError("Cannot process numeric score") + if data == "safe output": return data - elif isinstance(data, str): - if "secret" in data.lower(): - # Another type of error - result = 1 / 0 # ZeroDivisionError - return data - elif isinstance(data, list): - # Try to access non-existent index - if len(data) > 0: - _ = data[100] # IndexError - return data - return data + raise TypeError("private exception detail") - # Set the broken masking function braintrust.set_masking_function(broken_masking_function) - - # Create test experiment - from braintrust.logger import Experiment, ObjectMetadata, ProjectExperimentMetadata - - project_metadata = ObjectMetadata(id="test_project", name="test_project", full_info=dict()) - experiment_metadata = ObjectMetadata(id="test_experiment", name="test_experiment", full_info=dict()) - metadata = ProjectExperimentMetadata(project=project_metadata, experiment=experiment_metadata) - lazy_metadata = LazyValue(lambda: metadata, use_mutex=False) - experiment = Experiment(lazy_metadata=lazy_metadata) - - # Log data that will trigger various errors - experiment.log( - input={"password": "my-password", "user": "test"}, - output="This contains SECRET information", - expected=["item1", "item2"], - metadata={"safe": "data"}, - scores={"score": 1.0}, # Add a safe score that won't trigger error - ) - - experiment.flush() - - # Check the logged data - logs = with_memory_logger.pop() - assert len(logs) == 1 - log = logs[0] - - # Verify error handling - # The input should have an error message because of the password field - assert log["input"] == "ERROR: Failed to mask field 'input' - ValueError" - - # The output should have an error message because of division by zero - assert log["output"] == "ERROR: Failed to mask field 'output' - ZeroDivisionError" - - # The expected should have an error message because of index error - assert log["expected"] == "ERROR: Failed to mask field 'expected' - IndexError" - - # Metadata should be fine since it doesn't trigger any errors - assert log["metadata"] == {"safe": "data"} - - # Test with scores that triggers an error - experiment.log( - input={"data": "test"}, - output="result", - scores={"accuracy": 0.95}, # This will trigger an error - ) - - logs2 = with_memory_logger.pop() - assert len(logs2) == 1 - log2 = logs2[0] - - # Scores should be dropped and error should be logged - assert "scores" not in log2 - assert "error" in log2 - assert log2["error"] == "ERROR: Failed to mask field 'scores' - TypeError" - - # Test with metrics that triggers an error - experiment.log( - input={"data": "test2"}, - output="result2", - scores={"score": 1.0}, # Safe score - metrics={"accuracy": 0.95}, # This will trigger an error - ) - - logs3 = with_memory_logger.pop() - assert len(logs3) == 1 - log3 = logs3[0] - - # Metrics should be dropped and error should be logged - assert "metrics" not in log3 - assert "error" in log3 - assert log3["error"] == "ERROR: Failed to mask field 'metrics' - TypeError" - - # Test with both scores and metrics failing - experiment.log( - input={"data": "test3"}, - output="result3", - scores={"accuracy": 0.85}, # This will trigger an error - metrics={"accuracy": 0.95}, # This will also trigger an error - ) - - logs4 = with_memory_logger.pop() - assert len(logs4) == 1 - log4 = logs4[0] - - # Both should be dropped and errors should be concatenated - assert "scores" not in log4 - assert "metrics" not in log4 - assert "error" in log4 - assert "ERROR: Failed to mask field 'scores' - TypeError" in log4["error"] - assert "ERROR: Failed to mask field 'metrics' - TypeError" in log4["error"] - assert "; " in log4["error"] # Check that errors are joined with semicolon - - # Test with logger and nested spans test_logger = init_test_logger("test_masking_errors_logger") + test_logger.log( + input={"password": "private input"}, + output="safe output", + expected="private expected", + metadata={"token": "private metadata"}, + scores={"accuracy": 0.85}, + metrics={"accuracy": 0.95}, + error="existing application error", + tags=["untouched"], + ) - with test_logger.start_span("parent") as parent: - parent.log(input={"api_key": "key123", "password": "secret"}, metadata={"request_id": "req-123"}) - - with parent.start_span("child") as child: - child.log(output="Result with secret data", expected=[1, 2, 3]) - - test_logger.flush() - - # Check nested span logs - logs = with_memory_logger.pop() - assert len(logs) == 2 # parent and child - - # Find parent and child by span_attributes - parent_log = next(log for log in logs if log.get("span_attributes", {}).get("name") == "parent") - child_log = next(log for log in logs if log.get("span_attributes", {}).get("name") == "child") - - # Parent should have error in input - assert parent_log["input"] == "ERROR: Failed to mask field 'input' - ValueError" - - # Child should have errors in output and expected - assert child_log["output"] == "ERROR: Failed to mask field 'output' - ZeroDivisionError" - assert child_log["expected"] == "ERROR: Failed to mask field 'expected' - IndexError" - - # Clean up - braintrust.set_masking_function(None) + [record] = with_memory_logger.pop() + assert isinstance(record["input"], str) + assert isinstance(record["expected"], str) + assert isinstance(record["metadata"]["error"], str) + assert record["output"] == "safe output" + assert record["tags"] == ["untouched"] + assert "scores" not in record + assert "metrics" not in record + assert "existing application error" in record["error"] + assert "scores" in record["error"] + assert "metrics" in record["error"] + assert "private" not in json.dumps(record) def test_attachment_unreadable_path_logs_warning(caplog): diff --git a/py/src/braintrust/test_span_customizer.py b/py/src/braintrust/test_span_customizer.py new file mode 100644 index 000000000..82294d39d --- /dev/null +++ b/py/src/braintrust/test_span_customizer.py @@ -0,0 +1,447 @@ +import inspect +import json +from copy import deepcopy +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from braintrust import Attachment, SpanCustomizer, SpanExportData, auto_instrument, logger, set_span_customizers +from braintrust.functions.stream import BraintrustJsonChunk, BraintrustStream +from braintrust.span_customizer import _customize_span_export +from braintrust.test_helpers import init_test_exp, with_memory_logger # noqa: F401 # type: ignore[reportUnusedImport] +from braintrust.util import LazyValue + + +@pytest.fixture(autouse=True) +def reset_customizers(): + set_span_customizers(None) + try: + yield + finally: + set_span_customizers(None) + + +@pytest.fixture +def test_logger(): + metadata = logger.OrgProjectMetadata( + org_id="org", project=logger.ObjectMetadata(id="project", name="project", full_info={}) + ) + return logger.Logger(LazyValue(lambda: metadata, use_mutex=False)) + + +def test_order_replacement_and_registration_snapshot(): + observed = [] + + class Replace(SpanCustomizer): + def on_span_export(self, data: SpanExportData) -> SpanExportData: + # Reconfiguration during export applies only to subsequent records. + set_span_customizers([]) + return {"output": "redacted"} + + class Observe(SpanCustomizer): + def on_span_export(self, data: SpanExportData) -> SpanExportData: + observed.append(dict(data)) + data["metadata"] = {"safe": True} + return data + + customizers = [Replace(), Observe()] + set_span_customizers(customizers) + customizers.clear() + result = _customize_span_export({"id": "row", "input": "secret", "output": "secret"}) + assert observed == [{"id": "row", "output": "redacted"}] + assert result == {"id": "row", "output": "redacted", "metadata": {"safe": True}} + assert _customize_span_export({"output": "next"}) == {"output": "next"} + + +class _Valid(SpanCustomizer): + def on_span_export(self, data): + return {"output": "valid"} + + +class _Async(SpanCustomizer): + async def on_span_export(self, data): # type: ignore[override] + return data + + +@pytest.mark.parametrize( + "invalid", + [ + _Valid, # the class, not an instance + lambda data: data, + SimpleNamespace(on_span_export="not callable"), + object(), + _Async(), + ], + ids=["class", "function", "non-callable-hook", "no-hook", "async-hook"], +) +def test_registration_rejects_misconfigured_customizers(invalid): + valid = _Valid() + set_span_customizers([valid]) + with pytest.raises(TypeError): + set_span_customizers([valid, invalid]) + with pytest.raises(TypeError): + auto_instrument(span_customizers=[invalid]) + # A rejected registration leaves the previous configuration in place. + assert _customize_span_export({"id": "row", "output": "secret"}) == {"id": "row", "output": "valid"} + + +def test_duck_typed_customizers_are_accepted(): + class DuckTyped: + def on_span_export(self, data): + return {"output": "duck"} + + set_span_customizers([DuckTyped()]) # type: ignore[list-item] + assert _customize_span_export({"id": "row", "output": "secret"}) == {"id": "row", "output": "duck"} + + +def test_explicit_customizers_override_without_changing_global_configuration(): + class Global(SpanCustomizer): + def on_span_export(self, data): + return {"output": "global"} + + class Local(SpanCustomizer): + def on_span_export(self, data): + return {"output": "local"} + + set_span_customizers([Global()]) + record = {"id": "row", "output": "original"} + assert _customize_span_export(record, [Local()]) == {"id": "row", "output": "local"} + assert _customize_span_export(record, []) == record + assert _customize_span_export(record) == {"id": "row", "output": "global"} + + +def test_protected_protocol_fields_restored_between_hooks(): + protected = { + "id": "row", + "span_id": "span", + "root_span_id": "root", + "span_parents": ["parent"], + "org_id": "org", + "project_id": "project", + "log_id": "g", + "function_data": {"nested": ["routing"]}, + "_is_merge": True, + "_merge_paths": [["metadata", "nested"]], + "_parent_id": "parent-row", + "_object_delete": False, + "_array_delete": [["tags", "private"]], + "_xact_id": "transaction", + } + original = deepcopy(protected) + observed = [] + + class Corrupt(SpanCustomizer): + def on_span_export(self, data: SpanExportData) -> SpanExportData: + data["span_parents"].append("wrong") + data["_merge_paths"][0].append("wrong") + data["_array_delete"][0].clear() + data["function_data"]["nested"].clear() + for key in protected: + data[key] = "wrong" + data["experiment_id"] = "wrong-destination" + data["dataset_id"] = "wrong-destination" + data["prompt_session_id"] = "wrong-destination" + return data + + class Observe(SpanCustomizer): + def on_span_export(self, data: SpanExportData) -> SpanExportData: + observed.append(deepcopy(data)) + # A second mutation must not poison the stored snapshot either. + data["_merge_paths"][0].clear() + data["span_parents"].clear() + return {"output": "safe"} + + set_span_customizers([Corrupt(), Observe()]) + result = _customize_span_export(protected) + assert observed == [original] + assert result == {**original, "output": "safe"} + assert protected == original + + +@pytest.mark.parametrize("failure", ["exception", "none", "list", "coroutine", "awaitable"]) +def test_fail_open_keeps_mutations_and_continues(failure): + observed = [] + coroutines = [] + + class AwaitableRecord(dict): + def __await__(self): + raise AssertionError("Hooks must not be awaited") + yield + + async def asynchronous_result(): + raise AssertionError("Async hook must not execute") + + class Broken(SpanCustomizer): + def on_span_export(self, data): + data["output"] = "redacted" + data["id"] = "wrong" + if failure == "exception": + raise ValueError("private exception text") + if failure == "none": + return None + if failure == "list": + return [] + if failure == "awaitable": + return AwaitableRecord(output="wrong") + coroutine = asynchronous_result() + coroutines.append(coroutine) + return coroutine + + class Next(SpanCustomizer): + def on_span_export(self, data): + observed.append(dict(data)) + return data + + set_span_customizers([Broken(), Next()]) + result = _customize_span_export({"id": "row", "output": "secret"}) + assert observed == [{"id": "row", "output": "redacted"}] + assert result == observed[0] + assert all(inspect.getcoroutinestate(c) == inspect.CORO_CLOSED for c in coroutines) + + +def test_native_scope_incremental_redaction_and_provider_isolation(with_memory_logger, test_logger): + seen = [] + + class Redact(SpanCustomizer): + def on_span_export(self, data): + seen.append((data["id"], set(data))) + for key in ("input", "output"): + if key in data: + data[key]["secret"] = "redacted" + data.pop("error", None) + return data + + set_span_customizers([Redact()]) + application_input = {"secret": "input"} + with test_logger.start_span(name="manual", input=application_input) as manual: + with manual.start_span( + name="instrumented", input=application_input, internal={"instrumentation": "test-auto"} + ) as span: + # Flush before completion to expose the incremental lifecycle. + first = with_memory_logger.pop() + assert next(row for row in first if row["id"] == span.id)["input"] == {"secret": "redacted"} + assert next(row for row in first if row["id"] == manual.id)["input"] == {"secret": "input"} + stream = BraintrustStream([BraintrustJsonChunk(data='{"secret":"output"}')]) + span.log(output=stream, error="private error") + span.set_attributes(name="renamed") + with span.start_span(name="manual child", input=application_input) as manual_child: + pass + span.log_feedback(expected="feedback value") + + metadata = logger.ProjectDatasetMetadata( + project=logger.ObjectMetadata(id="project", name="project", full_info={}), + dataset=logger.ObjectMetadata(id="dataset", name="dataset", full_info={}), + ) + dataset = logger.Dataset(LazyValue(lambda: metadata, use_mutex=False), legacy=False) + dataset_id = dataset.insert(input={"secret": "dataset"}) + rows = with_memory_logger.pop() + exported = next(row for row in rows if row["id"] == span.id) + assert exported["output"] == {"secret": "redacted"} + assert "error" not in exported + assert exported["expected"] == "feedback value" + assert exported["span_attributes"]["name"] == "renamed" + assert next(row for row in rows if row["id"] == manual_child.id)["input"] == application_input + assert next(row for row in rows if row["id"] == dataset_id)["input"] == {"secret": "dataset"} + assert application_input == {"secret": "input"} + assert stream.final_value() == {"secret": "output"} + assert [row_id for row_id, _ in seen] == [span.id] * 4 + assert ["input" in keys for _, keys in seen] == [True, False, False, False] + assert ["output" in keys for _, keys in seen] == [False, True, False, False] + assert not any("expected" in keys for _, keys in seen) + + +def test_customization_precedes_attachments_masking_and_reuses_records_on_retry( + monkeypatch, with_memory_logger, test_logger +): + attachment = Attachment(data=b"private", filename="private.txt", content_type="text/plain") + seen_attachments = [] + invocations = [] + + class Redact(SpanCustomizer): + def on_span_export(self, data): + invocations.append(data["id"]) + if "input" in data: + seen_attachments.append(data["input"]) + data["input"] = "redacted" + if "output" in data: + data["output"] = "redacted output" + return data + + set_span_customizers([Redact()]) + span = test_logger.start_span(input=attachment, internal={"instrumentation": "test-auto"}) + span.log(output="private output") + span.end() + pending = list(with_memory_logger.logs) + with_memory_logger.logs.clear() + + # A later lazy record fails once, after earlier records have been customized. + attempts = 0 + + def resolve_later_record(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("retry lazy resolution") + return {"id": "other", "project_id": "project", "log_id": "g", "output": "other"} + + pending.append(LazyValue(resolve_later_record, use_mutex=False)) + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + monkeypatch.setattr(logger.time, "sleep", lambda _: None) + connection = MagicMock() + payloads = [] + + def send(_path, *, data): + payloads.append(data) + if len(payloads) == 1: + raise ConnectionError("retry transport") + return SimpleNamespace(ok=True) + + connection.post.side_effect = send + background = logger._HTTPBackgroundLogger(LazyValue(lambda: connection, use_mutex=False)) + background.num_tries = 2 + background.sync_flush = True + masked = [] + + def mask(value): + masked.append(value) + return value + + background.set_masking_function(mask) + background._max_request_size_result = {"max_request_size": 6_000_000, "can_use_overflow": False} + for item in pending: + background.queue.put(item) + background.flush() + + assert attempts == 2 + assert invocations == [span.id] * 3 + assert seen_attachments == [attachment] + assert seen_attachments[0] is attachment + assert "redacted" in masked + assert "redacted output" in masked + assert payloads[0] == payloads[1] + exported = next(row for row in json.loads(payloads[1])["rows"] if row["id"] == span.id) + assert exported["input"] == "redacted" + assert exported["output"] == "redacted output" + + +@pytest.mark.parametrize("backend", ["memory", "http"]) +def test_masking_remains_logger_local_and_runs_on_merged_manual_records(monkeypatch, backend): + class InstrumentationOnly(SpanCustomizer): + def on_span_export(self, data): + return {"output": "must not run on manual records"} + + set_span_customizers([InstrumentationOnly()]) + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "1") + background = ( + logger._MemoryBackgroundLogger() + if backend == "memory" + else logger._HTTPBackgroundLogger(LazyValue(lambda: MagicMock(), use_mutex=False)) + ) + other_background = logger._MemoryBackgroundLogger() + + def export(target): + records = [ + {"id": "row", "project_id": "project", "input": None, "metadata": {"secret": "private"}}, + {"id": "row", "project_id": "project", "_is_merge": True, "metadata": {"redact": True}}, + ] + pending = [LazyValue(lambda record=record: record, use_mutex=False) for record in records] + if isinstance(target, logger._MemoryBackgroundLogger): + target.log(*pending) + return target.pop()[0] + rows, _ = target._unwrap_lazy_values(pending) + return rows[0] + + def mask(value): + if value is None: + return "masked null" + if isinstance(value, dict) and value.get("redact"): + return {"secret": "redacted"} + return value + + background.set_masking_function(mask) + masked = export(background) + assert masked["metadata"] == {"secret": "redacted"} + assert masked["input"] == "masked null" + assert "output" not in masked + assert export(other_background)["metadata"] == {"secret": "private", "redact": True} + + background.set_masking_function(lambda _: "replacement") + assert export(background)["input"] == "replacement" + background.set_masking_function(None) + unmasked = export(background) + assert unmasked["metadata"] == {"secret": "private", "redact": True} + assert unmasked["input"] is None + assert "output" not in unmasked + + +def test_auto_instrument_registration_and_disable(monkeypatch, with_memory_logger, test_logger): + # Registration should work without importing optional provider libraries. + monkeypatch.setattr("braintrust.auto._instrument_integration", lambda _: False) + + class Redact(SpanCustomizer): + def on_span_export(self, data): + if "input" in data: + data["input"] = "redacted" + return data + + auto_instrument(span_customizers=[Redact()]) + auto_instrument() # Omitted configuration must not reset customizers. + with test_logger.start_span(input="private", internal={"instrumentation": "test-auto"}): + pass + assert with_memory_logger.pop()[0]["input"] == "redacted" + auto_instrument(span_customizers=[]) + with test_logger.start_span(input="untouched", internal={"instrumentation": "test-auto"}): + pass + assert with_memory_logger.pop()[0]["input"] == "untouched" + + +def _cached(experiment, span): + return {cached.span_id: cached for cached in experiment.state.span_cache.get_by_root_span_id(span.root_span_id)} + + +@pytest.fixture +def span_cache_experiment(with_memory_logger): + experiment = init_test_exp("span-customizer-cache") + experiment.state.span_cache.start() + try: + yield experiment + finally: + experiment.state.span_cache.dispose() + experiment.state.span_cache.stop() + + +def test_span_cache_stores_customized_content(with_memory_logger, span_cache_experiment): + class Redact(SpanCustomizer): + def on_span_export(self, data): + # Nested in-place edits and replacement must both reach the cache. + if "metadata" in data: + data["metadata"]["secret"] = "redacted" + if "output" in data: + data["output"] = "redacted" + return data + + set_span_customizers([Redact()]) + experiment = span_cache_experiment + with experiment.start_span(name="manual", metadata={"secret": "manual"}) as manual: + with manual.start_span( + name="instrumented", metadata={"secret": "private"}, internal={"instrumentation": "test-auto"} + ) as span: + span.log(output="private output") + + # Customized records are cached at export resolution, never with raw content. + assert span.span_id not in _cached(experiment, manual) + rows = with_memory_logger.pop() + cached = _cached(experiment, span) + assert cached[span.span_id].metadata == {"secret": "redacted"} + assert cached[span.span_id].output == "redacted" + assert cached[manual.span_id].metadata == {"secret": "manual"} + exported = next(row for row in rows if row["id"] == span.id) + assert exported["metadata"] == {"secret": "redacted"} + assert exported["output"] == "redacted" + + +def test_span_cache_written_eagerly_without_customizers(with_memory_logger, span_cache_experiment): + experiment = span_cache_experiment + with experiment.start_span(name="instrumented", internal={"instrumentation": "test-auto"}) as span: + span.log(output="output") + assert _cached(experiment, span)[span.span_id].output == "output" diff --git a/py/src/braintrust/wrappers/langchain.py b/py/src/braintrust/wrappers/langchain.py index c9eff1fad..b8656715f 100644 --- a/py/src/braintrust/wrappers/langchain.py +++ b/py/src/braintrust/wrappers/langchain.py @@ -54,6 +54,9 @@ def _start_span(self, parent_run_id, run_id, name: str | None, **kwargs: Any) -> else: parent_span = braintrust + internal = dict(kwargs.get("internal") or {}) + internal.setdefault("instrumentation", "langchain-auto") + kwargs["internal"] = internal span = parent_span.start_span(name=name, **kwargs) langchain_parent.set(span) self.spans[run_id] = span