Skip to content

Commit 2847ea5

Browse files
thpiercethpierce
authored
Merge pull request #826 from zzhlogin/enablement_bedrock
Add bedrock and bedrockRuntime support.
2 parents dd65f52 + 3678a6c commit 2847ea5

File tree

4 files changed

+115
-0
lines changed

4 files changed

+115
-0
lines changed

‎awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,4 +59,11 @@ private AwsAttributeKeys() {}
5959
static final AttributeKey<String> AWS_QUEUE_NAME = AttributeKey.stringKey("aws.queue.name");
6060
static final AttributeKey<String> AWS_STREAM_NAME = AttributeKey.stringKey("aws.stream.name");
6161
static final AttributeKey<String> AWS_TABLE_NAME = AttributeKey.stringKey("aws.table.name");
62+
static final AttributeKey<String> AWS_AGENT_ID = AttributeKey.stringKey("aws.bedrock.agent.id");
63+
static final AttributeKey<String> AWS_KNOWLEDGE_BASE_ID =
64+
AttributeKey.stringKey("aws.bedrock.knowledge_base.id");
65+
static final AttributeKey<String> AWS_DATA_SOURCE_ID =
66+
AttributeKey.stringKey("aws.bedrock.data_source.id");
67+
static final AttributeKey<String> AWS_GUARDRAIL_ID =
68+
AttributeKey.stringKey("aws.bedrock.guardrail.id");
6269
}

‎awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGenerator.java

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,11 @@
4141
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_PORT;
4242
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_ADDRESS;
4343
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_PORT;
44+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_AGENT_ID;
4445
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
46+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATA_SOURCE_ID;
47+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
48+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGE_BASE_ID;
4549
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
4650
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
4751
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_NAME;
@@ -54,6 +58,7 @@
5458
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_SPAN_KIND;
5559
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_STREAM_NAME;
5660
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_TABLE_NAME;
61+
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL;
5762
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.MAX_KEYWORD_LENGTH;
5863
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.SQL_DIALECT_PATTERN;
5964
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.UNKNOWN_OPERATION;
@@ -106,6 +111,8 @@ final class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
106111
private static final String NORMALIZED_KINESIS_SERVICE_NAME = "AWS::Kinesis";
107112
private static final String NORMALIZED_S3_SERVICE_NAME = "AWS::S3";
108113
private static final String NORMALIZED_SQS_SERVICE_NAME = "AWS::SQS";
114+
private static final String NORMALIZED_BEDROCK_SERVICE_NAME = "AWS::Bedrock";
115+
private static final String NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME = "AWS::BedrockRuntime";
109116

110117
// Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
111118
private static final String GRAPHQL = "graphql";
@@ -363,6 +370,20 @@ private static String normalizeRemoteServiceName(SpanData span, String serviceNa
363370
case "AmazonSQS": // AWS SDK v1
364371
case "Sqs": // AWS SDK v2
365372
return NORMALIZED_SQS_SERVICE_NAME;
373+
// For Bedrock, Bedrock Agent, and Bedrock Agent Runtime, we can align with AWS Cloud
374+
// Control and use AWS::Bedrock for RemoteService.
375+
case "AmazonBedrock": // AWS SDK v1
376+
case "Bedrock": // AWS SDK v2
377+
case "AWSBedrockAgentRuntime": // AWS SDK v1
378+
case "BedrockAgentRuntime": // AWS SDK v2
379+
case "AWSBedrockAgent": // AWS SDK v1
380+
case "BedrockAgent": // AWS SDK v2
381+
return NORMALIZED_BEDROCK_SERVICE_NAME;
382+
// For BedrockRuntime, we are using AWS::BedrockRuntime as the associated remote resource
383+
// (Model) is not listed in Cloud Control.
384+
case "AmazonBedrockRuntime": // AWS SDK v1
385+
case "BedrockRuntime": // AWS SDK v2
386+
return NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME;
366387
default:
367388
return "AWS::" + serviceName;
368389
}
@@ -406,6 +427,26 @@ private static void setRemoteResourceTypeAndIdentifier(SpanData span, Attributes
406427
remoteResourceType = Optional.of(NORMALIZED_SQS_SERVICE_NAME + "::Queue");
407428
remoteResourceIdentifier =
408429
SqsUrlParser.getQueueName(escapeDelimiters(span.getAttributes().get(AWS_QUEUE_URL)));
430+
} else if (isKeyPresent(span, AWS_AGENT_ID)) {
431+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Agent");
432+
remoteResourceIdentifier =
433+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_AGENT_ID)));
434+
} else if (isKeyPresent(span, AWS_KNOWLEDGE_BASE_ID)) {
435+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::KnowledgeBase");
436+
remoteResourceIdentifier =
437+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_KNOWLEDGE_BASE_ID)));
438+
} else if (isKeyPresent(span, AWS_DATA_SOURCE_ID)) {
439+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::DataSource");
440+
remoteResourceIdentifier =
441+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_DATA_SOURCE_ID)));
442+
} else if (isKeyPresent(span, AWS_GUARDRAIL_ID)) {
443+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Guardrail");
444+
remoteResourceIdentifier =
445+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_GUARDRAIL_ID)));
446+
} else if (isKeyPresent(span, GEN_AI_REQUEST_MODEL)) {
447+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Model");
448+
remoteResourceIdentifier =
449+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(GEN_AI_REQUEST_MODEL)));
409450
}
410451
} else if (isDBSpan(span)) {
411452
remoteResourceType = Optional.of(DB_CONNECTION_RESOURCE_TYPE);

‎awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsSpanProcessingUtil.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ final class AwsSpanProcessingUtil {
5656
// The current longest command word is DATETIME_INTERVAL_PRECISION at 27 characters.
5757
// If we add a longer keyword to the sql dialect keyword list, need to update the constant below.
5858
static final int MAX_KEYWORD_LENGTH = 27;
59+
// TODO: Use Semantic Conventions once upgrade once upgrade to v1.26.0
60+
static final AttributeKey<String> GEN_AI_REQUEST_MODEL =
61+
AttributeKey.stringKey("gen_ai.request.model");
5962
static final Pattern SQL_DIALECT_PATTERN =
6063
Pattern.compile("^(?:" + String.join("|", getDialectKeywords()) + ")\\b");
6164

‎awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGeneratorTest.java

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,11 @@
2121
import static org.assertj.core.api.Assertions.assertThat;
2222
import static org.mockito.Mockito.mock;
2323
import static org.mockito.Mockito.when;
24+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_AGENT_ID;
2425
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
26+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATA_SOURCE_ID;
27+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
28+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGE_BASE_ID;
2529
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
2630
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
2731
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_NAME;
@@ -34,6 +38,7 @@
3438
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_SPAN_KIND;
3539
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_STREAM_NAME;
3640
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_TABLE_NAME;
41+
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL;
3742
import static software.amazon.opentelemetry.javaagent.providers.MetricAttributeGenerator.DEPENDENCY_METRIC;
3843
import static software.amazon.opentelemetry.javaagent.providers.MetricAttributeGenerator.SERVICE_METRIC;
3944

@@ -701,6 +706,57 @@ public void testSdkClientSpanWithRemoteResourceAttributes() {
701706
validateRemoteResourceAttributes("AWS::DynamoDB::Table", "aws_table^^name");
702707
mockAttribute(AWS_TABLE_NAME, null);
703708

709+
// Validate behaviour of AWS_BEDROCK_AGENT_ID attribute, then remove it.
710+
mockAttribute(AWS_AGENT_ID, "test_agent_id");
711+
validateRemoteResourceAttributes("AWS::Bedrock::Agent", "test_agent_id");
712+
mockAttribute(AWS_AGENT_ID, null);
713+
714+
// Validate behaviour of AWS_BEDROCK_AGENT_ID attribute with special chars(^), then remove it.
715+
mockAttribute(AWS_AGENT_ID, "test_agent_^id");
716+
validateRemoteResourceAttributes("AWS::Bedrock::Agent", "test_agent_^^id");
717+
mockAttribute(AWS_AGENT_ID, null);
718+
719+
// Validate behaviour of AWS_KNOWLEDGE_BASE_ID attribute, then remove it.
720+
mockAttribute(AWS_KNOWLEDGE_BASE_ID, "test_knowledgeBase_id");
721+
validateRemoteResourceAttributes("AWS::Bedrock::KnowledgeBase", "test_knowledgeBase_id");
722+
mockAttribute(AWS_KNOWLEDGE_BASE_ID, null);
723+
724+
// Validate behaviour of AWS_KNOWLEDGE_BASE_ID attribute with special chars(^), then remove it.
725+
mockAttribute(AWS_KNOWLEDGE_BASE_ID, "test_knowledgeBase_^id");
726+
validateRemoteResourceAttributes("AWS::Bedrock::KnowledgeBase", "test_knowledgeBase_^^id");
727+
mockAttribute(AWS_KNOWLEDGE_BASE_ID, null);
728+
729+
// Validate behaviour of AWS_DATA_SOURCE_ID attribute, then remove it.
730+
mockAttribute(AWS_DATA_SOURCE_ID, "test_datasource_id");
731+
validateRemoteResourceAttributes("AWS::Bedrock::DataSource", "test_datasource_id");
732+
mockAttribute(AWS_DATA_SOURCE_ID, null);
733+
734+
// Validate behaviour of AWS_DATA_SOURCE_ID attribute with special chars(^), then remove
735+
// it.
736+
mockAttribute(AWS_DATA_SOURCE_ID, "test_datasource_^id");
737+
validateRemoteResourceAttributes("AWS::Bedrock::DataSource", "test_datasource_^^id");
738+
mockAttribute(AWS_DATA_SOURCE_ID, null);
739+
740+
// Validate behaviour of AWS_GUARDRAIL_ID attribute, then remove it.
741+
mockAttribute(AWS_GUARDRAIL_ID, "test_guardrail_id");
742+
validateRemoteResourceAttributes("AWS::Bedrock::Guardrail", "test_guardrail_id");
743+
mockAttribute(AWS_GUARDRAIL_ID, null);
744+
745+
// Validate behaviour of AWS_GUARDRAIL_ID attribute with special chars(^), then remove it.
746+
mockAttribute(AWS_GUARDRAIL_ID, "test_guardrail_^id");
747+
validateRemoteResourceAttributes("AWS::Bedrock::Guardrail", "test_guardrail_^^id");
748+
mockAttribute(AWS_GUARDRAIL_ID, null);
749+
750+
// Validate behaviour of AWS_BEDROCK_RUNTIME_MODEL_ID attribute, then remove it.
751+
mockAttribute(GEN_AI_REQUEST_MODEL, "test.service_id");
752+
validateRemoteResourceAttributes("AWS::Bedrock::Model", "test.service_id");
753+
mockAttribute(GEN_AI_REQUEST_MODEL, null);
754+
755+
// Validate behaviour of AWS_BEDROCK_RUNTIME_MODEL_ID attribute with special chars(^), then
756+
// remove it.
757+
mockAttribute(GEN_AI_REQUEST_MODEL, "test.service_^id");
758+
validateRemoteResourceAttributes("AWS::Bedrock::Model", "test.service_^^id");
759+
mockAttribute(GEN_AI_REQUEST_MODEL, null);
704760
mockAttribute(RPC_SYSTEM, "null");
705761
}
706762

@@ -1102,12 +1158,20 @@ public void testNormalizeRemoteServiceName_AwsSdk() {
11021158
testAwsSdkServiceNormalization("AmazonKinesis", "AWS::Kinesis");
11031159
testAwsSdkServiceNormalization("Amazon S3", "AWS::S3");
11041160
testAwsSdkServiceNormalization("AmazonSQS", "AWS::SQS");
1161+
testAwsSdkServiceNormalization("Bedrock", "AWS::Bedrock");
1162+
testAwsSdkServiceNormalization("AWSBedrockAgentRuntime", "AWS::Bedrock");
1163+
testAwsSdkServiceNormalization("AWSBedrockAgent", "AWS::Bedrock");
1164+
testAwsSdkServiceNormalization("AmazonBedrockRuntime", "AWS::BedrockRuntime");
11051165

11061166
// AWS SDK V2
11071167
testAwsSdkServiceNormalization("DynamoDb", "AWS::DynamoDB");
11081168
testAwsSdkServiceNormalization("Kinesis", "AWS::Kinesis");
11091169
testAwsSdkServiceNormalization("S3", "AWS::S3");
11101170
testAwsSdkServiceNormalization("Sqs", "AWS::SQS");
1171+
testAwsSdkServiceNormalization("Bedrock", "AWS::Bedrock");
1172+
testAwsSdkServiceNormalization("BedrockAgentRuntime", "AWS::Bedrock");
1173+
testAwsSdkServiceNormalization("BedrockAgent", "AWS::Bedrock");
1174+
testAwsSdkServiceNormalization("BedrockRuntime", "AWS::BedrockRuntime");
11111175
}
11121176

11131177
private void testAwsSdkServiceNormalization(String serviceName, String expectedRemoteService) {

0 commit comments

Comments
 (0)