Skip to content

Commit

Permalink
Merge pull request #826 from zzhlogin/enablement_bedrock
Browse files Browse the repository at this point in the history 
Add bedrock and bedrockRuntime support.
  • Loading branch information
@thpierce
thpierce authored Jul 9, 2024
2 parents dd65f52 + 3678a6c commit 2847ea5
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 0 deletions.
7 changes: 7 additions & 0 deletions ...der/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,4 +59,11 @@ private AwsAttributeKeys() {}
static final AttributeKey<String> AWS_QUEUE_NAME = AttributeKey.stringKey("aws.queue.name");
static final AttributeKey<String> AWS_STREAM_NAME = AttributeKey.stringKey("aws.stream.name");
static final AttributeKey<String> AWS_TABLE_NAME = AttributeKey.stringKey("aws.table.name");
static final AttributeKey<String> AWS_AGENT_ID = AttributeKey.stringKey("aws.bedrock.agent.id");
static final AttributeKey<String> AWS_KNOWLEDGE_BASE_ID =
AttributeKey.stringKey("aws.bedrock.knowledge_base.id");
static final AttributeKey<String> AWS_DATA_SOURCE_ID =
AttributeKey.stringKey("aws.bedrock.data_source.id");
static final AttributeKey<String> AWS_GUARDRAIL_ID =
AttributeKey.stringKey("aws.bedrock.guardrail.id");
}
41 changes: 41 additions & 0 deletions ...n/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGenerator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,11 @@
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_PORT;
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_ADDRESS;
import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_PORT;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_AGENT_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATA_SOURCE_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGE_BASE_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_NAME;
Expand All @@ -54,6 +58,7 @@
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_SPAN_KIND;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_STREAM_NAME;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_TABLE_NAME;
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL;
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.MAX_KEYWORD_LENGTH;
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.SQL_DIALECT_PATTERN;
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.UNKNOWN_OPERATION;
Expand Down Expand Up @@ -106,6 +111,8 @@ final class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
private static final String NORMALIZED_KINESIS_SERVICE_NAME = "AWS::Kinesis";
private static final String NORMALIZED_S3_SERVICE_NAME = "AWS::S3";
private static final String NORMALIZED_SQS_SERVICE_NAME = "AWS::SQS";
private static final String NORMALIZED_BEDROCK_SERVICE_NAME = "AWS::Bedrock";
private static final String NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME = "AWS::BedrockRuntime";

// Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
private static final String GRAPHQL = "graphql";
Expand Down Expand Up @@ -363,6 +370,20 @@ private static String normalizeRemoteServiceName(SpanData span, String serviceNa
case "AmazonSQS": // AWS SDK v1
case "Sqs": // AWS SDK v2
return NORMALIZED_SQS_SERVICE_NAME;
// For Bedrock, Bedrock Agent, and Bedrock Agent Runtime, we can align with AWS Cloud
// Control and use AWS::Bedrock for RemoteService.
case "AmazonBedrock": // AWS SDK v1
case "Bedrock": // AWS SDK v2
case "AWSBedrockAgentRuntime": // AWS SDK v1
case "BedrockAgentRuntime": // AWS SDK v2
case "AWSBedrockAgent": // AWS SDK v1
case "BedrockAgent": // AWS SDK v2
return NORMALIZED_BEDROCK_SERVICE_NAME;
// For BedrockRuntime, we are using AWS::BedrockRuntime as the associated remote resource
// (Model) is not listed in Cloud Control.
case "AmazonBedrockRuntime": // AWS SDK v1
case "BedrockRuntime": // AWS SDK v2
return NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME;
default:
return "AWS::" + serviceName;
}
Expand Down Expand Up @@ -406,6 +427,26 @@ private static void setRemoteResourceTypeAndIdentifier(SpanData span, Attributes
remoteResourceType = Optional.of(NORMALIZED_SQS_SERVICE_NAME + "::Queue");
remoteResourceIdentifier =
SqsUrlParser.getQueueName(escapeDelimiters(span.getAttributes().get(AWS_QUEUE_URL)));
} else if (isKeyPresent(span, AWS_AGENT_ID)) {
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Agent");
remoteResourceIdentifier =
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_AGENT_ID)));
} else if (isKeyPresent(span, AWS_KNOWLEDGE_BASE_ID)) {
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::KnowledgeBase");
remoteResourceIdentifier =
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_KNOWLEDGE_BASE_ID)));
} else if (isKeyPresent(span, AWS_DATA_SOURCE_ID)) {
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::DataSource");
remoteResourceIdentifier =
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_DATA_SOURCE_ID)));
} else if (isKeyPresent(span, AWS_GUARDRAIL_ID)) {
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Guardrail");
remoteResourceIdentifier =
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_GUARDRAIL_ID)));
} else if (isKeyPresent(span, GEN_AI_REQUEST_MODEL)) {
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Model");
remoteResourceIdentifier =
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(GEN_AI_REQUEST_MODEL)));
}
} else if (isDBSpan(span)) {
remoteResourceType = Optional.of(DB_CONNECTION_RESOURCE_TYPE);
Expand Down
3 changes: 3 additions & 0 deletions ...rc/main/java/software/amazon/opentelemetry/javaagent/providers/AwsSpanProcessingUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,9 @@ final class AwsSpanProcessingUtil {
// The current longest command word is DATETIME_INTERVAL_PRECISION at 27 characters.
// If we add a longer keyword to the sql dialect keyword list, need to update the constant below.
static final int MAX_KEYWORD_LENGTH = 27;
// TODO: Use Semantic Conventions once upgrade once upgrade to v1.26.0
static final AttributeKey<String> GEN_AI_REQUEST_MODEL =
AttributeKey.stringKey("gen_ai.request.model");
static final Pattern SQL_DIALECT_PATTERN =
Pattern.compile("^(?:" + String.join("|", getDialectKeywords()) + ")\\b");

Expand Down
64 changes: 64 additions & 0 deletions ...va/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGeneratorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,11 @@
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_AGENT_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATA_SOURCE_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGE_BASE_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_NAME;
Expand All @@ -34,6 +38,7 @@
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_SPAN_KIND;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_STREAM_NAME;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_TABLE_NAME;
import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL;
import static software.amazon.opentelemetry.javaagent.providers.MetricAttributeGenerator.DEPENDENCY_METRIC;
import static software.amazon.opentelemetry.javaagent.providers.MetricAttributeGenerator.SERVICE_METRIC;

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

// Validate behaviour of AWS_BEDROCK_AGENT_ID attribute, then remove it.
mockAttribute(AWS_AGENT_ID, "test_agent_id");
validateRemoteResourceAttributes("AWS::Bedrock::Agent", "test_agent_id");
mockAttribute(AWS_AGENT_ID, null);

// Validate behaviour of AWS_BEDROCK_AGENT_ID attribute with special chars(^), then remove it.
mockAttribute(AWS_AGENT_ID, "test_agent_^id");
validateRemoteResourceAttributes("AWS::Bedrock::Agent", "test_agent_^^id");
mockAttribute(AWS_AGENT_ID, null);

// Validate behaviour of AWS_KNOWLEDGE_BASE_ID attribute, then remove it.
mockAttribute(AWS_KNOWLEDGE_BASE_ID, "test_knowledgeBase_id");
validateRemoteResourceAttributes("AWS::Bedrock::KnowledgeBase", "test_knowledgeBase_id");
mockAttribute(AWS_KNOWLEDGE_BASE_ID, null);

// Validate behaviour of AWS_KNOWLEDGE_BASE_ID attribute with special chars(^), then remove it.
mockAttribute(AWS_KNOWLEDGE_BASE_ID, "test_knowledgeBase_^id");
validateRemoteResourceAttributes("AWS::Bedrock::KnowledgeBase", "test_knowledgeBase_^^id");
mockAttribute(AWS_KNOWLEDGE_BASE_ID, null);

// Validate behaviour of AWS_DATA_SOURCE_ID attribute, then remove it.
mockAttribute(AWS_DATA_SOURCE_ID, "test_datasource_id");
validateRemoteResourceAttributes("AWS::Bedrock::DataSource", "test_datasource_id");
mockAttribute(AWS_DATA_SOURCE_ID, null);

// Validate behaviour of AWS_DATA_SOURCE_ID attribute with special chars(^), then remove
// it.
mockAttribute(AWS_DATA_SOURCE_ID, "test_datasource_^id");
validateRemoteResourceAttributes("AWS::Bedrock::DataSource", "test_datasource_^^id");
mockAttribute(AWS_DATA_SOURCE_ID, null);

// Validate behaviour of AWS_GUARDRAIL_ID attribute, then remove it.
mockAttribute(AWS_GUARDRAIL_ID, "test_guardrail_id");
validateRemoteResourceAttributes("AWS::Bedrock::Guardrail", "test_guardrail_id");
mockAttribute(AWS_GUARDRAIL_ID, null);

// Validate behaviour of AWS_GUARDRAIL_ID attribute with special chars(^), then remove it.
mockAttribute(AWS_GUARDRAIL_ID, "test_guardrail_^id");
validateRemoteResourceAttributes("AWS::Bedrock::Guardrail", "test_guardrail_^^id");
mockAttribute(AWS_GUARDRAIL_ID, null);

// Validate behaviour of AWS_BEDROCK_RUNTIME_MODEL_ID attribute, then remove it.
mockAttribute(GEN_AI_REQUEST_MODEL, "test.service_id");
validateRemoteResourceAttributes("AWS::Bedrock::Model", "test.service_id");
mockAttribute(GEN_AI_REQUEST_MODEL, null);

// Validate behaviour of AWS_BEDROCK_RUNTIME_MODEL_ID attribute with special chars(^), then
// remove it.
mockAttribute(GEN_AI_REQUEST_MODEL, "test.service_^id");
validateRemoteResourceAttributes("AWS::Bedrock::Model", "test.service_^^id");
mockAttribute(GEN_AI_REQUEST_MODEL, null);
mockAttribute(RPC_SYSTEM, "null");
}

Expand Down Expand Up @@ -1102,12 +1158,20 @@ public void testNormalizeRemoteServiceName_AwsSdk() {
testAwsSdkServiceNormalization("AmazonKinesis", "AWS::Kinesis");
testAwsSdkServiceNormalization("Amazon S3", "AWS::S3");
testAwsSdkServiceNormalization("AmazonSQS", "AWS::SQS");
testAwsSdkServiceNormalization("Bedrock", "AWS::Bedrock");
testAwsSdkServiceNormalization("AWSBedrockAgentRuntime", "AWS::Bedrock");
testAwsSdkServiceNormalization("AWSBedrockAgent", "AWS::Bedrock");
testAwsSdkServiceNormalization("AmazonBedrockRuntime", "AWS::BedrockRuntime");

// AWS SDK V2
testAwsSdkServiceNormalization("DynamoDb", "AWS::DynamoDB");
testAwsSdkServiceNormalization("Kinesis", "AWS::Kinesis");
testAwsSdkServiceNormalization("S3", "AWS::S3");
testAwsSdkServiceNormalization("Sqs", "AWS::SQS");
testAwsSdkServiceNormalization("Bedrock", "AWS::Bedrock");
testAwsSdkServiceNormalization("BedrockAgentRuntime", "AWS::Bedrock");
testAwsSdkServiceNormalization("BedrockAgent", "AWS::Bedrock");
testAwsSdkServiceNormalization("BedrockRuntime", "AWS::BedrockRuntime");
}

private void testAwsSdkServiceNormalization(String serviceName, String expectedRemoteService) {
Expand Down

0 comments on commit 2847ea5

Please sign in to comment.