diff --git a/src/strands/telemetry/config.py b/src/strands/telemetry/config.py index 928bc0e87..0509c7440 100644 --- a/src/strands/telemetry/config.py +++ b/src/strands/telemetry/config.py @@ -6,6 +6,7 @@ import logging from importlib.metadata import version +from typing import Any import opentelemetry.metrics as metrics_api import opentelemetry.sdk.metrics as metrics_sdk @@ -118,22 +119,46 @@ def _initialize_tracer(self) -> None: ) ) - def setup_console_exporter(self) -> "StrandsTelemetry": - """Set up console exporter for the tracer provider.""" + def setup_console_exporter(self, **kwargs: Any) -> "StrandsTelemetry": + """Set up console exporter for the tracer provider. + + Args: + **kwargs: Optional keyword arguments passed directly to + OpenTelemetry's ConsoleSpanExporter initializer. + + Returns: + self: Enables method chaining. + + This method configures a SimpleSpanProcessor with a ConsoleSpanExporter, + allowing trace data to be output to the console. Any additional keyword + arguments provided will be forwarded to the ConsoleSpanExporter. + """ try: logger.info("Enabling console export") - console_processor = SimpleSpanProcessor(ConsoleSpanExporter()) + console_processor = SimpleSpanProcessor(ConsoleSpanExporter(**kwargs)) 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) -> "StrandsTelemetry": - """Set up OTLP exporter for the tracer provider.""" + def setup_otlp_exporter(self, **kwargs: Any) -> "StrandsTelemetry": + """Set up OTLP exporter for the tracer provider. + + Args: + **kwargs: Optional keyword arguments passed directly to + OpenTelemetry's OTLPSpanExporter initializer. + + Returns: + self: Enables method chaining. + + This method configures a BatchSpanProcessor with an OTLPSpanExporter, + allowing trace data to be exported to an OTLP endpoint. Any additional + keyword arguments provided will be forwarded to the OTLPSpanExporter. + """ from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter try: - otlp_exporter = OTLPSpanExporter() + otlp_exporter = OTLPSpanExporter(**kwargs) batch_processor = BatchSpanProcessor(otlp_exporter) self.tracer_provider.add_span_processor(batch_processor) logger.info("OTLP exporter configured") diff --git a/tests/strands/telemetry/test_config.py b/tests/strands/telemetry/test_config.py index 0a81d5e2d..658d4d08a 100644 --- a/tests/strands/telemetry/test_config.py +++ b/tests/strands/telemetry/test_config.py @@ -168,9 +168,9 @@ def test_setup_console_exporter(mock_resource, mock_tracer_provider, mock_consol telemetry = StrandsTelemetry() # Set the tracer_provider directly telemetry.tracer_provider = mock_tracer_provider.return_value - telemetry.setup_console_exporter() + telemetry.setup_console_exporter(foo="bar") - mock_console_exporter.assert_called_once() + mock_console_exporter.assert_called_once_with(foo="bar") mock_simple_processor.assert_called_once_with(mock_console_exporter.return_value) mock_tracer_provider.return_value.add_span_processor.assert_called() @@ -182,9 +182,9 @@ def test_setup_otlp_exporter(mock_resource, mock_tracer_provider, mock_otlp_expo telemetry = StrandsTelemetry() # Set the tracer_provider directly telemetry.tracer_provider = mock_tracer_provider.return_value - telemetry.setup_otlp_exporter() + telemetry.setup_otlp_exporter(foo="bar") - mock_otlp_exporter.assert_called_once() + mock_otlp_exporter.assert_called_once_with(foo="bar") mock_batch_processor.assert_called_once_with(mock_otlp_exporter.return_value) mock_tracer_provider.return_value.add_span_processor.assert_called()