Configure Span Pipeline for BedrockAgentCore

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_utils.py

@@ -1,13 +1,16 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
 import sys
 from logging import Logger, getLogger
 
 import pkg_resources
 
 _logger: Logger = getLogger(__name__)
 
+AGENT_OBSERVABILITY_ENABLED = "AGENT_OBSERVABILITY_ENABLED"
+
 
 def is_installed(req: str) -> bool:
     """Is the given required package installed?"""
@@ -21,3 +24,8 @@ def is_installed(req: str) -> bool:
         _logger.debug("Skipping instrumentation patch: package %s, exception: %s", req, exc)
         return False
     return True
+
+
+def is_agent_observability_enabled() -> bool:
+    """Is the Agentic AI monitoring flag set to true?"""
+    return os.environ.get(AGENT_OBSERVABILITY_ENABLED, "false").lower() == "true"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py

@@ -11,6 +11,7 @@
 
 from amazon.opentelemetry.distro._aws_attribute_keys import AWS_LOCAL_SERVICE
 from amazon.opentelemetry.distro._aws_resource_attribute_configurator import get_service_attribute
+from amazon.opentelemetry.distro._utils import is_agent_observability_enabled
 from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler
 from amazon.opentelemetry.distro.attribute_propagating_span_processor_builder import (
     AttributePropagatingSpanProcessorBuilder,
@@ -27,7 +28,7 @@
 from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from amazon.opentelemetry.distro.scope_based_exporter import ScopeBasedPeriodicExportingMetricReader
 from amazon.opentelemetry.distro.scope_based_filtering_view import ScopeBasedRetainingView
-from opentelemetry._logs import set_logger_provider
+from opentelemetry._logs import get_logger_provider, set_logger_provider
 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
@@ -359,7 +360,15 @@ def _customize_span_exporter(span_exporter: SpanExporter, resource: Resource) ->
         _logger.info("Detected using AWS OTLP Traces Endpoint.")
 
         if isinstance(span_exporter, OTLPSpanExporter):
-            span_exporter = OTLPAwsSpanExporter(endpoint=traces_endpoint)
+            if is_agent_observability_enabled():
+                # Span exporter needs an instance of logger provider in ai agent
+                # observability case because we need to split input/output prompts
+                # from span attributes and send them to the logs pipeline per
+                # the new Gen AI semantic convention from OTel
+                # ref: https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/
+                span_exporter = OTLPAwsSpanExporter(endpoint=traces_endpoint, logger_provider=get_logger_provider())
+            else:
+                span_exporter = OTLPAwsSpanExporter(endpoint=traces_endpoint)
 
         else:
             _logger.warning(

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/traces/otlp_aws_span_exporter.py

@@ -6,6 +6,7 @@
 from amazon.opentelemetry.distro.exporter.otlp.aws.common.aws_auth_session import AwsAuthSession
 from opentelemetry.exporter.otlp.proto.http import Compression
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk._logs import LoggerProvider
 
 
 class OTLPAwsSpanExporter(OTLPSpanExporter):
@@ -18,8 +19,10 @@ def __init__(
         headers: Optional[Dict[str, str]] = None,
         timeout: Optional[int] = None,
         compression: Optional[Compression] = None,
+        logger_provider: Optional[LoggerProvider] = None,
     ):
         self._aws_region = None
+        self._logger_provider = logger_provider
 
         if endpoint:
             self._aws_region = endpoint.split(".")[1]

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py

@@ -311,6 +311,35 @@ def test_customize_span_exporter(self):
         self.assertIsInstance(customized_exporter._delegate, OTLPUdpSpanExporter)
         os.environ.pop("AWS_LAMBDA_FUNCTION_NAME", None)
 
+    def test_customize_span_exporter_with_agent_observability(self):
+        # Test that logger_provider is passed when agent observability is enabled
+        os.environ["AGENT_OBSERVABIILTY_ENABLED"] = "true"
+        os.environ[OTEL_EXPORTER_OTLP_TRACES_ENDPOINT] = "https://xray.us-east-1.amazonaws.com/v1/traces"
+
+        mock_logger_provider = MagicMock()
+        with patch(
+            "amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_logger_provider",
+            return_value=mock_logger_provider,
+        ):
+            mock_exporter = MagicMock(spec=OTLPSpanExporter)
+            result = _customize_span_exporter(mock_exporter, Resource.get_empty())
+
+            self.assertIsInstance(result, OTLPAwsSpanExporter)
+            self.assertEqual(result._logger_provider, mock_logger_provider)
+
+        # Test that logger_provider is not passed when agent observability is disabled
+        os.environ["AGENT_OBSERVABIILTY_ENABLED"] = "false"
+
+        mock_exporter = MagicMock(spec=OTLPSpanExporter)
+        result = _customize_span_exporter(mock_exporter, Resource.get_empty())
+
+        self.assertIsInstance(result, OTLPAwsSpanExporter)
+        self.assertIsNone(result._logger_provider)
+
+        # Clean up
+        os.environ.pop("AGENT_OBSERVABIILTY_ENABLED", None)
+        os.environ.pop(OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, None)
+
     def test_customize_span_exporter_sigv4(self):
 
         traces_good_endpoints = [

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_aws_span_exporter.py

@@ -0,0 +1,29 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from unittest import TestCase
+from unittest.mock import MagicMock
+
+from amazon.opentelemetry.distro.exporter.otlp.aws.traces.otlp_aws_span_exporter import OTLPAwsSpanExporter
+from opentelemetry.sdk._logs import LoggerProvider
+
+
+class TestOTLPAwsSpanExporter(TestCase):
+    def test_init_with_logger_provider(self):
+        # Test initialization with logger_provider
+        mock_logger_provider = MagicMock(spec=LoggerProvider)
+        endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
+
+        exporter = OTLPAwsSpanExporter(endpoint=endpoint, logger_provider=mock_logger_provider)
+
+        self.assertEqual(exporter._logger_provider, mock_logger_provider)
+        self.assertEqual(exporter._aws_region, "us-east-1")
+
+    def test_init_without_logger_provider(self):
+        # Test initialization without logger_provider (default behavior)
+        endpoint = "https://xray.us-west-2.amazonaws.com/v1/traces"
+
+        exporter = OTLPAwsSpanExporter(endpoint=endpoint)
+
+        self.assertIsNone(exporter._logger_provider)
+        self.assertEqual(exporter._aws_region, "us-west-2")