From 1142b565e134ad06c34fb2bd14045d028b268a18 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Tue, 31 Mar 2026 16:32:59 -0600 Subject: [PATCH] feature: filterAISpans --- .../braintrust/config/BraintrustConfig.java | 6 +++ .../braintrust/trace/BraintrustSampler.java | 36 +++++++++++++ .../trace/BraintrustSpanProcessor.java | 22 +++++++- .../trace/BraintrustTracingTest.java | 48 ++++++++++++++++++ .../braintrust/trace/FilterAISpansTest.java | 50 +++++++++++++++++++ .../java/dev/braintrust/TestHarness.java | 7 +++ 6 files changed, 168 insertions(+), 1 deletion(-) create mode 100644 braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java create mode 100644 braintrust-sdk/src/test/java/dev/braintrust/trace/FilterAISpansTest.java diff --git a/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java b/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java index 158cad51..c87f4055 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/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; diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSampler.java new file mode 100644 index 00000000..aa0efb04 --- /dev/null +++ b/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 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; + } + } +} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java index 0ff98c89..532193bd 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java +++ b/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 samplers; private final ConcurrentMap parentContexts = new ConcurrentHashMap<>(); BraintrustSpanProcessor(BraintrustConfig config, SpanProcessor delegate) { this.config = config; this.delegate = delegate; + this.samplers = buildSamplers(config); + } + + private static List buildSamplers(BraintrustConfig config) { + var samplers = new java.util.ArrayList(); + 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 diff --git a/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java b/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java index eff7ccf2..6b3bda06 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/trace/BraintrustTracingTest.java +++ b/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(); diff --git a/braintrust-sdk/src/test/java/dev/braintrust/trace/FilterAISpansTest.java b/braintrust-sdk/src/test/java/dev/braintrust/trace/FilterAISpansTest.java new file mode 100644 index 00000000..993e3e30 --- /dev/null +++ b/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)); + } +} diff --git a/test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java b/test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java index f1096d06..f45bd8ba 100644 --- a/test-harness/src/testFixtures/java/dev/braintrust/TestHarness.java +++ b/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 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()); }