Change recording to create a span of type "record_root"

examples/experimental/otel_exporter.ipynb

@@ -109,6 +109,7 @@
     "session = TruSession()\n",
     "session.experimental_enable_feature(\"otel_tracing\")\n",
     "session.reset_database()\n",
+    "\n",
     "init(session, debug=True)"
    ]
   },

src/core/trulens/core/app.py

@@ -3,6 +3,7 @@
 from abc import ABC
 from abc import ABCMeta
 from abc import abstractmethod
+import contextlib
 import contextvars
 import datetime
 import inspect
@@ -421,6 +422,18 @@ def tru(self) -> core_connector.DBConnector:
         pydantic.PrivateAttr(default_factory=dict)
     )
 
+    tokens: List[object] = []
+    """
+    OTEL context tokens for the current context manager. These tokens are how the OTEL
+    context api keeps track of what is changed in the context, and used to undo the changes.
+    """
+
+    span_context: Optional[contextlib.AbstractContextManager] = None
+    """
+    Span context manager. Required to help keep track of the appropriate span context
+    to enter/exit.
+    """
+
     def __init__(
         self,
         connector: Optional[core_connector.DBConnector] = None,
@@ -1048,27 +1061,29 @@ def __enter__(self):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return _App.__enter__(self)
+            return OTELApp.__enter__(self)
 
         ctx = core_instruments._RecordingContext(app=self)
 
         token = self.recording_contexts.set(ctx)
         ctx.token = token
 
-        # self._set_context_vars()
-
         return ctx
 
     # For use as a context manager.
     def __exit__(self, exc_type, exc_value, exc_tb):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return _App.__exit__(self, exc_type, exc_value, exc_tb)
+            return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
 
         ctx = self.recording_contexts.get()
         self.recording_contexts.reset(ctx.token)
@@ -1085,9 +1100,11 @@ async def __aenter__(self):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return await _App.__aenter__(self)
+            return OTELApp.__enter__(self)
 
         ctx = core_instruments._RecordingContext(app=self)
 
@@ -1103,9 +1120,11 @@ async def __aexit__(self, exc_type, exc_value, exc_tb):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.app import _App
+            from trulens.experimental.otel_tracing.core.instrument import (
+                App as OTELApp,
+            )
 
-            return await _App.__aexit__(self, exc_type, exc_value, exc_tb)
+            return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
 
         ctx = self.recording_contexts.get()
         self.recording_contexts.reset(ctx.token)

src/core/trulens/experimental/otel_tracing/core/instrument.py

@@ -1,8 +1,14 @@
 from functools import wraps
 import logging
 from typing import Any, Callable, Dict, Optional, Union
+import uuid
 
 from opentelemetry import trace
+from opentelemetry.baggage import get_baggage
+from opentelemetry.baggage import remove_baggage
+from opentelemetry.baggage import set_baggage
+import opentelemetry.context as context_api
+from trulens.core import app as core_app
 from trulens.experimental.otel_tracing.core.init import TRULENS_SERVICE_NAME
 from trulens.otel.semconv.trace import SpanAttributes
 
@@ -84,6 +90,20 @@ def wrapper(*args, **kwargs):
                     # It's on the user to deal with None as a return value.
                     func_exception = e
 
+                span.set_attribute("name", func.__name__)
+                span.set_attribute("kind", "SPAN_KIND_TRULENS")
+                span.set_attribute(
+                    "parent_span_id", span.get_span_context().span_id
+                )
+                span.set_attribute(
+                    SpanAttributes.RECORD_ID,
+                    str(get_baggage(SpanAttributes.RECORD_ID)),
+                )
+                span.set_attribute(
+                    SpanAttributes.APP_ID,
+                    str(get_baggage(SpanAttributes.APP_ID)),
+                )
+
                 try:
                     attributes_to_add = {}
 
@@ -137,3 +157,69 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return inner_decorator
+
+
+class App(core_app.App):
+    # For use as a context manager.
+    def __enter__(self):
+        logger.debug("Entering the OTEL app context.")
+
+        # Note: This is not the same as the record_id in the core app since the OTEL
+        # tracing is currently separate from the old records behavior
+        otel_record_id = str(uuid.uuid4())
+
+        tracer = trace.get_tracer_provider().get_tracer(TRULENS_SERVICE_NAME)
+
+        # Calling set_baggage does not actually add the baggage to the current context, but returns a new one
+        # To avoid issues with remembering to add/remove the baggage, we attach it to the runtime context.
+        self.tokens.append(
+            context_api.attach(
+                set_baggage(SpanAttributes.RECORD_ID, otel_record_id)
+            )
+        )
+        self.tokens.append(
+            context_api.attach(set_baggage(SpanAttributes.APP_ID, self.app_id))
+        )
+
+        # Use start_as_current_span as a context manager
+        self.span_context = tracer.start_as_current_span("root")
+        root_span = self.span_context.__enter__()
+
+        logger.debug(str(get_baggage(SpanAttributes.RECORD_ID)))
+
+        # Set general span attributes
+        root_span.set_attribute("kind", "SPAN_KIND_TRULENS")
+        root_span.set_attribute("name", "root")
+        root_span.set_attribute(
+            SpanAttributes.SPAN_TYPE, SpanAttributes.SpanType.RECORD_ROOT
+        )
+        root_span.set_attribute(SpanAttributes.APP_ID, self.app_id)
+        root_span.set_attribute(SpanAttributes.RECORD_ID, otel_record_id)
+
+        # Set record root specific attributes
+        root_span.set_attribute(
+            SpanAttributes.RECORD_ROOT.APP_NAME, self.app_name
+        )
+        root_span.set_attribute(
+            SpanAttributes.RECORD_ROOT.APP_VERSION, self.app_version
+        )
+        root_span.set_attribute(SpanAttributes.RECORD_ROOT.APP_ID, self.app_id)
+        root_span.set_attribute(
+            SpanAttributes.RECORD_ROOT.RECORD_ID, otel_record_id
+        )
+
+        return root_span
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        remove_baggage(SpanAttributes.RECORD_ID)
+        remove_baggage(SpanAttributes.APP_ID)
+
+        logger.debug("Exiting the OTEL app context.")
+
+        while self.tokens:
+            # Clearing the context once we're done with this root span.
+            # See https://github.com/open-telemetry/opentelemetry-python/issues/2432#issuecomment-1593458684
+            context_api.detach(self.tokens.pop())
+
+        if self.span_context:
+            self.span_context.__exit__(exc_type, exc_value, exc_tb)

src/otel/semconv/trulens/otel/semconv/trace.py

@@ -56,6 +56,12 @@ class SpanAttributes:
     User-defined selector name for the current span.
     """
 
+    RECORD_ID = BASE + "record_id"
+    """ID of the record that the span belongs to."""
+
+    APP_ID = BASE + "app_id"
+    """ID of the app that the span belongs to."""
+
     class SpanType(str, Enum):
         """Span type attribute values.