support filterAISpans

braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java

@@ -39,6 +39,7 @@ public final class BraintrustConfig extends BaseConfig {
     private final boolean debug = getConfig("BRAINTRUST_DEBUG", false);
     private final Duration requestTimeout =
             Duration.ofSeconds(getConfig("BRAINTRUST_REQUEST_TIMEOUT", 30));
+    private final Boolean filterAISpans = getConfig("BRAINTRUST_FILTER_AI_SPANS", false);
 
     /** Custom SSL context for OTLP exporter. Builder-only field, not backed by envars. */
     private final SSLContext sslContext;
@@ -187,6 +188,11 @@ public Builder requestTimeout(Duration value) {
             return this;
         }
 
+        public Builder filterAISpans(boolean value) {
+            envOverrides.put("BRAINTRUST_FILTER_AI_SPANS", String.valueOf(value));
+            return this;
+        }
+
         public Builder sslContext(SSLContext value) {
             this.sslContext = value;
             return this;

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

@@ -0,0 +1,36 @@
+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) {
+            for (var key : span.getAttributes().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 !samplers.isEmpty() || delegate.isEndRequired();
     }
 
     @Override

braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.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;
 
@@ -64,6 +69,49 @@ void spanProcessorAddsParentFromConfig() {
                 "braintrust.parent should be set from the config's default project name");
     }
 
+    @Test
+    void filterAISpansDropsNonAISpans() {
+        // 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();
+
+        // Span with no attributes — should be filtered out
+        var emptySpan = tracer.spanBuilder("empty-span").startSpan();
+        emptySpan.end();
+
+        // Substring non-match — should be filtered out
+        var substringSpan = tracer.spanBuilder("not-ai").startSpan();
+        substringSpan.setAttribute("zbraintrust.foo", "bar");
+        substringSpan.end();
+
+        // 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();
+
+        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());
+    }
+
     private void doSimpleOtelTrace(Tracer tracer) {
         // use tracer to create a simple trace with a root span and a child span
         var span = tracer.spanBuilder("unit-test-root").startSpan();

braintrust-sdk/src/test/java/dev/braintrust/trace/FilterAISpansTest.java

@@ -0,0 +1,50 @@
+package dev.braintrust.trace;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import org.junit.jupiter.api.Test;
+
+public class FilterAISpansTest {
+
+    private final BraintrustSampler.FilterAISpans filter = new BraintrustSampler.FilterAISpans();
+    private final SdkTracerProvider tracerProvider = SdkTracerProvider.builder().build();
+
+    private ReadableSpan createSpan(String name, String attrKey, String attrValue) {
+        var tracer = tracerProvider.get("test");
+        var spanBuilder = tracer.spanBuilder(name);
+        if (attrKey != null) {
+            spanBuilder.setAttribute(AttributeKey.stringKey(attrKey), attrValue);
+        }
+        var span = spanBuilder.startSpan();
+        span.end();
+        return (ReadableSpan) span;
+    }
+
+    @Test
+    public void shouldKeepSpanWithMatchingPrefix() {
+        var span = createSpan("llm-call", "gen_ai.model", "gpt-4");
+        assertTrue(filter.sample(span));
+    }
+
+    @Test
+    public void shouldDropSpanWithNoMatchingAttributes() {
+        var span = createSpan("http-call", "http.method", "GET");
+        assertFalse(filter.sample(span));
+    }
+
+    @Test
+    public void shouldDropSpanWithEmptyAttributes() {
+        var span = createSpan("empty-span", null, null);
+        assertFalse(filter.sample(span));
+    }
+
+    @Test
+    public void shouldNotMatchSubstring() {
+        var span = createSpan("not-ai", "zbraintrust.foo", "bar");
+        assertFalse(filter.sample(span));
+    }
+}

test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java

@@ -19,6 +19,7 @@
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 import javax.annotation.Nonnull;
 import lombok.Getter;
 import lombok.SneakyThrows;
@@ -49,6 +50,11 @@ public class TestHarness {
     }
 
     public static TestHarness setup() {
+        return setup(cfg -> cfg);
+    }
+
+    public static TestHarness setup(
+            Function<BraintrustConfig.Builder, BraintrustConfig.Builder> configCustomizer) {
         var configBuilder =
                 BraintrustConfig.builder()
                         .apiUrl(vcr.getUrlForTargetBase("https://api.braintrust.dev"))
@@ -60,6 +66,7 @@ public static TestHarness setup() {
                             "BRAINTRUST_API_KEY",
                             "sk-000000000000000000000000000000000000000000000000"));
         }
+        configBuilder = configCustomizer.apply(configBuilder);
         return setup(configBuilder.build());
     }