fix(typesafe): gate image capture behind env var

py/src/braintrust/env.py

@@ -228,6 +228,7 @@ class BraintrustEnv:
     FAILED_PUBLISH_PAYLOADS_DIR = EnvVar("BRAINTRUST_FAILED_PUBLISH_PAYLOADS_DIR", EnvParser.STRING)
     ALL_PUBLISH_PAYLOADS_DIR = EnvVar("BRAINTRUST_ALL_PUBLISH_PAYLOADS_DIR", EnvParser.STRING)
     DISABLE_ATEXIT_FLUSH = EnvVar("BRAINTRUST_DISABLE_ATEXIT_FLUSH", EnvParser.BOOL)
+    CAPTURE_ATTACHMENTS = EnvVar("BRAINTRUST_CAPTURE_ATTACHMENTS", EnvParser.BOOL)
     CAPTURE_USER_AUDIO_ATTACHMENTS = EnvVar("BRAINTRUST_CAPTURE_USER_AUDIO_ATTACHMENTS", EnvParser.BOOL)
     CAPTURE_AGENT_AUDIO_ATTACHMENTS = EnvVar("BRAINTRUST_CAPTURE_AGENT_AUDIO_ATTACHMENTS", EnvParser.BOOL)
     MAX_GENERATOR_ITEMS = EnvVar("BRAINTRUST_MAX_GENERATOR_ITEMS", EnvParser.INT)

py/src/braintrust/integrations/test_decisions.py

@@ -254,7 +254,8 @@ def test_decisions_nested_media_and_existing_references():
     assert questions["bounds"]["instructions"]["image"] == image
 
 
-def test_decisions_media_transform_handles_serialized_images():
+def test_decisions_media_transform_handles_serialized_images(monkeypatch):
+    monkeypatch.setenv("BRAINTRUST_CAPTURE_ATTACHMENTS", "true")
     data_url = decisions_image_input()[0]["content"][1]["image_url"]
     image = {"type": "image", "format": "png", "data_base64": data_url.split(",", 1)[1]}
     questions = {"bounds": {"type": "box", "instructions": "The target"}}

py/src/braintrust/integrations/typesafe/test_typesafe.py

@@ -6,6 +6,7 @@
 from braintrust import Attachment, logger
 from braintrust.integrations.test_utils import assert_metrics_are_valid, verify_autoinstrument_script
 from braintrust.integrations.typesafe import setup_typesafe, wrap_typesafe
+from braintrust.integrations.typesafe.tracing import _span_input
 from braintrust.span_types import SpanTypeAttribute
 from braintrust.test_helpers import init_test_logger
 from typesafe_sdk import AsyncTypeSafeClient, Choice, Noul, Score, TypeSafeClient, TypeSafeError
@@ -133,7 +134,8 @@ def test_auto_instrument_typesafe_subprocess():
 @pytest.mark.asyncio
 @pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
 @pytest.mark.vcr
-async def test_typesafe_boxes_and_images(memory_logger, is_async):
+async def test_typesafe_boxes_and_images(memory_logger, is_async, monkeypatch):
+    monkeypatch.setenv("BRAINTRUST_CAPTURE_ATTACHMENTS", "true")
     if not hasattr(typesafe_sdk, "Box"):
         pytest.skip("Box and Image require TypeSafe SDK >=0.7.5")
 
@@ -213,8 +215,21 @@ class CustomResponse(BaseModel):
     assert "classifications" not in span
 
 
+def test_typesafe_images_are_omitted_by_default(monkeypatch):
+    if not hasattr(typesafe_sdk, "Image"):
+        pytest.skip("Image requires TypeSafe SDK >=0.7.5")
+
+    monkeypatch.delenv("BRAINTRUST_CAPTURE_ATTACHMENTS", raising=False)
+    image = typesafe_sdk.Image(Path(__file__).parents[2] / "fixtures" / "decision-box.png")
+
+    span_input = _span_input({"photo": image}, {}, {})
+
+    assert span_input["state"]["photo"] == "<omitted>"
+
+
 @pytest.mark.vcr
-def test_typesafe_nested_image_assets(memory_logger):
+def test_typesafe_nested_image_assets(memory_logger, monkeypatch):
+    monkeypatch.setenv("BRAINTRUST_CAPTURE_ATTACHMENTS", "true")
     if not hasattr(typesafe_sdk, "Box"):
         pytest.skip("Box and Image require TypeSafe SDK >=0.7.5")
 

py/src/braintrust/integrations/typesafe/tracing.py

@@ -6,6 +6,7 @@
 from collections.abc import Mapping
 from typing import Any
 
+from braintrust.env import BraintrustEnv
 from braintrust.integrations.question import (
     _ANSWER_FIELDS,
     _QUESTION_FIELDS,
@@ -24,6 +25,7 @@
 
 _INSTRUMENTATION = "typesafe"
 _MISSING = object()
+_OMITTED = "<omitted>"
 
 
 def start_span(*args, **kwargs):
@@ -83,7 +85,18 @@ def _image_payload(value: Any) -> Any:
     # Only interpret TypeSafe's Image objects as files/buffers. Arbitrary JSON
     # strings in the request must never be treated as local paths.
     image_type = getattr(sys.modules.get("typesafe_sdk"), "Image", ())
-    if isinstance(value, image_type):
+    is_typesafe_image = isinstance(value, image_type)
+    is_encoded_image = (
+        isinstance(value, Mapping)
+        and "data_base64" in value
+        and value.get("format") in ("png", "jpg", "jpeg")
+        and isinstance(value["data_base64"], str)
+    )
+    is_data_url = isinstance(value, str) and value.startswith("data:image/")
+    if (is_typesafe_image or is_encoded_image or is_data_url) and not BraintrustEnv.CAPTURE_ATTACHMENTS.get(False):
+        return _OMITTED
+
+    if is_typesafe_image:
         data = value.data
         if isinstance(data, os.PathLike):
             with open(data, "rb") as image_file:
@@ -101,14 +114,13 @@ def _image_payload(value: Any) -> Any:
             resolved = _materialize_attachment(data, mime_type=mime_type)
             if resolved is not None:
                 return resolved.attachment
-    if isinstance(value, Mapping) and "data_base64" in value:
+    if is_encoded_image:
         image_format = value.get("format")
-        if image_format in ("png", "jpg", "jpeg") and isinstance(value["data_base64"], str):
-            resolved = _materialize_attachment(
-                value["data_base64"], mime_type="image/png" if image_format == "png" else "image/jpeg"
-            )
-            if resolved is not None:
-                return resolved.attachment
+        resolved = _materialize_attachment(
+            value["data_base64"], mime_type="image/png" if image_format == "png" else "image/jpeg"
+        )
+        if resolved is not None:
+            return resolved.attachment
     return value