Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add bedrock and bedrockRuntime support. #826

Merged
merged 7 commits into from
Jul 9, 2024
Merged

Add bedrock and bedrockRuntime support. #826

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
2ce0833
Add bedrock and bedrockRuntime support.
zzhlogin Jun 28, 2024
06d874e
Align with Python
zzhlogin Jul 8, 2024
01eb413
Apply spotless.
zzhlogin Jul 8, 2024
a9ab2d8
Fix naming.
zzhlogin Jul 9, 2024
c7d0c49
Refine keys name.
zzhlogin Jul 9, 2024
53e87ce
update comment.
zzhlogin Jul 9, 2024
3678a6c
Trigger workflow: no-op commit
zzhlogin Jul 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions ...der/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,4 +56,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_KNOWLEDGEBASE_ID =
AttributeKey.stringKey("aws.bedrock.knowledgebase.id");
static final AttributeKey<String> AWS_DATASOURCE_ID =
AttributeKey.stringKey("aws.bedrock.datasource.id");
zzhlogin marked this conversation as resolved.
Show resolved Hide resolved
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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,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_DATASOURCE_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGEBASE_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 @@ -52,6 +56,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 @@ -104,6 +109,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 @@ -360,6 +367,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 @@ -403,6 +424,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_KNOWLEDGEBASE_ID)) {
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::KnowledgeBase");
remoteResourceIdentifier =
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_KNOWLEDGEBASE_ID)));
} else if (isKeyPresent(span, AWS_DATASOURCE_ID)) {
remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::DataSource");
remoteResourceIdentifier =
Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_DATASOURCE_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
Open in desktop
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
Open in desktop
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_DATASOURCE_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGEBASE_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 @@ -33,6 +37,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 @@ -700,6 +705,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_KNOWLEDGEBASE_ID attribute, then remove it.
mockAttribute(AWS_KNOWLEDGEBASE_ID, "test_knowledgeBase_id");
validateRemoteResourceAttributes("AWS::Bedrock::KnowledgeBase", "test_knowledgeBase_id");
mockAttribute(AWS_KNOWLEDGEBASE_ID, null);

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

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

// Validate behaviour of AWS_BEDROCK_DATASOURCE_ID attribute with special chars(^), then remove
// it.
mockAttribute(AWS_DATASOURCE_ID, "test_datasource_^id");
validateRemoteResourceAttributes("AWS::Bedrock::DataSource", "test_datasource_^^id");
mockAttribute(AWS_DATASOURCE_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 @@ -1047,12 +1103,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
Loading