Add instrumentation of AWS Bedrock to use gen_ai conventions (#13355)
This commit is contained in:
Anuraag (Rag) Agrawal
GitHub
parent
f0d80b2e55
commit
17634e2d19
|
|
@ -356,9 +356,10 @@ tasks.withType<Test>().configureEach {
|
|||
val trustStore = project(":testing-common").file("src/misc/testing-keystore.p12")
|
||||
// Work around payara not working when this is set for some reason.
|
||||
// Don't set for:
|
||||
// - aws-sdk as we have tests that interact with AWS and need normal trustStore
|
||||
// - camel as we have tests that interact with AWS and need normal trustStore
|
||||
// - vaadin as tests need to be able to download nodejs when not cached in ~/.vaadin/
|
||||
if (project.name != "jaxrs-2.0-payara-testing" && !project.path.contains("vaadin") && project.description != "camel-2-20") {
|
||||
if (project.name != "jaxrs-2.0-payara-testing" && !project.path.contains("vaadin") && project.description != "camel-2-20" && !project.path.contains("aws-sdk")) {
|
||||
jvmArgumentProviders.add(KeystoreArgumentsProvider(trustStore))
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,113 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.api.incubator.semconv.genai;
|
||||
|
||||
import static io.opentelemetry.api.common.AttributeKey.doubleKey;
|
||||
import static io.opentelemetry.api.common.AttributeKey.longKey;
|
||||
import static io.opentelemetry.api.common.AttributeKey.stringArrayKey;
|
||||
import static io.opentelemetry.api.common.AttributeKey.stringKey;
|
||||
import static io.opentelemetry.instrumentation.api.internal.AttributesExtractorUtil.internalSet;
|
||||
|
||||
import io.opentelemetry.api.common.AttributeKey;
|
||||
import io.opentelemetry.api.common.AttributesBuilder;
|
||||
import io.opentelemetry.context.Context;
|
||||
import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
|
||||
import java.util.List;
|
||||
import javax.annotation.Nullable;
|
||||
|
||||
/**
|
||||
* Extractor of <a href="https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/">GenAI
|
||||
* attributes</a>.
|
||||
*
|
||||
* <p>This class delegates to a type-specific {@link GenAiAttributesGetter} for individual attribute
|
||||
* extraction from request/response objects.
|
||||
*/
|
||||
public final class GenAiAttributesExtractor<REQUEST, RESPONSE>
|
||||
implements AttributesExtractor<REQUEST, RESPONSE> {
|
||||
|
||||
// copied from GenAiIncubatingAttributes
|
||||
private static final AttributeKey<String> GEN_AI_OPERATION_NAME =
|
||||
stringKey("gen_ai.operation.name");
|
||||
private static final AttributeKey<List<String>> GEN_AI_REQUEST_ENCODING_FORMATS =
|
||||
stringArrayKey("gen_ai.request.encoding_formats");
|
||||
private static final AttributeKey<Double> GEN_AI_REQUEST_FREQUENCY_PENALTY =
|
||||
doubleKey("gen_ai.request.frequency_penalty");
|
||||
private static final AttributeKey<Long> GEN_AI_REQUEST_MAX_TOKENS =
|
||||
longKey("gen_ai.request.max_tokens");
|
||||
private static final AttributeKey<String> GEN_AI_REQUEST_MODEL =
|
||||
stringKey("gen_ai.request.model");
|
||||
private static final AttributeKey<Double> GEN_AI_REQUEST_PRESENCE_PENALTY =
|
||||
doubleKey("gen_ai.request.presence_penalty");
|
||||
private static final AttributeKey<Long> GEN_AI_REQUEST_SEED = longKey("gen_ai.request.seed");
|
||||
private static final AttributeKey<List<String>> GEN_AI_REQUEST_STOP_SEQUENCES =
|
||||
stringArrayKey("gen_ai.request.stop_sequences");
|
||||
private static final AttributeKey<Double> GEN_AI_REQUEST_TEMPERATURE =
|
||||
doubleKey("gen_ai.request.temperature");
|
||||
private static final AttributeKey<Double> GEN_AI_REQUEST_TOP_K =
|
||||
doubleKey("gen_ai.request.top_k");
|
||||
private static final AttributeKey<Double> GEN_AI_REQUEST_TOP_P =
|
||||
doubleKey("gen_ai.request.top_p");
|
||||
private static final AttributeKey<List<String>> GEN_AI_RESPONSE_FINISH_REASONS =
|
||||
stringArrayKey("gen_ai.response.finish_reasons");
|
||||
private static final AttributeKey<String> GEN_AI_RESPONSE_ID = stringKey("gen_ai.response.id");
|
||||
private static final AttributeKey<String> GEN_AI_RESPONSE_MODEL =
|
||||
stringKey("gen_ai.response.model");
|
||||
private static final AttributeKey<String> GEN_AI_SYSTEM = stringKey("gen_ai.system");
|
||||
private static final AttributeKey<Long> GEN_AI_USAGE_INPUT_TOKENS =
|
||||
longKey("gen_ai.usage.input_tokens");
|
||||
private static final AttributeKey<Long> GEN_AI_USAGE_OUTPUT_TOKENS =
|
||||
longKey("gen_ai.usage.output_tokens");
|
||||
|
||||
/** Creates the GenAI attributes extractor. */
|
||||
public static <REQUEST, RESPONSE> AttributesExtractor<REQUEST, RESPONSE> create(
|
||||
GenAiAttributesGetter<REQUEST, RESPONSE> attributesGetter) {
|
||||
return new GenAiAttributesExtractor<>(attributesGetter);
|
||||
}
|
||||
|
||||
private final GenAiAttributesGetter<REQUEST, RESPONSE> getter;
|
||||
|
||||
private GenAiAttributesExtractor(GenAiAttributesGetter<REQUEST, RESPONSE> getter) {
|
||||
this.getter = getter;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onStart(AttributesBuilder attributes, Context parentContext, REQUEST request) {
|
||||
internalSet(attributes, GEN_AI_OPERATION_NAME, getter.getOperationName(request));
|
||||
internalSet(attributes, GEN_AI_SYSTEM, getter.getSystem(request));
|
||||
internalSet(attributes, GEN_AI_REQUEST_MODEL, getter.getRequestModel(request));
|
||||
internalSet(attributes, GEN_AI_REQUEST_SEED, getter.getRequestSeed(request));
|
||||
internalSet(
|
||||
attributes, GEN_AI_REQUEST_ENCODING_FORMATS, getter.getRequestEncodingFormats(request));
|
||||
internalSet(
|
||||
attributes, GEN_AI_REQUEST_FREQUENCY_PENALTY, getter.getRequestFrequencyPenalty(request));
|
||||
internalSet(attributes, GEN_AI_REQUEST_MAX_TOKENS, getter.getRequestMaxTokens(request));
|
||||
internalSet(
|
||||
attributes, GEN_AI_REQUEST_PRESENCE_PENALTY, getter.getRequestPresencePenalty(request));
|
||||
internalSet(attributes, GEN_AI_REQUEST_STOP_SEQUENCES, getter.getRequestStopSequences(request));
|
||||
internalSet(attributes, GEN_AI_REQUEST_TEMPERATURE, getter.getRequestTemperature(request));
|
||||
internalSet(attributes, GEN_AI_REQUEST_TOP_K, getter.getRequestTopK(request));
|
||||
internalSet(attributes, GEN_AI_REQUEST_TOP_P, getter.getRequestTopP(request));
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onEnd(
|
||||
AttributesBuilder attributes,
|
||||
Context context,
|
||||
REQUEST request,
|
||||
@Nullable RESPONSE response,
|
||||
@Nullable Throwable error) {
|
||||
internalSet(
|
||||
attributes,
|
||||
GEN_AI_RESPONSE_FINISH_REASONS,
|
||||
getter.getResponseFinishReasons(request, response));
|
||||
internalSet(attributes, GEN_AI_RESPONSE_ID, getter.getResponseId(request, response));
|
||||
internalSet(attributes, GEN_AI_RESPONSE_MODEL, getter.getResponseModel(request, response));
|
||||
internalSet(
|
||||
attributes, GEN_AI_USAGE_INPUT_TOKENS, getter.getUsageInputTokens(request, response));
|
||||
internalSet(
|
||||
attributes, GEN_AI_USAGE_OUTPUT_TOKENS, getter.getUsageOutputTokens(request, response));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,67 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.api.incubator.semconv.genai;
|
||||
|
||||
import java.util.List;
|
||||
import javax.annotation.Nullable;
|
||||
|
||||
/**
|
||||
* An interface for getting GenAI attributes.
|
||||
*
|
||||
* <p>Instrumentation authors will create implementations of this interface for their specific
|
||||
* library/framework. It will be used by the {@link GenAiAttributesExtractor} to obtain the various
|
||||
* GenAI attributes in a type-generic way.
|
||||
*/
|
||||
public interface GenAiAttributesGetter<REQUEST, RESPONSE> {
|
||||
String getOperationName(REQUEST request);
|
||||
|
||||
String getSystem(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
String getRequestModel(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
Long getRequestSeed(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
List<String> getRequestEncodingFormats(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
Double getRequestFrequencyPenalty(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
Long getRequestMaxTokens(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
Double getRequestPresencePenalty(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
List<String> getRequestStopSequences(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
Double getRequestTemperature(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
Double getRequestTopK(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
Double getRequestTopP(REQUEST request);
|
||||
|
||||
@Nullable
|
||||
List<String> getResponseFinishReasons(REQUEST request, RESPONSE response);
|
||||
|
||||
@Nullable
|
||||
String getResponseId(REQUEST request, RESPONSE response);
|
||||
|
||||
@Nullable
|
||||
String getResponseModel(REQUEST request, RESPONSE response);
|
||||
|
||||
@Nullable
|
||||
Long getUsageInputTokens(REQUEST request, RESPONSE response);
|
||||
|
||||
@Nullable
|
||||
Long getUsageOutputTokens(REQUEST request, RESPONSE response);
|
||||
}
|
||||
|
|
@ -0,0 +1,37 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.api.incubator.semconv.genai;
|
||||
|
||||
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
|
||||
|
||||
/** A {@link SpanNameExtractor} for GenAI requests. */
|
||||
public final class GenAiSpanNameExtractor<REQUEST> implements SpanNameExtractor<REQUEST> {
|
||||
|
||||
/**
|
||||
* Returns a {@link SpanNameExtractor} that constructs the span name according to GenAI semantic
|
||||
* conventions: {@code <gen_ai.operation.name> <gen_ai.request.model>}.
|
||||
*/
|
||||
public static <REQUEST> SpanNameExtractor<REQUEST> create(
|
||||
GenAiAttributesGetter<REQUEST, ?> attributesGetter) {
|
||||
return new GenAiSpanNameExtractor<>(attributesGetter);
|
||||
}
|
||||
|
||||
private final GenAiAttributesGetter<REQUEST, ?> getter;
|
||||
|
||||
private GenAiSpanNameExtractor(GenAiAttributesGetter<REQUEST, ?> getter) {
|
||||
this.getter = getter;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String extract(REQUEST request) {
|
||||
String operation = getter.getOperationName(request);
|
||||
String model = getter.getRequestModel(request);
|
||||
if (model == null) {
|
||||
return operation;
|
||||
}
|
||||
return operation + ' ' + model;
|
||||
}
|
||||
}
|
||||
|
|
@ -11,6 +11,7 @@ muzzle {
|
|||
// client, which is not target of instrumentation anyways.
|
||||
extraDependency("software.amazon.awssdk:protocol-core")
|
||||
|
||||
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sqs")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sns")
|
||||
excludeInstrumentationName("aws-sdk-2.2-lambda")
|
||||
|
|
@ -43,6 +44,7 @@ muzzle {
|
|||
// client, which is not target of instrumentation anyways.
|
||||
extraDependency("software.amazon.awssdk:protocol-core")
|
||||
|
||||
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sns")
|
||||
excludeInstrumentationName("aws-sdk-2.2-lambda")
|
||||
|
||||
|
|
@ -58,6 +60,7 @@ muzzle {
|
|||
// client, which is not target of instrumentation anyways.
|
||||
extraDependency("software.amazon.awssdk:protocol-core")
|
||||
|
||||
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sqs")
|
||||
excludeInstrumentationName("aws-sdk-2.2-lambda")
|
||||
|
||||
|
|
@ -72,12 +75,25 @@ muzzle {
|
|||
// client, which is not target of instrumentation anyways.
|
||||
extraDependency("software.amazon.awssdk:protocol-core")
|
||||
|
||||
excludeInstrumentationName("aws-sdk-2.2-bedrock-runtime")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sqs")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sns")
|
||||
|
||||
// several software.amazon.awssdk artifacts are missing for this version
|
||||
skip("2.17.200")
|
||||
}
|
||||
pass {
|
||||
group.set("software.amazon.awssdk")
|
||||
module.set("bedrock-runtime")
|
||||
versions.set("[2.25.63,)")
|
||||
// Used by all SDK services, the only case it isn't is an SDK extension such as a custom HTTP
|
||||
// client, which is not target of instrumentation anyways.
|
||||
extraDependency("software.amazon.awssdk:protocol-core")
|
||||
|
||||
excludeInstrumentationName("aws-sdk-2.2-lambda")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sqs")
|
||||
excludeInstrumentationName("aws-sdk-2.2-sns")
|
||||
}
|
||||
}
|
||||
|
||||
dependencies {
|
||||
|
|
@ -120,6 +136,18 @@ testing {
|
|||
implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:library"))
|
||||
}
|
||||
}
|
||||
|
||||
val testBedrockRuntime by registering(JvmTestSuite::class) {
|
||||
dependencies {
|
||||
implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing"))
|
||||
if (findProperty("testLatestDeps") as Boolean) {
|
||||
implementation("software.amazon.awssdk:bedrockruntime:+")
|
||||
} else {
|
||||
// First release with Converse API
|
||||
implementation("software.amazon.awssdk:bedrockruntime:2.25.63")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,20 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
|
||||
|
||||
/**
|
||||
* This class is internal and is hence not for public use. Its APIs are unstable and can change at
|
||||
* any time.
|
||||
*/
|
||||
public final class BedrockRuntimeAdviceBridge {
|
||||
private BedrockRuntimeAdviceBridge() {}
|
||||
|
||||
public static void referenceForMuzzleOnly() {
|
||||
throw new UnsupportedOperationException(
|
||||
BedrockRuntimeImpl.class.getName()
|
||||
+ " referencing for muzzle, should never be actually called");
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,45 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.javaagent.instrumentation.awssdk.v2_2;
|
||||
|
||||
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
|
||||
import static net.bytebuddy.matcher.ElementMatchers.none;
|
||||
|
||||
import com.google.auto.service.AutoService;
|
||||
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.BedrockRuntimeAdviceBridge;
|
||||
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
|
||||
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
|
||||
import net.bytebuddy.asm.Advice;
|
||||
import net.bytebuddy.matcher.ElementMatcher;
|
||||
|
||||
@AutoService(InstrumentationModule.class)
|
||||
public class BedrockRuntimeInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
|
||||
|
||||
public BedrockRuntimeInstrumentationModule() {
|
||||
super("aws-sdk-2.2-bedrock-runtime");
|
||||
}
|
||||
|
||||
@Override
|
||||
public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
|
||||
return hasClassesNamed("software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient");
|
||||
}
|
||||
|
||||
@Override
|
||||
public void doTransform(TypeTransformer transformer) {
|
||||
transformer.applyAdviceToMethod(
|
||||
none(), BedrockRuntimeInstrumentationModule.class.getName() + "$RegisterAdvice");
|
||||
}
|
||||
|
||||
@SuppressWarnings("unused")
|
||||
public static class RegisterAdvice {
|
||||
@Advice.OnMethodExit(suppress = Throwable.class)
|
||||
public static void onExit() {
|
||||
// (indirectly) using BedrockRuntimeImpl class here to make sure it is available from
|
||||
// BedrockRuntimeAccess (injected into app classloader) and checked by Muzzle
|
||||
BedrockRuntimeAdviceBridge.referenceForMuzzleOnly();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
|
||||
|
||||
import io.opentelemetry.instrumentation.awssdk.v2_2.AbstractAws2BedrockRuntimeTest;
|
||||
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
|
||||
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
|
||||
import org.junit.jupiter.api.extension.RegisterExtension;
|
||||
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
|
||||
|
||||
class Aws2BedrockRuntimeTest extends AbstractAws2BedrockRuntimeTest {
|
||||
@RegisterExtension
|
||||
static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
|
||||
|
||||
@Override
|
||||
protected InstrumentationExtension getTesting() {
|
||||
return testing;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected ClientOverrideConfiguration.Builder createOverrideConfigurationBuilder() {
|
||||
return ClientOverrideConfiguration.builder();
|
||||
}
|
||||
}
|
||||
|
|
@ -51,3 +51,12 @@ run over API limitations set by AWS.
|
|||
If this does not fulfill your use case, perhaps because you are
|
||||
using the same SDK with a different non-AWS managed service, let us know so we can provide
|
||||
configuration for this behavior.
|
||||
|
||||
## Development
|
||||
|
||||
### Testing
|
||||
|
||||
Some tests use recorded API responses to run through instrumentation. By default, recordings
|
||||
are used, but if needing to add new tests/recordings or update existing ones, run the tests with
|
||||
the `RECORD_WITH_REAL_API` environment variable set. AWS credentials will need to be correctly
|
||||
configured to work.
|
||||
|
|
|
|||
|
|
@ -14,6 +14,11 @@ dependencies {
|
|||
compileOnly("software.amazon.awssdk:json-utils:2.17.0")
|
||||
compileOnly(project(":muzzle")) // For @NoMuzzle
|
||||
|
||||
// Don't use library to make sure base test is run with the floor version.
|
||||
// bedrock runtime is tested separately in testBedrockRuntime.
|
||||
// First release with Converse API
|
||||
compileOnly("software.amazon.awssdk:bedrockruntime:2.25.63")
|
||||
|
||||
testImplementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing"))
|
||||
|
||||
testLibrary("software.amazon.awssdk:dynamodb:2.2.0")
|
||||
|
|
@ -56,6 +61,18 @@ testing {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
val testBedrockRuntime by registering(JvmTestSuite::class) {
|
||||
dependencies {
|
||||
implementation(project())
|
||||
implementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing"))
|
||||
if (findProperty("testLatestDeps") as Boolean) {
|
||||
implementation("software.amazon.awssdk:bedrockruntime:+")
|
||||
} else {
|
||||
implementation("software.amazon.awssdk:bedrockruntime:2.25.63")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -72,6 +89,7 @@ tasks {
|
|||
}
|
||||
|
||||
check {
|
||||
dependsOn(testing.suites)
|
||||
dependsOn(testStableSemconv)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ public class AwsSdkTelemetry {
|
|||
private final Instrumenter<SqsProcessRequest, Response> consumerProcessInstrumenter;
|
||||
private final Instrumenter<ExecutionAttributes, Response> producerInstrumenter;
|
||||
private final Instrumenter<ExecutionAttributes, Response> dynamoDbInstrumenter;
|
||||
private final Instrumenter<ExecutionAttributes, Response> bedrockRuntimeInstrumenter;
|
||||
private final boolean captureExperimentalSpanAttributes;
|
||||
@Nullable private final TextMapPropagator messagingPropagator;
|
||||
private final boolean useXrayPropagator;
|
||||
|
|
@ -86,6 +87,7 @@ public class AwsSdkTelemetry {
|
|||
this.consumerProcessInstrumenter = instrumenterFactory.consumerProcessInstrumenter();
|
||||
this.producerInstrumenter = instrumenterFactory.producerInstrumenter();
|
||||
this.dynamoDbInstrumenter = instrumenterFactory.dynamoDbInstrumenter();
|
||||
this.bedrockRuntimeInstrumenter = instrumenterFactory.bedrockRuntimeInstrumenter();
|
||||
this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
|
||||
this.recordIndividualHttpError = recordIndividualHttpError;
|
||||
}
|
||||
|
|
@ -101,6 +103,7 @@ public class AwsSdkTelemetry {
|
|||
consumerProcessInstrumenter,
|
||||
producerInstrumenter,
|
||||
dynamoDbInstrumenter,
|
||||
bedrockRuntimeInstrumenter,
|
||||
captureExperimentalSpanAttributes,
|
||||
messagingPropagator,
|
||||
useXrayPropagator,
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ import io.opentelemetry.api.trace.Span;
|
|||
import io.opentelemetry.context.Context;
|
||||
import io.opentelemetry.context.propagation.TextMapPropagator;
|
||||
import io.opentelemetry.instrumentation.api.incubator.semconv.db.DbClientMetrics;
|
||||
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesExtractor;
|
||||
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiSpanNameExtractor;
|
||||
import io.opentelemetry.instrumentation.api.incubator.semconv.messaging.MessageOperation;
|
||||
import io.opentelemetry.instrumentation.api.incubator.semconv.messaging.MessagingAttributesExtractor;
|
||||
import io.opentelemetry.instrumentation.api.incubator.semconv.messaging.MessagingAttributesGetter;
|
||||
|
|
@ -219,6 +221,18 @@ public final class AwsSdkInstrumenterFactory {
|
|||
true);
|
||||
}
|
||||
|
||||
public Instrumenter<ExecutionAttributes, Response> bedrockRuntimeInstrumenter() {
|
||||
return createInstrumenter(
|
||||
openTelemetry,
|
||||
GenAiSpanNameExtractor.create(BedrockRuntimeAttributesGetter.INSTANCE),
|
||||
SpanKindExtractor.alwaysClient(),
|
||||
attributesExtractors(),
|
||||
builder ->
|
||||
builder.addAttributesExtractor(
|
||||
GenAiAttributesExtractor.create(BedrockRuntimeAttributesGetter.INSTANCE)),
|
||||
true);
|
||||
}
|
||||
|
||||
private static <REQUEST, RESPONSE> Instrumenter<REQUEST, RESPONSE> createInstrumenter(
|
||||
OpenTelemetry openTelemetry,
|
||||
SpanNameExtractor<REQUEST> spanNameExtractor,
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
|
||||
|
||||
import static io.opentelemetry.instrumentation.awssdk.v2_2.internal.AwsSdkRequestType.BEDROCK_RUNTIME;
|
||||
import static io.opentelemetry.instrumentation.awssdk.v2_2.internal.AwsSdkRequestType.DYNAMODB;
|
||||
import static io.opentelemetry.instrumentation.awssdk.v2_2.internal.AwsSdkRequestType.KINESIS;
|
||||
import static io.opentelemetry.instrumentation.awssdk.v2_2.internal.AwsSdkRequestType.S3;
|
||||
|
|
@ -119,7 +120,8 @@ enum AwsSdkRequest {
|
|||
"ProvisionedThroughput.ReadCapacityUnits"),
|
||||
request(
|
||||
"aws.dynamodb.provisioned_throughput.write_capacity_units",
|
||||
"ProvisionedThroughput.WriteCapacityUnits"));
|
||||
"ProvisionedThroughput.WriteCapacityUnits")),
|
||||
ConverseRequest(BEDROCK_RUNTIME, "ConverseRequest", request("gen_ai.request.model", "modelId"));
|
||||
|
||||
private final AwsSdkRequestType type;
|
||||
private final String requestClass;
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ enum AwsSdkRequestType {
|
|||
SQS(request("aws.queue.url", "QueueUrl"), request("aws.queue.name", "QueueName")),
|
||||
KINESIS(request("aws.stream.name", "StreamName")),
|
||||
DYNAMODB(request("aws.table.name", "TableName")),
|
||||
BEDROCK_RUNTIME(),
|
||||
SNS(
|
||||
/*
|
||||
* Only one of TopicArn and TargetArn are permitted on an SNS request.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,87 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
|
||||
|
||||
import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
|
||||
import java.util.List;
|
||||
import javax.annotation.Nullable;
|
||||
import software.amazon.awssdk.core.SdkRequest;
|
||||
import software.amazon.awssdk.core.SdkResponse;
|
||||
|
||||
final class BedrockRuntimeAccess {
|
||||
private BedrockRuntimeAccess() {}
|
||||
|
||||
private static final boolean enabled;
|
||||
|
||||
static {
|
||||
boolean isEnabled = true;
|
||||
if (!PluginImplUtil.isImplPresent("BedrockRuntimeImpl")) {
|
||||
// Muzzle disabled the instrumentation.
|
||||
isEnabled = false;
|
||||
} else {
|
||||
try {
|
||||
Class.forName("software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest");
|
||||
} catch (ClassNotFoundException e) {
|
||||
// Application does not include library
|
||||
isEnabled = false;
|
||||
}
|
||||
}
|
||||
enabled = isEnabled;
|
||||
}
|
||||
|
||||
@NoMuzzle
|
||||
static boolean isBedrockRuntimeRequest(SdkRequest request) {
|
||||
return enabled && BedrockRuntimeImpl.isBedrockRuntimeRequest(request);
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static String getModelId(SdkRequest request) {
|
||||
return enabled ? BedrockRuntimeImpl.getModelId(request) : null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static Long getMaxTokens(SdkRequest request) {
|
||||
return enabled ? BedrockRuntimeImpl.getMaxTokens(request) : null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static Double getTemperature(SdkRequest request) {
|
||||
return enabled ? BedrockRuntimeImpl.getTemperature(request) : null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static Double getTopP(SdkRequest request) {
|
||||
return enabled ? BedrockRuntimeImpl.getTopP(request) : null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static List<String> getStopSequences(SdkRequest request) {
|
||||
return enabled ? BedrockRuntimeImpl.getStopSequences(request) : null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static String getStopReason(SdkResponse response) {
|
||||
return enabled ? BedrockRuntimeImpl.getStopReason(response) : null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static Long getUsageInputTokens(SdkResponse response) {
|
||||
return enabled ? BedrockRuntimeImpl.getUsageInputTokens(response) : null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@NoMuzzle
|
||||
static Long getUsageOutputTokens(SdkResponse response) {
|
||||
return enabled ? BedrockRuntimeImpl.getUsageOutputTokens(response) : null;
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,150 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
|
||||
|
||||
import static io.opentelemetry.instrumentation.awssdk.v2_2.internal.TracingExecutionInterceptor.SDK_REQUEST_ATTRIBUTE;
|
||||
|
||||
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesGetter;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import javax.annotation.Nullable;
|
||||
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
|
||||
import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
|
||||
|
||||
enum BedrockRuntimeAttributesGetter
|
||||
implements GenAiAttributesGetter<ExecutionAttributes, Response> {
|
||||
INSTANCE;
|
||||
|
||||
// copied from GenAiIncubatingAttributes
|
||||
private static final class GenAiOperationNameIncubatingValues {
|
||||
static final String CHAT = "chat";
|
||||
|
||||
private GenAiOperationNameIncubatingValues() {}
|
||||
}
|
||||
|
||||
private static final class GenAiSystemIncubatingValues {
|
||||
static final String AWS_BEDROCK = "aws.bedrock";
|
||||
|
||||
private GenAiSystemIncubatingValues() {}
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getOperationName(ExecutionAttributes executionAttributes) {
|
||||
String operation = executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
|
||||
if (operation != null) {
|
||||
switch (operation) {
|
||||
case "Converse":
|
||||
return GenAiOperationNameIncubatingValues.CHAT;
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getSystem(ExecutionAttributes executionAttributes) {
|
||||
return GenAiSystemIncubatingValues.AWS_BEDROCK;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public String getRequestModel(ExecutionAttributes executionAttributes) {
|
||||
return BedrockRuntimeAccess.getModelId(executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Long getRequestSeed(ExecutionAttributes executionAttributes) {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public List<String> getRequestEncodingFormats(ExecutionAttributes executionAttributes) {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Double getRequestFrequencyPenalty(ExecutionAttributes executionAttributes) {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Long getRequestMaxTokens(ExecutionAttributes executionAttributes) {
|
||||
return BedrockRuntimeAccess.getMaxTokens(
|
||||
executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Double getRequestPresencePenalty(ExecutionAttributes executionAttributes) {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public List<String> getRequestStopSequences(ExecutionAttributes executionAttributes) {
|
||||
return BedrockRuntimeAccess.getStopSequences(
|
||||
executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Double getRequestTemperature(ExecutionAttributes executionAttributes) {
|
||||
return BedrockRuntimeAccess.getTemperature(
|
||||
executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Double getRequestTopK(ExecutionAttributes executionAttributes) {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Double getRequestTopP(ExecutionAttributes executionAttributes) {
|
||||
return BedrockRuntimeAccess.getTopP(executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public List<String> getResponseFinishReasons(
|
||||
ExecutionAttributes executionAttributes, Response response) {
|
||||
String stopReason = BedrockRuntimeAccess.getStopReason(response.getSdkResponse());
|
||||
if (stopReason == null) {
|
||||
return null;
|
||||
}
|
||||
return Arrays.asList(stopReason);
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public String getResponseId(ExecutionAttributes executionAttributes, Response response) {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public String getResponseModel(ExecutionAttributes executionAttributes, Response response) {
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Long getUsageInputTokens(ExecutionAttributes executionAttributes, Response response) {
|
||||
return BedrockRuntimeAccess.getUsageInputTokens(response.getSdkResponse());
|
||||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public Long getUsageOutputTokens(ExecutionAttributes executionAttributes, Response response) {
|
||||
return BedrockRuntimeAccess.getUsageOutputTokens(response.getSdkResponse());
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
|
||||
|
||||
import java.util.List;
|
||||
import javax.annotation.Nullable;
|
||||
import software.amazon.awssdk.core.SdkRequest;
|
||||
import software.amazon.awssdk.core.SdkResponse;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.InferenceConfiguration;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.StopReason;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.TokenUsage;
|
||||
|
||||
final class BedrockRuntimeImpl {
|
||||
private BedrockRuntimeImpl() {}
|
||||
|
||||
static boolean isBedrockRuntimeRequest(SdkRequest request) {
|
||||
if (request instanceof ConverseRequest) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static String getModelId(SdkRequest request) {
|
||||
if (request instanceof ConverseRequest) {
|
||||
return ((ConverseRequest) request).modelId();
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static Long getMaxTokens(SdkRequest request) {
|
||||
if (request instanceof ConverseRequest) {
|
||||
InferenceConfiguration config = ((ConverseRequest) request).inferenceConfig();
|
||||
if (config != null) {
|
||||
return integerToLong(config.maxTokens());
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static Double getTemperature(SdkRequest request) {
|
||||
if (request instanceof ConverseRequest) {
|
||||
InferenceConfiguration config = ((ConverseRequest) request).inferenceConfig();
|
||||
if (config != null) {
|
||||
return floatToDouble(config.temperature());
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static Double getTopP(SdkRequest request) {
|
||||
if (request instanceof ConverseRequest) {
|
||||
InferenceConfiguration config = ((ConverseRequest) request).inferenceConfig();
|
||||
if (config != null) {
|
||||
return floatToDouble(config.topP());
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static List<String> getStopSequences(SdkRequest request) {
|
||||
if (request instanceof ConverseRequest) {
|
||||
InferenceConfiguration config = ((ConverseRequest) request).inferenceConfig();
|
||||
if (config != null) {
|
||||
return config.stopSequences();
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static String getStopReason(SdkResponse response) {
|
||||
if (response instanceof ConverseResponse) {
|
||||
StopReason reason = ((ConverseResponse) response).stopReason();
|
||||
if (reason != null) {
|
||||
return reason.toString();
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static Long getUsageInputTokens(SdkResponse response) {
|
||||
if (response instanceof ConverseResponse) {
|
||||
TokenUsage usage = ((ConverseResponse) response).usage();
|
||||
if (usage != null) {
|
||||
return integerToLong(usage.inputTokens());
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
static Long getUsageOutputTokens(SdkResponse response) {
|
||||
if (response instanceof ConverseResponse) {
|
||||
TokenUsage usage = ((ConverseResponse) response).usage();
|
||||
if (usage != null) {
|
||||
return integerToLong(usage.outputTokens());
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
@Nullable
|
||||
private static Long integerToLong(Integer value) {
|
||||
if (value == null) {
|
||||
return null;
|
||||
}
|
||||
return Long.valueOf(value);
|
||||
}
|
||||
|
||||
@Nullable
|
||||
private static Double floatToDouble(Float value) {
|
||||
if (value == null) {
|
||||
return null;
|
||||
}
|
||||
return Double.valueOf(value);
|
||||
}
|
||||
}
|
||||
|
|
@ -27,7 +27,6 @@ import java.time.Instant;
|
|||
import java.util.Optional;
|
||||
import java.util.stream.Collectors;
|
||||
import javax.annotation.Nullable;
|
||||
import software.amazon.awssdk.auth.signer.AwsSignerExecutionAttribute;
|
||||
import software.amazon.awssdk.awscore.AwsResponse;
|
||||
import software.amazon.awssdk.core.ClientType;
|
||||
import software.amazon.awssdk.core.SdkRequest;
|
||||
|
|
@ -77,6 +76,7 @@ public final class TracingExecutionInterceptor implements ExecutionInterceptor {
|
|||
private final Instrumenter<SqsProcessRequest, Response> consumerProcessInstrumenter;
|
||||
private final Instrumenter<ExecutionAttributes, Response> producerInstrumenter;
|
||||
private final Instrumenter<ExecutionAttributes, Response> dynamoDbInstrumenter;
|
||||
private final Instrumenter<ExecutionAttributes, Response> bedrockRuntimeInstrumenter;
|
||||
private final boolean captureExperimentalSpanAttributes;
|
||||
|
||||
static final AttributeKey<String> HTTP_ERROR_MSG =
|
||||
|
|
@ -105,12 +105,14 @@ public final class TracingExecutionInterceptor implements ExecutionInterceptor {
|
|||
private final boolean recordIndividualHttpError;
|
||||
private final FieldMapper fieldMapper;
|
||||
|
||||
@SuppressWarnings("TooManyParameters") // internal method
|
||||
public TracingExecutionInterceptor(
|
||||
Instrumenter<ExecutionAttributes, Response> requestInstrumenter,
|
||||
Instrumenter<SqsReceiveRequest, Response> consumerReceiveInstrumenter,
|
||||
Instrumenter<SqsProcessRequest, Response> consumerProcessInstrumenter,
|
||||
Instrumenter<ExecutionAttributes, Response> producerInstrumenter,
|
||||
Instrumenter<ExecutionAttributes, Response> dynamoDbInstrumenter,
|
||||
Instrumenter<ExecutionAttributes, Response> bedrockRuntimeInstrumenter,
|
||||
boolean captureExperimentalSpanAttributes,
|
||||
TextMapPropagator messagingPropagator,
|
||||
boolean useXrayPropagator,
|
||||
|
|
@ -120,6 +122,7 @@ public final class TracingExecutionInterceptor implements ExecutionInterceptor {
|
|||
this.consumerProcessInstrumenter = consumerProcessInstrumenter;
|
||||
this.producerInstrumenter = producerInstrumenter;
|
||||
this.dynamoDbInstrumenter = dynamoDbInstrumenter;
|
||||
this.bedrockRuntimeInstrumenter = bedrockRuntimeInstrumenter;
|
||||
this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
|
||||
this.messagingPropagator = messagingPropagator;
|
||||
this.useXrayPropagator = useXrayPropagator;
|
||||
|
|
@ -128,6 +131,7 @@ public final class TracingExecutionInterceptor implements ExecutionInterceptor {
|
|||
}
|
||||
|
||||
@Override
|
||||
@SuppressWarnings("deprecation") // need to access deprecated signer
|
||||
public SdkRequest modifyRequest(
|
||||
Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
|
||||
|
||||
|
|
@ -144,7 +148,8 @@ public final class TracingExecutionInterceptor implements ExecutionInterceptor {
|
|||
|
||||
// Ignore presign request. These requests don't run all interceptor methods and the span created
|
||||
// here would never be ended and scope closed.
|
||||
if (executionAttributes.getAttribute(AwsSignerExecutionAttribute.PRESIGNER_EXPIRATION)
|
||||
if (executionAttributes.getAttribute(
|
||||
software.amazon.awssdk.auth.signer.AwsSignerExecutionAttribute.PRESIGNER_EXPIRATION)
|
||||
!= null) {
|
||||
return request;
|
||||
}
|
||||
|
|
@ -451,6 +456,9 @@ public final class TracingExecutionInterceptor implements ExecutionInterceptor {
|
|||
if (SqsAccess.isSqsProducerRequest(request)) {
|
||||
return producerInstrumenter;
|
||||
}
|
||||
if (BedrockRuntimeAccess.isBedrockRuntimeRequest(request)) {
|
||||
return bedrockRuntimeInstrumenter;
|
||||
}
|
||||
if (awsSdkRequest != null && awsSdkRequest.type() == DYNAMODB) {
|
||||
return dynamoDbInstrumenter;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,39 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
|
||||
|
||||
import io.opentelemetry.instrumentation.awssdk.v2_2.AbstractAws2BedrockRuntimeTest;
|
||||
import io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkTelemetry;
|
||||
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
|
||||
import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.extension.RegisterExtension;
|
||||
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
|
||||
|
||||
class Aws2BedrockRuntimeTest extends AbstractAws2BedrockRuntimeTest {
|
||||
|
||||
@RegisterExtension
|
||||
private static final LibraryInstrumentationExtension testing =
|
||||
LibraryInstrumentationExtension.create();
|
||||
|
||||
@Override
|
||||
protected InstrumentationExtension getTesting() {
|
||||
return testing;
|
||||
}
|
||||
|
||||
private static AwsSdkTelemetry telemetry;
|
||||
|
||||
@BeforeAll
|
||||
static void setup() {
|
||||
telemetry = AwsSdkTelemetry.create(testing.getOpenTelemetry());
|
||||
}
|
||||
|
||||
@Override
|
||||
protected ClientOverrideConfiguration.Builder createOverrideConfigurationBuilder() {
|
||||
return ClientOverrideConfiguration.builder()
|
||||
.addExecutionInterceptor(telemetry.newExecutionInterceptor());
|
||||
}
|
||||
}
|
||||
|
|
@ -11,6 +11,7 @@ dependencies {
|
|||
|
||||
// compileOnly because we never want to pin the low version implicitly; need to add dependencies
|
||||
// explicitly in user projects, e.g. using testLatestDeps.
|
||||
compileOnly("software.amazon.awssdk:bedrockruntime:2.25.63")
|
||||
compileOnly("software.amazon.awssdk:dynamodb:2.2.0")
|
||||
compileOnly("software.amazon.awssdk:ec2:2.2.0")
|
||||
compileOnly("software.amazon.awssdk:kinesis:2.2.0")
|
||||
|
|
@ -24,6 +25,11 @@ dependencies {
|
|||
// needed for SQS - using emq directly as localstack references emq v0.15.7 ie WITHOUT AWS trace header propagation
|
||||
implementation("org.elasticmq:elasticmq-rest-sqs_2.13")
|
||||
|
||||
// used to record LLM responses in bedrock tests
|
||||
implementation("com.github.tomakehurst:wiremock-jre8:2.35.2")
|
||||
implementation("com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.18.2")
|
||||
implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.18.2")
|
||||
|
||||
implementation("com.google.guava:guava")
|
||||
|
||||
implementation("io.opentelemetry:opentelemetry-api")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,165 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2;
|
||||
|
||||
import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
|
||||
import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.satisfies;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_OPERATION_NAME;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_REQUEST_MAX_TOKENS;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_REQUEST_MODEL;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_REQUEST_STOP_SEQUENCES;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_REQUEST_TEMPERATURE;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_REQUEST_TOP_P;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_RESPONSE_FINISH_REASONS;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_SYSTEM;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_USAGE_INPUT_TOKENS;
|
||||
import static io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes.GEN_AI_USAGE_OUTPUT_TOKENS;
|
||||
import static java.util.Arrays.asList;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.within;
|
||||
|
||||
import io.opentelemetry.api.trace.SpanKind;
|
||||
import io.opentelemetry.instrumentation.awssdk.v2_2.recording.RecordingExtension;
|
||||
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
|
||||
import io.opentelemetry.semconv.incubating.GenAiIncubatingAttributes;
|
||||
import java.net.URI;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.RegisterExtension;
|
||||
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
|
||||
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
|
||||
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient;
|
||||
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClientBuilder;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.InferenceConfiguration;
|
||||
import software.amazon.awssdk.services.bedrockruntime.model.Message;
|
||||
|
||||
public abstract class AbstractAws2BedrockRuntimeTest {
|
||||
|
||||
private static final String API_URL = "https://bedrock-runtime.us-east-1.amazonaws.com";
|
||||
|
||||
@RegisterExtension static final RecordingExtension recording = new RecordingExtension(API_URL);
|
||||
|
||||
protected abstract InstrumentationExtension getTesting();
|
||||
|
||||
protected abstract ClientOverrideConfiguration.Builder createOverrideConfigurationBuilder();
|
||||
|
||||
private static void configureClient(BedrockRuntimeClientBuilder builder) {
|
||||
builder
|
||||
.region(Region.US_EAST_1)
|
||||
.endpointOverride(URI.create("http://localhost:" + recording.getPort()));
|
||||
if (recording.isRecording()) {
|
||||
builder.putAuthScheme(new FixedHostAwsV4AuthScheme(API_URL));
|
||||
} else {
|
||||
builder.credentialsProvider(
|
||||
StaticCredentialsProvider.create(AwsBasicCredentials.create("testing", "testing")));
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void testConverseBasic() {
|
||||
BedrockRuntimeClientBuilder builder = BedrockRuntimeClient.builder();
|
||||
builder.overrideConfiguration(createOverrideConfigurationBuilder().build());
|
||||
configureClient(builder);
|
||||
BedrockRuntimeClient client = builder.build();
|
||||
|
||||
String modelId = "amazon.titan-text-lite-v1";
|
||||
ConverseResponse response =
|
||||
client.converse(
|
||||
ConverseRequest.builder()
|
||||
.modelId(modelId)
|
||||
.messages(
|
||||
Message.builder()
|
||||
.role(ConversationRole.USER)
|
||||
.content(ContentBlock.fromText("Say this is a test"))
|
||||
.build())
|
||||
.build());
|
||||
|
||||
assertThat(response.output().message().content().get(0).text())
|
||||
.isEqualTo("Hi there! How can I help you today?");
|
||||
|
||||
getTesting()
|
||||
.waitAndAssertTraces(
|
||||
trace ->
|
||||
trace.hasSpansSatisfyingExactly(
|
||||
span ->
|
||||
span.hasName("chat amazon.titan-text-lite-v1")
|
||||
.hasKind(SpanKind.CLIENT)
|
||||
.hasAttributesSatisfying(
|
||||
equalTo(
|
||||
GEN_AI_SYSTEM,
|
||||
GenAiIncubatingAttributes.GenAiSystemIncubatingValues
|
||||
.AWS_BEDROCK),
|
||||
equalTo(
|
||||
GEN_AI_OPERATION_NAME,
|
||||
GenAiIncubatingAttributes.GenAiOperationNameIncubatingValues
|
||||
.CHAT),
|
||||
equalTo(GEN_AI_REQUEST_MODEL, modelId),
|
||||
equalTo(GEN_AI_USAGE_INPUT_TOKENS, 8),
|
||||
equalTo(GEN_AI_USAGE_OUTPUT_TOKENS, 14),
|
||||
equalTo(GEN_AI_RESPONSE_FINISH_REASONS, asList("end_turn")))));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testConverseOptions() {
|
||||
BedrockRuntimeClientBuilder builder = BedrockRuntimeClient.builder();
|
||||
builder.overrideConfiguration(createOverrideConfigurationBuilder().build());
|
||||
configureClient(builder);
|
||||
BedrockRuntimeClient client = builder.build();
|
||||
|
||||
String modelId = "amazon.titan-text-lite-v1";
|
||||
ConverseResponse response =
|
||||
client.converse(
|
||||
ConverseRequest.builder()
|
||||
.modelId(modelId)
|
||||
.messages(
|
||||
Message.builder()
|
||||
.role(ConversationRole.USER)
|
||||
.content(ContentBlock.fromText("Say this is a test"))
|
||||
.build())
|
||||
.inferenceConfig(
|
||||
InferenceConfiguration.builder()
|
||||
.maxTokens(10)
|
||||
.temperature(0.8f)
|
||||
.topP(1f)
|
||||
.stopSequences("|")
|
||||
.build())
|
||||
.build());
|
||||
|
||||
assertThat(response.output().message().content().get(0).text()).isEqualTo("This is an LLM (");
|
||||
|
||||
getTesting()
|
||||
.waitAndAssertTraces(
|
||||
trace ->
|
||||
trace.hasSpansSatisfyingExactly(
|
||||
span ->
|
||||
span.hasName("chat amazon.titan-text-lite-v1")
|
||||
.hasKind(SpanKind.CLIENT)
|
||||
.hasAttributesSatisfying(
|
||||
equalTo(
|
||||
GEN_AI_SYSTEM,
|
||||
GenAiIncubatingAttributes.GenAiSystemIncubatingValues
|
||||
.AWS_BEDROCK),
|
||||
equalTo(
|
||||
GEN_AI_OPERATION_NAME,
|
||||
GenAiIncubatingAttributes.GenAiOperationNameIncubatingValues
|
||||
.CHAT),
|
||||
equalTo(GEN_AI_REQUEST_MODEL, modelId),
|
||||
equalTo(GEN_AI_REQUEST_MAX_TOKENS, 10),
|
||||
satisfies(
|
||||
GEN_AI_REQUEST_TEMPERATURE,
|
||||
temp -> temp.isCloseTo(0.8, within(0.0001))),
|
||||
equalTo(GEN_AI_REQUEST_TOP_P, 1.0),
|
||||
equalTo(GEN_AI_REQUEST_STOP_SEQUENCES, asList("|")),
|
||||
equalTo(GEN_AI_USAGE_INPUT_TOKENS, 8),
|
||||
equalTo(GEN_AI_USAGE_OUTPUT_TOKENS, 10),
|
||||
equalTo(GEN_AI_RESPONSE_FINISH_REASONS, asList("max_tokens")))));
|
||||
}
|
||||
}
|
||||
|
|
@ -59,7 +59,6 @@ import software.amazon.awssdk.core.ResponseInputStream;
|
|||
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
|
||||
import software.amazon.awssdk.core.exception.SdkClientException;
|
||||
import software.amazon.awssdk.core.exception.SdkException;
|
||||
import software.amazon.awssdk.core.retry.RetryPolicy;
|
||||
import software.amazon.awssdk.http.apache.ApacheHttpClient;
|
||||
import software.amazon.awssdk.regions.Region;
|
||||
import software.amazon.awssdk.services.ec2.Ec2AsyncClient;
|
||||
|
|
@ -92,6 +91,7 @@ import software.amazon.awssdk.services.sqs.SqsClientBuilder;
|
|||
import software.amazon.awssdk.services.sqs.model.CreateQueueRequest;
|
||||
import software.amazon.awssdk.services.sqs.model.SendMessageRequest;
|
||||
|
||||
@SuppressWarnings("deprecation") // We need to use deprecated APIs for testing older versions.
|
||||
public abstract class AbstractAws2ClientTest extends AbstractAws2ClientCoreTest {
|
||||
private static final String QUEUE_URL = "http://xxx/somequeue";
|
||||
|
||||
|
|
@ -638,7 +638,10 @@ public abstract class AbstractAws2ClientTest extends AbstractAws2ClientCoreTest
|
|||
S3Client.builder()
|
||||
.overrideConfiguration(
|
||||
createOverrideConfigurationBuilder()
|
||||
.retryPolicy(RetryPolicy.builder().numRetries(1).build())
|
||||
.retryPolicy(
|
||||
software.amazon.awssdk.core.retry.RetryPolicy.builder()
|
||||
.numRetries(1)
|
||||
.build())
|
||||
.build())
|
||||
.endpointOverride(clientUri)
|
||||
.region(Region.AP_NORTHEAST_1)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,81 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2;
|
||||
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import software.amazon.awssdk.http.SdkHttpRequest;
|
||||
import software.amazon.awssdk.http.auth.aws.internal.scheme.DefaultAwsV4AuthScheme;
|
||||
import software.amazon.awssdk.http.auth.aws.internal.signer.DefaultAwsV4HttpSigner;
|
||||
import software.amazon.awssdk.http.auth.aws.scheme.AwsV4AuthScheme;
|
||||
import software.amazon.awssdk.http.auth.aws.signer.AwsV4HttpSigner;
|
||||
import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest;
|
||||
import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest;
|
||||
import software.amazon.awssdk.http.auth.spi.signer.SignRequest;
|
||||
import software.amazon.awssdk.http.auth.spi.signer.SignedRequest;
|
||||
import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
|
||||
import software.amazon.awssdk.identity.spi.IdentityProvider;
|
||||
import software.amazon.awssdk.identity.spi.IdentityProviders;
|
||||
|
||||
final class FixedHostAwsV4AuthScheme implements AwsV4AuthScheme {
|
||||
|
||||
private final FixedHostAwsV4HttpSigner signer;
|
||||
|
||||
public FixedHostAwsV4AuthScheme(String apiUrl) {
|
||||
this.signer = new FixedHostAwsV4HttpSigner(apiUrl);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String schemeId() {
|
||||
return AwsV4AuthScheme.SCHEME_ID;
|
||||
}
|
||||
|
||||
@Override
|
||||
public IdentityProvider<AwsCredentialsIdentity> identityProvider(
|
||||
IdentityProviders identityProviders) {
|
||||
return DefaultAwsV4AuthScheme.create().identityProvider(identityProviders);
|
||||
}
|
||||
|
||||
@Override
|
||||
public AwsV4HttpSigner signer() {
|
||||
return signer;
|
||||
}
|
||||
|
||||
private static class FixedHostAwsV4HttpSigner implements AwsV4HttpSigner {
|
||||
private static final AwsV4HttpSigner DEFAULT = new DefaultAwsV4HttpSigner();
|
||||
|
||||
private final String apiUrl;
|
||||
|
||||
FixedHostAwsV4HttpSigner(String apiUrl) {
|
||||
this.apiUrl = apiUrl;
|
||||
}
|
||||
|
||||
@Override
|
||||
public SignedRequest sign(SignRequest<? extends AwsCredentialsIdentity> request) {
|
||||
SdkHttpRequest original = request.request();
|
||||
SignRequest<? extends AwsCredentialsIdentity> override =
|
||||
request.toBuilder()
|
||||
.request(
|
||||
request.request().toBuilder().port(443).protocol("https").host(apiUrl).build())
|
||||
.build();
|
||||
SignedRequest signed = DEFAULT.sign(override);
|
||||
return signed.toBuilder()
|
||||
.request(
|
||||
signed.request().toBuilder()
|
||||
.protocol(original.protocol())
|
||||
.host(original.host())
|
||||
.port(original.port())
|
||||
.build())
|
||||
.build();
|
||||
}
|
||||
|
||||
@Override
|
||||
public CompletableFuture<AsyncSignedRequest> signAsync(
|
||||
AsyncSignRequest<? extends AwsCredentialsIdentity> request) {
|
||||
// TODO: Implement
|
||||
return null;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.recording;
|
||||
|
||||
import com.github.tomakehurst.wiremock.common.FileSource;
|
||||
import com.github.tomakehurst.wiremock.extension.Parameters;
|
||||
import com.github.tomakehurst.wiremock.extension.StubMappingTransformer;
|
||||
import com.github.tomakehurst.wiremock.matching.ContentPattern;
|
||||
import com.github.tomakehurst.wiremock.matching.EqualToJsonPattern;
|
||||
import com.github.tomakehurst.wiremock.stubbing.StubMapping;
|
||||
import java.util.List;
|
||||
|
||||
public final class PrettyPrintEqualToJsonStubMappingTransformer extends StubMappingTransformer {
|
||||
@Override
|
||||
public String getName() {
|
||||
return "pretty-print-equal-to-json";
|
||||
}
|
||||
|
||||
@Override
|
||||
public StubMapping transform(StubMapping stubMapping, FileSource files, Parameters parameters) {
|
||||
List<ContentPattern<?>> patterns = stubMapping.getRequest().getBodyPatterns();
|
||||
if (patterns != null) {
|
||||
for (int i = 0; i < patterns.size(); i++) {
|
||||
ContentPattern<?> pattern = patterns.get(i);
|
||||
if (!(pattern instanceof EqualToJsonPattern)) {
|
||||
continue;
|
||||
}
|
||||
EqualToJsonPattern equalToJsonPattern = (EqualToJsonPattern) pattern;
|
||||
patterns.set(
|
||||
i,
|
||||
new EqualToJsonPattern(
|
||||
equalToJsonPattern.getExpected(), // pretty printed,
|
||||
// We exact match the request unlike the default.
|
||||
false,
|
||||
false));
|
||||
}
|
||||
}
|
||||
return stubMapping;
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.recording;
|
||||
|
||||
import static com.github.tomakehurst.wiremock.client.WireMock.proxyAllTo;
|
||||
import static com.github.tomakehurst.wiremock.client.WireMock.recordSpec;
|
||||
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
|
||||
|
||||
import com.github.tomakehurst.wiremock.common.SingleRootFileSource;
|
||||
import com.github.tomakehurst.wiremock.junit5.WireMockExtension;
|
||||
import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo;
|
||||
import org.junit.jupiter.api.extension.AfterTestExecutionCallback;
|
||||
import org.junit.jupiter.api.extension.ExtensionContext;
|
||||
|
||||
public final class RecordingExtension extends WireMockExtension
|
||||
implements AfterTestExecutionCallback {
|
||||
/**
|
||||
* Setting this to true will make the tests call the real API instead and record the responses.
|
||||
* You'll have to setup credentials for this to work.
|
||||
*/
|
||||
private static final boolean RECORD_WITH_REAL_API = System.getenv("RECORD_WITH_REAL_API") != null;
|
||||
|
||||
private final String apiUrl;
|
||||
|
||||
@SuppressWarnings({"unchecked", "varargs"})
|
||||
public RecordingExtension(String apiUrl) {
|
||||
super(
|
||||
WireMockExtension.newInstance()
|
||||
.options(
|
||||
options()
|
||||
.extensions(
|
||||
ResponseHeaderScrubber.class,
|
||||
PrettyPrintEqualToJsonStubMappingTransformer.class)
|
||||
.mappingSource(
|
||||
new YamlFileMappingsSource(
|
||||
new SingleRootFileSource("../testing/src/main/resources")
|
||||
.child("mappings")))));
|
||||
this.apiUrl = apiUrl;
|
||||
}
|
||||
|
||||
public boolean isRecording() {
|
||||
return RECORD_WITH_REAL_API;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onBeforeEach(WireMockRuntimeInfo wireMock) {
|
||||
// Set a low priority so recordings are used when available
|
||||
if (RECORD_WITH_REAL_API) {
|
||||
stubFor(proxyAllTo(apiUrl).atPriority(Integer.MAX_VALUE));
|
||||
startRecording(
|
||||
recordSpec()
|
||||
.forTarget(apiUrl)
|
||||
.transformers("scrub-response-header", "pretty-print-equal-to-json")
|
||||
// Include all bodies inline.
|
||||
.extractTextBodiesOver(Long.MAX_VALUE)
|
||||
.extractBinaryBodiesOver(Long.MAX_VALUE));
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onAfterEach(WireMockRuntimeInfo wireMockRuntimeInfo) {
|
||||
if (RECORD_WITH_REAL_API) {
|
||||
stopRecording();
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void afterTestExecution(ExtensionContext context) {
|
||||
YamlFileMappingsSource.setCurrentTest(context);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,38 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.recording;
|
||||
|
||||
import com.github.tomakehurst.wiremock.common.FileSource;
|
||||
import com.github.tomakehurst.wiremock.extension.Parameters;
|
||||
import com.github.tomakehurst.wiremock.extension.ResponseTransformer;
|
||||
import com.github.tomakehurst.wiremock.http.HttpHeader;
|
||||
import com.github.tomakehurst.wiremock.http.HttpHeaders;
|
||||
import com.github.tomakehurst.wiremock.http.Request;
|
||||
import com.github.tomakehurst.wiremock.http.Response;
|
||||
|
||||
public final class ResponseHeaderScrubber extends ResponseTransformer {
|
||||
@Override
|
||||
public String getName() {
|
||||
return "scrub-response-header";
|
||||
}
|
||||
|
||||
@Override
|
||||
public Response transform(
|
||||
Request request, Response response, FileSource fileSource, Parameters parameters) {
|
||||
HttpHeaders scrubbed = HttpHeaders.noHeaders();
|
||||
for (HttpHeader header : response.getHeaders().all()) {
|
||||
switch (header.key()) {
|
||||
case "Set-Cookie":
|
||||
scrubbed = scrubbed.plus(HttpHeader.httpHeader("Set-Cookie", "test_set_cookie"));
|
||||
break;
|
||||
default:
|
||||
scrubbed = scrubbed.plus(header);
|
||||
break;
|
||||
}
|
||||
}
|
||||
return Response.Builder.like(response).but().headers(scrubbed).build();
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,203 @@
|
|||
/*
|
||||
* Copyright The OpenTelemetry Authors
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
package io.opentelemetry.instrumentation.awssdk.v2_2.recording;
|
||||
|
||||
import static com.github.tomakehurst.wiremock.common.AbstractFileSource.byFileExtension;
|
||||
import static com.github.tomakehurst.wiremock.common.Exceptions.throwUnchecked;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonInclude.Include;
|
||||
import com.fasterxml.jackson.core.JsonGenerator;
|
||||
import com.fasterxml.jackson.core.JsonParser;
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.DeserializationFeature;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.fasterxml.jackson.databind.ObjectWriter;
|
||||
import com.fasterxml.jackson.databind.SerializationFeature;
|
||||
import com.fasterxml.jackson.databind.cfg.JsonNodeFeature;
|
||||
import com.fasterxml.jackson.dataformat.yaml.YAMLGenerator;
|
||||
import com.fasterxml.jackson.dataformat.yaml.YAMLMapper;
|
||||
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
|
||||
import com.github.tomakehurst.wiremock.common.FileSource;
|
||||
import com.github.tomakehurst.wiremock.common.Json;
|
||||
import com.github.tomakehurst.wiremock.common.JsonException;
|
||||
import com.github.tomakehurst.wiremock.common.NotWritableException;
|
||||
import com.github.tomakehurst.wiremock.common.TextFile;
|
||||
import com.github.tomakehurst.wiremock.standalone.MappingFileException;
|
||||
import com.github.tomakehurst.wiremock.standalone.MappingsSource;
|
||||
import com.github.tomakehurst.wiremock.stubbing.StubMapping;
|
||||
import com.github.tomakehurst.wiremock.stubbing.StubMappings;
|
||||
import java.io.IOException;
|
||||
import java.lang.reflect.Method;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Paths;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
import java.util.UUID;
|
||||
import java.util.regex.Pattern;
|
||||
import java.util.stream.Collectors;
|
||||
import org.junit.jupiter.api.extension.ExtensionContext;
|
||||
|
||||
// Mostly the same as
|
||||
// https://github.com/wiremock/wiremock/blob/master/src/main/java/com/github/tomakehurst/wiremock/standalone/JsonFileMappingsSource.java
|
||||
// replacing Json with Yaml.
|
||||
final class YamlFileMappingsSource implements MappingsSource {
|
||||
|
||||
private static final ObjectMapper yamlMapper =
|
||||
new YAMLMapper()
|
||||
.enable(YAMLGenerator.Feature.MINIMIZE_QUOTES)
|
||||
.enable(YAMLGenerator.Feature.ALWAYS_QUOTE_NUMBERS_AS_STRINGS)
|
||||
// For non-YAML, follow
|
||||
// https://github.com/wiremock/wiremock/blob/master/src/main/java/com/github/tomakehurst/wiremock/common/Json.java#L41
|
||||
.setSerializationInclusion(Include.NON_NULL)
|
||||
.configure(JsonNodeFeature.STRIP_TRAILING_BIGDECIMAL_ZEROES, false)
|
||||
.configure(JsonParser.Feature.ALLOW_COMMENTS, true)
|
||||
.configure(JsonParser.Feature.ALLOW_SINGLE_QUOTES, true)
|
||||
.configure(JsonParser.Feature.IGNORE_UNDEFINED, true)
|
||||
.configure(JsonGenerator.Feature.WRITE_BIGDECIMAL_AS_PLAIN, true)
|
||||
.configure(DeserializationFeature.USE_BIG_DECIMAL_FOR_FLOATS, true)
|
||||
.registerModule(new JavaTimeModule())
|
||||
.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS)
|
||||
.enable(JsonParser.Feature.INCLUDE_SOURCE_IN_LOCATION);
|
||||
|
||||
private static final ThreadLocal<ExtensionContext> currentTest = new ThreadLocal<>();
|
||||
|
||||
private static final Pattern NON_ALPHANUMERIC = Pattern.compile("[^\\w-.]");
|
||||
|
||||
static void setCurrentTest(ExtensionContext context) {
|
||||
currentTest.set(context);
|
||||
}
|
||||
|
||||
private final FileSource mappingsFileSource;
|
||||
private final Map<UUID, StubMappingFileMetadata> fileNameMap;
|
||||
|
||||
YamlFileMappingsSource(FileSource mappingsFileSource) {
|
||||
this.mappingsFileSource = mappingsFileSource;
|
||||
fileNameMap = new HashMap<>();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void save(List<StubMapping> stubMappings) {
|
||||
for (StubMapping mapping : stubMappings) {
|
||||
if (mapping != null && mapping.isDirty()) {
|
||||
save(mapping);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void save(StubMapping stubMapping) {
|
||||
StubMappingFileMetadata fileMetadata = fileNameMap.get(stubMapping.getId());
|
||||
if (fileMetadata == null) {
|
||||
ExtensionContext test = currentTest.get();
|
||||
Method method = test.getTestMethod().get();
|
||||
String filename = method.getDeclaringClass().getName() + "." + method.getName();
|
||||
fileMetadata = new StubMappingFileMetadata(sanitize(filename) + ".yaml", false);
|
||||
}
|
||||
|
||||
if (fileMetadata.multi) {
|
||||
throw new NotWritableException(
|
||||
"Stubs loaded from multi-mapping files are read-only, and therefore cannot be saved");
|
||||
}
|
||||
|
||||
String yaml = "";
|
||||
try {
|
||||
ObjectWriter objectWriter =
|
||||
yamlMapper.writerWithDefaultPrettyPrinter().withView(Json.PrivateView.class);
|
||||
yaml = objectWriter.writeValueAsString(stubMapping);
|
||||
} catch (IOException ioe) {
|
||||
throwUnchecked(ioe, String.class);
|
||||
}
|
||||
TextFile outFile = mappingsFileSource.getTextFileNamed(fileMetadata.path);
|
||||
// For multiple requests from the same test, we append as a multi-file yaml.
|
||||
if (Files.exists(Paths.get(outFile.getPath()))) {
|
||||
String existing = outFile.readContentsAsString();
|
||||
yaml = existing + yaml;
|
||||
}
|
||||
mappingsFileSource.writeTextFile(fileMetadata.path, yaml);
|
||||
|
||||
fileNameMap.put(stubMapping.getId(), fileMetadata);
|
||||
stubMapping.setDirty(false);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void remove(StubMapping stubMapping) {
|
||||
StubMappingFileMetadata fileMetadata = fileNameMap.get(stubMapping.getId());
|
||||
if (fileMetadata.multi) {
|
||||
throw new NotWritableException(
|
||||
"Stubs loaded from multi-mapping files are read-only, and therefore cannot be removed");
|
||||
}
|
||||
|
||||
mappingsFileSource.deleteFile(fileMetadata.path);
|
||||
fileNameMap.remove(stubMapping.getId());
|
||||
}
|
||||
|
||||
@Override
|
||||
public void removeAll() {
|
||||
if (anyFilesAreMultiMapping()) {
|
||||
throw new NotWritableException(
|
||||
"Some stubs were loaded from multi-mapping files which are read-only, so remove all cannot be performed");
|
||||
}
|
||||
|
||||
for (StubMappingFileMetadata fileMetadata : fileNameMap.values()) {
|
||||
mappingsFileSource.deleteFile(fileMetadata.path);
|
||||
}
|
||||
fileNameMap.clear();
|
||||
}
|
||||
|
||||
private boolean anyFilesAreMultiMapping() {
|
||||
return fileNameMap.values().stream().anyMatch(input -> input.multi);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void loadMappingsInto(StubMappings stubMappings) {
|
||||
if (!mappingsFileSource.exists()) {
|
||||
return;
|
||||
}
|
||||
|
||||
List<TextFile> mappingFiles =
|
||||
mappingsFileSource.listFilesRecursively().stream()
|
||||
.filter(byFileExtension("yaml"))
|
||||
.collect(Collectors.toList());
|
||||
for (TextFile mappingFile : mappingFiles) {
|
||||
try {
|
||||
List<StubMapping> mappings =
|
||||
yamlMapper
|
||||
.readValues(
|
||||
yamlMapper.createParser(mappingFile.readContentsAsString()), StubMapping.class)
|
||||
.readAll();
|
||||
for (StubMapping mapping : mappings) {
|
||||
mapping.setDirty(false);
|
||||
stubMappings.addMapping(mapping);
|
||||
StubMappingFileMetadata fileMetadata =
|
||||
new StubMappingFileMetadata(mappingFile.getPath(), mappings.size() > 1);
|
||||
fileNameMap.put(mapping.getId(), fileMetadata);
|
||||
}
|
||||
} catch (JsonProcessingException e) {
|
||||
throw new MappingFileException(
|
||||
mappingFile.getPath(), JsonException.fromJackson(e).getErrors().first().getDetail());
|
||||
} catch (IOException e) {
|
||||
throwUnchecked(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static String sanitize(String s) {
|
||||
String decoratedString = String.join("-", s.split(" "));
|
||||
return NON_ALPHANUMERIC.matcher(decoratedString).replaceAll("").toLowerCase(Locale.ROOT);
|
||||
}
|
||||
|
||||
private static class StubMappingFileMetadata {
|
||||
final String path;
|
||||
final boolean multi;
|
||||
|
||||
public StubMappingFileMetadata(String path, boolean multi) {
|
||||
this.path = path;
|
||||
this.multi = multi;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
---
|
||||
id: af4ed37f-021f-43eb-beac-519f4da59b49
|
||||
name: model_amazontitan-text-lite-v1_converse
|
||||
request:
|
||||
url: /model/amazon.titan-text-lite-v1/converse
|
||||
method: POST
|
||||
bodyPatterns:
|
||||
- equalToJson: |-
|
||||
{
|
||||
"messages" : [ {
|
||||
"role" : "user",
|
||||
"content" : [ {
|
||||
"text" : "Say this is a test"
|
||||
} ]
|
||||
} ]
|
||||
}
|
||||
ignoreArrayOrder: false
|
||||
ignoreExtraElements: false
|
||||
response:
|
||||
status: 200
|
||||
body: "{\"metrics\":{\"latencyMs\":743},\"output\":{\"message\":{\"content\":[{\"\
|
||||
text\":\"Hi there! How can I help you today?\"}],\"role\":\"assistant\"}},\"stopReason\"\
|
||||
:\"end_turn\",\"usage\":{\"inputTokens\":8,\"outputTokens\":14,\"totalTokens\"\
|
||||
:22}}"
|
||||
headers:
|
||||
Date: "Thu, 20 Feb 2025 04:36:27 GMT"
|
||||
Content-Type: application/json
|
||||
x-amzn-RequestId: 2c52dcff-377c-40eb-8367-df9d1a40b746
|
||||
uuid: af4ed37f-021f-43eb-beac-519f4da59b49
|
||||
persistent: true
|
||||
insertionIndex: 4
|
||||
|
|
@ -0,0 +1,36 @@
|
|||
---
|
||||
id: d50f134a-048f-41e2-ad78-f702880b78d6
|
||||
name: model_amazontitan-text-lite-v1_converse
|
||||
request:
|
||||
url: /model/amazon.titan-text-lite-v1/converse
|
||||
method: POST
|
||||
bodyPatterns:
|
||||
- equalToJson: |-
|
||||
{
|
||||
"messages" : [ {
|
||||
"role" : "user",
|
||||
"content" : [ {
|
||||
"text" : "Say this is a test"
|
||||
} ]
|
||||
} ],
|
||||
"inferenceConfig" : {
|
||||
"maxTokens" : 10,
|
||||
"temperature" : 0.8,
|
||||
"topP" : 1,
|
||||
"stopSequences" : [ "|" ]
|
||||
}
|
||||
}
|
||||
ignoreArrayOrder: false
|
||||
ignoreExtraElements: false
|
||||
response:
|
||||
status: 200
|
||||
body: "{\"metrics\":{\"latencyMs\":659},\"output\":{\"message\":{\"content\":[{\"\
|
||||
text\":\"This is an LLM (\"}],\"role\":\"assistant\"}},\"stopReason\":\"max_tokens\"\
|
||||
,\"usage\":{\"inputTokens\":8,\"outputTokens\":10,\"totalTokens\":18}}"
|
||||
headers:
|
||||
Date: "Fri, 14 Feb 2025 06:31:20 GMT"
|
||||
Content-Type: application/json
|
||||
x-amzn-RequestId: ea2aad1f-b9a2-4c38-ac72-fb0215291d6a
|
||||
uuid: d50f134a-048f-41e2-ad78-f702880b78d6
|
||||
persistent: true
|
||||
insertionIndex: 2
|
||||
Loading…
Reference in New Issue