Skip to content

Commit 17634e2

Browse files
anuraagaanuraaga
authored
Add instrumentation of AWS Bedrock to use gen_ai conventions (#13355)
1 parent f0d80b2 commit 17634e2

File tree

29 files changed

+1483
-6
lines changed

29 files changed

+1483
-6
lines changed

‎conventions/src/main/kotlin/otel.java-conventions.gradle.kts

+2-1
Original file line numberDiff line numberDiff line change
@@ -356,9 +356,10 @@ tasks.withType<Test>().configureEach {
356356
val trustStore = project(":testing-common").file("src/misc/testing-keystore.p12")
357357
// Work around payara not working when this is set for some reason.
358358
// Don't set for:
359+
// - aws-sdk as we have tests that interact with AWS and need normal trustStore
359360
// - camel as we have tests that interact with AWS and need normal trustStore
360361
// - vaadin as tests need to be able to download nodejs when not cached in ~/.vaadin/
361-
if (project.name != "jaxrs-2.0-payara-testing" && !project.path.contains("vaadin") && project.description != "camel-2-20") {
362+
if (project.name != "jaxrs-2.0-payara-testing" && !project.path.contains("vaadin") && project.description != "camel-2-20" && !project.path.contains("aws-sdk")) {
362363
jvmArgumentProviders.add(KeystoreArgumentsProvider(trustStore))
363364
}
364365

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/genai/GenAiAttributesExtractor.java

+113
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.api.incubator.semconv.genai;
7+
8+
import static io.opentelemetry.api.common.AttributeKey.doubleKey;
9+
import static io.opentelemetry.api.common.AttributeKey.longKey;
10+
import static io.opentelemetry.api.common.AttributeKey.stringArrayKey;
11+
import static io.opentelemetry.api.common.AttributeKey.stringKey;
12+
import static io.opentelemetry.instrumentation.api.internal.AttributesExtractorUtil.internalSet;
13+
14+
import io.opentelemetry.api.common.AttributeKey;
15+
import io.opentelemetry.api.common.AttributesBuilder;
16+
import io.opentelemetry.context.Context;
17+
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
18+
import java.util.List;
19+
import javax.annotation.Nullable;
20+
21+
/**
22+
* Extractor of <a href="https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/">GenAI
23+
* attributes</a>.
24+
*
25+
* <p>This class delegates to a type-specific {@link GenAiAttributesGetter} for individual attribute
26+
* extraction from request/response objects.
27+
*/
28+
public final class GenAiAttributesExtractor<REQUEST, RESPONSE>
29+
implements AttributesExtractor<REQUEST, RESPONSE> {
30+
31+
// copied from GenAiIncubatingAttributes
32+
private static final AttributeKey<String> GEN_AI_OPERATION_NAME =
33+
stringKey("gen_ai.operation.name");
34+
private static final AttributeKey<List<String>> GEN_AI_REQUEST_ENCODING_FORMATS =
35+
stringArrayKey("gen_ai.request.encoding_formats");
36+
private static final AttributeKey<Double> GEN_AI_REQUEST_FREQUENCY_PENALTY =
37+
doubleKey("gen_ai.request.frequency_penalty");
38+
private static final AttributeKey<Long> GEN_AI_REQUEST_MAX_TOKENS =
39+
longKey("gen_ai.request.max_tokens");
40+
private static final AttributeKey<String> GEN_AI_REQUEST_MODEL =
41+
stringKey("gen_ai.request.model");
42+
private static final AttributeKey<Double> GEN_AI_REQUEST_PRESENCE_PENALTY =
43+
doubleKey("gen_ai.request.presence_penalty");
44+
private static final AttributeKey<Long> GEN_AI_REQUEST_SEED = longKey("gen_ai.request.seed");
45+
private static final AttributeKey<List<String>> GEN_AI_REQUEST_STOP_SEQUENCES =
46+
stringArrayKey("gen_ai.request.stop_sequences");
47+
private static final AttributeKey<Double> GEN_AI_REQUEST_TEMPERATURE =
48+
doubleKey("gen_ai.request.temperature");
49+
private static final AttributeKey<Double> GEN_AI_REQUEST_TOP_K =
50+
doubleKey("gen_ai.request.top_k");
51+
private static final AttributeKey<Double> GEN_AI_REQUEST_TOP_P =
52+
doubleKey("gen_ai.request.top_p");
53+
private static final AttributeKey<List<String>> GEN_AI_RESPONSE_FINISH_REASONS =
54+
stringArrayKey("gen_ai.response.finish_reasons");
55+
private static final AttributeKey<String> GEN_AI_RESPONSE_ID = stringKey("gen_ai.response.id");
56+
private static final AttributeKey<String> GEN_AI_RESPONSE_MODEL =
57+
stringKey("gen_ai.response.model");
58+
private static final AttributeKey<String> GEN_AI_SYSTEM = stringKey("gen_ai.system");
59+
private static final AttributeKey<Long> GEN_AI_USAGE_INPUT_TOKENS =
60+
longKey("gen_ai.usage.input_tokens");
61+
private static final AttributeKey<Long> GEN_AI_USAGE_OUTPUT_TOKENS =
62+
longKey("gen_ai.usage.output_tokens");
63+
64+
/** Creates the GenAI attributes extractor. */
65+
public static <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> create(
66+
GenAiAttributesGetter<REQUEST, RESPONSE> attributesGetter) {
67+
return new GenAiAttributesExtractor<>(attributesGetter);
68+
}
69+
70+
private final GenAiAttributesGetter<REQUEST, RESPONSE> getter;
71+
72+
private GenAiAttributesExtractor(GenAiAttributesGetter<REQUEST, RESPONSE> getter) {
73+
this.getter = getter;
74+
}
75+
76+
@Override
77+
public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST request) {
78+
internalSet(attributes, GEN_AI_OPERATION_NAME, getter.getOperationName(request));
79+
internalSet(attributes, GEN_AI_SYSTEM, getter.getSystem(request));
80+
internalSet(attributes, GEN_AI_REQUEST_MODEL, getter.getRequestModel(request));
81+
internalSet(attributes, GEN_AI_REQUEST_SEED, getter.getRequestSeed(request));
82+
internalSet(
83+
attributes, GEN_AI_REQUEST_ENCODING_FORMATS, getter.getRequestEncodingFormats(request));
84+
internalSet(
85+
attributes, GEN_AI_REQUEST_FREQUENCY_PENALTY, getter.getRequestFrequencyPenalty(request));
86+
internalSet(attributes, GEN_AI_REQUEST_MAX_TOKENS, getter.getRequestMaxTokens(request));
87+
internalSet(
88+
attributes, GEN_AI_REQUEST_PRESENCE_PENALTY, getter.getRequestPresencePenalty(request));
89+
internalSet(attributes, GEN_AI_REQUEST_STOP_SEQUENCES, getter.getRequestStopSequences(request));
90+
internalSet(attributes, GEN_AI_REQUEST_TEMPERATURE, getter.getRequestTemperature(request));
91+
internalSet(attributes, GEN_AI_REQUEST_TOP_K, getter.getRequestTopK(request));
92+
internalSet(attributes, GEN_AI_REQUEST_TOP_P, getter.getRequestTopP(request));
93+
}
94+
95+
@Override
96+
public void onEnd(
97+
AttributesBuilder attributes,
98+
Context context,
99+
REQUEST request,
100+
@Nullable RESPONSE response,
101+
@Nullable Throwable error) {
102+
internalSet(
103+
attributes,
104+
GEN_AI_RESPONSE_FINISH_REASONS,
105+
getter.getResponseFinishReasons(request, response));
106+
internalSet(attributes, GEN_AI_RESPONSE_ID, getter.getResponseId(request, response));
107+
internalSet(attributes, GEN_AI_RESPONSE_MODEL, getter.getResponseModel(request, response));
108+
internalSet(
109+
attributes, GEN_AI_USAGE_INPUT_TOKENS, getter.getUsageInputTokens(request, response));
110+
internalSet(
111+
attributes, GEN_AI_USAGE_OUTPUT_TOKENS, getter.getUsageOutputTokens(request, response));
112+
}
113+
}

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/genai/GenAiAttributesGetter.java

+67
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.api.incubator.semconv.genai;
7+
8+
import java.util.List;
9+
import javax.annotation.Nullable;
10+
11+
/**
12+
* An interface for getting GenAI attributes.
13+
*
14+
* <p>Instrumentation authors will create implementations of this interface for their specific
15+
* library/framework. It will be used by the {@link GenAiAttributesExtractor} to obtain the various
16+
* GenAI attributes in a type-generic way.
17+
*/
18+
public interface GenAiAttributesGetter<REQUEST, RESPONSE> {
19+
String getOperationName(REQUEST request);
20+
21+
String getSystem(REQUEST request);
22+
23+
@Nullable
24+
String getRequestModel(REQUEST request);
25+
26+
@Nullable
27+
Long getRequestSeed(REQUEST request);
28+
29+
@Nullable
30+
List<String> getRequestEncodingFormats(REQUEST request);
31+
32+
@Nullable
33+
Double getRequestFrequencyPenalty(REQUEST request);
34+
35+
@Nullable
36+
Long getRequestMaxTokens(REQUEST request);
37+
38+
@Nullable
39+
Double getRequestPresencePenalty(REQUEST request);
40+
41+
@Nullable
42+
List<String> getRequestStopSequences(REQUEST request);
43+
44+
@Nullable
45+
Double getRequestTemperature(REQUEST request);
46+
47+
@Nullable
48+
Double getRequestTopK(REQUEST request);
49+
50+
@Nullable
51+
Double getRequestTopP(REQUEST request);
52+
53+
@Nullable
54+
List<String> getResponseFinishReasons(REQUEST request, RESPONSE response);
55+
56+
@Nullable
57+
String getResponseId(REQUEST request, RESPONSE response);
58+
59+
@Nullable
60+
String getResponseModel(REQUEST request, RESPONSE response);
61+
62+
@Nullable
63+
Long getUsageInputTokens(REQUEST request, RESPONSE response);
64+
65+
@Nullable
66+
Long getUsageOutputTokens(REQUEST request, RESPONSE response);
67+
}

‎instrumentation-api-incubator/src/main/java/io/opentelemetry/instrumentation/api/incubator/semconv/genai/GenAiSpanNameExtractor.java

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.api.incubator.semconv.genai;
7+
8+
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
9+
10+
/** A {@link SpanNameExtractor} for GenAI requests. */
11+
public final class GenAiSpanNameExtractor<REQUEST> implements SpanNameExtractor<REQUEST> {
12+
13+
/**
14+
* Returns a {@link SpanNameExtractor} that constructs the span name according to GenAI semantic
15+
* conventions: {@code <gen_ai.operation.name> <gen_ai.request.model>}.
16+
*/
17+
public static <REQUEST> SpanNameExtractor<REQUEST> create(
18+
GenAiAttributesGetter<REQUEST, ?> attributesGetter) {
19+
return new GenAiSpanNameExtractor<>(attributesGetter);
20+
}
21+
22+
private final GenAiAttributesGetter<REQUEST, ?> getter;
23+
24+
private GenAiSpanNameExtractor(GenAiAttributesGetter<REQUEST, ?> getter) {
25+
this.getter = getter;
26+
}
27+
28+
@Override
29+
public String extract(REQUEST request) {
30+
String operation = getter.getOperationName(request);
31+
String model = getter.getRequestModel(request);
32+
if (model == null) {
33+
return operation;
34+
}
35+
return operation + ' ' + model;
36+
}
37+
}

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

+28
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ muzzle {
1111
// client, which is not target of instrumentation anyways.
1212
extraDependency("software.amazon.awssdk:protocol-core")
1313

14+
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
1415
excludeInstrumentationName("aws-sdk-2.2-sqs")
1516
excludeInstrumentationName("aws-sdk-2.2-sns")
1617
excludeInstrumentationName("aws-sdk-2.2-lambda")
@@ -43,6 +44,7 @@ muzzle {
4344
// client, which is not target of instrumentation anyways.
4445
extraDependency("software.amazon.awssdk:protocol-core")
4546

47+
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
4648
excludeInstrumentationName("aws-sdk-2.2-sns")
4749
excludeInstrumentationName("aws-sdk-2.2-lambda")
4850

@@ -58,6 +60,7 @@ muzzle {
5860
// client, which is not target of instrumentation anyways.
5961
extraDependency("software.amazon.awssdk:protocol-core")
6062

63+
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
6164
excludeInstrumentationName("aws-sdk-2.2-sqs")
6265
excludeInstrumentationName("aws-sdk-2.2-lambda")
6366

@@ -72,12 +75,25 @@ muzzle {
7275
// client, which is not target of instrumentation anyways.
7376
extraDependency("software.amazon.awssdk:protocol-core")
7477

78+
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
7579
excludeInstrumentationName("aws-sdk-2.2-sqs")
7680
excludeInstrumentationName("aws-sdk-2.2-sns")
7781

7882
// several software.amazon.awssdk artifacts are missing for this version
7983
skip("2.17.200")
8084
}
85+
pass {
86+
group.set("software.amazon.awssdk")
87+
module.set("bedrock-runtime")
88+
versions.set("[2.25.63,)")
89+
// Used by all SDK services, the only case it isn't is an SDK extension such as a custom HTTP
90+
// client, which is not target of instrumentation anyways.
91+
extraDependency("software.amazon.awssdk:protocol-core")
92+
93+
excludeInstrumentationName("aws-sdk-2.2-lambda")
94+
excludeInstrumentationName("aws-sdk-2.2-sqs")
95+
excludeInstrumentationName("aws-sdk-2.2-sns")
96+
}
8197
}
8298

8399
dependencies {
@@ -120,6 +136,18 @@ testing {
120136
implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:library"))
121137
}
122138
}
139+
140+
val testBedrockRuntime by registering(JvmTestSuite::class) {
141+
dependencies {
142+
implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing"))
143+
if (findProperty("testLatestDeps") as Boolean) {
144+
implementation("software.amazon.awssdk:bedrockruntime:+")
145+
} else {
146+
// First release with Converse API
147+
implementation("software.amazon.awssdk:bedrockruntime:2.25.63")
148+
}
149+
}
150+
}
123151
}
124152
}
125153

‎instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/BedrockRuntimeAdviceBridge.java

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
7+
8+
/**
9+
* This class is internal and is hence not for public use. Its APIs are unstable and can change at
10+
* any time.
11+
*/
12+
public final class BedrockRuntimeAdviceBridge {
13+
private BedrockRuntimeAdviceBridge() {}
14+
15+
public static void referenceForMuzzleOnly() {
16+
throw new UnsupportedOperationException(
17+
BedrockRuntimeImpl.class.getName()
18+
+ " referencing for muzzle, should never be actually called");
19+
}
20+
}

‎instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/BedrockRuntimeInstrumentationModule.java

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.awssdk.v2_2;
7+
8+
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
9+
import static net.bytebuddy.matcher.ElementMatchers.none;
10+
11+
import com.google.auto.service.AutoService;
12+
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.BedrockRuntimeAdviceBridge;
13+
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
14+
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
15+
import net.bytebuddy.asm.Advice;
16+
import net.bytebuddy.matcher.ElementMatcher;
17+
18+
@AutoService(InstrumentationModule.class)
19+
public class BedrockRuntimeInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
20+
21+
public BedrockRuntimeInstrumentationModule() {
22+
super("aws-sdk-2.2-bedrock-runtime");
23+
}
24+
25+
@Override
26+
public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
27+
return hasClassesNamed("software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient");
28+
}
29+
30+
@Override
31+
public void doTransform(TypeTransformer transformer) {
32+
transformer.applyAdviceToMethod(
33+
none(), BedrockRuntimeInstrumentationModule.class.getName() + "$RegisterAdvice");
34+
}
35+
36+
@SuppressWarnings("unused")
37+
public static class RegisterAdvice {
38+
@Advice.OnMethodExit(suppress = Throwable.class)
39+
public static void onExit() {
40+
// (indirectly) using BedrockRuntimeImpl class here to make sure it is available from
41+
// BedrockRuntimeAccess (injected into app classloader) and checked by Muzzle
42+
BedrockRuntimeAdviceBridge.referenceForMuzzleOnly();
43+
}
44+
}
45+
}

‎instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/testBedrockRuntime/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/Aws2BedrockRuntimeTest.java

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
7+
8+
import io.opentelemetry.instrumentation.awssdk.v2_2.AbstractAws2BedrockRuntimeTest;
9+
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
10+
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
11+
import org.junit.jupiter.api.extension.RegisterExtension;
12+
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
13+
14+
class Aws2BedrockRuntimeTest extends AbstractAws2BedrockRuntimeTest {
15+
@RegisterExtension
16+
static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
17+
18+
@Override
19+
protected InstrumentationExtension getTesting() {
20+
return testing;
21+
}
22+
23+
@Override
24+
protected ClientOverrideConfiguration.Builder createOverrideConfigurationBuilder() {
25+
return ClientOverrideConfiguration.builder();
26+
}
27+
}

0 commit comments

Comments
 (0)