feat(sdk): allow disabling metrics programmatically

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)."""

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: