Skip to content

Commit 359d60f

Browse files
zzhloginzzhlogin
authored
Add bedrock runtime Java SDK V1
This PR add AWS Bedrock and BedrockRuntime support for AWS Java SDK V1 with the following change: A. **Bedrock**: Extract `guardrailId` from API response, and add into `"aws.bedrock.guardrail.id"` span attribute. B. **Bedrock Agent**: Extract `agentId` from both API request and response, and add into `"aws.bedrock.agent.id"` span attribute. Extract `knowledgeBaseId` from API request, and add into `"aws.bedrock.knowledgebase.id"` span attribute. Extract `dataSourceId` from both API request and response,, and add into `"aws.bedrock.datasource.id"` span attribute. The instrumentation is on API operation level, we make sure only one attribute is extracted per API call, there will be no overlap/conflict to identify the resource. C. **Bedrock Agent Runtime**: Extract `agentId` from API request, and add into `"aws.bedrock.agent.id"` span attribute. Extract `knowledgeBaseId` from API request, and add into `"aws.bedrock.knowledgebase.id"` span attribute. D. **Bedrock Runtime**: Extract the following attributes and add into span according to [ Gen AI semantic-conventions](https://github.com/open-telemetry/semantic-conventions/blob/main/docs/attributes-registry/gen-ai.md): ``` gen_ai.request.model gen_ai.system ```
1 parent 305f32e commit 359d60f

File tree

9 files changed

+372
-19
lines changed

9 files changed

+372
-19
lines changed

‎instrumentation/aws-sdk/aws-sdk-1.11/javaagent/build.gradle.kts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,10 @@ dependencies {
4747
testLibrary("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
4848
testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
4949
testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
50+
testLibrary("com.amazonaws:aws-java-sdk-bedrock:1.12.744")
51+
testLibrary("com.amazonaws:aws-java-sdk-bedrockagent:1.12.744")
52+
testLibrary("com.amazonaws:aws-java-sdk-bedrockagentruntime:1.12.744")
53+
testLibrary("com.amazonaws:aws-java-sdk-bedrockruntime:1.12.744")
5054

5155
testImplementation(project(":instrumentation:aws-sdk:aws-sdk-1.11:testing"))
5256

‎instrumentation/aws-sdk/aws-sdk-1.11/library-autoconfigure/build.gradle.kts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@ dependencies {
1818
testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
1919
testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
2020
testLibrary("com.amazonaws:aws-java-sdk-sqs:1.11.106")
21+
testLibrary("com.amazonaws:aws-java-sdk-bedrock:1.12.744")
22+
testLibrary("com.amazonaws:aws-java-sdk-bedrockagent:1.12.744")
23+
testLibrary("com.amazonaws:aws-java-sdk-bedrockagentruntime:1.12.744")
24+
testLibrary("com.amazonaws:aws-java-sdk-bedrockruntime:1.12.744")
2125

2226
// last version that does not use json protocol
2327
latestDepTestLibrary("com.amazonaws:aws-java-sdk-sqs:1.12.583")

‎instrumentation/aws-sdk/aws-sdk-1.11/library/build.gradle.kts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ dependencies {
1717
testLibrary("com.amazonaws:aws-java-sdk-kinesis:1.11.106")
1818
testLibrary("com.amazonaws:aws-java-sdk-dynamodb:1.11.106")
1919
testLibrary("com.amazonaws:aws-java-sdk-sns:1.11.106")
20+
testLibrary("com.amazonaws:aws-java-sdk-bedrock:1.12.744")
21+
testLibrary("com.amazonaws:aws-java-sdk-bedrockagent:1.12.744")
22+
testLibrary("com.amazonaws:aws-java-sdk-bedrockagentruntime:1.12.744")
23+
testLibrary("com.amazonaws:aws-java-sdk-bedrockruntime:1.12.744")
2024

2125
// last version that does not use json protocol
2226
latestDepTestLibrary("com.amazonaws:aws-java-sdk-sqs:1.12.583")

‎instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsBedrockResourceType.java‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.awssdk.v1_11;
7+
8+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_AGENT_ID;
9+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_DATASOURCE_ID;
10+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGEBASE_ID;
11+
12+
import io.opentelemetry.api.common.AttributeKey;
13+
import java.util.Arrays;
14+
import java.util.HashMap;
15+
import java.util.List;
16+
import java.util.Map;
17+
import java.util.function.Function;
18+
19+
enum AwsBedrockResourceType {
20+
AGENT_TYPE(AWS_AGENT_ID, RequestAccess::getAgentId),
21+
DATA_SOURCE_TYPE(AWS_DATASOURCE_ID, RequestAccess::getDataSourceId),
22+
KNOWLEDGE_BASE_TYPE(AWS_KNOWLEDGEBASE_ID, RequestAccess::getKnowledgeBaseId);
23+
24+
@SuppressWarnings("ImmutableEnumChecker")
25+
private final AttributeKey<String> keyAttribute;
26+
27+
@SuppressWarnings("ImmutableEnumChecker")
28+
private final Function<Object, String> attributeValueAccessor;
29+
30+
AwsBedrockResourceType(
31+
AttributeKey<String> keyAttribute, Function<Object, String> attributeValueAccessor) {
32+
this.keyAttribute = keyAttribute;
33+
this.attributeValueAccessor = attributeValueAccessor;
34+
}
35+
36+
public AttributeKey<String> getKeyAttribute() {
37+
return keyAttribute;
38+
}
39+
40+
public Function<Object, String> getAttributeValueAccessor() {
41+
return attributeValueAccessor;
42+
}
43+
44+
public static AwsBedrockResourceType getRequestType(String requestClass) {
45+
return AwsBedrockResourceTypeMap.BEDROCK_REQUEST_MAP.get(requestClass);
46+
}
47+
48+
public static AwsBedrockResourceType getResponseType(String responseClass) {
49+
return AwsBedrockResourceTypeMap.BEDROCK_RESPONSE_MAP.get(responseClass);
50+
}
51+
52+
private static class AwsBedrockResourceTypeMap {
53+
private static final Map<String, AwsBedrockResourceType> BEDROCK_REQUEST_MAP = new HashMap<>();
54+
private static final Map<String, AwsBedrockResourceType> BEDROCK_RESPONSE_MAP = new HashMap<>();
55+
56+
// Bedrock request/response mapping
57+
// We only support operations that are related to the resource and where the context contains
58+
// the AgentID/DataSourceID/KnowledgeBaseID.
59+
// AgentID
60+
private static final List<String> agentRequestClasses =
61+
Arrays.asList(
62+
"CreateAgentActionGroupRequest",
63+
"CreateAgentAliasRequest",
64+
"DeleteAgentActionGroupRequest",
65+
"DeleteAgentAliasRequest",
66+
"DeleteAgentRequest",
67+
"DeleteAgentVersionRequest",
68+
"GetAgentActionGroupRequest",
69+
"GetAgentAliasRequest",
70+
"GetAgentRequest",
71+
"GetAgentVersionRequest",
72+
"ListAgentActionGroupsRequest",
73+
"ListAgentAliasesRequest",
74+
"ListAgentKnowledgeBasesRequest",
75+
"ListAgentVersionsRequest",
76+
"PrepareAgentRequest",
77+
"UpdateAgentActionGroupRequest",
78+
"UpdateAgentAliasRequest",
79+
"UpdateAgentRequest");
80+
private static final List<String> agentResponseClasses =
81+
Arrays.asList(
82+
"DeleteAgentAliasResult",
83+
"DeleteAgentResult",
84+
"DeleteAgentVersionResult",
85+
"PrepareAgentResult");
86+
// DataSourceID
87+
private static final List<String> dataSourceRequestClasses =
88+
Arrays.asList("DeleteDataSourceRequest", "GetDataSourceRequest", "UpdateDataSourceRequest");
89+
private static final List<String> dataSourceResponseClasses =
90+
Arrays.asList("DeleteDataSourceResult");
91+
// KnowledgeBaseID
92+
private static final List<String> knowledgeBaseRequestClasses =
93+
Arrays.asList(
94+
"AssociateAgentKnowledgeBaseRequest",
95+
"CreateDataSourceRequest",
96+
"DeleteKnowledgeBaseRequest",
97+
"DisassociateAgentKnowledgeBaseRequest",
98+
"GetAgentKnowledgeBaseRequest",
99+
"GetKnowledgeBaseRequest",
100+
"ListDataSourcesRequest",
101+
"UpdateAgentKnowledgeBaseRequest");
102+
private static final List<String> knowledgeBaseResponseClasses =
103+
Arrays.asList("DeleteKnowledgeBaseResult");
104+
105+
private AwsBedrockResourceTypeMap() {}
106+
107+
static {
108+
// Populate the BEDROCK_REQUEST_MAP
109+
for (String agentRequestClass : agentRequestClasses) {
110+
BEDROCK_REQUEST_MAP.put(agentRequestClass, AwsBedrockResourceType.AGENT_TYPE);
111+
}
112+
for (String dataSourceRequestClass : dataSourceRequestClasses) {
113+
BEDROCK_REQUEST_MAP.put(dataSourceRequestClass, AwsBedrockResourceType.DATA_SOURCE_TYPE);
114+
}
115+
for (String knowledgeBaseRequestClass : knowledgeBaseRequestClasses) {
116+
BEDROCK_REQUEST_MAP.put(
117+
knowledgeBaseRequestClass, AwsBedrockResourceType.KNOWLEDGE_BASE_TYPE);
118+
}
119+
120+
// Populate the BEDROCK_RESPONSE_MAP
121+
for (String agentResponseClass : agentResponseClasses) {
122+
BEDROCK_REQUEST_MAP.put(agentResponseClass, AwsBedrockResourceType.AGENT_TYPE);
123+
}
124+
for (String dataSourceResponseClass : dataSourceResponseClasses) {
125+
BEDROCK_REQUEST_MAP.put(dataSourceResponseClass, AwsBedrockResourceType.DATA_SOURCE_TYPE);
126+
}
127+
for (String knowledgeBaseResponseClass : knowledgeBaseResponseClasses) {
128+
BEDROCK_REQUEST_MAP.put(
129+
knowledgeBaseResponseClass, AwsBedrockResourceType.KNOWLEDGE_BASE_TYPE);
130+
}
131+
}
132+
}
133+
}

‎instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsExperimentalAttributes.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,16 @@ final class AwsExperimentalAttributes {
1818
static final AttributeKey<String> AWS_STREAM_NAME = stringKey("aws.stream.name");
1919
static final AttributeKey<String> AWS_TABLE_NAME = stringKey("aws.table.name");
2020
static final AttributeKey<String> AWS_REQUEST_ID = stringKey("aws.requestId");
21+
static final AttributeKey<String> AWS_AGENT_ID = stringKey("aws.bedrock.agent.id");
22+
static final AttributeKey<String> AWS_KNOWLEDGEBASE_ID =
23+
stringKey("aws.bedrock.knowledgebase.id");
24+
static final AttributeKey<String> AWS_DATASOURCE_ID = stringKey("aws.bedrock.datasource.id");
25+
static final AttributeKey<String> AWS_GUARDRAIL_ID = stringKey("aws.bedrock.guardrail.id");
26+
27+
// TODO: Merge in gen_ai attributes in opentelemetry-semconv-incubating once upgrade to v1.26.0
28+
static final AttributeKey<String> AWS_BEDROCK_RUNTIME_MODEL_ID =
29+
stringKey("gen_ai.request.model");
30+
static final AttributeKey<String> AWS_BEDROCK_SYSTEM = stringKey("gen_ai.system");
2131

2232
private AwsExperimentalAttributes() {}
2333
}

‎instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkExperimentalAttributesExtractor.java‎

Lines changed: 122 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,13 @@
66
package io.opentelemetry.instrumentation.awssdk.v1_11;
77

88
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_AGENT;
9+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_AGENT_ID;
10+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_BEDROCK_RUNTIME_MODEL_ID;
11+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_BEDROCK_SYSTEM;
912
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_BUCKET_NAME;
1013
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_ENDPOINT;
14+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_GUARDRAIL_ID;
15+
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGEBASE_ID;
1116
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_QUEUE_NAME;
1217
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_QUEUE_URL;
1318
import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_REQUEST_ID;
@@ -21,34 +26,36 @@
2126
import io.opentelemetry.api.common.AttributesBuilder;
2227
import io.opentelemetry.context.Context;
2328
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
29+
import java.util.Objects;
2430
import java.util.function.Function;
2531
import javax.annotation.Nullable;
2632

2733
class AwsSdkExperimentalAttributesExtractor
2834
implements AttributesExtractor<Request<?>, Response<?>> {
2935
private static final String COMPONENT_NAME = "java-aws-sdk";
36+
private static final String BEDROCK_SERVICE = "AmazonBedrock";
37+
private static final String BEDROCK_AGENT_SERVICE = "AWSBedrockAgent";
38+
private static final String BEDROCK_AGENT_RUNTIME_SERVICE = "AWSBedrockAgentRuntime";
39+
private static final String BEDROCK_RUNTIME_SERVICE = "AmazonBedrockRuntime";
3040

3141
@Override
3242
public void onStart(AttributesBuilder attributes, Context parentContext, Request<?> request) {
3343
attributes.put(AWS_AGENT, COMPONENT_NAME);
3444
attributes.put(AWS_ENDPOINT, request.getEndpoint().toString());
3545

3646
Object originalRequest = request.getOriginalRequest();
37-
setRequestAttribute(attributes, AWS_BUCKET_NAME, originalRequest, RequestAccess::getBucketName);
38-
setRequestAttribute(attributes, AWS_QUEUE_URL, originalRequest, RequestAccess::getQueueUrl);
39-
setRequestAttribute(attributes, AWS_QUEUE_NAME, originalRequest, RequestAccess::getQueueName);
40-
setRequestAttribute(attributes, AWS_STREAM_NAME, originalRequest, RequestAccess::getStreamName);
41-
setRequestAttribute(attributes, AWS_TABLE_NAME, originalRequest, RequestAccess::getTableName);
42-
}
47+
String requestClassName = originalRequest.getClass().getSimpleName();
48+
setAttribute(attributes, AWS_BUCKET_NAME, originalRequest, RequestAccess::getBucketName);
49+
setAttribute(attributes, AWS_QUEUE_URL, originalRequest, RequestAccess::getQueueUrl);
50+
setAttribute(attributes, AWS_QUEUE_NAME, originalRequest, RequestAccess::getQueueName);
51+
setAttribute(attributes, AWS_STREAM_NAME, originalRequest, RequestAccess::getStreamName);
52+
setAttribute(attributes, AWS_TABLE_NAME, originalRequest, RequestAccess::getTableName);
4353

44-
private static void setRequestAttribute(
45-
AttributesBuilder attributes,
46-
AttributeKey<String> key,
47-
Object request,
48-
Function<Object, String> getter) {
49-
String value = getter.apply(request);
50-
if (value != null) {
51-
attributes.put(key, value);
54+
// Get serviceName defined in the AWS Java SDK V1 Request class.
55+
String serviceName = request.getServiceName();
56+
// Extract request attributes only for Bedrock services.
57+
if (isBedrockService(serviceName)) {
58+
bedrockOnStart(attributes, originalRequest, requestClassName, serviceName);
5259
}
5360
}
5461

@@ -59,12 +66,108 @@ public void onEnd(
5966
Request<?> request,
6067
@Nullable Response<?> response,
6168
@Nullable Throwable error) {
62-
if (response != null && response.getAwsResponse() instanceof AmazonWebServiceResponse) {
63-
AmazonWebServiceResponse<?> awsResp = (AmazonWebServiceResponse<?>) response.getAwsResponse();
64-
String requestId = awsResp.getRequestId();
65-
if (requestId != null) {
66-
attributes.put(AWS_REQUEST_ID, requestId);
69+
if (response != null) {
70+
Object awsResp = response.getAwsResponse();
71+
if (awsResp instanceof AmazonWebServiceResponse) {
72+
AmazonWebServiceResponse<?> awsWebServiceResponse = (AmazonWebServiceResponse<?>) awsResp;
73+
String requestId = awsWebServiceResponse.getRequestId();
74+
if (requestId != null) {
75+
attributes.put(AWS_REQUEST_ID, requestId);
76+
}
77+
}
78+
// Get serviceName defined in the AWS Java SDK V1 Request class.
79+
String serviceName = request.getServiceName();
80+
// Extract response attributes for Bedrock services
81+
if (awsResp != null && isBedrockService(serviceName)) {
82+
bedrockOnEnd(attributes, awsResp, serviceName);
6783
}
6884
}
6985
}
86+
87+
private static void bedrockOnStart(
88+
AttributesBuilder attributes,
89+
Object originalRequest,
90+
String requestClassName,
91+
String serviceName) {
92+
switch (serviceName) {
93+
case BEDROCK_SERVICE:
94+
setAttribute(attributes, AWS_GUARDRAIL_ID, originalRequest, RequestAccess::getGuardrailId);
95+
break;
96+
case BEDROCK_AGENT_SERVICE:
97+
AwsBedrockResourceType resourceType =
98+
AwsBedrockResourceType.getRequestType(requestClassName);
99+
if (resourceType != null) {
100+
setAttribute(
101+
attributes,
102+
resourceType.getKeyAttribute(),
103+
originalRequest,
104+
resourceType.getAttributeValueAccessor());
105+
}
106+
break;
107+
case BEDROCK_AGENT_RUNTIME_SERVICE:
108+
setAttribute(attributes, AWS_AGENT_ID, originalRequest, RequestAccess::getAgentId);
109+
setAttribute(
110+
attributes, AWS_KNOWLEDGEBASE_ID, originalRequest, RequestAccess::getKnowledgeBaseId);
111+
break;
112+
case BEDROCK_RUNTIME_SERVICE:
113+
if (!Objects.equals(requestClassName, "InvokeModelRequest")) {
114+
break;
115+
}
116+
attributes.put(AWS_BEDROCK_SYSTEM, "aws_bedrock");
117+
Function<Object, String> getter = RequestAccess::getModelId;
118+
String modelId = getter.apply(originalRequest);
119+
attributes.put(AWS_BEDROCK_RUNTIME_MODEL_ID, modelId);
120+
break;
121+
default:
122+
break;
123+
}
124+
}
125+
126+
private static void bedrockOnEnd(
127+
AttributesBuilder attributes, Object awsResp, String serviceName) {
128+
switch (serviceName) {
129+
case BEDROCK_SERVICE:
130+
setAttribute(attributes, AWS_GUARDRAIL_ID, awsResp, RequestAccess::getGuardrailId);
131+
break;
132+
case BEDROCK_AGENT_SERVICE:
133+
String responseClassName = awsResp.getClass().getSimpleName();
134+
AwsBedrockResourceType resourceType =
135+
AwsBedrockResourceType.getResponseType(responseClassName);
136+
if (resourceType != null) {
137+
setAttribute(
138+
attributes,
139+
resourceType.getKeyAttribute(),
140+
awsResp,
141+
resourceType.getAttributeValueAccessor());
142+
}
143+
break;
144+
case BEDROCK_AGENT_RUNTIME_SERVICE:
145+
setAttribute(attributes, AWS_AGENT_ID, awsResp, RequestAccess::getAgentId);
146+
setAttribute(attributes, AWS_KNOWLEDGEBASE_ID, awsResp, RequestAccess::getKnowledgeBaseId);
147+
break;
148+
default:
149+
break;
150+
}
151+
}
152+
153+
private static boolean isBedrockService(String serviceName) {
154+
// Check if the serviceName belongs to Bedrock Services defined in AWS Java SDK V1.
155+
// For example <a
156+
// href="https://github.com/aws/aws-sdk-java/blob/38031248a696468e19a4670c0c4585637d5e7cc6/aws-java-sdk-bedrock/src/main/java/com/amazonaws/services/bedrock/AmazonBedrock.java#L34">AmazonBedrock</a>
157+
return serviceName.equals(BEDROCK_SERVICE)
158+
|| serviceName.equals(BEDROCK_AGENT_SERVICE)
159+
|| serviceName.equals(BEDROCK_AGENT_RUNTIME_SERVICE)
160+
|| serviceName.equals(BEDROCK_RUNTIME_SERVICE);
161+
}
162+
163+
private static void setAttribute(
164+
AttributesBuilder attributes,
165+
AttributeKey<String> key,
166+
Object request,
167+
Function<Object, String> getter) {
168+
String value = getter.apply(request);
169+
if (value != null) {
170+
attributes.put(key, value);
171+
}
172+
}
70173
}

0 commit comments

Comments
 (0)