fix: allow user override of id generator (#1315)

Arize-ai · Feb 27, 2025 · 1916749 · 1916749
1 parent b59d70a
commit 1916749
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py
@@ -39,7 +39,7 @@
     set_value,
 )
 from opentelemetry.sdk.trace import TracerProvider as OTelTracerProvider
-from opentelemetry.sdk.trace.id_generator import IdGenerator
+from opentelemetry.sdk.trace.id_generator import IdGenerator, RandomIdGenerator
 from opentelemetry.trace import (
     INVALID_SPAN,
     INVALID_SPAN_ID,
@@ -612,7 +612,10 @@ def __init__(self, wrapped: Tracer, config: TraceConfig) -> None:
 
     @property
     def id_generator(self) -> IdGenerator:
-        return self._self_id_generator
+        ans = getattr(self.__wrapped__, "id_generator", None)
+        if ans and ans.__class__ is RandomIdGenerator:
+            return self._self_id_generator
+        return cast(IdGenerator, ans)
 
     @contextmanager
     def start_as_current_span(

diff --git a/python/openinference-instrumentation/tests/test_id_generator.py b/python/openinference-instrumentation/tests/test_id_generator.py
@@ -1,15 +1,28 @@
 from random import seed
 
+import pytest
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 
 from openinference.instrumentation import OITracer, TraceConfig
 
 
-def test_id_generator_is_unaffected_by_seed() -> None:
+@pytest.mark.parametrize(
+    "override",
+    [
+        pytest.param(True, id="should-allow-user-override"),
+        pytest.param(False, id="should-be-unaffected-by-seed"),
+    ],
+)
+def test_id_generator(override: bool) -> None:
     in_memory_span_exporter = InMemorySpanExporter()
-    tracer_provider = TracerProvider()
+    tracer_provider = (
+        TracerProvider(id_generator=type("Foo", (RandomIdGenerator,), {})())
+        if override
+        else TracerProvider()
+    )
     tracer_provider.add_span_processor(SimpleSpanProcessor(in_memory_span_exporter))
     tracer = tracer_provider.get_tracer(__name__)
     oi_tracer = OITracer(tracer, TraceConfig())
@@ -21,5 +34,5 @@ def test_id_generator_is_unaffected_by_seed() -> None:
                 tr.start_span("child").end()
     spans = in_memory_span_exporter.get_finished_spans()
     assert len(spans) == n * 2 * 2
-    assert len(set(span.context.trace_id for span in spans)) == (n + 1)
-    assert len(set(span.context.span_id for span in spans)) == (n + 1) * 2
+    assert len(set(span.context.trace_id for span in spans)) == 1 if override else (n + 1)
+    assert len(set(span.context.span_id for span in spans)) == 2 if override else (n + 1) * 2