wip

Author: realark

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

@@ -0,0 +1,37 @@
+package dev.braintrust.trace;
+
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import java.util.List;
+
+/**
+ * A filter that decides whether a finalized span should be exported. Samplers are evaluated in
+ * {@link BraintrustSpanProcessor#onEnd} after all attributes have been set on the span.
+ */
+interface BraintrustSampler {
+    /** Returns {@code true} if the span should be exported, {@code false} to discard it. */
+    boolean sample(ReadableSpan span);
+
+    /** Keeps only spans that have at least one attribute with a known AI instrumentation prefix. */
+    class FilterAISpans implements BraintrustSampler {
+        static final List<String> AI_ATTR_PREFIXES =
+                List.of("gen_ai.", "braintrust.", "llm.", "ai.", "traceloop.");
+
+        @Override
+        public boolean sample(ReadableSpan span) {
+            var attributes = span.toSpanData().getAttributes();
+            for (var key : attributes.asMap().keySet()) {
+                var keyName = key.getKey();
+                // Skip internal attributes injected by the span processor itself
+                if (keyName.equals(BraintrustTracing.PARENT_KEY)) {
+                    continue;
+                }
+                for (var prefix : AI_ATTR_PREFIXES) {
+                    if (keyName.startsWith(prefix)) {
+                        return true;
+                    }
+                }
+            }
+            return false;
+        }
+    }
+}

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

@@ -7,6 +7,7 @@
 import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SpanProcessor;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -26,11 +27,21 @@ public class BraintrustSpanProcessor implements SpanProcessor {
 
     private final BraintrustConfig config;
     private final SpanProcessor delegate;
+    private final List<BraintrustSampler> samplers;
     private final ConcurrentMap<String, ParentContext> parentContexts = new ConcurrentHashMap<>();
 
     BraintrustSpanProcessor(BraintrustConfig config, SpanProcessor delegate) {
         this.config = config;
         this.delegate = delegate;
+        this.samplers = buildSamplers(config);
+    }
+
+    private static List<BraintrustSampler> buildSamplers(BraintrustConfig config) {
+        var samplers = new java.util.ArrayList<BraintrustSampler>();
+        if (config.filterAISpans()) {
+            samplers.add(new BraintrustSampler.FilterAISpans());
+        }
+        return List.copyOf(samplers);
     }
 
     @Override
@@ -93,12 +104,21 @@ public void onEnd(ReadableSpan span) {
         if (config.debug()) {
             logSpanDetails(span);
         }
+        for (var sampler : samplers) {
+            if (!sampler.sample(span)) {
+                log.debug(
+                        "Span filtered by {}: {}",
+                        sampler.getClass().getSimpleName(),
+                        span.getName());
+                return;
+            }
+        }
         delegate.onEnd(span);
     }
 
     @Override
     public boolean isEndRequired() {
-        return delegate.isEndRequired();
+        return true;
     }
 
     @Override

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

@@ -117,9 +117,6 @@ public static void enable(
                                 .setMaxExportBatchSize(maxExportBatchSize)
                                 .build());
         tracerProviderBuilder.addResource(resource).addSpanProcessor(spanProcessor);
-        if (config.filterAISpans()) {
-            tracerProviderBuilder.setSampler(Samplers.SAMPLE_AI_SPANS_ONLY);
-        }
         // logs
         var logProcessor =
                 BatchLogRecordProcessor.builder(new BraintrustLogExporter(config))

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

@@ -3,8 +3,13 @@
 import static org.junit.jupiter.api.Assertions.*;
 
 import dev.braintrust.TestHarness;
+import dev.braintrust.UnitTestSpanExporter;
+import dev.braintrust.config.BraintrustConfig;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -66,25 +71,42 @@ void spanProcessorAddsParentFromConfig() {
 
     @Test
     void filterAISpansDropsNonAISpans() {
-        testHarness = TestHarness.setup(configBuilder -> configBuilder.filterAISpans(true));
+        // Build a tracer where the BraintrustSpanProcessor wraps a test exporter,
+        // so the filter in onEnd() gates what the exporter receives.
+        var config =
+                BraintrustConfig.builder()
+                        .apiKey("test-key")
+                        .apiUrl("http://localhost:1234")
+                        .filterAISpans(true)
+                        .build();
+
+        var spanExporter = new UnitTestSpanExporter();
+        var processor =
+                new BraintrustSpanProcessor(config, SimpleSpanProcessor.create(spanExporter));
+        var tracerProvider = SdkTracerProvider.builder().addSpanProcessor(processor).build();
+        var tracer = tracerProvider.get("test");
+
+        // Non-AI span — should be filtered out
+        var nonAiSpan = tracer.spanBuilder("http-call").startSpan();
+        nonAiSpan.setAttribute("http.method", "GET");
+        nonAiSpan.end();
 
-        var tracer = testHarness.openTelemetry().getTracer("unit test");
+        // Span with no attributes — should be filtered out
+        var emptySpan = tracer.spanBuilder("empty-span").startSpan();
+        emptySpan.end();
 
-        // Emit a non-AI span — should be dropped by the sampler
-        var nonAiSpan =
-                tracer.spanBuilder("http-call")
-                        .setAttribute(AttributeKey.stringKey("http.method"), "GET")
-                        .startSpan();
-        nonAiSpan.end();
+        // Substring non-match — should be filtered out
+        var substringSpan = tracer.spanBuilder("not-ai").startSpan();
+        substringSpan.setAttribute("zbraintrust.foo", "bar");
+        substringSpan.end();
 
-        // Emit an AI span — should pass the sampler
-        var aiSpan =
-                tracer.spanBuilder("llm-call")
-                        .setAttribute(AttributeKey.stringKey("gen_ai.model"), "gpt-4")
-                        .startSpan();
+        // AI span with attribute set AFTER startSpan — should pass the filter
+        var aiSpan = tracer.spanBuilder("llm-call").startSpan();
+        aiSpan.setAttribute("gen_ai.model", "gpt-4");
         aiSpan.end();
 
-        var spans = testHarness.awaitExportedSpans();
+        tracerProvider.forceFlush().join(5, TimeUnit.SECONDS);
+        var spans = spanExporter.getFinishedSpanItems();
 
         assertEquals(1, spans.size(), "Only the AI span should be exported");
         assertEquals("llm-call", spans.get(0).getName());