diff --git a/py/src/braintrust/bt_json.py b/py/src/braintrust/bt_json.py index e6035d6b..120e8e4e 100644 --- a/py/src/braintrust/bt_json.py +++ b/py/src/braintrust/bt_json.py @@ -3,6 +3,7 @@ import math import warnings from collections.abc import Callable, Mapping +from functools import cache from typing import Any, NamedTuple, cast, overload @@ -16,9 +17,30 @@ _HAS_ORJSON = False -def _to_bt_safe(v: Any) -> Any: +class _LoggerTypes(NamedTuple): + Span: type + Experiment: type + Dataset: type + Logger: type + BaseAttachment: type + ReadonlyAttachment: type + + +@cache +def _logger_types() -> _LoggerTypes: + # Imported lazily to avoid circular imports, and cached because `_to_bt_safe` runs per logged value. + from braintrust.logger import BaseAttachment, Dataset, Experiment, Logger, ReadonlyAttachment, Span + + return _LoggerTypes(Span, Experiment, Dataset, Logger, BaseAttachment, ReadonlyAttachment) + + +def _to_bt_safe(v: Any, copy_dataclass_field: Callable[[Any, int], Any] | None = None, depth: int = 0) -> Any: """ Converts the object to a Braintrust-safe representation (i.e. Attachment objects are safe (specially handled by background logger)). + + `copy_dataclass_field(value, depth)` converts dataclass field values. `bt_safe_deep_copy` passes its own + walker and the fields' depth so they share its cycle and depth tracking; without it, the dataclass + starts a new `bt_safe_deep_copy`. """ # Fast path: check primitives via type identity before hitting # isinstance checks against abstract classes or Pydantic model_dump. @@ -37,8 +59,7 @@ def _to_bt_safe(v: Any) -> Any: if isinstance(v, (int, str)): return v - # avoid circular imports - from braintrust.logger import BaseAttachment, Dataset, Experiment, Logger, ReadonlyAttachment, Span + Span, Experiment, Dataset, Logger, BaseAttachment, ReadonlyAttachment = _logger_types() if isinstance(v, Span): return "" @@ -59,10 +80,9 @@ def _to_bt_safe(v: Any) -> Any: 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)} + if copy_dataclass_field is None: + return bt_safe_deep_copy(v) + return _dataclass_to_dict(v, copy_dataclass_field, depth) # 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): @@ -100,6 +120,20 @@ def _to_bt_safe(v: Any) -> Any: return bt_loads(bt_dumps(v, encoder=_str_encoder)) +def _dataclass_to_dict(v: Any, copy_field: Callable[[Any, int], Any], depth: int) -> dict[str, Any]: + # 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.) + result = {} + for f in dataclasses.fields(v): + try: + result[f.name] = copy_field(getattr(v, f.name), depth) + except Exception: + # Keep the other fields (and any attachments in them) when one can't be read. + result[f.name] = f"" + return result + + @overload def bt_safe_deep_copy( obj: Mapping[str, Any], @@ -155,43 +189,14 @@ def _deep_copy_object(v: Any, depth: int = 0) -> Any: if depth >= max_depth: return "" - # Fast path for dict (the most common container in log data). - # Uses type identity instead of isinstance(v, Mapping) which is slow. - if tv is dict: - obj_id = id(v) - if obj_id in visited: - return "" - visited_add(obj_id) - try: - result = {} - for k in v: - if type(k) is str: - key_str = k - else: - try: - key_str = str(k) - except Exception: - key_str = f"" - result[key_str] = _deep_copy_object(v[k], depth + 1) - return result - finally: - visited_discard(obj_id) - elif tv is list or tv is tuple: - obj_id = id(v) - if obj_id in visited: - return "" - visited_add(obj_id) - try: - return [_deep_copy_object(x, depth + 1) for x in v] - finally: - visited_discard(obj_id) - # Slow path for non-builtin Mapping/set types. - elif isinstance(v, Mapping): - obj_id = id(v) - if obj_id in visited: - return "" - visited_add(obj_id) - try: + obj_id = id(v) + if obj_id in visited: + return "" + visited_add(obj_id) + try: + # Fast paths for builtin containers (the most common values in log data). + # Type identity checks avoid the slow isinstance(v, Mapping) for them. + if tv is dict or (tv is not list and tv is not tuple and isinstance(v, Mapping)): result = {} for k in v: if type(k) is str: @@ -203,22 +208,17 @@ def _deep_copy_object(v: Any, depth: int = 0) -> Any: key_str = f"" result[key_str] = _deep_copy_object(v[k], depth + 1) return result - finally: - visited_discard(obj_id) - elif isinstance(v, set): - obj_id = id(v) - if obj_id in visited: - return "" - visited_add(obj_id) - try: + if tv is list or tv is tuple or isinstance(v, set): return [_deep_copy_object(x, depth + 1) for x in v] - finally: - visited_discard(obj_id) - try: - return _to_bt_safe(v) - except Exception: - return f"" + try: + # Dataclass fields go back through this walker so nested attachments keep + # their identity and share its cycle/depth tracking. + return _to_bt_safe(v, _deep_copy_object, depth + 1) + except Exception: + return f"" + finally: + visited_discard(obj_id) return _deep_copy_object(obj) @@ -234,15 +234,12 @@ 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): + if isinstance(v, _logger_types().BaseAttachment): v = v.reference return v diff --git a/py/src/braintrust/test_bt_json.py b/py/src/braintrust/test_bt_json.py index a1d4b711..b8a6821f 100644 --- a/py/src/braintrust/test_bt_json.py +++ b/py/src/braintrust/test_bt_json.py @@ -1,14 +1,120 @@ # pyright: reportUnknownVariableType=false # pyright: reportUnknownArgumentType=false # pyright: reportPrivateUsage=false +import dataclasses import json import warnings from typing import Any from unittest import TestCase +import pytest from braintrust.bt_json import bt_dumps, bt_safe_deep_copy -from braintrust.logger import Attachment, ExternalAttachment -from braintrust.test_helpers import with_memory_logger # noqa: F401 +from braintrust.logger import Attachment, ExternalAttachment, ReadonlyAttachment +from braintrust.test_helpers import DataclassAttachment, with_memory_logger # noqa: F401 + + +@dataclasses.dataclass +class _Payload: + value: Any + + +@dataclasses.dataclass +class _DataclassReadonlyAttachment(ReadonlyAttachment): + reference: dict[str, Any] + + def __post_init__(self): + super().__init__(self.reference) + + +def test_dataclass_nested_attachments(): + attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain") + payload = _Payload( + { + "direct": attachment, + "list": [attachment], + "tuple": (attachment,), + "dataclass": _Payload(attachment), + "dataclasses": [_Payload(attachment)], + } + ) + + assert bt_safe_deep_copy(payload) == { + "value": { + "direct": attachment, + "list": [attachment], + "tuple": [attachment], + "dataclass": {"value": attachment}, + "dataclasses": [{"value": attachment}], + } + } + + +@pytest.mark.parametrize( + "base_name,placeholder", + [("_NoopSpan", ""), ("Experiment", ""), ("Dataset", ""), ("Logger", "")], +) +def test_dataclass_sdk_object_subclasses_keep_placeholder(base_name, placeholder): + from braintrust import logger + + # The `secret` field must never be copied into logged events. + obj = dataclasses.make_dataclass( + "DataclassObject", [("secret", str, "token")], bases=(getattr(logger, base_name),) + )() + + assert bt_safe_deep_copy({"direct": obj, "nested": _Payload([obj])}) == { + "direct": placeholder, + "nested": {"value": [placeholder]}, + } + + +def test_dataclass_attachment_subclasses_keep_special_handling(): + attachment = DataclassAttachment("test-key") + reference = { + "type": "braintrust_attachment", + "key": "readonly-key", + "filename": "a.txt", + "content_type": "text/plain", + } + readonly = _DataclassReadonlyAttachment(reference) + + assert bt_safe_deep_copy({"direct": [attachment, readonly], "nested": _Payload([attachment, readonly])}) == { + "direct": [attachment, reference], + "nested": {"value": [attachment, reference]}, + } + + +def test_dataclass_cycles_and_depth_share_walker_state(): + attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain") + payload = _Payload([attachment]) + payload.value.append(payload) + + # Shared (non-circular) references are copied each time; only true cycles are replaced. + assert bt_safe_deep_copy([payload, payload]) == [ + {"value": [attachment, ""]}, + {"value": [attachment, ""]}, + ] + assert bt_safe_deep_copy(_Payload([_Payload("leaf")]), max_depth=2) == {"value": [""]} + + +def test_dataclass_unreadable_field_keeps_readable_fields(): + @dataclasses.dataclass + class Payload: + files: list[Any] + value: int = 0 + unset: int = dataclasses.field(init=False) + + def __getattribute__(self, name): + if name == "value": + raise RuntimeError("Cannot read field") + return super().__getattribute__(name) + + attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain") + + assert bt_safe_deep_copy(Payload([attachment])) == { + "files": [attachment], + "value": "", + "unset": "", + } class TestBTJson(TestCase): @@ -586,36 +692,6 @@ class ModelWithAttachment(BaseModel): except ImportError: self.skipTest("Pydantic not available") - def test_bt_safe_deep_copy_dataclass_with_attachments(self): - """Test that dataclasses with attachments are handled correctly.""" - from dataclasses import dataclass - - attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain") - - @dataclass - class DataWithAttachment: - name: str - value: int - file: Attachment - - data = DataWithAttachment(name="test", value=42, file=attachment) - - result = bt_safe_deep_copy({"data": data}) - - # Dataclasses with Attachment fields are now properly converted by - # recursively applying _to_bt_safe to each field instead of using - # dataclasses.asdict() which would try to deepcopy the Attachment. - self.assertIsInstance(result["data"], dict) - self.assertEqual(result["data"]["name"], "test") - self.assertEqual(result["data"]["value"], 42) - self.assertIs(result["data"]["file"], attachment) - - # Attachments in regular dicts also work fine - dict_with_attachment = {"name": "test", "value": 42, "file": attachment} - result2 = bt_safe_deep_copy({"data": dict_with_attachment}) - self.assertIsInstance(result2["data"], dict) - self.assertIs(result2["data"]["file"], attachment) - def test_bt_safe_deep_copy_circular_in_pydantic_deferred_to_bt_dumps(self): """Test that circular references inside Pydantic model results bypass bt_safe_deep_copy detection. diff --git a/py/src/braintrust/test_helpers.py b/py/src/braintrust/test_helpers.py index 8b0086b7..f592677a 100644 --- a/py/src/braintrust/test_helpers.py +++ b/py/src/braintrust/test_helpers.py @@ -1,9 +1,10 @@ +import dataclasses import os from contextlib import contextmanager import pytest from braintrust import logger -from braintrust.logger import ObjectMetadata, OrgProjectMetadata, ProjectExperimentMetadata +from braintrust.logger import BaseAttachment, ObjectMetadata, OrgProjectMetadata, ProjectExperimentMetadata from braintrust.util import LazyValue @@ -129,6 +130,32 @@ def find_span_by_name(spans, name): raise AssertionError(f"Expected span named {name!r}. Available spans: {available_names}") +@dataclasses.dataclass +class DataclassAttachment(BaseAttachment): + """A `BaseAttachment` subclass that is also a dataclass, for serialization tests.""" + + key: str + + @property + def reference(self): + return { + "type": "braintrust_attachment", + "key": self.key, + "filename": "file.txt", + "content_type": "text/plain", + } + + @property + def data(self): + return b"data" + + def upload(self): + raise AssertionError("Tests must not upload attachments") + + def debug_info(self): + return {"key": self.key} + + def init_test_exp(experiment_name: str, project_name: str = None): """ Initialize an experiment for testing with fake project and experiment metadata. diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 19dc3add..48e4d416 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -48,6 +48,7 @@ from braintrust.prompt_cache.parameters_cache import ParametersCache from braintrust.prompt_cache.prompt_cache import PromptCache from braintrust.test_helpers import ( + DataclassAttachment, assert_dict_matches, assert_logged_out, init_test_exp, @@ -3426,6 +3427,30 @@ def test_attachment_upload_tracked_on_flush(with_memory_logger, with_simulate_lo assert with_memory_logger.upload_attempts[0] is attachment +def test_dataclass_attachments_reach_flush(with_memory_logger, with_simulate_login): + from dataclasses import dataclass + + @dataclass + class Payload: + files: list[BaseAttachment] + + attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain") + external = ExternalAttachment(url="s3://bucket/key", filename="file.pdf", content_type="application/pdf") + dataclass_attachment = DataclassAttachment("test-key") + payload = Payload([attachment, external]) + test_logger = init_test_logger(__name__) + with test_logger.start_span(name="dataclass-attachments") as span: + span.log(input=payload, output=dataclass_attachment) + test_logger.flush() + + # The memory logger exercises extraction without making upload requests. + assert with_memory_logger.upload_attempts == [attachment, external, dataclass_attachment] + logged = with_memory_logger.pop()[0] + assert logged["input"] == {"files": [attachment.reference, external.reference]} + assert logged["output"] == dataclass_attachment.reference + assert payload.files == [attachment, external] + + def test_same_attachment_logged_twice_tracked_twice(with_memory_logger, with_simulate_login): """Test that same attachment logged twice appears twice in upload attempts.""" attachment = Attachment(data=b"data", filename="file.txt", content_type="text/plain")