Enable exporting spans to snowflake stage if a TruLensSnowflakeSpanExporter is provided

Makefile

@@ -32,12 +32,13 @@ env-tests:
 		nbconvert \
 		nbformat \
 		opentelemetry-sdk \
+		opentelemetry-proto \
 		pre-commit \
 		pytest \
 		pytest-azurepipelines \
 		pytest-cov \
 		pytest-subtests \
-		ruff \
+		ruff
 
 env-tests-required:
 	poetry install --only required \

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)"
    ]
@@ -100,17 +91,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import os\n",
+    "\n",
     "import dotenv\n",
+    "from trulens.connectors.snowflake import SnowflakeConnector\n",
     "from trulens.core.session import TruSession\n",
-    "from trulens.experimental.otel_tracing.core.init import init\n",
+    "from trulens.experimental.otel_tracing.core.exporter import (\n",
+    "    TruLensSnowflakeSpanExporter,\n",
+    ")\n",
     "\n",
     "dotenv.load_dotenv()\n",
     "\n",
-    "session = TruSession()\n",
-    "session.experimental_enable_feature(\"otel_tracing\")\n",
-    "session.reset_database()\n",
+    "connection_params = {\n",
+    "    \"account\": os.environ[\"SNOWFLAKE_ACCOUNT\"],\n",
+    "    \"user\": os.environ[\"SNOWFLAKE_USER\"],\n",
+    "    \"password\": os.environ[\"SNOWFLAKE_USER_PASSWORD\"],\n",
+    "    \"role\": os.environ.get(\"SNOWFLAKE_ROLE\", \"ENGINEER\"),\n",
+    "    \"database\": os.environ.get(\"SNOWFLAKE_DATABASE\"),\n",
+    "    \"schema\": os.environ.get(\"SNOWFLAKE_SCHEMA\"),\n",
+    "    \"warehouse\": os.environ.get(\"SNOWFLAKE_WAREHOUSE\"),\n",
+    "}\n",
+    "\n",
+    "connector = SnowflakeConnector(\n",
+    "    **connection_params,\n",
+    ")\n",
+    "\n",
+    "exporter = TruLensSnowflakeSpanExporter(connector=connector)\n",
     "\n",
-    "init(session, debug=True)"
+    "# Note that the connector for the session can be different from the connector for the exporter.\n",
+    "session = TruSession(_experimental_otel_exporter=exporter)"
    ]
   },
   {
@@ -122,14 +131,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=\"123\") as recording:\n",
     "    test_app.respond_to_query(\"test\")\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(\"throw\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "session.experimental_force_flush()"
+   ]
   }
  ],
  "metadata": {

src/core/pyproject.toml

@@ -49,6 +49,7 @@ importlib-resources = "^6.0"
 trulens-otel-semconv = { version = "^1.0.0", optional = true }
 opentelemetry-api = { version = "^1.0.0", optional = true }
 opentelemetry-sdk = { version = "^1.0.0", optional = true }
+opentelemetry-proto = { version = "^1.0.0", optional = true }
 
 [tool.poetry.group.tqdm]
 optional = true

src/core/trulens/core/app.py

@@ -898,7 +898,7 @@ def __enter__(self):
             core_experimental.Feature.OTEL_TRACING
         ):
             from trulens.experimental.otel_tracing.core.instrument import (
-                App as OTELApp,
+                OTELRecordingContext as OTELApp,
             )
 
             return OTELApp.__enter__(self)
@@ -916,16 +916,14 @@ def __exit__(self, exc_type, exc_value, exc_tb):
             core_experimental.Feature.OTEL_TRACING
         ):
             from trulens.experimental.otel_tracing.core.instrument import (
-                App as OTELApp,
+                OTELRecordingContext as OTELApp,
             )
 
             return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
 
         ctx = self.recording_contexts.get()
         self.recording_contexts.reset(ctx.token)
 
-        # self._reset_context_vars()
-
         if exc_type is not None:
             raise exc_value
 
@@ -937,7 +935,7 @@ async def __aenter__(self):
             core_experimental.Feature.OTEL_TRACING
         ):
             from trulens.experimental.otel_tracing.core.instrument import (
-                App as OTELApp,
+                OTELRecordingContext as OTELApp,
             )
 
             return OTELApp.__enter__(self)
@@ -957,7 +955,7 @@ async def __aexit__(self, exc_type, exc_value, exc_tb):
             core_experimental.Feature.OTEL_TRACING
         ):
             from trulens.experimental.otel_tracing.core.instrument import (
-                App as OTELApp,
+                OTELRecordingContext as OTELApp,
             )
 
             return OTELApp.__exit__(self, exc_type, exc_value, exc_tb)
@@ -972,6 +970,22 @@ 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.model_construct(
+            **self.model_dump(), 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.
 
@@ -177,15 +181,38 @@ 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
+    def experimental_otel_exporter(self, value: Optional[SpanExporter]):
         otel_tracing_feature._FeatureSetup.assert_optionals_installed()
 
         from trulens.experimental.otel_tracing.core.session import _TruSession
 
         _TruSession._setup_otel_exporter(self, value)
 
+    def experimental_force_flush(
+        self, timeout_millis: Optional[int] = None
+    ) -> 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.
+        """
+        timeout_millis = timeout_millis or 300000
+
+        if (
+            not self.experimental_feature(
+                core_experimental.Feature.OTEL_TRACING
+            )
+            or self._experimental_tracer_provider is None
+        ):
+            return False
+
+        return self._experimental_tracer_provider.force_flush(timeout_millis)
+
     def __str__(self) -> str:
         return f"TruSession({self.connector})"
 
@@ -205,9 +232,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 +261,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

src/core/trulens/core/utils/requirements.txt

@@ -9,3 +9,4 @@ trulens-dashboard >= 1.0.0
 # talk about these packages as required.
 opentelemetry-api >= 1.0.0
 opentelemetry-sdk >= 1.0.0
+opentelemetry-proto >= 1.0.0

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

@@ -8,7 +8,7 @@
 """Feature controlling the use of this module."""
 
 REQUIREMENT = import_utils.format_import_errors(
-    ["opentelemetry-api", "opentelemetry-sdk"],
+    ["opentelemetry-api", "opentelemetry-sdk", "opentelemetry-proto"],
     purpose="otel_tracing experimental feature",
 )
 """Optional modules required for the otel_tracing experimental feature."""