Skip to content

Commit 2ce0833

Browse files
zzhloginzzhlogin
committed
Add bedrock and bedrockRuntime support.
1 parent ff295b7 commit 2ce0833

File tree

3 files changed

+83
-0
lines changed

3 files changed

+83
-0
lines changed

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,4 +56,13 @@ private AwsAttributeKeys() {}
5656
static final AttributeKey<String> AWS_QUEUE_NAME = AttributeKey.stringKey("aws.queue.name");
5757
static final AttributeKey<String> AWS_STREAM_NAME = AttributeKey.stringKey("aws.stream.name");
5858
static final AttributeKey<String> AWS_TABLE_NAME = AttributeKey.stringKey("aws.table.name");
59+
static final AttributeKey<String> AWS_AGENT_ID = AttributeKey.stringKey("aws.bedrock.agent_id");
60+
static final AttributeKey<String> AWS_KNOWLEDGEBASE_ID =
61+
AttributeKey.stringKey("aws.bedrock.knowledgebase_id");
62+
static final AttributeKey<String> AWS_DATASOURCE_ID =
63+
AttributeKey.stringKey("aws.bedrock.datasource_id");
64+
static final AttributeKey<String> AWS_GUARDRAIL_ID =
65+
AttributeKey.stringKey("aws.bedrock.guardrail_id");
66+
static final AttributeKey<String> AWS_BEDROCK_RUNTIME_MODEL_ID =
67+
AttributeKey.stringKey("gen_ai.request.model");
5968
}

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,11 @@
4141
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_ADDRESS;
4242
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_PORT;
4343
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
44+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_AGENT_ID;
45+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BEDROCK_RUNTIME_MODEL_ID;
46+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATASOURCE_ID;
47+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
48+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGEBASE_ID;
4449
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
4550
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
4651
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_NAME;
@@ -104,6 +109,8 @@ final class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
104109
private static final String NORMALIZED_KINESIS_SERVICE_NAME = "AWS::Kinesis";
105110
private static final String NORMALIZED_S3_SERVICE_NAME = "AWS::S3";
106111
private static final String NORMALIZED_SQS_SERVICE_NAME = "AWS::SQS";
112+
private static final String NORMALIZED_BEDROCK_SERVICE_NAME = "AWS::Bedrock";
113+
private static final String NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME = "AWS::BedrockRuntime";
107114

108115
// Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
109116
private static final String GRAPHQL = "graphql";
@@ -360,6 +367,15 @@ private static String normalizeRemoteServiceName(SpanData span, String serviceNa
360367
case "AmazonSQS": // AWS SDK v1
361368
case "Sqs": // AWS SDK v2
362369
return NORMALIZED_SQS_SERVICE_NAME;
370+
case "Bedrock": // AWS SDK v2 & v1
371+
case "AWSBedrockAgentRuntime": // AWS SDK v1
372+
case "BedrockAgentRuntime": // AWS SDK v2
373+
case "AWSBedrockAgent": // AWS SDK v1
374+
case "BedrockAgent": // AWS SDK v2
375+
return NORMALIZED_BEDROCK_SERVICE_NAME;
376+
case "AmazonBedrockRuntime": // AWS SDK v1
377+
case "BedrockRuntime": // AWS SDK v2
378+
return NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME;
363379
default:
364380
return "AWS::" + serviceName;
365381
}
@@ -403,6 +419,27 @@ private static void setRemoteResourceTypeAndIdentifier(SpanData span, Attributes
403419
remoteResourceType = Optional.of(NORMALIZED_SQS_SERVICE_NAME + "::Queue");
404420
remoteResourceIdentifier =
405421
SqsUrlParser.getQueueName(escapeDelimiters(span.getAttributes().get(AWS_QUEUE_URL)));
422+
} else if (isKeyPresent(span, AWS_AGENT_ID)) {
423+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Agent");
424+
remoteResourceIdentifier =
425+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_AGENT_ID)));
426+
} else if (isKeyPresent(span, AWS_KNOWLEDGEBASE_ID)) {
427+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::KnowledgeBase");
428+
remoteResourceIdentifier =
429+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_KNOWLEDGEBASE_ID)));
430+
} else if (isKeyPresent(span, AWS_DATASOURCE_ID)) {
431+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::DataSource");
432+
remoteResourceIdentifier =
433+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_DATASOURCE_ID)));
434+
} else if (isKeyPresent(span, AWS_GUARDRAIL_ID)) {
435+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Guardrail");
436+
remoteResourceIdentifier =
437+
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_GUARDRAIL_ID)));
438+
} else if (isKeyPresent(span, AWS_BEDROCK_RUNTIME_MODEL_ID)) {
439+
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Model");
440+
remoteResourceIdentifier =
441+
Optional.ofNullable(
442+
escapeDelimiters(span.getAttributes().get(AWS_BEDROCK_RUNTIME_MODEL_ID)));
406443
}
407444
} else if (isDBSpan(span)) {
408445
remoteResourceType = Optional.of(DB_CONNECTION_RESOURCE_TYPE);

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,12 @@
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;
25+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BEDROCK_RUNTIME_MODEL_ID;
2426
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
27+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATASOURCE_ID;
28+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
29+
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGEBASE_ID;
2530
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
2631
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
2732
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_NAME;
@@ -700,6 +705,30 @@ public void testSdkClientSpanWithRemoteResourceAttributes() {
700705
validateRemoteResourceAttributes("AWS::DynamoDB::Table", "aws_table^^name");
701706
mockAttribute(AWS_TABLE_NAME, null);
702707

708+
// Validate behaviour of AWS_BEDROCK_AGENT_ID attribute, then remove it.
709+
mockAttribute(AWS_AGENT_ID, "test_agent_id");
710+
validateRemoteResourceAttributes("AWS::Bedrock::Agent", "test_agent_id");
711+
mockAttribute(AWS_AGENT_ID, null);
712+
713+
// Validate behaviour of AWS_KNOWLEDGEBASE_ID attribute, then remove it.
714+
mockAttribute(AWS_KNOWLEDGEBASE_ID, "test_knowledgeBase_id");
715+
validateRemoteResourceAttributes("AWS::Bedrock::KnowledgeBase", "test_knowledgeBase_id");
716+
mockAttribute(AWS_KNOWLEDGEBASE_ID, null);
717+
718+
// Validate behaviour of AWS_BEDROCK_DATASOURCE_ID attribute, then remove it.
719+
mockAttribute(AWS_DATASOURCE_ID, "test_datasource_id");
720+
validateRemoteResourceAttributes("AWS::Bedrock::DataSource", "test_datasource_id");
721+
mockAttribute(AWS_DATASOURCE_ID, null);
722+
723+
// Validate behaviour of AWS_GUARDRAIL_ID attribute, then remove it.
724+
mockAttribute(AWS_GUARDRAIL_ID, "test_guardrail_id");
725+
validateRemoteResourceAttributes("AWS::Bedrock::Guardrail", "test_guardrail_id");
726+
mockAttribute(AWS_GUARDRAIL_ID, null);
727+
728+
// Validate behaviour of AWS_BEDROCK_RUNTIME_MODEL_ID attribute, then remove it.
729+
mockAttribute(AWS_BEDROCK_RUNTIME_MODEL_ID, "test.service-id");
730+
validateRemoteResourceAttributes("AWS::Bedrock::Model", "test.service-id");
731+
mockAttribute(AWS_BEDROCK_RUNTIME_MODEL_ID, null);
703732
mockAttribute(RPC_SYSTEM, "null");
704733
}
705734

@@ -1047,12 +1076,20 @@ public void testNormalizeRemoteServiceName_AwsSdk() {
10471076
testAwsSdkServiceNormalization("AmazonKinesis", "AWS::Kinesis");
10481077
testAwsSdkServiceNormalization("Amazon S3", "AWS::S3");
10491078
testAwsSdkServiceNormalization("AmazonSQS", "AWS::SQS");
1079+
testAwsSdkServiceNormalization("Bedrock", "AWS::Bedrock");
1080+
testAwsSdkServiceNormalization("AWSBedrockAgentRuntime", "AWS::Bedrock");
1081+
testAwsSdkServiceNormalization("AWSBedrockAgent", "AWS::Bedrock");
1082+
testAwsSdkServiceNormalization("AmazonBedrockRuntime", "AWS::BedrockRuntime");
10501083

10511084
// AWS SDK V2
10521085
testAwsSdkServiceNormalization("DynamoDb", "AWS::DynamoDB");
10531086
testAwsSdkServiceNormalization("Kinesis", "AWS::Kinesis");
10541087
testAwsSdkServiceNormalization("S3", "AWS::S3");
10551088
testAwsSdkServiceNormalization("Sqs", "AWS::SQS");
1089+
testAwsSdkServiceNormalization("Bedrock", "AWS::Bedrock");
1090+
testAwsSdkServiceNormalization("BedrockAgentRuntime", "AWS::Bedrock");
1091+
testAwsSdkServiceNormalization("BedrockAgent", "AWS::Bedrock");
1092+
testAwsSdkServiceNormalization("BedrockRuntime", "AWS::BedrockRuntime");
10561093
}
10571094

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

0 commit comments

Comments
 (0)