diff --git a/py/src/braintrust/otel/__init__.py b/py/src/braintrust/otel/__init__.py index f8f5c121..84946a41 100644 --- a/py/src/braintrust/otel/__init__.py +++ b/py/src/braintrust/otel/__init__.py @@ -92,6 +92,29 @@ def _with_span_origin_attributes(span, environment): return _SpanWithAttributes(span, attributes) +class _SpanOriginProcessor: + """Add Braintrust span origin attributes after any upstream filtering.""" + + def __init__(self, processor, environment): + self._processor = processor + self._environment = environment + + def on_start(self, span, parent_context=None): + self._processor.on_start(span, parent_context) + + def on_end(self, span): + self._processor.on_end(_with_span_origin_attributes(span, self._environment)) + + def _on_ending(self, span): + _forward_on_ending(self._processor, span) + + def shutdown(self): + self._processor.shutdown() + + def force_flush(self, timeout_millis=30000): + return self._processor.force_flush(timeout_millis) + + class AISpanProcessor: """ A span processor that filters spans to only export filtered telemetry. @@ -355,14 +378,13 @@ def __init__( if SpanProcessor is None: SpanProcessor = BatchSpanProcessor - # Always create a BatchSpanProcessor first - processor = SpanProcessor(self._exporter) + # Add provenance downstream so filtering decisions use the span's original attributes. + processor = _SpanOriginProcessor(SpanProcessor(self._exporter), self._environment) if filter_ai_spans: - # Wrap the BatchSpanProcessor with filtering + # Wrap the origin-injecting processor so filtering runs first. self._processor = AISpanProcessor(processor, custom_filter=custom_filter) else: - # Use BatchSpanProcessor directly self._processor = processor def on_start(self, span, parent_context=None): @@ -418,7 +440,7 @@ def _get_parent_otel_braintrust_parent(self, parent_context): def on_end(self, span): """Forward span end events to the inner processor.""" self._exporter.initialize() - self._processor.on_end(_with_span_origin_attributes(span, self._environment)) + self._processor.on_end(span) def _on_ending(self, span): """Forward pre-end hook when the wrapped processor supports it.""" diff --git a/py/src/braintrust/test_otel.py b/py/src/braintrust/test_otel.py index bb4adcbb..2ad206b3 100644 --- a/py/src/braintrust/test_otel.py +++ b/py/src/braintrust/test_otel.py @@ -124,8 +124,11 @@ def test_braintrust_span_processor_merges_span_origin_with_context_json_set_afte memory_exporter = InMemorySpanExporter() provider = TracerProvider() - processor = BraintrustSpanProcessor(api_key="test-api-key", parent="project_name:test") - processor._processor = SimpleSpanProcessor(memory_exporter) + processor = BraintrustSpanProcessor( + api_key="test-api-key", + parent="project_name:test", + SpanProcessor=lambda _exporter: SimpleSpanProcessor(memory_exporter), + ) provider.add_span_processor(processor) tracer = provider.get_tracer("test_tracer") @@ -146,6 +149,43 @@ def test_braintrust_span_processor_merges_span_origin_with_context_json_set_afte provider.shutdown() +def test_braintrust_span_processor_filters_before_adding_span_origin(): + if not _check_otel_installed(): + pytest.skip("OpenTelemetry SDK not fully installed, skipping test") + + from braintrust.otel import BraintrustSpanProcessor + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + memory_exporter = InMemorySpanExporter() + provider = TracerProvider() + processor = BraintrustSpanProcessor( + api_key="test-api-key", + parent="project_name:test", + filter_ai_spans=True, + SpanProcessor=lambda _exporter: SimpleSpanProcessor(memory_exporter), + ) + provider.add_span_processor(processor) + tracer = provider.get_tracer("test_tracer") + + try: + with tracer.start_as_current_span("GET /health") as span: + span.set_attribute("http.route", "/health") + with tracer.start_as_current_span("SELECT users") as span: + span.set_attribute("db.system", "postgresql") + with tracer.start_as_current_span("chat") as span: + span.set_attribute("gen_ai.operation.name", "chat") + + provider.force_flush() + spans = memory_exporter.get_finished_spans() + assert [span.name for span in spans] == ["chat"] + context = json.loads(spans[0].attributes["braintrust.context_json"]) + assert context["span_origin"]["instrumentation"]["name"] == "braintrust-python-otel" + finally: + provider.shutdown() + + def test_merge_span_origin_context_uses_passed_instrumentation_name(): from braintrust.span_origin import merge_span_origin_context