Enable exporting spans to snowflake stage if a TruLensSnowflakeSpanExporter is provided

examples/experimental/otel_exporter.ipynb

@@ -1,15 +1,5 @@
 {
  "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# !pip install opentelemetry-api\n",
-    "# !pip install opentelemetry-sdk"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -41,6 +31,7 @@
     "formatter = logging.Formatter(\n",
     "    \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"\n",
     ")\n",
+    "handler.addFilter(logging.Filter(\"trulens\"))\n",
     "handler.setFormatter(formatter)\n",
     "root.addHandler(handler)"
    ]
@@ -67,7 +58,7 @@
     "    @instrument(\n",
     "        attributes=lambda ret, exception, *args, **kwargs: {\n",
     "            \"nested2_ret\": ret,\n",
-    "            \"nested2_args[0]\": args[0],\n",
+    "            \"nested2_args[1]\": args[1],\n",
     "        }\n",
     "    )\n",
     "    def nested2(self, query: str) -> str:\n",
@@ -100,17 +91,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
+    "\n",
     "import dotenv\n",
+    "from trulens.core.experimental import Feature\n",
     "from trulens.core.session import TruSession\n",
-    "from trulens.experimental.otel_tracing.core.init import init\n",
     "\n",
     "dotenv.load_dotenv()\n",
-    "\n",
-    "session = TruSession()\n",
-    "session.experimental_enable_feature(\"otel_tracing\")\n",
-    "session.reset_database()\n",
-    "\n",
-    "init(session, debug=True)"
+    "os.environ[\"TRULENS_OTEL_TRACING\"] = \"1\"\n",
+    "session = TruSession(experimental_feature_flags={Feature.OTEL_TRACING: True})"
    ]
   },
   {
@@ -122,14 +111,28 @@
     "from trulens.apps.custom import TruCustomApp\n",
     "\n",
     "test_app = TestApp()\n",
-    "custom_app = TruCustomApp(test_app)\n",
+    "custom_app = TruCustomApp(\n",
+    "    test_app,\n",
+    "    app_name=\"database.schema.app\",\n",
+    "    app_version=\"1.0.0\",\n",
+    "    session=session,\n",
+    ")\n",
     "\n",
-    "with custom_app as recording:\n",
+    "with custom_app(run_name=\"test run\", input_id=\"456\") as recording:\n",
     "    test_app.respond_to_query(\"test\")\n",
     "\n",
-    "with custom_app as recording:\n",
+    "with custom_app() as recording:\n",
     "    test_app.respond_to_query(\"throw\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "session.experimental_force_flush()"
+   ]
   }
  ],
  "metadata": {

src/core/trulens/core/app.py

@@ -3,7 +3,6 @@
 from abc import ABC
 from abc import ABCMeta
 from abc import abstractmethod
-import contextlib
 import contextvars
 import datetime
 import inspect
@@ -423,18 +422,6 @@ 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,
@@ -889,20 +876,18 @@ def model_dump(self, *args, redact_keys: bool = False, **kwargs):
             **kwargs,
         )
 
-    # For use as a context manager.
-    def __enter__(self):
-        if not core_instruments.Instrument._have_context():
-            raise RuntimeError(core_endpoint._NO_CONTEXT_WARNING)
-
+    def _prevent_invalid_otel_syntax(self):
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
-            from trulens.experimental.otel_tracing.core.instrument import (
-                App as OTELApp,
-            )
+            raise RuntimeError("Invalid TruLens OTEL Tracing syntax.")
 
-            return OTELApp.__enter__(self)
+    # For use as a context manager.
+    def __enter__(self):
+        if not core_instruments.Instrument._have_context():
+            raise RuntimeError(core_endpoint._NO_CONTEXT_WARNING)
 
+        self._prevent_invalid_otel_syntax()
         ctx = core_instruments._RecordingContext(app=self)
 
         token = self.recording_contexts.set(ctx)
@@ -912,35 +897,19 @@ def __enter__(self):
 
     # 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.instrument import (
-                App as OTELApp,
-            )
-
-            return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
+        self._prevent_invalid_otel_syntax()
 
         ctx = self.recording_contexts.get()
         self.recording_contexts.reset(ctx.token)
 
-        # self._reset_context_vars()
-
         if exc_type is not None:
             raise exc_value
 
         return
 
     # For use as a context manager.
     async def __aenter__(self):
-        if self.session.experimental_feature(
-            core_experimental.Feature.OTEL_TRACING
-        ):
-            from trulens.experimental.otel_tracing.core.instrument import (
-                App as OTELApp,
-            )
-
-            return OTELApp.__enter__(self)
+        self._prevent_invalid_otel_syntax()
 
         ctx = core_instruments._RecordingContext(app=self)
 
@@ -953,14 +922,7 @@ async def __aenter__(self):
 
     # For use as a context manager.
     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.instrument import (
-                App as OTELApp,
-            )
-
-            return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
+        self._prevent_invalid_otel_syntax()
 
         ctx = self.recording_contexts.get()
         self.recording_contexts.reset(ctx.token)
@@ -972,6 +934,20 @@ async def __aexit__(self, exc_type, exc_value, exc_tb):
 
         return
 
+    def __call__(self, *, run_name: str = "", input_id: str = ""):
+        if not self.session.experimental_feature(
+            core_experimental.Feature.OTEL_TRACING
+        ):
+            raise RuntimeError("OTEL Tracing is not enabled for this session.")
+
+        from trulens.experimental.otel_tracing.core.instrument import (
+            OTELRecordingContext as OTELApp,
+        )
+
+        # Pylance shows an error here, but it is likely a false positive. due to the overriden
+        # model dump returning json instead of a dict.
+        return OTELApp(app=self, run_name=run_name, input_id=input_id)
+
     def _set_context_vars(self):
         # HACK: For debugging purposes, try setting/resetting all context vars
         # used in trulens around the app context managers due to bugs in trying

src/core/trulens/core/session.py

@@ -45,6 +45,8 @@
 from trulens.experimental.otel_tracing import _feature as otel_tracing_feature
 
 if TYPE_CHECKING:
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import SpanExporter
     from trulens.core import app as base_app
 
 tqdm = None
@@ -157,16 +159,18 @@ class TruSession(
     """Database Connector to use. If not provided, a default is created and
     used."""
 
-    _experimental_otel_exporter: Optional[
-        Any
-    ] = (  # Any = otel_export_sdk.SpanExporter
+    _experimental_otel_exporter: Optional[SpanExporter] = pydantic.PrivateAttr(
+        None
+    )
+
+    _experimental_tracer_provider: Optional[TracerProvider] = (
         pydantic.PrivateAttr(None)
     )
 
     @property
     def experimental_otel_exporter(
         self,
-    ) -> Any:  # Any = Optional[otel_export_sdk.SpanExporter]
+    ) -> Optional[SpanExporter]:
         """EXPERIMENTAL(otel_tracing): OpenTelemetry SpanExporter to send spans
         to.
 
@@ -176,15 +180,27 @@ def experimental_otel_exporter(
 
         return self._experimental_otel_exporter
 
-    @experimental_otel_exporter.setter
-    def experimental_otel_exporter(
-        self, value: Optional[Any]
-    ):  # Any = otel_export_sdk.SpanExporter
-        otel_tracing_feature._FeatureSetup.assert_optionals_installed()
+    def experimental_force_flush(self, timeout_millis: int = 300000) -> bool:
+        """
+        Force flush the OpenTelemetry exporters.
+
+        Args:
+            timeout_millis: The maximum amount of time to wait for spans to be
+                processed.
+
+        Returns:
+            False if the timeout is exceeded, feature is not enabled, or the provider doesn't exist, True otherwise.
+        """
 
-        from trulens.experimental.otel_tracing.core.session import _TruSession
+        if (
+            not self.experimental_feature(
+                core_experimental.Feature.OTEL_TRACING
+            )
+            or self._experimental_tracer_provider is None
+        ):
+            return False
 
-        _TruSession._setup_otel_exporter(self, value)
+        return self._experimental_tracer_provider.force_flush(timeout_millis)
 
     def __str__(self) -> str:
         return f"TruSession({self.connector})"
@@ -205,9 +221,7 @@ def __init__(
                 Iterable[core_experimental.Feature],
             ]
         ] = None,
-        _experimental_otel_exporter: Optional[
-            Any
-        ] = None,  # Any = otel_export_sdk.SpanExporter
+        _experimental_otel_exporter: Optional[SpanExporter] = None,
         **kwargs,
     ):
         if python_utils.safe_hasattr(self, "connector"):
@@ -236,24 +250,31 @@ def __init__(
                 f"Cannot provide both `connector` and connector argument(s) {extra_keys}."
             )
 
+        connector = connector or core_connector.DefaultDBConnector(
+            **connector_args
+        )
+
         super().__init__(
-            connector=connector
-            or core_connector.DefaultDBConnector(**connector_args),
+            connector=connector,
             **self_args,
         )
 
         # for WithExperimentalSettings mixin
         if experimental_feature_flags is not None:
             self.experimental_set_features(experimental_feature_flags)
 
-        if _experimental_otel_exporter is not None:
-            otel_tracing_feature.assert_optionals_installed()
+        if _experimental_otel_exporter is not None or self.experimental_feature(
+            core_experimental.Feature.OTEL_TRACING
+        ):
+            otel_tracing_feature._FeatureSetup.assert_optionals_installed()
 
             from trulens.experimental.otel_tracing.core.session import (
                 _TruSession,
             )
 
-            _TruSession._setup_otel_exporter(self, _experimental_otel_exporter)
+            _TruSession._setup_otel_exporter(
+                self, connector, _experimental_otel_exporter
+            )
 
     def App(self, *args, app: Optional[Any] = None, **kwargs) -> base_app.App:
         """Create an App from the given App constructor arguments by guessing