chore: allow custom tracer_provider and chain setup (strands-agents#316)

Author: poshinchen

src/strands/telemetry/__init__.py

@@ -3,7 +3,7 @@
 This module provides metrics and tracing functionality.
 """
 
-from .config import StrandsTelemetry, get_otel_resource
+from .config import StrandsTelemetry
 from .metrics import EventLoopMetrics, MetricsClient, Trace, metrics_to_string
 from .tracer import Tracer, get_tracer
 
@@ -16,8 +16,6 @@
     # Tracer
     "Tracer",
     "get_tracer",
-    # Resource
-    "get_otel_resource",
     # Telemetry Setup
     "StrandsTelemetry",
 ]

src/strands/telemetry/config.py

@@ -10,7 +10,6 @@
 import opentelemetry.trace as trace_api
 from opentelemetry import propagate
 from opentelemetry.baggage.propagation import W3CBaggagePropagator
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.propagators.composite import CompositePropagator
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider
@@ -41,46 +40,57 @@ def get_otel_resource() -> Resource:
 class StrandsTelemetry:
     """OpenTelemetry configuration and setup for Strands applications.
 
-    It automatically initializes a tracer provider with text map propagators.
-    Trace exporters (console, OTLP) can be set up individually using dedicated methods.
+    Automatically initializes a tracer provider with text map propagators.
+    Trace exporters (console, OTLP) can be set up individually using dedicated methods
+    that support method chaining for convenient configuration.
+
+    Args:
+        tracer_provider: Optional pre-configured SDKTracerProvider. If None,
+            a new one will be created and set as the global tracer provider.
 
     Environment Variables:
         Environment variables are handled by the underlying OpenTelemetry SDK:
         - OTEL_EXPORTER_OTLP_ENDPOINT: OTLP endpoint URL
         - OTEL_EXPORTER_OTLP_HEADERS: Headers for OTLP requests
 
-    Example:
-        Basic setup with console exporter:
-        >>> telemetry = StrandsTelemetry()
-        >>> telemetry.setup_console_exporter()
+    Examples:
+        Quick setup with method chaining:
+        >>> StrandsTelemetry().setup_console_exporter().setup_otlp_exporter()
 
-        Setup with OTLP exporter:
-        >>> telemetry = StrandsTelemetry()
-        >>> telemetry.setup_otlp_exporter()
+        Using a custom tracer provider:
+        >>> StrandsTelemetry(tracer_provider=my_provider).setup_console_exporter()
 
-        Setup with both exporters:
+        Step-by-step configuration:
         >>> telemetry = StrandsTelemetry()
         >>> telemetry.setup_console_exporter()
         >>> telemetry.setup_otlp_exporter()
 
     Note:
-        The tracer provider is automatically initialized upon instantiation.
-        Exporters must be explicitly configured using the setup methods.
-        Failed exporter configurations are logged but do not raise exceptions.
+        - The tracer provider is automatically initialized upon instantiation
+        - When no tracer_provider is provided, the instance sets itself as the global provider
+        - Exporters must be explicitly configured using the setup methods
+        - Failed exporter configurations are logged but do not raise exceptions
+        - All setup methods return self to enable method chaining
     """
 
     def __init__(
         self,
+        tracer_provider: SDKTracerProvider | None = None,
     ) -> None:
         """Initialize the StrandsTelemetry instance.
 
-        Automatically sets up the OpenTelemetry infrastructure.
+        Args:
+            tracer_provider: Optional pre-configured tracer provider.
+                If None, a new one will be created and set as global.
 
         The instance is ready to use immediately after initialization, though
         trace exporters must be configured separately using the setup methods.
         """
-        self.resource = get_otel_resource()
-        self._initialize_tracer()
+        if tracer_provider:
+            self.tracer_provider = tracer_provider
+        else:
+            self.resource = get_otel_resource()
+            self._initialize_tracer()
 
     def _initialize_tracer(self) -> None:
         """Initialize the OpenTelemetry tracer."""
@@ -102,21 +112,25 @@ def _initialize_tracer(self) -> None:
             )
         )
 
-    def setup_console_exporter(self) -> None:
+    def setup_console_exporter(self) -> "StrandsTelemetry":
         """Set up console exporter for the tracer provider."""
         try:
             logger.info("enabling console export")
             console_processor = SimpleSpanProcessor(ConsoleSpanExporter())
             self.tracer_provider.add_span_processor(console_processor)
         except Exception as e:
             logger.exception("error=<%s> | Failed to configure console exporter", e)
+        return self
 
-    def setup_otlp_exporter(self) -> None:
+    def setup_otlp_exporter(self) -> "StrandsTelemetry":
         """Set up OTLP exporter for the tracer provider."""
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+
         try:
             otlp_exporter = OTLPSpanExporter()
             batch_processor = BatchSpanProcessor(otlp_exporter)
             self.tracer_provider.add_span_processor(batch_processor)
             logger.info("OTLP exporter configured")
         except Exception as e:
             logger.exception("error=<%s> | Failed to configure OTLP exporter", e)
+        return self

tests/strands/telemetry/test_config.py

@@ -47,7 +47,7 @@ def mock_console_exporter():
 
 @pytest.fixture
 def mock_otlp_exporter():
-    with mock.patch("strands.telemetry.config.OTLPSpanExporter") as mock_otlp_exporter:
+    with mock.patch("opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter") as mock_otlp_exporter:
         yield mock_otlp_exporter