Skip to content

Commit

Permalink
draft
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-gtokernliang
sfc-gh-gtokernliang committed Dec 19, 2024
1 parent d5ce738 commit 712f97b
Show file tree
Hide file tree
Showing 4 changed files with 226 additions and 136 deletions.
267 changes: 144 additions & 123 deletions examples/experimental/otel_exporter.ipynb
View file
Original file line number Diff line number Diff line change
@@ -1,126 +1,147 @@
{
"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,
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"import sys\n",
"\n",
"# Add base dir to path to be able to access test folder.\n",
"base_dir = Path().cwd().parent.parent.resolve()\n",
"if str(base_dir) not in sys.path:\n",
" print(f\"Adding {base_dir} to sys.path\")\n",
" sys.path.append(str(base_dir))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from trulens.experimental.otel_tracing.core.instrument import instrument\n",
"\n",
"\n",
"class TestApp:\n",
" @instrument()\n",
" def respond_to_query(self, query: str) -> str:\n",
" return f\"answer: {self.nested(query)}\"\n",
"\n",
" @instrument(attributes={\"nested_attr1\": \"value1\"})\n",
" def nested(self, query: str) -> str:\n",
" return f\"nested: {self.nested2(query)}\"\n",
"\n",
" @instrument(\n",
" attributes=lambda ret, *args, **kwargs: {\n",
" \"nested2_ret\": ret,\n",
" \"nested2_args[0]\": args[0],\n",
" }\n",
" )\n",
" def nested2(self, query: str) -> str:\n",
" return f\"nested2: {query}\""
]
},
{
"cell_type": "code",
"execution_count": null,
"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",
"\n",
"dotenv.load_dotenv()\n",
"\n",
"connection_params = {\n",
" \"account\": os.environ[\"SNOWFLAKE_ACCOUNT\"],\n",
" \"user\": os.environ[\"SNOWFLAKE_USER\"],\n",
" \"password\": os.environ[\"SNOWFLAKE_USER_PASSWORD\"],\n",
" \"database\": os.environ[\"SNOWFLAKE_DATABASE\"],\n",
" \"schema\": os.environ[\"SNOWFLAKE_SCHEMA\"],\n",
" \"warehouse\": os.environ[\"SNOWFLAKE_WAREHOUSE\"],\n",
" \"role\": os.environ[\"SNOWFLAKE_ROLE\"],\n",
"}\n",
"\n",
"connector = SnowflakeConnector(\n",
" **connection_params, database_redact_keys=True, database_args=None\n",
")\n",
"session = TruSession(connector=connector)\n",
"session.reset_database()\n",
"init(session, debug=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from trulens.apps.custom import TruCustomApp\n",
"\n",
"test_app = TestApp()\n",
"custom_app = TruCustomApp(test_app)\n",
"\n",
"with custom_app as recording:\n",
" test_app.respond_to_query(\"test\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "trulens",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
}
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# !pip install opentelemetry-api\n",
"# !pip install opentelemetry-sdk"
]
},
"nbformat": 4,
"nbformat_minor": 2
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"import sys\n",
"\n",
"# Add base dir to path to be able to access test folder.\n",
"base_dir = Path().cwd().parent.parent.resolve()\n",
"if str(base_dir) not in sys.path:\n",
" print(f\"Adding {base_dir} to sys.path\")\n",
" sys.path.append(str(base_dir))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import logging\n",
"\n",
"root = logging.getLogger()\n",
"root.setLevel(logging.DEBUG)\n",
"handler = logging.StreamHandler(sys.stdout)\n",
"handler.setLevel(logging.DEBUG)\n",
"handler.addFilter(logging.Filter(\"trulens\"))\n",
"formatter = logging.Formatter(\n",
" \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"\n",
")\n",
"handler.setFormatter(formatter)\n",
"root.addHandler(handler)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from trulens.experimental.otel_tracing.core.instrument import instrument\n",
"\n",
"\n",
"class TestApp:\n",
" @instrument()\n",
" def respond_to_query(self, query: str) -> str:\n",
" return f\"answer: {self.nested(query)}\"\n",
"\n",
" @instrument(attributes={\"nested_attr1\": \"value1\"})\n",
" def nested(self, query: str) -> str:\n",
" return f\"nested: {self.nested2(query)}\"\n",
"\n",
" @instrument(\n",
" attributes=lambda ret, *args, **kwargs: {\n",
" \"nested2_ret\": ret,\n",
" \"nested2_args[0]\": args[0],\n",
" }\n",
" )\n",
" def nested2(self, query: str) -> str:\n",
" return f\"nested2: {query}\""
]
},
{
"cell_type": "code",
"execution_count": null,
"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",
"\n",
"dotenv.load_dotenv()\n",
"\n",
"connection_params = {\n",
" \"account\": os.environ[\"SNOWFLAKE_ACCOUNT\"],\n",
" \"user\": os.environ[\"SNOWFLAKE_USER\"],\n",
" \"password\": os.environ[\"SNOWFLAKE_USER_PASSWORD\"],\n",
" \"database\": os.environ[\"SNOWFLAKE_DATABASE\"],\n",
" \"schema\": os.environ[\"SNOWFLAKE_SCHEMA\"],\n",
" \"warehouse\": os.environ[\"SNOWFLAKE_WAREHOUSE\"],\n",
" \"role\": os.environ[\"SNOWFLAKE_ROLE\"],\n",
"}\n",
"\n",
"connector = SnowflakeConnector(\n",
" **connection_params, database_redact_keys=True, database_args=None\n",
")\n",
"session = TruSession(connector=connector)\n",
"session.experimental_enable_feature(\"otel_tracing\")\n",
"session.reset_database()\n",
"init(session, debug=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from trulens.apps.custom import TruCustomApp\n",
"\n",
"test_app = TestApp()\n",
"custom_app = TruCustomApp(test_app)\n",
"\n",
"with custom_app as recording:\n",
" test_app.respond_to_query(\"test\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "trulens",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
25 changes: 20 additions & 5 deletions src/core/trulens/core/app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from abc import ABC
from abc import ABCMeta
from abc import abstractmethod
import contextlib
import contextvars
import datetime
import inspect
Expand Down Expand Up @@ -421,6 +422,17 @@ def tru(self) -> core_connector.DBConnector:
pydantic.PrivateAttr(default_factory=dict)
)

token: Optional[object] = None
"""
OTEL context token for the current context manager.
"""

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,
Expand Down Expand Up @@ -1058,7 +1070,6 @@ def __enter__(self):

token = self.recording_contexts.set(ctx)
ctx.token = token
# self._set_context_vars()

return ctx

Expand Down Expand Up @@ -1088,9 +1099,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)

Expand All @@ -1106,9 +1119,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)
Expand Down
Loading

0 comments on commit 712f97b

Please sign in to comment.