Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Minor cleanup of some context manager stuff and better handling of GetItemOrAttribute. #1687

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 7 additions & 18 deletions src/core/trulens/core/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
Hashable,
Iterable,
List,
Literal,
Optional,
Sequence,
Set,
Expand Down Expand Up @@ -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
):
Expand All @@ -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):
Expand All @@ -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
):
Expand All @@ -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
Expand Down
13 changes: 7 additions & 6 deletions src/core/trulens/core/database/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion src/core/trulens/core/utils/serial.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
5 changes: 5 additions & 0 deletions src/core/trulens/experimental/otel_tracing/core/session.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down