Improve resource attributes handling

Author: mxiamxia

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

@@ -33,3 +33,4 @@
 AWS_LAMBDA_FUNCTION_NAME: str = "aws.lambda.function.name"
 AWS_LAMBDA_RESOURCEMAPPING_ID: str = "aws.lambda.resource_mapping.id"
 AWS_LAMBDA_FUNCTION_ARN: str = "aws.lambda.function.arn"
+AWS_AI_AGENT_TYPE: str = "aws.ai.agent.type"

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

@@ -9,7 +9,7 @@
 from importlib_metadata import version
 from typing_extensions import override
 
-from amazon.opentelemetry.distro._aws_attribute_keys import AWS_LOCAL_SERVICE
+from amazon.opentelemetry.distro._aws_attribute_keys import AWS_AI_AGENT_TYPE, 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, is_installed
 from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler
@@ -174,7 +174,7 @@ def _initialize_components():
             AwsEksResourceDetector(),
             AwsEcsResourceDetector(),
         ]
-        if not _is_lambda_environment()
+        if not (_is_lambda_environment() or is_agent_observability_enabled())
         else []
     )
 
@@ -555,7 +555,15 @@ def _customize_resource(resource: Resource) -> Resource:
     if is_unknown:
         _logger.debug("No valid service name found")
 
-    return resource.merge(Resource.create({AWS_LOCAL_SERVICE: service_name}))
+    custom_attributes = {AWS_LOCAL_SERVICE: service_name}
+
+    if is_agent_observability_enabled():
+        # Add aws.ai.agent.type if it doesn't exist in the resource
+        if resource and resource.attributes.get(AWS_AI_AGENT_TYPE) is None:
+            # Set a default agent type for AI agent observability
+            custom_attributes[AWS_AI_AGENT_TYPE] = "default"
+
+    return resource.merge(Resource.create(custom_attributes))
 
 
 def _is_application_signals_enabled():

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

@@ -70,8 +70,6 @@ def _configure(self, **kwargs):
 
         os.environ.setdefault(OTEL_EXPORTER_OTLP_PROTOCOL, "http/protobuf")
 
-        super(AwsOpenTelemetryDistro, self)._configure()
-
         os.environ.setdefault(OTEL_PROPAGATORS, "xray,tracecontext,b3,b3multi")
         os.environ.setdefault(OTEL_PYTHON_ID_GENERATOR, "xray")
         os.environ.setdefault(
@@ -117,5 +115,7 @@ def _configure(self, **kwargs):
             # Disable AWS Application Signals by default
             os.environ.setdefault(APPLICATION_SIGNALS_ENABLED_CONFIG, "false")
 
+        super(AwsOpenTelemetryDistro, self)._configure()
+
         if kwargs.get("apply_patches", True):
             apply_instrumentation_patches()

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

@@ -10,6 +10,7 @@
 
 from requests import Session
 
+from amazon.opentelemetry.distro._aws_attribute_keys import AWS_AI_AGENT_TYPE, AWS_LOCAL_SERVICE
 from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler
 from amazon.opentelemetry.distro.attribute_propagating_span_processor import AttributePropagatingSpanProcessor
 from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor
@@ -27,6 +28,7 @@
     _custom_import_sampler,
     _customize_logs_exporter,
     _customize_metric_exporters,
+    _customize_resource,
     _customize_sampler,
     _customize_span_exporter,
     _customize_span_processors,
@@ -67,6 +69,7 @@
 from opentelemetry.sdk.trace import Span, SpanProcessor, Tracer, TracerProvider
 from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.sdk.trace.sampling import DEFAULT_ON, Sampler
+from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.trace import get_tracer_provider
 
 
@@ -1091,6 +1094,53 @@ def test_customize_metric_exporters_with_emf(self):
             self.assertEqual(len(metric_readers), 1)
             self.assertIsInstance(metric_readers[0], PeriodicExportingMetricReader)
 
+    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")
+    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")
+    def test_customize_resource_without_agent_observability(self, mock_get_service_attribute, mock_is_agent_enabled):
+        """Test _customize_resource when agent observability is disabled"""
+        mock_is_agent_enabled.return_value = False
+        mock_get_service_attribute.return_value = ("test-service", False)
+
+        resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"})
+        result = _customize_resource(resource)
+
+        # Should only have AWS_LOCAL_SERVICE added
+        self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")
+        self.assertNotIn(AWS_AI_AGENT_TYPE, result.attributes)
+
+    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")
+    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")
+    def test_customize_resource_with_agent_observability_default(
+        self, mock_get_service_attribute, mock_is_agent_enabled
+    ):
+        """Test _customize_resource when agent observability is enabled with default agent type"""
+        mock_is_agent_enabled.return_value = True
+        mock_get_service_attribute.return_value = ("test-service", False)
+
+        resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"})
+        result = _customize_resource(resource)
+
+        # Should have both AWS_LOCAL_SERVICE and AWS_AI_AGENT_TYPE with default value
+        self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")
+        self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "default")
+
+    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")
+    @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")
+    def test_customize_resource_with_existing_agent_type(self, mock_get_service_attribute, mock_is_agent_enabled):
+        """Test _customize_resource when agent type already exists in resource"""
+        mock_is_agent_enabled.return_value = True
+        mock_get_service_attribute.return_value = ("test-service", False)
+
+        # Create resource with existing agent type
+        resource = Resource.create(
+            {ResourceAttributes.SERVICE_NAME: "test-service", AWS_AI_AGENT_TYPE: "existing-agent"}
+        )
+        result = _customize_resource(resource)
+
+        # Should preserve existing agent type and not override it
+        self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")
+        self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "existing-agent")
+
 
 def validate_distro_environ():
     tc: TestCase = TestCase()