braintrust 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

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/db_fields.py

@@ -5,6 +5,7 @@ ID_FIELD = "id"
 
 IS_MERGE_FIELD = "_is_merge"
 MERGE_PATHS_FIELD = "_merge_paths"
+ARRAY_DELETE_FIELD = "_array_delete"
 
 AUDIT_SOURCE_FIELD = "_audit_source"
 AUDIT_METADATA_FIELD = "_audit_metadata"

braintrust/framework.py

@@ -47,7 +47,7 @@ from .resource_manager import ResourceManager
 from .score import Score, is_score, is_scorer
 from .serializable_data_class import SerializableDataClass
 from .span_types import SpanTypeAttribute
-from .util import bt_iscoroutinefunction, eprint
+from .util import bt_iscoroutinefunction, eprint, merge_dicts
 
 Input = TypeVar("Input")
 Output = TypeVar("Output")
@@ -1284,8 +1284,17 @@ async def _run_evaluator_internal(
     event_loop = asyncio.get_event_loop()
 
     async def await_or_run_scorer(root_span, scorer, name, **kwargs):
+        # Merge purpose into parent's propagated_event rather than replacing it
+        parent_propagated = root_span.propagated_event or {}
+        merged_propagated = merge_dicts(
+            {**parent_propagated},
+            {"span_attributes": {"purpose": "scorer"}},
+        )
         with root_span.start_span(
-            name=name, span_attributes={"type": SpanTypeAttribute.SCORE}, input=dict(**kwargs)
+            name=name,
+            span_attributes={"type": SpanTypeAttribute.SCORE, "purpose": "scorer"},
+            propagated_event=merged_propagated,
+            input=dict(**kwargs),
         ) as span:
             score = scorer
             if hasattr(scorer, "eval_async"):
@@ -1550,9 +1559,9 @@ def build_local_summary(
     scores_by_name = defaultdict(lambda: (0, 0))
     for result in results:
         for name, score in result.scores.items():
-            curr = scores_by_name[name]
-            if curr is None:
+            if score is None:
                 continue
+            curr = scores_by_name[name]
             scores_by_name[name] = (curr[0] + score, curr[1] + 1)
     longest_score_name = max(len(name) for name in scores_by_name) if scores_by_name else 0
     avg_scores = {

braintrust/logger.py

@@ -9,7 +9,6 @@ import inspect
 import io
 import json
 import logging
-import math
 import os
 import sys
 import textwrap
@@ -46,7 +45,7 @@ from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
 from . import context, id_gen
-from .bt_json import bt_dumps, bt_loads
+from .bt_json import bt_dumps, bt_safe_deep_copy
 from .db_fields import (
     ASYNC_SCORING_CONTROL_FIELD,
     AUDIT_METADATA_FIELD,
@@ -271,6 +270,10 @@ class _NoopSpan(Span):
     def id(self):
         return ""
 
+    @property
+    def propagated_event(self):
+        return None
+
     def log(self, **event: Any):
         pass
 
@@ -739,13 +742,6 @@ def construct_logs3_data(items: Sequence[str]):
     return '{"rows": ' + rowsS + ', "api_version": ' + str(DATA_API_VERSION) + "}"
 
 
-def _check_json_serializable(event):
-    try:
-        return bt_dumps(event)
-    except (TypeError, ValueError) as e:
-        raise Exception(f"All logged values must be JSON-serializable: {event}") from e
-
-
 class _MaskingError:
     """Internal class to signal masking errors that need special handling."""
 
@@ -795,6 +791,7 @@ class _MemoryBackgroundLogger(_BackgroundLogger):
         self.lock = threading.Lock()
         self.logs = []
         self.masking_function: Callable[[Any], Any] | None = None
+        self.upload_attempts: list[BaseAttachment] = []  # Track upload attempts
 
     def enforce_queue_size_limit(self, enforce: bool) -> None:
         pass
@@ -808,7 +805,21 @@ class _MemoryBackgroundLogger(_BackgroundLogger):
         self.masking_function = masking_function
 
     def flush(self, batch_size: int | None = None):
-        pass
+        """Flush the memory logger, extracting attachments and tracking upload attempts."""
+        with self.lock:
+            if not self.logs:
+                return
+
+            # Unwrap lazy values and extract attachments
+            logs = [l.get() for l in self.logs]
+
+            # Extract attachments from all logs
+            attachments: list[BaseAttachment] = []
+            for log in logs:
+                _extract_attachments(log, attachments)
+
+            # Track upload attempts (don't actually call upload() in tests)
+            self.upload_attempts.extend(attachments)
 
     def pop(self):
         with self.lock:
@@ -1959,24 +1970,14 @@ def get_span_parent_object(
 
 def _try_log_input(span, f_sig, f_args, f_kwargs):
     if f_sig:
-        bound_args = f_sig.bind(*f_args, **f_kwargs).arguments
-        input_serializable = bound_args
+        input_data = f_sig.bind(*f_args, **f_kwargs).arguments
     else:
-        input_serializable = dict(args=f_args, kwargs=f_kwargs)
-    try:
-        _check_json_serializable(input_serializable)
-    except Exception as e:
-        input_serializable = "<input not json-serializable>: " + str(e)
-    span.log(input=input_serializable)
+        input_data = dict(args=f_args, kwargs=f_kwargs)
+    span.log(input=input_data)
 
 
 def _try_log_output(span, output):
-    output_serializable = output
-    try:
-        _check_json_serializable(output)
-    except Exception as e:
-        output_serializable = "<output not json-serializable>: " + str(e)
-    span.log(output=output_serializable)
+    span.log(output=output)
 
 
 F = TypeVar("F", bound=Callable[..., Any])
@@ -2426,91 +2427,6 @@ def _validate_and_sanitize_experiment_log_full_args(event: Mapping[str, Any], ha
     return event
 
 
-def _deep_copy_event(event: Mapping[str, Any]) -> dict[str, Any]:
-    """
-    Creates a deep copy of the given event. Replaces references to user objects
-    with placeholder strings to ensure serializability, except for `Attachment`
-    and `ExternalAttachment` objects, which are preserved and not deep-copied.
-
-    Handles circular references and excessive nesting depth to prevent
-    RecursionError during serialization.
-    """
-    # Maximum depth to prevent hitting Python's recursion limit
-    # Python's default limit is ~1000, we use a conservative limit
-    # to account for existing call stack usage from pytest, application code, etc.
-    MAX_DEPTH = 200
-
-    # 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)
-
-        if isinstance(v, Span):
-            return "<span>"
-        elif isinstance(v, Experiment):
-            return "<experiment>"
-        elif isinstance(v, Dataset):
-            return "<dataset>"
-        elif isinstance(v, Logger):
-            return "<logger>"
-        elif isinstance(v, BaseAttachment):
-            return v
-        elif isinstance(v, ReadonlyAttachment):
-            return v.reference
-        elif isinstance(v, float):
-            # Handle NaN and Infinity for JSON compatibility
-            if math.isnan(v):
-                return "NaN"
-            elif math.isinf(v):
-                return "Infinity" if v > 0 else "-Infinity"
-            return v
-        elif isinstance(v, (int, str, bool)) or v is None:
-            # Skip roundtrip for primitive types.
-            return v
-        else:
-            # 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.
-            return bt_loads(bt_dumps(v))
-
-    return _deep_copy_object(event)
-
-
 class ObjectIterator(Generic[T]):
     def __init__(self, refetch_fn: Callable[[], Sequence[T]]):
         self.refetch_fn = refetch_fn
@@ -3060,7 +2976,7 @@ def _log_feedback_impl(
     metadata = update_event.pop("metadata")
     update_event = {k: v for k, v in update_event.items() if v is not None}
 
-    update_event = _deep_copy_event(update_event)
+    update_event = bt_safe_deep_copy(update_event)
 
     def parent_ids():
         exporter = _get_exporter()
@@ -3116,7 +3032,7 @@ def _update_span_impl(
         event=event,
     )
 
-    update_event = _deep_copy_event(update_event)
+    update_event = bt_safe_deep_copy(update_event)
 
     def parent_ids():
         exporter = _get_exporter()
@@ -3936,14 +3852,10 @@ class SpanImpl(Span):
             **{IS_MERGE_FIELD: self._is_merge},
         )
 
-        serializable_partial_record = _deep_copy_event(partial_record)
-        _check_json_serializable(serializable_partial_record)
+        serializable_partial_record = bt_safe_deep_copy(partial_record)
         if serializable_partial_record.get("metrics", {}).get("end") is not None:
             self._logged_end_time = serializable_partial_record["metrics"]["end"]
 
-        if len(serializable_partial_record.get("tags", [])) > 0 and self.span_parents:
-            raise Exception("Tags can only be logged to the root span")
-
         def compute_record() -> dict[str, Any]:
             exporter = _get_exporter()
             return dict(
@@ -4304,8 +4216,7 @@ class Dataset(ObjectFetcher[DatasetEvent]):
             args[IS_MERGE_FIELD] = True
             args = _filter_none_args(args)  # If merging, then remove None values to prevent null value writes
 
-        _check_json_serializable(args)
-        args = _deep_copy_event(args)
+        args = bt_safe_deep_copy(args)
 
         def compute_args() -> dict[str, Any]:
             return dict(
@@ -4408,8 +4319,7 @@ class Dataset(ObjectFetcher[DatasetEvent]):
                 "_object_delete": True,  # XXX potentially place this in the logging endpoint
             },
         )
-        _check_json_serializable(partial_args)
-        partial_args = _deep_copy_event(partial_args)
+        partial_args = bt_safe_deep_copy(partial_args)
 
         def compute_args():
             return dict(

braintrust/otel/__init__.py

@@ -90,18 +90,13 @@ class AISpanProcessor:
     def _should_keep_filtered_span(self, span):
         """
         Keep spans if:
-        1. It's a root span (no parent)
-        2. Custom filter returns True/False (if provided)
-        3. Span name starts with 'gen_ai.', 'braintrust.', 'llm.', 'ai.', or 'traceloop.'
-        4. Any attribute name starts with those prefixes
+        1. Custom filter returns True/False (if provided)
+        2. Span name starts with 'gen_ai.', 'braintrust.', 'llm.', 'ai.', or 'traceloop.'
+        3. Any attribute name starts with those prefixes
         """
         if not span:
             return False
 
-        # Braintrust requires root spans, so always keep them
-        if span.parent is None:
-            return True
-
         # Apply custom filter if provided
         if self._custom_filter:
             custom_result = self._custom_filter(span)
@@ -384,6 +379,9 @@ def _get_braintrust_parent(object_type, object_id: str | None = None, compute_ar
 
     return None
 
+def is_root_span(span) -> bool:
+    """Returns True if the span is a root span (no parent span)."""
+    return getattr(span, "parent", None) is None
 
 def context_from_span_export(export_str: str):
     """
@@ -522,15 +520,17 @@ def add_span_parent_to_baggage(span, ctx=None):
     return add_parent_to_baggage(parent_value, ctx=ctx)
 
 
-def parent_from_headers(headers: dict[str, str]) -> str | None:
+def parent_from_headers(headers: dict[str, str], propagator=None) -> str | None:
     """
-    Extract a Braintrust-compatible parent string from W3C Trace Context headers.
+    Extract a Braintrust-compatible parent string from trace context headers.
 
-    This converts OTEL trace context headers (traceparent/baggage) into a format
-    that can be passed as the 'parent' parameter to Braintrust's start_span() method.
+    This converts OTEL trace context headers into a format that can be passed
+    as the 'parent' parameter to Braintrust's start_span() method.
 
     Args:
-        headers: Dictionary with 'traceparent' and optionally 'baggage' keys
+        headers: Dictionary with trace context headers (e.g., 'traceparent'/'baggage' for W3C)
+        propagator: Optional custom TextMapPropagator. If not provided, uses the
+                   globally registered propagator (W3C TraceContext by default).
 
     Returns:
         Braintrust V4 export string that can be used as parent parameter,
@@ -545,6 +545,12 @@ def parent_from_headers(headers: dict[str, str]) -> str | None:
         >>> parent = parent_from_headers(headers)
         >>> with project.start_span(name="service_c", parent=parent) as span:
         >>>     span.log(input="BT span as child of OTEL parent")
+
+        >>> # Using a custom propagator (e.g., B3 format)
+        >>> from opentelemetry.propagators.b3 import B3MultiFormat
+        >>> propagator = B3MultiFormat()
+        >>> headers = {'X-B3-TraceId': '...', 'X-B3-SpanId': '...', 'baggage': '...'}
+        >>> parent = parent_from_headers(headers, propagator=propagator)
     """
     if not OTEL_AVAILABLE:
         raise ImportError(INSTALL_ERR_MSG)
@@ -553,8 +559,11 @@ def parent_from_headers(headers: dict[str, str]) -> str | None:
     from opentelemetry import baggage, trace
     from opentelemetry.propagate import extract
 
-    # Extract context from headers using W3C Trace Context propagator
-    ctx = extract(headers)
+    # Extract context from headers using provided propagator or global propagator
+    if propagator is not None:
+        ctx = propagator.extract(headers)
+    else:
+        ctx = extract(headers)
 
     # Get span from context
     span = trace.get_current_span(ctx)