Skip to content

Commit c3e1190

Browse files
authored
Improve resource attributes handling (#414)
*Description of changes:* - Implement automatic injection of aws.ai.agent.type resource attribute when agent observability is enabled - Set default value to "default" if not already present in resource - Add comprehensive unit tests for _customize_resource function - Diable default resource detectors for GenAI cases By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
1 parent 12fabd6 commit c3e1190

File tree

4 files changed

+64
-5
lines changed

4 files changed

+64
-5
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,3 +33,4 @@
3333
AWS_LAMBDA_FUNCTION_NAME: str = "aws.lambda.function.name"
3434
AWS_LAMBDA_RESOURCEMAPPING_ID: str = "aws.lambda.resource_mapping.id"
3535
AWS_LAMBDA_FUNCTION_ARN: str = "aws.lambda.function.arn"
36+
AWS_AI_AGENT_TYPE: str = "aws.ai.agent.type"

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from importlib_metadata import version
1010
from typing_extensions import override
1111

12-
from amazon.opentelemetry.distro._aws_attribute_keys import AWS_LOCAL_SERVICE
12+
from amazon.opentelemetry.distro._aws_attribute_keys import AWS_AI_AGENT_TYPE, AWS_LOCAL_SERVICE
1313
from amazon.opentelemetry.distro._aws_resource_attribute_configurator import get_service_attribute
1414
from amazon.opentelemetry.distro._utils import is_agent_observability_enabled, is_installed
1515
from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler
@@ -174,7 +174,7 @@ def _initialize_components():
174174
AwsEksResourceDetector(),
175175
AwsEcsResourceDetector(),
176176
]
177-
if not _is_lambda_environment()
177+
if not (_is_lambda_environment() or is_agent_observability_enabled())
178178
else []
179179
)
180180

@@ -555,7 +555,15 @@ def _customize_resource(resource: Resource) -> Resource:
555555
if is_unknown:
556556
_logger.debug("No valid service name found")
557557

558-
return resource.merge(Resource.create({AWS_LOCAL_SERVICE: service_name}))
558+
custom_attributes = {AWS_LOCAL_SERVICE: service_name}
559+
560+
if is_agent_observability_enabled():
561+
# Add aws.ai.agent.type if it doesn't exist in the resource
562+
if resource and resource.attributes.get(AWS_AI_AGENT_TYPE) is None:
563+
# Set a default agent type for AI agent observability
564+
custom_attributes[AWS_AI_AGENT_TYPE] = "default"
565+
566+
return resource.merge(Resource.create(custom_attributes))
559567

560568

561569
def _is_application_signals_enabled():

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,6 @@ def _configure(self, **kwargs):
7070

7171
os.environ.setdefault(OTEL_EXPORTER_OTLP_PROTOCOL, "http/protobuf")
7272

73-
super(AwsOpenTelemetryDistro, self)._configure()
74-
7573
os.environ.setdefault(OTEL_PROPAGATORS, "xray,tracecontext,b3,b3multi")
7674
os.environ.setdefault(OTEL_PYTHON_ID_GENERATOR, "xray")
7775
os.environ.setdefault(
@@ -117,5 +115,7 @@ def _configure(self, **kwargs):
117115
# Disable AWS Application Signals by default
118116
os.environ.setdefault(APPLICATION_SIGNALS_ENABLED_CONFIG, "false")
119117

118+
super(AwsOpenTelemetryDistro, self)._configure()
119+
120120
if kwargs.get("apply_patches", True):
121121
apply_instrumentation_patches()

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
from requests import Session
1212

13+
from amazon.opentelemetry.distro._aws_attribute_keys import AWS_AI_AGENT_TYPE, AWS_LOCAL_SERVICE
1314
from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler
1415
from amazon.opentelemetry.distro.attribute_propagating_span_processor import AttributePropagatingSpanProcessor
1516
from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor
@@ -27,6 +28,7 @@
2728
_custom_import_sampler,
2829
_customize_logs_exporter,
2930
_customize_metric_exporters,
31+
_customize_resource,
3032
_customize_sampler,
3133
_customize_span_exporter,
3234
_customize_span_processors,
@@ -67,6 +69,7 @@
6769
from opentelemetry.sdk.trace import Span, SpanProcessor, Tracer, TracerProvider
6870
from opentelemetry.sdk.trace.export import SpanExporter
6971
from opentelemetry.sdk.trace.sampling import DEFAULT_ON, Sampler
72+
from opentelemetry.semconv.resource import ResourceAttributes
7073
from opentelemetry.trace import get_tracer_provider
7174

7275

@@ -1091,6 +1094,53 @@ def test_customize_metric_exporters_with_emf(self):
10911094
self.assertEqual(len(metric_readers), 1)
10921095
self.assertIsInstance(metric_readers[0], PeriodicExportingMetricReader)
10931096

1097+
@patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")
1098+
@patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")
1099+
def test_customize_resource_without_agent_observability(self, mock_get_service_attribute, mock_is_agent_enabled):
1100+
"""Test _customize_resource when agent observability is disabled"""
1101+
mock_is_agent_enabled.return_value = False
1102+
mock_get_service_attribute.return_value = ("test-service", False)
1103+
1104+
resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"})
1105+
result = _customize_resource(resource)
1106+
1107+
# Should only have AWS_LOCAL_SERVICE added
1108+
self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")
1109+
self.assertNotIn(AWS_AI_AGENT_TYPE, result.attributes)
1110+
1111+
@patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")
1112+
@patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")
1113+
def test_customize_resource_with_agent_observability_default(
1114+
self, mock_get_service_attribute, mock_is_agent_enabled
1115+
):
1116+
"""Test _customize_resource when agent observability is enabled with default agent type"""
1117+
mock_is_agent_enabled.return_value = True
1118+
mock_get_service_attribute.return_value = ("test-service", False)
1119+
1120+
resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"})
1121+
result = _customize_resource(resource)
1122+
1123+
# Should have both AWS_LOCAL_SERVICE and AWS_AI_AGENT_TYPE with default value
1124+
self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")
1125+
self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "default")
1126+
1127+
@patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled")
1128+
@patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute")
1129+
def test_customize_resource_with_existing_agent_type(self, mock_get_service_attribute, mock_is_agent_enabled):
1130+
"""Test _customize_resource when agent type already exists in resource"""
1131+
mock_is_agent_enabled.return_value = True
1132+
mock_get_service_attribute.return_value = ("test-service", False)
1133+
1134+
# Create resource with existing agent type
1135+
resource = Resource.create(
1136+
{ResourceAttributes.SERVICE_NAME: "test-service", AWS_AI_AGENT_TYPE: "existing-agent"}
1137+
)
1138+
result = _customize_resource(resource)
1139+
1140+
# Should preserve existing agent type and not override it
1141+
self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service")
1142+
self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "existing-agent")
1143+
10941144

10951145
def validate_distro_environ():
10961146
tc: TestCase = TestCase()

0 commit comments

Comments
 (0)