fix(sdk): support overriding the span processor on_end hook (#2947)

aryehklein · nirga · web-flow · commit 6e7a123bf62b · 2025-05-16T16:33:07.000+02:00
Co-authored-by: Nir Gazit &lt;nirga@users.noreply.github.com&gt;
diff --git a/packages/traceloop-sdk/tests/cassettes/test_sdk_initialization/test_span_postprocess_callback.yaml b/packages/traceloop-sdk/tests/cassettes/test_sdk_initialization/test_span_postprocess_callback.yaml
@@ -0,0 +1,199 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": "Tell me a joke about opentelemetry"}],
+      "model": "gpt-3.5-turbo"}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '107'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+      user-agent:
+      - OpenAI/Python 1.12.0
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - 'false'
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.12.0
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.9.5
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAAA1RRTUsDMRC9768Yc/HSyrbrWu1FVLQKfoKgIFLS7OxubJKJySxapP9d0tYWL3N4
+        L+/Nm5efDEDoSoxBqFayst70j78+n267x9Hrw81DPRnd39QXd9dXz5PPb3t2KXpJQbMPVPynOlBk
+        vUHW5Na0CigZk+tglJ/kRTksihVhqUKTZI3nfnFQ9rkLM+rng2G5UbakFUYxhrcMAOBnNVNGV+G3
+        GEPe+0MsxigbFOPtIwARyCREyBh1ZOlY9HakIsfoVrFf2gVUugJuEXygJkhrMcAsoJxD5+FLcwvk
+        0TEatMhhcQrnqGQXETSDos5Ubp8hIlpgAg5SIUhoiCpQ5ByqVMee2GxfbmMbanygWTrRdcZs8Vo7
+        HdtpQBnJpYiRya/lywzgfVVP9+9i4QNZz1OmObpkOCjXdmL3ITtyONyQTCzNDi9G2SafiIvIaKe1
+        dg0GH/S6q9pPj48G5ZE8OZS5yJbZLwAAAP//AwAUd8GRNQIAAA==
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 85c044a25bdc0d6e-MXP
+      Cache-Control:
+      - no-cache, must-revalidate
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Tue, 27 Feb 2024 12:00:34 GMT
+      Server:
+      - cloudflare
+      Set-Cookie:
+      - __cf_bm=J7CW3gXyUth9bXxP62KOXJBT9fqPtQCn6rOtIOYSOms-1709035234-1.0-Af3zRwkM02ElV8pGlA2ndZpn9K5kxgF0BOjGNlpNg3Dv/qKUKDqX5KjIeb/o2pyy3ZD0WS15+EM040L1eN/yQ4s=;
+        path=/; expires=Tue, 27-Feb-24 12:30:34 GMT; domain=.api.openai.com; HttpOnly;
+        Secure; SameSite=None
+      - _cfuvid=tr0R1O18cmcX2Q2PFjB12pv_1Mu4R.MQAHgSyTr7lfY-1709035234516-0.0-604800000;
+        path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
+      Transfer-Encoding:
+      - chunked
+      access-control-allow-origin:
+      - '*'
+      alt-svc:
+      - h3=":443"; ma=86400
+      openai-model:
+      - gpt-3.5-turbo-0125
+      openai-organization:
+      - traceloop
+      openai-processing-ms:
+      - '410'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=15724800; includeSubDomains
+      x-ratelimit-limit-requests:
+      - '5000'
+      x-ratelimit-limit-tokens:
+      - '160000'
+      x-ratelimit-remaining-requests:
+      - '4999'
+      x-ratelimit-remaining-tokens:
+      - '159974'
+      x-ratelimit-reset-requests:
+      - 12ms
+      x-ratelimit-reset-tokens:
+      - 9ms
+      x-request-id:
+      - req_008b3333db5c78b9ded9415f11929844
+    status:
+      code: 200
+      message: OK
+- request:
+    body: '{"messages": [{"role": "user", "content": "Tell me a joke about opentelemetry"}],
+      "model": "gpt-3.5-turbo"}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '107'
+      content-type:
+      - application/json
+      cookie:
+      - __cf_bm=J7CW3gXyUth9bXxP62KOXJBT9fqPtQCn6rOtIOYSOms-1709035234-1.0-Af3zRwkM02ElV8pGlA2ndZpn9K5kxgF0BOjGNlpNg3Dv/qKUKDqX5KjIeb/o2pyy3ZD0WS15+EM040L1eN/yQ4s=;
+        _cfuvid=tr0R1O18cmcX2Q2PFjB12pv_1Mu4R.MQAHgSyTr7lfY-1709035234516-0.0-604800000
+      host:
+      - api.openai.com
+      user-agent:
+      - OpenAI/Python 1.12.0
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - 'false'
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.12.0
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.9.5
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAAA1RRXWsbMRB8v1+x1bMdzj5favslEEpo6DcEUmiLkaX1nWKdVl2t45jg/15059j0
+        RYiZnWFm97UAUM6qJSjTajFd9OP5/u+PL3d384eX5uHD/fb7x2f++s192n9uqsefapQVtH5CI2+q
+        K0Nd9CiOwkAbRi2YXSfvy0VZ1dNq1hMdWfRZ1kQZV1f1WHa8pnE5mdYnZUvOYFJL+FUAALz2b84Y
+        LL6oJZSjN6TDlHSDankeAlBMPiNKp+SS6CBqdCENBcHQx35sD2CdBWkRKGIQ9Nih8AEsPqOniAxr
+        dqEBDV5biwxCsCfe3sDvcItG7xJm9QFa1GyHb0C0aPOksDY97xiSYEzv1CnH8VzAUxOZ1rls2Hl/
+        xjcuuNSuGHWikMMmoTjIjwXAn35Ru/+6q8jURVkJbTFkw0k92KnLaS7kdHYihUT7C14tilM+lQ5J
+        sFttXGiQI7tha5u4ml9P6mu9mOlSFcfiHwAAAP//AwB+6qFIPwIAAA==
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 85c044a7edf50d6e-MXP
+      Cache-Control:
+      - no-cache, must-revalidate
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Tue, 27 Feb 2024 12:00:35 GMT
+      Server:
+      - cloudflare
+      Transfer-Encoding:
+      - chunked
+      access-control-allow-origin:
+      - '*'
+      alt-svc:
+      - h3=":443"; ma=86400
+      openai-model:
+      - gpt-3.5-turbo-0125
+      openai-organization:
+      - traceloop
+      openai-processing-ms:
+      - '750'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=15724800; includeSubDomains
+      x-ratelimit-limit-requests:
+      - '5000'
+      x-ratelimit-limit-tokens:
+      - '160000'
+      x-ratelimit-remaining-requests:
+      - '4999'
+      x-ratelimit-remaining-tokens:
+      - '159974'
+      x-ratelimit-reset-requests:
+      - 12ms
+      x-ratelimit-reset-tokens:
+      - 9ms
+      x-request-id:
+      - req_bf1ec13b40c7a03dd17dcca80cbd9031
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/packages/traceloop-sdk/tests/conftest.py b/packages/traceloop-sdk/tests/conftest.py
@@ -1,13 +1,16 @@
 """Unit tests configuration module."""
 
 import os
+import re
 import pytest
 from traceloop.sdk import Traceloop
 from traceloop.sdk.instruments import Instruments
 from traceloop.sdk.tracing.tracing import TracerWrapper
-from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, BatchSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from opentelemetry.context import attach, Context
+from opentelemetry.sdk.trace import ReadableSpan
+
 pytest_plugins = []
 
 
@@ -70,6 +73,44 @@ def on_start(self, span, parent_context=None):
         TracerWrapper.instance = _trace_wrapper_instance
 
 
+@pytest.fixture(scope="function")
+def exporter_with_custom_span_postprocess_callback(exporter):
+
+    if hasattr(TracerWrapper, "instance"):
+        _trace_wrapper_instance = TracerWrapper.instance
+        del TracerWrapper.instance
+
+    def span_postprocess_callback(span: ReadableSpan) -> None:
+        prompt_pattern = re.compile(r"gen_ai\.prompt\.\d+\.content$")
+        completion_pattern = re.compile(r"gen_ai\.completion\.\d+\.content$")
+        if hasattr(span, "_attributes"):
+            attributes = span._attributes if span._attributes else {}
+            # Find and encode all matching attributes
+            for key, value in attributes.items():
+                if (prompt_pattern.match(key) or completion_pattern.match(key)) and isinstance(value, str):
+                    attributes[key] = "REDACTED"  # Modify the attributes directly
+
+    Traceloop.init(
+        exporter=exporter,
+        span_postprocess_callback=span_postprocess_callback,
+    )
+
+    yield exporter
+
+    if hasattr(TracerWrapper, "instance"):
+        # Get the span processor
+        if hasattr(TracerWrapper.instance, "_TracerWrapper__spans_processor"):
+            span_processor = TracerWrapper.instance._TracerWrapper__spans_processor
+            # Reset the on_end method to its original class implementation.
+            # This is needed to make this test run in isolation as SpanProcessor is a singleton.
+            if isinstance(span_processor, SimpleSpanProcessor):
+                span_processor.on_end = SimpleSpanProcessor.on_end.__get__(span_processor, SimpleSpanProcessor)
+            elif isinstance(span_processor, BatchSpanProcessor):
+                span_processor.on_end = BatchSpanProcessor.on_end.__get__(span_processor, BatchSpanProcessor)
+    if _trace_wrapper_instance:
+        TracerWrapper.instance = _trace_wrapper_instance
+
+
 @pytest.fixture
 def exporter_with_custom_instrumentations():
     # Clear singleton if existed
diff --git a/packages/traceloop-sdk/tests/test_sdk_initialization.py b/packages/traceloop-sdk/tests/test_sdk_initialization.py
@@ -33,6 +33,19 @@ def run_workflow():
     assert workflow_span.attributes["custom_span"] == "yes"
 
 
+@pytest.mark.vcr
+def test_span_postprocess_callback(exporter_with_custom_span_postprocess_callback, openai_client):
+    openai_client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "Tell me a joke about opentelemetry"}],
+    )
+
+    spans = exporter_with_custom_span_postprocess_callback.get_finished_spans()
+    open_ai_span = spans[0]
+    assert open_ai_span.attributes["gen_ai.prompt.0.content"] == "REDACTED"
+    assert open_ai_span.attributes["gen_ai.completion.0.content"] == "REDACTED"
+
+
 def test_instruments(exporter_with_custom_instrumentations):
     @workflow()
     def run_workflow():
diff --git a/packages/traceloop-sdk/traceloop/sdk/__init__.py b/packages/traceloop-sdk/traceloop/sdk/__init__.py
@@ -2,9 +2,9 @@
 import sys
 from pathlib import Path
 
-from typing import Optional, Set
+from typing import Callable, Optional, Set
 from colorama import Fore
-from opentelemetry.sdk.trace import SpanProcessor
+from opentelemetry.sdk.trace import SpanProcessor, ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.sdk.metrics.export import MetricExporter
 from opentelemetry.sdk._logs.export import LogExporter
@@ -66,6 +66,7 @@ def init(
         instruments: Optional[Set[Instruments]] = None,
         block_instruments: Optional[Set[Instruments]] = None,
         image_uploader: Optional[ImageUploader] = None,
+        span_postprocess_callback: Optional[Callable[[ReadableSpan], None]] = None,
     ) -> Optional[Client]:
         if not enabled:
             TracerWrapper.set_disabled(True)
@@ -147,6 +148,7 @@ def init(
             image_uploader=image_uploader or ImageUploader(api_endpoint, api_key),
             instruments=instruments,
             block_instruments=block_instruments,
+            span_postprocess_callback=span_postprocess_callback,
         )
 
         if not is_metrics_enabled() or not metrics_exporter and exporter:
diff --git a/packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py b/packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py
@@ -12,7 +12,7 @@
     OTLPSpanExporter as GRPCExporter,
 )
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace import TracerProvider, SpanProcessor
+from opentelemetry.sdk.trace import TracerProvider, SpanProcessor, ReadableSpan
 from opentelemetry.propagators.textmap import TextMapPropagator
 from opentelemetry.propagate import set_global_textmap
 from opentelemetry.sdk.trace.export import (
@@ -73,6 +73,7 @@ def __new__(
         instruments: Optional[Set[Instruments]] = None,
         block_instruments: Optional[Set[Instruments]] = None,
         image_uploader: ImageUploader = None,
+        span_postprocess_callback: Optional[Callable[[ReadableSpan], None]] = None,
     ) -> "TracerWrapper":
         if not hasattr(cls, "instance"):
             obj = cls.instance = super(TracerWrapper, cls).__new__(cls)
@@ -120,6 +121,16 @@ def __new__(
                         obj.__spans_exporter
                     )
                 obj.__spans_processor_original_on_start = None
+                if span_postprocess_callback:
+                    # Create a wrapper that calls both the custom and original methods
+                    original_on_end = obj.__spans_processor.on_end
+
+                    def wrapped_on_end(span):
+                        # Call the custom on_end first
+                        span_postprocess_callback(span)
+                        # Then call the original to ensure normal processing
+                        original_on_end(span)
+                    obj.__spans_processor.on_end = wrapped_on_end
 
             obj.__spans_processor.on_start = obj._span_processor_on_start
             obj.__tracer_provider.add_span_processor(obj.__spans_processor)
