From 3764b350e5b17b44b1124d3e51d39daf25beadaf Mon Sep 17 00:00:00 2001
From: David Kurokawa <david.kurokawa@snowflake.com>
Date: Fri, 6 Dec 2024 11:12:18 -0800
Subject: [PATCH] Minor fixes from Piotr's big PR that AFAICT aren't really
 related to the main thrust of the OTEL work.

---
 src/core/trulens/core/app.py                  | 25 ++++++-------------
 src/core/trulens/core/database/base.py        | 13 +++++-----
 src/core/trulens/core/utils/serial.py         |  2 +-
 .../experimental/otel_tracing/core/session.py |  5 ++++
 4 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/src/core/trulens/core/app.py b/src/core/trulens/core/app.py
index f27ac2cdd..d864fbd8a 100644
--- a/src/core/trulens/core/app.py
+++ b/src/core/trulens/core/app.py
@@ -19,6 +19,7 @@
     Hashable,
     Iterable,
     List,
+    Literal,
     Optional,
     Sequence,
     Set,
@@ -1057,12 +1058,10 @@ def __enter__(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):
+    def __exit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
@@ -1073,12 +1072,8 @@ def __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
-
-        return
+        # False means dont suppress exceptions.
+        return False
 
     # For use as a context manager.
     async def __aenter__(self):
@@ -1094,12 +1089,10 @@ async def __aenter__(self):
         token = self.recording_contexts.set(ctx)
         ctx.token = token
 
-        # self._set_context_vars()
-
         return ctx
 
     # For use as a context manager.
-    async def __aexit__(self, exc_type, exc_value, exc_tb):
+    async def __aexit__(self, exc_type, exc_value, exc_tb) -> Literal[False]:
         if self.session.experimental_feature(
             core_experimental.Feature.OTEL_TRACING
         ):
@@ -1110,12 +1103,8 @@ async def __aexit__(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
-
-        return
+        # False means dont suppress exceptions.
+        return False
 
     def _set_context_vars(self):
         # HACK: For debugging purposes, try setting/resetting all context vars
diff --git a/src/core/trulens/core/database/base.py b/src/core/trulens/core/database/base.py
index 6999baf1f..e05700cf3 100644
--- a/src/core/trulens/core/database/base.py
+++ b/src/core/trulens/core/database/base.py
@@ -138,8 +138,8 @@ def insert_app(self, app: app_schema.AppDefinition) -> types_schema.AppID:
 
         Args:
             app: The app to insert or update. Note that only the
-                [AppDefinition][trulens.core.schema.app.AppDefinition] parts are serialized
-                hence the type hint.
+                [AppDefinition][trulens.core.schema.app.AppDefinition] parts
+                are serialized hence the type hint.
 
         Returns:
             The id of the given app.
@@ -366,8 +366,9 @@ def get_records_and_feedback(
     def insert_ground_truth(
         self, ground_truth: groundtruth_schema.GroundTruth
     ) -> types_schema.GroundTruthID:
-        """Insert a ground truth entry into the database. The ground truth id is generated
-        based on the ground truth content, so re-inserting is idempotent.
+        """Insert a ground truth entry into the database. The ground truth id
+        is generated based on the ground truth content, so re-inserting is
+        idempotent.
 
         Args:
             ground_truth: The ground truth entry to insert.
@@ -413,8 +414,8 @@ def get_ground_truths_by_dataset(self, dataset_name: str) -> pd.DataFrame:
     def insert_dataset(
         self, dataset: dataset_schema.Dataset
     ) -> types_schema.DatasetID:
-        """Insert a dataset into the database. The dataset id is generated based on the
-        dataset content, so re-inserting is idempotent.
+        """Insert a dataset into the database. The dataset id is generated
+        based on the dataset content, so re-inserting is idempotent.
 
         Args:
             dataset: The dataset to insert.
diff --git a/src/core/trulens/core/utils/serial.py b/src/core/trulens/core/utils/serial.py
index b4a062f3e..0fe4b258a 100644
--- a/src/core/trulens/core/utils/serial.py
+++ b/src/core/trulens/core/utils/serial.py
@@ -363,7 +363,7 @@ def get(self, obj: Dict[str, T]) -> Iterable[T]:
                         pass
 
         # Otherwise handle a dict or object with the named attribute.
-        elif isinstance(obj, Dict):
+        elif hasattr(obj, "__getitem__"):
             if self.item_or_attribute in obj:
                 yield obj[self.item_or_attribute]
             else:
diff --git a/src/core/trulens/experimental/otel_tracing/core/session.py b/src/core/trulens/experimental/otel_tracing/core/session.py
index a87116742..fc218f889 100644
--- a/src/core/trulens/experimental/otel_tracing/core/session.py
+++ b/src/core/trulens/experimental/otel_tracing/core/session.py
@@ -1,3 +1,8 @@
+"""OTEL tracing additions on top of TruSession.
+
+Adds the code for setting up an otel exporter.
+"""
+
 from typing import Any, Optional
 
 from trulens.core import experimental as core_experimental