braintrust 0.4.0__tar.gz → 0.4.1__tar.gz

{braintrust-0.4.0 → braintrust-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrust
-Version: 0.4.0
+Version: 0.4.1
 Summary: SDK for integrating Braintrust
 Home-page: https://www.braintrust.dev
 Author: Braintrust

braintrust-0.4.1/src/braintrust/bt_json.py

@@ -0,0 +1,275 @@
+import dataclasses
+import json
+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
+try:
+    import orjson
+
+    _HAS_ORJSON = True
+except ImportError:
+    _HAS_ORJSON = False
+
+
+
+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)).
+    """
+    # 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, v).model_dump(exclude_none=True)
+    except (AttributeError, TypeError):
+        pass
+
+    # Attempt to dump a Pydantic v1 `BaseModel`.
+    try:
+        return cast(Any, v).dict(exclude_none=True)
+    except (AttributeError, TypeError):
+        pass
+
+    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:
+        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.
+
+    This is used as a fallback when orjson is not available or fails.
+    """
+
+    def default(self, o: Any):
+        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: Any, encoder: Encoder | None=_json_encoder, **kwargs: Any) -> str:
+    """
+    Serialize obj to a JSON-formatted string.
+
+    Automatically uses orjson if available for better performance (3-5x faster),
+    with fallback to standard json library if orjson is not installed or fails.
+
+    Args:
+        obj: Object to serialize
+        encoder: Encoder to use, defaults to `_default_encoder`
+        **kwargs: Additional arguments (passed to json.dumps in fallback path)
+
+    Returns:
+        JSON string representation of obj
+    """
+    if _HAS_ORJSON:
+        # Try orjson first for better performance
+        try:
+            # pylint: disable=no-member  # orjson is a C extension, pylint can't introspect it
+            return orjson.dumps(  # type: ignore[possibly-unbound]
+                obj,
+                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")
+        except Exception:
+            # If orjson fails, fall back to standard json
+            pass
+
+    # 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=encoder.native if encoder else None, allow_nan=False, sort_keys=True, **kwargs)
+
+
+def bt_loads(s: str, **kwargs) -> Any:
+    """
+    Deserialize s (a str containing a JSON document) to a Python object.
+
+    Automatically uses orjson if available for better performance (2-3x faster),
+    with fallback to standard json library if orjson is not installed or fails.
+
+    Args:
+        s: JSON string to deserialize
+        **kwargs: Additional arguments (passed to json.loads in fallback path)
+
+    Returns:
+        Python object representation of JSON string
+    """
+    if _HAS_ORJSON:
+        # Try orjson first for better performance
+        try:
+            # pylint: disable=no-member  # orjson is a C extension, pylint can't introspect it
+            return orjson.loads(s)  # type: ignore[possibly-unbound]
+        except Exception:
+            # If orjson fails, fall back to standard json
+            pass
+
+    # Use standard json (either orjson not available or it failed)
+    return json.loads(s, **kwargs)

{braintrust-0.4.0 → braintrust-0.4.1}/src/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"):

{braintrust-0.4.0 → braintrust-0.4.1}/src/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,8 +3852,7 @@ 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"]
 
@@ -4304,8 +4219,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 +4322,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(