From 5837f522a520e52803adaaca68667b92c67df5b4 Mon Sep 17 00:00:00 2001 From: LuisFigueroaG Date: Mon, 27 Jul 2026 21:12:51 -0400 Subject: [PATCH] feat(sdk): allow disabling metrics programmatically --- .../tests/test_sdk_initialization.py | 107 ++++++++++++++++++ .../traceloop-sdk/traceloop/sdk/__init__.py | 47 +++++--- 2 files changed, 139 insertions(+), 15 deletions(-) diff --git a/packages/traceloop-sdk/tests/test_sdk_initialization.py b/packages/traceloop-sdk/tests/test_sdk_initialization.py index 18ef5d7d0f..d39daca90f 100644 --- a/packages/traceloop-sdk/tests/test_sdk_initialization.py +++ b/packages/traceloop-sdk/tests/test_sdk_initialization.py @@ -1,10 +1,12 @@ import json +import os import warnings import pytest from unittest.mock import patch from openai import OpenAI from traceloop.sdk import Traceloop from traceloop.sdk.decorators import workflow +from traceloop.sdk.instruments import Instruments from traceloop.sdk.tracing.tracing import TracerWrapper from opentelemetry.sdk.trace.export import SimpleSpanProcessor, BatchSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -309,6 +311,111 @@ def test_passing_both_kwargs_raises_type_error(isolated_tracer_wrapper): use_legacy_attributes=True, ) + +@pytest.mark.parametrize( + ("metrics_enabled", "environment_value", "expected_metrics_enabled"), + [ + (None, None, True), + (None, "true", True), + (None, "false", False), + (True, "false", True), + (False, "true", False), + ], +) +def test_metrics_enabled_precedence( + monkeypatch, + metrics_enabled, + environment_value, + expected_metrics_enabled, +): + """An explicit init value overrides the environment; None preserves it.""" + if environment_value is None: + monkeypatch.delenv("TRACELOOP_METRICS_ENABLED", raising=False) + else: + monkeypatch.setenv("TRACELOOP_METRICS_ENABLED", environment_value) + + metrics_exporter = object() + with ( + patch("traceloop.sdk.TracerWrapper"), + patch("traceloop.sdk.ImageUploader"), + patch("traceloop.sdk.MetricsWrapper") as metrics_wrapper, + ): + Traceloop.init( + exporter=InMemorySpanExporter(), + metrics_exporter=metrics_exporter, + metrics_enabled=metrics_enabled, + ) + + assert metrics_wrapper.called is expected_metrics_enabled + if expected_metrics_enabled: + metrics_wrapper.assert_called_once_with(exporter=metrics_exporter) + + +@pytest.mark.parametrize( + ("metrics_enabled", "environment_value"), + [(False, "true"), (True, "false")], +) +def test_metrics_enabled_is_visible_during_instrumentation( + monkeypatch, + isolated_tracer_wrapper, + metrics_enabled, + environment_value, +): + """Instrumentors must observe the explicit override while they initialize.""" + from opentelemetry.instrumentation.openai.utils import is_metrics_enabled + + monkeypatch.setenv("TRACELOOP_METRICS_ENABLED", environment_value) + observed_metrics_enabled = [] + metrics_exporter = object() + + def observe_metrics_enabled(*args, **kwargs): + observed_metrics_enabled.append(is_metrics_enabled()) + return True + + with ( + patch( + "traceloop.sdk.tracing.tracing.init_openai_instrumentor", + side_effect=observe_metrics_enabled, + ), + patch("traceloop.sdk.MetricsWrapper") as metrics_wrapper, + ): + Traceloop.init( + exporter=InMemorySpanExporter(), + metrics_exporter=metrics_exporter, + metrics_enabled=metrics_enabled, + instruments={Instruments.OPENAI}, + disable_batch=True, + ) + + assert observed_metrics_enabled == [metrics_enabled] + assert metrics_wrapper.called is metrics_enabled + if metrics_enabled: + metrics_wrapper.assert_called_once_with(exporter=metrics_exporter) + assert os.environ["TRACELOOP_METRICS_ENABLED"] == environment_value + + +@pytest.mark.parametrize("environment_value", [None, "true"]) +def test_metrics_enabled_restores_environment_when_tracer_initialization_fails( + monkeypatch, + environment_value, +): + """The process environment must be restored even if tracer setup fails.""" + if environment_value is None: + monkeypatch.delenv("TRACELOOP_METRICS_ENABLED", raising=False) + else: + monkeypatch.setenv("TRACELOOP_METRICS_ENABLED", environment_value) + + with patch("traceloop.sdk.TracerWrapper", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError, match="boom"): + Traceloop.init( + exporter=InMemorySpanExporter(), + metrics_enabled=False, + disable_batch=True, + ) + + assert os.environ.get("TRACELOOP_METRICS_ENABLED") == environment_value + + def test_use_attributes_defaults_to_true(isolated_tracer_wrapper): """When use_attributes is not passed, instrumentors keep the default spec-compliant behavior (emit prompts/completions as span attributes).""" diff --git a/packages/traceloop-sdk/traceloop/sdk/__init__.py b/packages/traceloop-sdk/traceloop/sdk/__init__.py index 6663429d99..7c7f463eee 100644 --- a/packages/traceloop-sdk/traceloop/sdk/__init__.py +++ b/packages/traceloop-sdk/traceloop/sdk/__init__.py @@ -73,6 +73,7 @@ def init( endpoint_is_traceloop: Optional[bool] = False, use_attributes: Optional[bool] = None, use_legacy_attributes: Optional[bool] = None, + metrics_enabled: Optional[bool] = None, ) -> Optional[Client]: """Initialize Traceloop tracing, metrics, and instrumentation. @@ -88,6 +89,9 @@ def init( events have nowhere to go and no prompt/completion data will be recorded. use_legacy_attributes: Deprecated alias for ``use_attributes``. Will be removed in a future release. + metrics_enabled: Enables or disables metric exporting. An explicit value + takes precedence over ``TRACELOOP_METRICS_ENABLED``. If ``None``, the + environment variable is used and metrics default to enabled. """ if use_attributes is not None and use_legacy_attributes is not None: raise TypeError( @@ -181,22 +185,35 @@ def init( TracerWrapper.set_static_params( resource_attributes, enable_content_tracing, api_endpoint, headers ) - Traceloop.__tracer_wrapper = TracerWrapper( - disable_batch=disable_batch, - processor=processor, - propagator=propagator, - exporter=exporter, - sampler=sampler, - should_enrich_metrics=should_enrich_metrics, - image_uploader=image_uploader or ImageUploader(api_endpoint, api_key), - instruments=instruments, - block_instruments=block_instruments, - span_postprocess_callback=span_postprocess_callback, - use_attributes=use_attributes, - ) - - metrics_disabled_by_config = not is_metrics_enabled() + metrics_enabled_by_config = is_metrics_enabled() if metrics_enabled is None else metrics_enabled + previous_metrics_enabled = os.environ.get("TRACELOOP_METRICS_ENABLED") + if metrics_enabled is not None: + os.environ["TRACELOOP_METRICS_ENABLED"] = str(metrics_enabled).lower() + try: + Traceloop.__tracer_wrapper = TracerWrapper( + disable_batch=disable_batch, + processor=processor, + propagator=propagator, + exporter=exporter, + sampler=sampler, + should_enrich_metrics=should_enrich_metrics, + image_uploader=image_uploader or ImageUploader(api_endpoint, api_key), + instruments=instruments, + block_instruments=block_instruments, + span_postprocess_callback=span_postprocess_callback, + use_attributes=use_attributes, + ) + finally: + if metrics_enabled is not None: + if previous_metrics_enabled is None: + os.environ.pop("TRACELOOP_METRICS_ENABLED", None) + else: + os.environ["TRACELOOP_METRICS_ENABLED"] = previous_metrics_enabled + + metrics_disabled_by_config = not metrics_enabled_by_config has_custom_spans_pipeline = processor or exporter + # A custom trace pipeline still needs a custom metrics exporter, regardless + # of metrics_enabled, to avoid sending metrics to an unintended endpoint. custom_trace_without_custom_metrics = has_custom_spans_pipeline and not metrics_exporter if metrics_disabled_by_config or custom_trace_without_custom_metrics: