convert inline base64 attachments to s3 pointers via span processor

braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentProcessor.java

@@ -0,0 +1,168 @@
+package dev.braintrust.trace;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.fasterxml.jackson.databind.node.TextNode;
+import dev.braintrust.json.BraintrustJsonMapper;
+import java.util.Base64;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import lombok.SneakyThrows;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Scans JSON content for base64 data URI attachments and replaces them with attachment references
+ * after uploading to S3.
+ *
+ * <p>Package-private; not exposed in the public API.
+ */
+@Slf4j
+class AttachmentProcessor {
+    /**
+     * quick heuristic to determine if the json payload contains a base64 encoded file
+     *
+     * <p>This is used for performance reasons as a fail-fast to avoid doing a json parse.
+     */
+    static final Pattern BASE64_DATA_URI_PATTERN =
+            Pattern.compile("data:([\\w/\\-.+]+);base64,([A-Za-z0-9+/=]{20,})");
+
+    private final AttachmentUploader uploader;
+
+    AttachmentProcessor(AttachmentUploader uploader) {
+        this.uploader = uploader;
+    }
+
+    /**
+     * Scans a JSON string for base64 data URIs, uploads them, and returns the modified JSON with
+     * attachment references.
+     *
+     * @param json the JSON string to scan
+     * @return the modified JSON with base64 data replaced by attachment references, or the original
+     *     JSON if no base64 data was found
+     */
+    String processAndUpload(String json) {
+        if (json == null || !BASE64_DATA_URI_PATTERN.matcher(json).find()) {
+            return json;
+        }
+
+        try {
+            JsonNode root = BraintrustJsonMapper.get().readTree(json);
+            AtomicBoolean modified = new AtomicBoolean(false);
+            JsonNode result = replaceBase64Attachments(root, modified);
+            return modified.get() ? BraintrustJsonMapper.get().writeValueAsString(result) : json;
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to process attachments in JSON", e);
+        }
+    }
+
+    // NOTE: not concerned with recursion blowing the stack because we're mutating AI vendor
+    // messages which are not deep enough for this to be a concern
+    private JsonNode replaceBase64Attachments(JsonNode node, AtomicBoolean modified) {
+        if (node.isTextual()) {
+            return replaceInText((TextNode) node, modified);
+        } else if (node.isObject()) {
+            ObjectNode objectNode = (ObjectNode) node;
+            ObjectNode result = objectNode.deepCopy();
+            var fieldNames = objectNode.fieldNames();
+            while (fieldNames.hasNext()) {
+                String fieldName = fieldNames.next();
+                JsonNode child = objectNode.get(fieldName);
+                result.set(fieldName, replaceBase64Attachments(child, modified));
+            }
+            return result;
+        } else if (node.isArray()) {
+            ArrayNode arrayNode = (ArrayNode) node;
+            ArrayNode result = arrayNode.deepCopy();
+            for (int i = 0; i < arrayNode.size(); i++) {
+                result.set(i, replaceBase64Attachments(arrayNode.get(i), modified));
+            }
+            return result;
+        }
+        return node;
+    }
+
+    @SneakyThrows
+    private JsonNode replaceInText(TextNode textNode, AtomicBoolean modified) {
+        String value = textNode.asText();
+        Matcher matcher = BASE64_DATA_URI_PATTERN.matcher(value);
+        if (!matcher.find()) {
+            return textNode;
+        }
+        if (!isEntirelyDataUri(value)) {
+            log.debug("found base64 string but text contained extra content {}", value);
+            return textNode;
+        }
+
+        matcher.reset();
+        StringBuilder sb = new StringBuilder();
+        while (matcher.find()) {
+            String contentType = matcher.group(1);
+            String base64Data = matcher.group(2);
+            byte[] data = Base64.getDecoder().decode(base64Data);
+
+            String extension = contentTypeToExtension(contentType);
+            String filename = "attachment" + extension;
+            AttachmentReference ref = AttachmentReference.create(filename, contentType);
+
+            try {
+                uploader.enqueue(ref, data);
+            } catch (IllegalStateException e) {
+                throw new RuntimeException("Failed to enqueue attachment upload", e);
+            }
+
+            String replacement =
+                    "{\"type\":\"braintrust_attachment\",\"content_type\":\""
+                            + contentType
+                            + "\",\"filename\":\""
+                            + filename
+                            + "\",\"key\":\""
+                            + ref.key()
+                            + "\"}";
+
+            matcher.appendReplacement(sb, Matcher.quoteReplacement(replacement));
+        }
+        matcher.appendTail(sb);
+
+        modified.set(true);
+
+        return BraintrustJsonMapper.get().readTree(sb.toString());
+    }
+
+    static boolean isEntirelyDataUri(String value) {
+        String trimmed = value.trim();
+        return trimmed.startsWith("data:")
+                && !trimmed.contains("\"")
+                && !trimmed.contains("\\")
+                && !trimmed.contains(" ");
+    }
+
+    private static String contentTypeToExtension(String contentType) {
+        switch (contentType.toLowerCase()) {
+            case "image/png":
+                return ".png";
+            case "image/jpeg":
+            case "image/jpg":
+                return ".jpg";
+            case "image/gif":
+                return ".gif";
+            case "image/webp":
+                return ".webp";
+            case "image/svg+xml":
+                return ".svg";
+            case "application/pdf":
+                return ".pdf";
+            case "text/plain":
+                return ".txt";
+            case "application/json":
+                return ".json";
+            default:
+                String[] parts = contentType.split("/");
+                if (parts.length == 2) {
+                    return "." + parts[1].split("[;\\-]")[0];
+                }
+                return "";
+        }
+    }
+}

braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentReference.java

@@ -0,0 +1,33 @@
+package dev.braintrust.trace;
+
+import java.util.Objects;
+import java.util.UUID;
+import javax.annotation.Nonnull;
+
+/**
+ * Represents an attachment reference stored on a span in place of uploaded attachment data.
+ *
+ * <p>Its shape intentionally matches the cross-SDK Braintrust attachment reference format.
+ */
+record AttachmentReference(
+        @Nonnull String type,
+        @Nonnull String filename,
+        @Nonnull String contentType,
+        @Nonnull String key) {
+
+    private static final String DEFAULT_TYPE = "braintrust_attachment";
+
+    /**
+     * Creates an attachment reference with a generated UUID key.
+     *
+     * @param filename the display filename for the attachment
+     * @param contentType the MIME type of the attachment content
+     * @return a new AttachmentReference with a unique key
+     */
+    static AttachmentReference create(@Nonnull String filename, @Nonnull String contentType) {
+        Objects.requireNonNull(filename, "filename cannot be null");
+        Objects.requireNonNull(contentType, "contentType cannot be null");
+        return new AttachmentReference(
+                DEFAULT_TYPE, filename, contentType, UUID.randomUUID().toString());
+    }
+}