diff --git a/pyproject.toml b/pyproject.toml index 0138c3ee558..f79d62bf21a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,7 @@ dependencies = [ "opentelemetry-proto", "opentelemetry-exporter-otlp", "openinference-semantic-conventions>=0.1.5", + "openinference-instrumentation", "openinference-instrumentation-langchain>=0.1.12", "openinference-instrumentation-llama-index>=1.2.0", "openinference-instrumentation-openai>=0.1.4", diff --git a/src/phoenix/trace/__init__.py b/src/phoenix/trace/__init__.py index 097e1fd86f3..eec0f3ed00d 100644 --- a/src/phoenix/trace/__init__.py +++ b/src/phoenix/trace/__init__.py @@ -1,5 +1,4 @@ -import contextlib -from typing import Iterator +from openinference.instrumentation import suppress_tracing from .projects import using_project from .span_evaluations import DocumentEvaluations, Evaluations, SpanEvaluations, TraceEvaluations @@ -12,24 +11,5 @@ "DocumentEvaluations", "TraceEvaluations", "using_project", + "suppress_tracing", ] - - -@contextlib.contextmanager -def suppress_tracing() -> Iterator[None]: - """ - Context manager to pause OpenTelemetry instrumentation. - - Examples: - with suppress_tracing(): - # No tracing will occur within this block - ... - """ - try: - from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY, attach, detach, set_value - except ImportError: - yield - return - token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) - yield - detach(token) diff --git a/tests/trace/test_suppress_tracing.py b/tests/trace/test_suppress_tracing.py deleted file mode 100644 index 6d98672b7b1..00000000000 --- a/tests/trace/test_suppress_tracing.py +++ /dev/null @@ -1,29 +0,0 @@ -from typing import Iterator - -import pytest -from opentelemetry.context import ( - _SUPPRESS_INSTRUMENTATION_KEY, - attach, - detach, - get_value, - set_value, -) -from phoenix.trace import suppress_tracing - - -def test_suppress_tracing(obj: object): - with suppress_tracing(): - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) is True - assert get_value(_SUPPRESS_INSTRUMENTATION_KEY) is obj - - -@pytest.fixture(autouse=True) -def instrument(obj: object) -> Iterator[None]: - token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, obj)) - yield - detach(token) - - -@pytest.fixture -def obj() -> object: - return object()