From c6bbbf7caca6278f57ee490ad1bbb61802f561d1 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 20 Oct 2026 15:58:38 -0700 Subject: [PATCH] fix(typesafe): gate image capture behind env var --- py/src/braintrust/env.py | 1 + .../braintrust/integrations/test_decisions.py | 3 +- .../integrations/typesafe/test_typesafe.py | 19 +++++++++++-- .../integrations/typesafe/tracing.py | 28 +++++++++++++------ 4 files changed, 40 insertions(+), 11 deletions(-) diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index f08afb668..3214b290c 100644 --- a/py/src/braintrust/env.py +++ b/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) diff --git a/py/src/braintrust/integrations/test_decisions.py b/py/src/braintrust/integrations/test_decisions.py index d3a6ee256..80fa4f2d8 100644 --- a/py/src/braintrust/integrations/test_decisions.py +++ b/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"}} diff --git a/py/src/braintrust/integrations/typesafe/test_typesafe.py b/py/src/braintrust/integrations/typesafe/test_typesafe.py index bc8c89dbc..2910c5b92 100644 --- a/py/src/braintrust/integrations/typesafe/test_typesafe.py +++ b/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"] == "" + + @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") diff --git a/py/src/braintrust/integrations/typesafe/tracing.py b/py/src/braintrust/integrations/typesafe/tracing.py index 4d478a31d..7766936ca 100644 --- a/py/src/braintrust/integrations/typesafe/tracing.py +++ b/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 = "" 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