diff --git a/instrumentation/aws-sdk/README.md b/instrumentation/aws-sdk/README.md index 9d2bffb8b7..1b4d677d3e 100644 --- a/instrumentation/aws-sdk/README.md +++ b/instrumentation/aws-sdk/README.md @@ -27,6 +27,22 @@ The aws-sdk instrumentation is an SPI-based implementation that extends the upst - Finds ADOT’s **AdotAwsSdkInstrumentationModule** - Registers **AdotAwsSdkTracingExecutionInterceptor** (order > 0) +#### _Note on Attribute Collection:_ +AWS SDK v1.11 and v2.2 handle attribute collection differently: + +**V1.11:** +- Maintains a separate AttributesBuilder during request/response lifecycle +- Collects ADOT-specific attributes alongside upstream processing without interference +- Injects collected attributes into span at the end of the request and response lifecycle hooks + + +**V2.2:** +- FieldMapper directly modifies spans during request/response processing +- Attributes are added to spans immediately when discovered +- Direct integration with span lifecycle + +This architectural difference exists due to upstream AWS SDK injecting attributes into spans differently for v1.11 and v2.2 + ### AWS SDK v1 Instrumentation Summary The AdotAwsSdkInstrumentationModule uses the instrumentation (specified in AdotAwsClientInstrumentation) to register the AdotAwsSdkTracingRequestHandler through `typeInstrumentations`. @@ -61,6 +77,28 @@ _**Important Notes:**_ - The upstream interceptor closes the span in [afterResponse](https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L116) and/or [afterError](https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L131). These hooks are inaccessible for span modification. `afterAttempt` is our final hook point, giving us access to both the fully processed response and active span. +**High-Level Sequence Diagram:** +![img.png](sequence-diagram-1.11.png) + +_Class Functionalities:_ +- `AdotAwsSdkTracingRequestHandler` + - Hooks into AWS SDK request/response lifecycle + - Adds ADOT-specific attributes to spans extracted by AwsSdkExperimentalAttributesExtractor +- `AwsSdkExperimentalAttributesExtractor` + - Extracts attributes from AWS requests/responses and enriches spans + - Uses RequestAccess to get field values + - Special handling for Bedrock services +- `RequestAccess` + - Provides access to AWS SDK object fields + - Caches method handles for performance + - Uses BedrockJsonParser for parsing LLM payloads +- `BedrockJsonParser` + - Custom JSON parser for Bedrock payloads + - Handles different LLM model formats +- `AwsBedrockResourceType` + - Maps Bedrock class names to resource types + - Provides attribute keys and accessors for each type + ### AWS SDK v2 Instrumentation Summary **AdotAwsSdkInstrumentationModule** @@ -87,7 +125,6 @@ _**Important Notes:**_ `modifyResponse` is our final hook point, giving us access to both the fully processed response and active span. **High-Level Sequence Diagram:** - ![img.png](sequence-diagram-2.2.png) _Class Functionalities:_ @@ -111,4 +148,16 @@ _Class Functionalities:_ - Uses reflection to access internal SDK classes - Caches method handles for performance - `BedrockJasonParser` - - Parses and extracts specific attributes from Bedrock LLM responses for GenAI telemetry \ No newline at end of file + - Parses and extracts specific attributes from Bedrock LLM responses for GenAI telemetry + +### Commands for Running Groovy Tests + +To run the BedrockJsonParserTest for aws-sdk v1.11: +```` +./gradlew :instrumentation:aws-sdk:test --tests "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParserTest" +```` + +To run the BedrockJsonParserTest for aws-sdk v2.2: +```` +./gradlew :instrumentation:aws-sdk:test --tests "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v2_2.BedrockJsonParserTest" +```` \ No newline at end of file diff --git a/instrumentation/aws-sdk/build.gradle.kts b/instrumentation/aws-sdk/build.gradle.kts index e2cd769150..5863df2a10 100644 --- a/instrumentation/aws-sdk/build.gradle.kts +++ b/instrumentation/aws-sdk/build.gradle.kts @@ -22,6 +22,7 @@ plugins { base.archivesBaseName = "aws-instrumentation-aws-sdk" dependencies { + compileOnly("com.google.code.findbugs:jsr305:3.0.2") compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api") compileOnly("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api") compileOnly("com.amazonaws:aws-java-sdk-core:1.11.0") @@ -37,4 +38,6 @@ dependencies { testImplementation("com.amazonaws:aws-java-sdk-core:1.11.0") testImplementation("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api") testImplementation("org.mockito:mockito-core:5.14.2") + testImplementation("com.google.guava:guava") + testImplementation("io.opentelemetry.javaagent:opentelemetry-testing-common") } diff --git a/instrumentation/aws-sdk/sequence-diagram-1.11.png b/instrumentation/aws-sdk/sequence-diagram-1.11.png new file mode 100644 index 0000000000..891b07f384 Binary files /dev/null and b/instrumentation/aws-sdk/sequence-diagram-1.11.png differ diff --git a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkInstrumentationModule.java b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkInstrumentationModule.java index 677ea4612f..628997dfa4 100644 --- a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkInstrumentationModule.java +++ b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkInstrumentationModule.java @@ -16,6 +16,7 @@ package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11; import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed; +import static net.bytebuddy.matcher.ElementMatchers.*; import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule; import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; @@ -46,7 +47,17 @@ public int order() { @Override public List getAdditionalHelperClassNames() { return Arrays.asList( - "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AdotAwsSdkTracingRequestHandler"); + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AdotAwsSdkTracingRequestHandler", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsSdkExperimentalAttributesExtractor", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsBedrockResourceType", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsBedrockResourceType$AwsBedrockResourceTypeMap", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser$JsonParser", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser$LlmJson", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.BedrockJsonParser$JsonPathResolver", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.RequestAccess", + "software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.RequestAccess$1"); } @Override diff --git a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkTracingRequestHandler.java b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkTracingRequestHandler.java index 4ead53f5b1..edd0d14564 100644 --- a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkTracingRequestHandler.java +++ b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AdotAwsSdkTracingRequestHandler.java @@ -15,11 +15,37 @@ package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11; +import static io.opentelemetry.instrumentation.api.internal.ConfigPropertiesUtil.getBoolean; + import com.amazonaws.Request; +import com.amazonaws.Response; import com.amazonaws.handlers.HandlerAfterAttemptContext; import com.amazonaws.handlers.RequestHandler2; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.common.AttributesBuilder; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.context.Context; /** + * This handler extends the AWS SDK v1.11 request handling chain to add ADOT-specific span + * attributes. It operates at two key points in the request lifecycle: + * + *

1. Request Phase (beforeRequest): + * + *

+ * + *

2. Response/Error Phase (afterAttempt): + * + *

+ * * Based on OpenTelemetry Java Instrumentation's AWS SDK v1.11 TracingRequestHandler * (release/v2.11.x). Adapts the base instrumentation pattern to add ADOT-specific functionality. * @@ -27,8 +53,13 @@ * href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/release/v2.11.x/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java">... */ public class AdotAwsSdkTracingRequestHandler extends RequestHandler2 { + private final AwsSdkExperimentalAttributesExtractor experimentalAttributesExtractor; + private final boolean captureExperimentalSpanAttributes = + getBoolean("otel.instrumentation.aws-sdk.experimental-span-attributes", true); - public AdotAwsSdkTracingRequestHandler() {} + public AdotAwsSdkTracingRequestHandler() { + this.experimentalAttributesExtractor = new AwsSdkExperimentalAttributesExtractor(); + } /** * This is the latest point we can obtain the Sdk Request after it is modified by the upstream @@ -38,7 +69,20 @@ public AdotAwsSdkTracingRequestHandler() {} * href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L58">reference */ @Override - public void beforeRequest(Request request) {} + public void beforeRequest(Request request) { + Span currentSpan = Span.current(); + + if (captureExperimentalSpanAttributes + && currentSpan != null + && currentSpan.getSpanContext().isValid()) { + AttributesBuilder attributes = Attributes.builder(); + experimentalAttributesExtractor.onStart(attributes, Context.current(), request); + + attributes + .build() + .forEach((key, value) -> currentSpan.setAttribute(key.getKey(), value.toString())); + } + } /** * This is the latest point to access the sdk response before the span closes in the upstream @@ -52,5 +96,23 @@ public void beforeRequest(Request request) {} * href="https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java#L131">reference */ @Override - public void afterAttempt(HandlerAfterAttemptContext context) {} + public void afterAttempt(HandlerAfterAttemptContext context) { + Span currentSpan = Span.current(); + + if (captureExperimentalSpanAttributes + && currentSpan != null + && currentSpan.getSpanContext().isValid()) { + Request request = context.getRequest(); + Response response = context.getResponse(); + Exception exception = context.getException(); + + AttributesBuilder attributes = Attributes.builder(); + experimentalAttributesExtractor.onEnd( + attributes, Context.current(), request, response, exception); + + attributes + .build() + .forEach((key, value) -> currentSpan.setAttribute(key.getKey(), value.toString())); + } + } } diff --git a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsBedrockResourceType.java b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsBedrockResourceType.java new file mode 100644 index 0000000000..d006bc365d --- /dev/null +++ b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsBedrockResourceType.java @@ -0,0 +1,143 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11; + +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_AGENT_ID; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_DATA_SOURCE_ID; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGE_BASE_ID; + +import io.opentelemetry.api.common.AttributeKey; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.Function; + +enum AwsBedrockResourceType { + AGENT_TYPE(AWS_AGENT_ID, RequestAccess::getAgentId), + DATA_SOURCE_TYPE(AWS_DATA_SOURCE_ID, RequestAccess::getDataSourceId), + KNOWLEDGE_BASE_TYPE(AWS_KNOWLEDGE_BASE_ID, RequestAccess::getKnowledgeBaseId); + + @SuppressWarnings("ImmutableEnumChecker") + private final AttributeKey keyAttribute; + + @SuppressWarnings("ImmutableEnumChecker") + private final Function attributeValueAccessor; + + AwsBedrockResourceType( + AttributeKey keyAttribute, Function attributeValueAccessor) { + this.keyAttribute = keyAttribute; + this.attributeValueAccessor = attributeValueAccessor; + } + + public AttributeKey getKeyAttribute() { + return keyAttribute; + } + + public Function getAttributeValueAccessor() { + return attributeValueAccessor; + } + + public static AwsBedrockResourceType getRequestType(String requestClass) { + return AwsBedrockResourceTypeMap.BEDROCK_REQUEST_MAP.get(requestClass); + } + + public static AwsBedrockResourceType getResponseType(String responseClass) { + return AwsBedrockResourceTypeMap.BEDROCK_RESPONSE_MAP.get(responseClass); + } + + private static class AwsBedrockResourceTypeMap { + private static final Map BEDROCK_REQUEST_MAP = new HashMap<>(); + private static final Map BEDROCK_RESPONSE_MAP = new HashMap<>(); + + // Bedrock request/response mapping + // We only support operations that are related to the resource and where the context contains + // the AgentID/DataSourceID/KnowledgeBaseID. + // AgentID + private static final List agentRequestClasses = + Arrays.asList( + "CreateAgentActionGroupRequest", + "CreateAgentAliasRequest", + "DeleteAgentActionGroupRequest", + "DeleteAgentAliasRequest", + "DeleteAgentRequest", + "DeleteAgentVersionRequest", + "GetAgentActionGroupRequest", + "GetAgentAliasRequest", + "GetAgentRequest", + "GetAgentVersionRequest", + "ListAgentActionGroupsRequest", + "ListAgentAliasesRequest", + "ListAgentKnowledgeBasesRequest", + "ListAgentVersionsRequest", + "PrepareAgentRequest", + "UpdateAgentActionGroupRequest", + "UpdateAgentAliasRequest", + "UpdateAgentRequest"); + private static final List agentResponseClasses = + Arrays.asList( + "DeleteAgentAliasResult", + "DeleteAgentResult", + "DeleteAgentVersionResult", + "PrepareAgentResult"); + // DataSourceID + private static final List dataSourceRequestClasses = + Arrays.asList("DeleteDataSourceRequest", "GetDataSourceRequest", "UpdateDataSourceRequest"); + private static final List dataSourceResponseClasses = + Arrays.asList("DeleteDataSourceResult"); + // KnowledgeBaseID + private static final List knowledgeBaseRequestClasses = + Arrays.asList( + "AssociateAgentKnowledgeBaseRequest", + "CreateDataSourceRequest", + "DeleteKnowledgeBaseRequest", + "DisassociateAgentKnowledgeBaseRequest", + "GetAgentKnowledgeBaseRequest", + "GetKnowledgeBaseRequest", + "ListDataSourcesRequest", + "UpdateAgentKnowledgeBaseRequest"); + private static final List knowledgeBaseResponseClasses = + Arrays.asList("DeleteKnowledgeBaseResult"); + + private AwsBedrockResourceTypeMap() {} + + static { + // Populate the BEDROCK_REQUEST_MAP + for (String agentRequestClass : agentRequestClasses) { + BEDROCK_REQUEST_MAP.put(agentRequestClass, AwsBedrockResourceType.AGENT_TYPE); + } + for (String dataSourceRequestClass : dataSourceRequestClasses) { + BEDROCK_REQUEST_MAP.put(dataSourceRequestClass, AwsBedrockResourceType.DATA_SOURCE_TYPE); + } + for (String knowledgeBaseRequestClass : knowledgeBaseRequestClasses) { + BEDROCK_REQUEST_MAP.put( + knowledgeBaseRequestClass, AwsBedrockResourceType.KNOWLEDGE_BASE_TYPE); + } + + // Populate the BEDROCK_RESPONSE_MAP + for (String agentResponseClass : agentResponseClasses) { + BEDROCK_REQUEST_MAP.put(agentResponseClass, AwsBedrockResourceType.AGENT_TYPE); + } + for (String dataSourceResponseClass : dataSourceResponseClasses) { + BEDROCK_REQUEST_MAP.put(dataSourceResponseClass, AwsBedrockResourceType.DATA_SOURCE_TYPE); + } + for (String knowledgeBaseResponseClass : knowledgeBaseResponseClasses) { + BEDROCK_REQUEST_MAP.put( + knowledgeBaseResponseClass, AwsBedrockResourceType.KNOWLEDGE_BASE_TYPE); + } + } + } +} diff --git a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsExperimentalAttributes.java b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsExperimentalAttributes.java new file mode 100644 index 0000000000..f1870caa1c --- /dev/null +++ b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsExperimentalAttributes.java @@ -0,0 +1,70 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11; + +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + * + * Modifications Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + */ + +import static io.opentelemetry.api.common.AttributeKey.stringKey; + +import io.opentelemetry.api.common.AttributeKey; + +final class AwsExperimentalAttributes { + + // 2025-07-22: Amazon addition + static final AttributeKey AWS_STREAM_ARN = stringKey("aws.stream.arn"); + static final AttributeKey AWS_TABLE_ARN = stringKey("aws.table.arn"); + static final AttributeKey AWS_AGENT_ID = stringKey("aws.bedrock.agent.id"); + static final AttributeKey AWS_KNOWLEDGE_BASE_ID = + stringKey("aws.bedrock.knowledge_base.id"); + static final AttributeKey AWS_DATA_SOURCE_ID = stringKey("aws.bedrock.data_source.id"); + static final AttributeKey AWS_GUARDRAIL_ID = stringKey("aws.bedrock.guardrail.id"); + static final AttributeKey AWS_GUARDRAIL_ARN = stringKey("aws.bedrock.guardrail.arn"); + // TODO: Merge in gen_ai attributes in opentelemetry-semconv-incubating once upgrade to v1.26.0 + static final AttributeKey AWS_BEDROCK_RUNTIME_MODEL_ID = + stringKey("gen_ai.request.model"); + static final AttributeKey AWS_BEDROCK_SYSTEM = stringKey("gen_ai.system"); + static final AttributeKey GEN_AI_REQUEST_MAX_TOKENS = + stringKey("gen_ai.request.max_tokens"); + static final AttributeKey GEN_AI_REQUEST_TEMPERATURE = + stringKey("gen_ai.request.temperature"); + static final AttributeKey GEN_AI_REQUEST_TOP_P = stringKey("gen_ai.request.top_p"); + static final AttributeKey GEN_AI_RESPONSE_FINISH_REASONS = + stringKey("gen_ai.response.finish_reasons"); + static final AttributeKey GEN_AI_USAGE_INPUT_TOKENS = + stringKey("gen_ai.usage.input_tokens"); + static final AttributeKey GEN_AI_USAGE_OUTPUT_TOKENS = + stringKey("gen_ai.usage.output_tokens"); + static final AttributeKey AWS_STATE_MACHINE_ARN = + stringKey("aws.stepfunctions.state_machine.arn"); + static final AttributeKey AWS_STEP_FUNCTIONS_ACTIVITY_ARN = + stringKey("aws.stepfunctions.activity.arn"); + static final AttributeKey AWS_SNS_TOPIC_ARN = stringKey("aws.sns.topic.arn"); + static final AttributeKey AWS_SECRET_ARN = stringKey("aws.secretsmanager.secret.arn"); + static final AttributeKey AWS_LAMBDA_NAME = stringKey("aws.lambda.function.name"); + static final AttributeKey AWS_LAMBDA_ARN = stringKey("aws.lambda.function.arn"); + static final AttributeKey AWS_LAMBDA_RESOURCE_ID = + stringKey("aws.lambda.resource_mapping.id"); + static final AttributeKey AWS_AUTH_ACCESS_KEY = stringKey("aws.auth.account.access_key"); + + // End of Amazon addition + + private AwsExperimentalAttributes() {} +} diff --git a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsSdkExperimentalAttributesExtractor.java b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsSdkExperimentalAttributesExtractor.java new file mode 100644 index 0000000000..5aa3d39d78 --- /dev/null +++ b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/AwsSdkExperimentalAttributesExtractor.java @@ -0,0 +1,243 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11; + +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + * + * Modifications Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + */ + +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_AGENT_ID; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_AUTH_ACCESS_KEY; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_BEDROCK_RUNTIME_MODEL_ID; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_BEDROCK_SYSTEM; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_GUARDRAIL_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_GUARDRAIL_ID; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGE_BASE_ID; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_LAMBDA_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_LAMBDA_NAME; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_LAMBDA_RESOURCE_ID; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_SECRET_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_SNS_TOPIC_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_STATE_MACHINE_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_STEP_FUNCTIONS_ACTIVITY_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_STREAM_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.AWS_TABLE_ARN; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.GEN_AI_REQUEST_MAX_TOKENS; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.GEN_AI_REQUEST_TEMPERATURE; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.GEN_AI_REQUEST_TOP_P; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.GEN_AI_RESPONSE_FINISH_REASONS; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.GEN_AI_USAGE_INPUT_TOKENS; +import static software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11.AwsExperimentalAttributes.GEN_AI_USAGE_OUTPUT_TOKENS; + +import com.amazonaws.Request; +import com.amazonaws.Response; +import com.amazonaws.auth.AWSCredentials; +import com.amazonaws.handlers.HandlerContextKey; +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.Objects; +import java.util.function.Function; +import javax.annotation.Nullable; + +class AwsSdkExperimentalAttributesExtractor + implements AttributesExtractor, Response> { + // 2025-07-22: Amazon addition + private static final String BEDROCK_SERVICE = "AmazonBedrock"; + private static final String BEDROCK_AGENT_SERVICE = "AWSBedrockAgent"; + private static final String BEDROCK_AGENT_RUNTIME_SERVICE = "AWSBedrockAgentRuntime"; + private static final String BEDROCK_RUNTIME_SERVICE = "AmazonBedrockRuntime"; + private static final HandlerContextKey AWS_CREDENTIALS = + new HandlerContextKey("AWSCredentials"); + + AwsSdkExperimentalAttributesExtractor() {} + + @Override + public void onStart(AttributesBuilder attributes, Context parentContext, Request request) { + + Object originalRequest = request.getOriginalRequest(); + String requestClassName = originalRequest.getClass().getSimpleName(); + + AWSCredentials credentials = request.getHandlerContext(AWS_CREDENTIALS); + if (credentials != null) { + String accessKeyId = credentials.getAWSAccessKeyId(); + if (accessKeyId != null) { + attributes.put(AWS_AUTH_ACCESS_KEY, accessKeyId); + } + } + + setAttribute(attributes, AWS_STREAM_ARN, originalRequest, RequestAccess::getStreamArn); + setAttribute( + attributes, AWS_STATE_MACHINE_ARN, originalRequest, RequestAccess::getStateMachineArn); + setAttribute( + attributes, + AWS_STEP_FUNCTIONS_ACTIVITY_ARN, + originalRequest, + RequestAccess::getStepFunctionsActivityArn); + setAttribute(attributes, AWS_SNS_TOPIC_ARN, originalRequest, RequestAccess::getSnsTopicArn); + setAttribute(attributes, AWS_SECRET_ARN, originalRequest, RequestAccess::getSecretArn); + setAttribute(attributes, AWS_LAMBDA_NAME, originalRequest, RequestAccess::getLambdaName); + setAttribute( + attributes, AWS_LAMBDA_RESOURCE_ID, originalRequest, RequestAccess::getLambdaResourceId); + // Get serviceName defined in the AWS Java SDK V1 Request class. + String serviceName = request.getServiceName(); + // Extract request attributes only for Bedrock services. + if (isBedrockService(serviceName)) { + bedrockOnStart(attributes, originalRequest, requestClassName, serviceName); + } + } + + @Override + public void onEnd( + AttributesBuilder attributes, + Context context, + Request request, + @Nullable Response response, + @Nullable Throwable error) { + if (response != null) { + Object awsResp = response.getAwsResponse(); + setAttribute(attributes, AWS_TABLE_ARN, awsResp, RequestAccess::getTableArn); + setAttribute(attributes, AWS_LAMBDA_ARN, awsResp, RequestAccess::getLambdaArn); + setAttribute(attributes, AWS_STATE_MACHINE_ARN, awsResp, RequestAccess::getStateMachineArn); + setAttribute( + attributes, + AWS_STEP_FUNCTIONS_ACTIVITY_ARN, + awsResp, + RequestAccess::getStepFunctionsActivityArn); + setAttribute(attributes, AWS_SNS_TOPIC_ARN, awsResp, RequestAccess::getSnsTopicArn); + setAttribute(attributes, AWS_SECRET_ARN, awsResp, RequestAccess::getSecretArn); + // Get serviceName defined in the AWS Java SDK V1 Request class. + String serviceName = request.getServiceName(); + // Extract response attributes for Bedrock services + if (awsResp != null && isBedrockService(serviceName)) { + bedrockOnEnd(attributes, awsResp, serviceName); + } + } + } + + private static void bedrockOnStart( + AttributesBuilder attributes, + Object originalRequest, + String requestClassName, + String serviceName) { + switch (serviceName) { + case BEDROCK_SERVICE: + setAttribute(attributes, AWS_GUARDRAIL_ID, originalRequest, RequestAccess::getGuardrailId); + break; + case BEDROCK_AGENT_SERVICE: + AwsBedrockResourceType resourceType = + AwsBedrockResourceType.getRequestType(requestClassName); + if (resourceType != null) { + setAttribute( + attributes, + resourceType.getKeyAttribute(), + originalRequest, + resourceType.getAttributeValueAccessor()); + } + break; + case BEDROCK_AGENT_RUNTIME_SERVICE: + setAttribute(attributes, AWS_AGENT_ID, originalRequest, RequestAccess::getAgentId); + setAttribute( + attributes, AWS_KNOWLEDGE_BASE_ID, originalRequest, RequestAccess::getKnowledgeBaseId); + break; + case BEDROCK_RUNTIME_SERVICE: + if (!Objects.equals(requestClassName, "InvokeModelRequest")) { + break; + } + attributes.put(AWS_BEDROCK_SYSTEM, "aws.bedrock"); + Function getter = RequestAccess::getModelId; + String modelId = getter.apply(originalRequest); + attributes.put(AWS_BEDROCK_RUNTIME_MODEL_ID, modelId); + + setAttribute( + attributes, GEN_AI_REQUEST_MAX_TOKENS, originalRequest, RequestAccess::getMaxTokens); + setAttribute( + attributes, GEN_AI_REQUEST_TEMPERATURE, originalRequest, RequestAccess::getTemperature); + setAttribute(attributes, GEN_AI_REQUEST_TOP_P, originalRequest, RequestAccess::getTopP); + setAttribute( + attributes, GEN_AI_USAGE_INPUT_TOKENS, originalRequest, RequestAccess::getInputTokens); + break; + default: + break; + } + } + + private static void bedrockOnEnd( + AttributesBuilder attributes, Object awsResp, String serviceName) { + switch (serviceName) { + case BEDROCK_SERVICE: + setAttribute(attributes, AWS_GUARDRAIL_ID, awsResp, RequestAccess::getGuardrailId); + setAttribute(attributes, AWS_GUARDRAIL_ARN, awsResp, RequestAccess::getGuardrailArn); + break; + case BEDROCK_AGENT_SERVICE: + String responseClassName = awsResp.getClass().getSimpleName(); + AwsBedrockResourceType resourceType = + AwsBedrockResourceType.getResponseType(responseClassName); + if (resourceType != null) { + setAttribute( + attributes, + resourceType.getKeyAttribute(), + awsResp, + resourceType.getAttributeValueAccessor()); + } + break; + case BEDROCK_AGENT_RUNTIME_SERVICE: + setAttribute(attributes, AWS_AGENT_ID, awsResp, RequestAccess::getAgentId); + setAttribute(attributes, AWS_KNOWLEDGE_BASE_ID, awsResp, RequestAccess::getKnowledgeBaseId); + break; + case BEDROCK_RUNTIME_SERVICE: + if (!Objects.equals(awsResp.getClass().getSimpleName(), "InvokeModelResult")) { + break; + } + + setAttribute(attributes, GEN_AI_USAGE_INPUT_TOKENS, awsResp, RequestAccess::getInputTokens); + setAttribute( + attributes, GEN_AI_USAGE_OUTPUT_TOKENS, awsResp, RequestAccess::getOutputTokens); + setAttribute( + attributes, GEN_AI_RESPONSE_FINISH_REASONS, awsResp, RequestAccess::getFinishReasons); + break; + default: + break; + } + } + + private static boolean isBedrockService(String serviceName) { + // Check if the serviceName belongs to Bedrock Services defined in AWS Java SDK V1. + // For example AmazonBedrock + return serviceName.equals(BEDROCK_SERVICE) + || serviceName.equals(BEDROCK_AGENT_SERVICE) + || serviceName.equals(BEDROCK_AGENT_RUNTIME_SERVICE) + || serviceName.equals(BEDROCK_RUNTIME_SERVICE); + } + + // End of Amazon addition + + private static void setAttribute( + AttributesBuilder attributes, + AttributeKey key, + Object request, + Function getter) { + String value = getter.apply(request); + if (value != null) { + attributes.put(key, value); + } + } +} diff --git a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/BedrockJsonParser.java b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/BedrockJsonParser.java new file mode 100644 index 0000000000..60297e4948 --- /dev/null +++ b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/BedrockJsonParser.java @@ -0,0 +1,277 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class BedrockJsonParser { + + // Prevent instantiation + private BedrockJsonParser() { + throw new UnsupportedOperationException("Utility class"); + } + + public static LlmJson parse(String jsonString) { + JsonParser parser = new JsonParser(jsonString); + Map jsonBody = parser.parse(); + return new LlmJson(jsonBody); + } + + static class JsonParser { + private final String json; + private int position; + + public JsonParser(String json) { + this.json = json.trim(); + this.position = 0; + } + + private void skipWhitespace() { + while (position < json.length() && Character.isWhitespace(json.charAt(position))) { + position++; + } + } + + private char currentChar() { + return json.charAt(position); + } + + private static boolean isHexDigit(char c) { + return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'); + } + + private void expect(char c) { + skipWhitespace(); + if (currentChar() != c) { + throw new IllegalArgumentException( + "Expected '" + c + "' but found '" + currentChar() + "'"); + } + position++; + } + + private String readString() { + skipWhitespace(); + expect('"'); // Ensure the string starts with a quote + StringBuilder result = new StringBuilder(); + while (currentChar() != '"') { + // Handle escape sequences + if (currentChar() == '\\') { + position++; // Move past the backslash + if (position >= json.length()) { + throw new IllegalArgumentException("Unexpected end of input in string escape sequence"); + } + char escapeChar = currentChar(); + switch (escapeChar) { + case '"': + case '\\': + case '/': + result.append(escapeChar); + break; + case 'b': + result.append('\b'); + break; + case 'f': + result.append('\f'); + break; + case 'n': + result.append('\n'); + break; + case 'r': + result.append('\r'); + break; + case 't': + result.append('\t'); + break; + case 'u': // Unicode escape sequence + if (position + 4 >= json.length()) { + throw new IllegalArgumentException("Invalid unicode escape sequence in string"); + } + char[] hexChars = new char[4]; + for (int i = 0; i < 4; i++) { + position++; // Move to the next character + char hexChar = json.charAt(position); + if (!isHexDigit(hexChar)) { + throw new IllegalArgumentException( + "Invalid hexadecimal digit in unicode escape sequence"); + } + hexChars[i] = hexChar; + } + int unicodeValue = Integer.parseInt(new String(hexChars), 16); + result.append((char) unicodeValue); + break; + default: + throw new IllegalArgumentException("Invalid escape character: \\" + escapeChar); + } + position++; + } else { + result.append(currentChar()); + position++; + } + } + position++; // Skip closing quote + return result.toString(); + } + + private Object readValue() { + skipWhitespace(); + char c = currentChar(); + + if (c == '"') { + return readString(); + } else if (Character.isDigit(c)) { + return readScopedNumber(); + } else if (c == '{') { + return readObject(); // JSON Objects + } else if (c == '[') { + return readArray(); // JSON Arrays + } else if (json.startsWith("true", position)) { + position += 4; + return true; + } else if (json.startsWith("false", position)) { + position += 5; + return false; + } else if (json.startsWith("null", position)) { + position += 4; + return null; // JSON null + } else { + throw new IllegalArgumentException("Unexpected character: " + c); + } + } + + private Number readScopedNumber() { + int start = position; + + // Consume digits and the optional decimal point + while (position < json.length() + && (Character.isDigit(json.charAt(position)) || json.charAt(position) == '.')) { + position++; + } + + String number = json.substring(start, position); + + if (number.contains(".")) { + double value = Double.parseDouble(number); + if (value < 0.0 || value > 1.0) { + throw new IllegalArgumentException( + "Value out of bounds for Bedrock Floating Point Attribute: " + number); + } + return value; + } else { + return Integer.parseInt(number); + } + } + + private Map readObject() { + Map map = new HashMap<>(); + expect('{'); + skipWhitespace(); + while (currentChar() != '}') { + String key = readString(); + expect(':'); + Object value = readValue(); + map.put(key, value); + skipWhitespace(); + if (currentChar() == ',') { + position++; + } + } + position++; // Skip closing brace + return map; + } + + private List readArray() { + List list = new ArrayList<>(); + expect('['); + skipWhitespace(); + while (currentChar() != ']') { + list.add(readValue()); + skipWhitespace(); + if (currentChar() == ',') { + position++; + } + } + position++; + return list; + } + + public Map parse() { + return readObject(); + } + } + + // Resolves paths in a JSON structure + static class JsonPathResolver { + + // Private constructor to prevent instantiation + private JsonPathResolver() { + throw new UnsupportedOperationException("Utility class"); + } + + public static Object resolvePath(LlmJson llmJson, String... paths) { + for (String path : paths) { + Object value = resolvePath(llmJson.getJsonBody(), path); + if (value != null) { + return value; + } + } + return null; + } + + private static Object resolvePath(Map json, String path) { + String[] keys = path.split("/"); + Object current = json; + + for (String key : keys) { + if (key.isEmpty()) { + continue; + } + + if (current instanceof Map) { + current = ((Map) current).get(key); + } else if (current instanceof List) { + try { + int index = Integer.parseInt(key); + current = ((List) current).get(index); + } catch (NumberFormatException | IndexOutOfBoundsException e) { + return null; + } + } else { + return null; + } + + if (current == null) { + return null; + } + } + return current; + } + } + + public static class LlmJson { + private final Map jsonBody; + + public LlmJson(Map jsonBody) { + this.jsonBody = jsonBody; + } + + public Map getJsonBody() { + return jsonBody; + } + } +} diff --git a/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/RequestAccess.java b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/RequestAccess.java new file mode 100644 index 0000000000..7232c7d3c8 --- /dev/null +++ b/instrumentation/aws-sdk/src/main/java/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/RequestAccess.java @@ -0,0 +1,508 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11; + +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + * + * Modifications Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + */ + +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.lang.reflect.Method; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Objects; +import java.util.stream.Stream; +import javax.annotation.Nullable; + +final class RequestAccess { + + private static final ClassValue REQUEST_ACCESSORS = + new ClassValue() { + @Override + protected RequestAccess computeValue(Class type) { + return new RequestAccess(type); + } + }; + + // 2025-07-22: Amazon addition + @Nullable + private static BedrockJsonParser.LlmJson parseTargetBody(ByteBuffer buffer) { + try { + byte[] bytes; + // Create duplicate to avoid mutating the original buffer position + ByteBuffer duplicate = buffer.duplicate(); + if (buffer.hasArray()) { + bytes = + Arrays.copyOfRange( + duplicate.array(), + duplicate.arrayOffset(), + duplicate.arrayOffset() + duplicate.remaining()); + } else { + bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + } + String jsonString = new String(bytes, StandardCharsets.UTF_8); // Convert to String + return BedrockJsonParser.parse(jsonString); + } catch (RuntimeException e) { + return null; + } + } + + @Nullable + private static BedrockJsonParser.LlmJson getJsonBody(Object target) { + if (target == null) { + return null; + } + + RequestAccess access = REQUEST_ACCESSORS.get(target.getClass()); + ByteBuffer bodyBuffer = invokeOrNullGeneric(access.getBody, target, ByteBuffer.class); + if (bodyBuffer == null) { + return null; + } + + return parseTargetBody(bodyBuffer); + } + + @Nullable + private static String findFirstMatchingPath(BedrockJsonParser.LlmJson jsonBody, String... paths) { + if (jsonBody == null) { + return null; + } + + return Stream.of(paths) + .map(path -> BedrockJsonParser.JsonPathResolver.resolvePath(jsonBody, path)) + .filter(Objects::nonNull) + .map(Object::toString) + .findFirst() + .orElse(null); + } + + @Nullable + private static String approximateTokenCount( + BedrockJsonParser.LlmJson jsonBody, String... textPaths) { + if (jsonBody == null) { + return null; + } + + return Stream.of(textPaths) + .map(path -> BedrockJsonParser.JsonPathResolver.resolvePath(jsonBody, path)) + .filter(value -> value instanceof String) + .map(value -> Integer.toString((int) Math.ceil(((String) value).length() / 6.0))) + .findFirst() + .orElse(null); + } + + // Model -> Path Mapping: + // Amazon Nova -> "/inferenceConfig/max_new_tokens" + // Amazon Titan -> "/textGenerationConfig/maxTokenCount" + // Anthropic Claude -> "/max_tokens" + // Cohere Command -> "/max_tokens" + // Cohere Command R -> "/max_tokens" + // AI21 Jamba -> "/max_tokens" + // Meta Llama -> "/max_gen_len" + // Mistral AI -> "/max_tokens" + @Nullable + static String getMaxTokens(Object target) { + BedrockJsonParser.LlmJson jsonBody = getJsonBody(target); + return findFirstMatchingPath( + jsonBody, + "/max_tokens", + "/max_gen_len", + "/textGenerationConfig/maxTokenCount", + "/inferenceConfig/max_new_tokens"); + } + + // Model -> Path Mapping: + // Amazon Nova -> "/inferenceConfig/temperature" + // Amazon Titan -> "/textGenerationConfig/temperature" + // Anthropic Claude -> "/temperature" + // Cohere Command -> "/temperature" + // Cohere Command R -> "/temperature" + // AI21 Jamba -> "/temperature" + // Meta Llama -> "/temperature" + // Mistral AI -> "/temperature" + @Nullable + static String getTemperature(Object target) { + BedrockJsonParser.LlmJson jsonBody = getJsonBody(target); + return findFirstMatchingPath( + jsonBody, + "/temperature", + "/textGenerationConfig/temperature", + "inferenceConfig/temperature"); + } + + // Model -> Path Mapping: + // Amazon Nova -> "/inferenceConfig/top_p" + // Amazon Titan -> "/textGenerationConfig/topP" + // Anthropic Claude -> "/top_p" + // Cohere Command -> "/p" + // Cohere Command R -> "/p" + // AI21 Jamba -> "/top_p" + // Meta Llama -> "/top_p" + // Mistral AI -> "/top_p" + @Nullable + static String getTopP(Object target) { + BedrockJsonParser.LlmJson jsonBody = getJsonBody(target); + return findFirstMatchingPath( + jsonBody, "/top_p", "/p", "/textGenerationConfig/topP", "/inferenceConfig/top_p"); + } + + // Model -> Path Mapping: + // Amazon Nova -> "/usage/inputTokens" + // Amazon Titan -> "/inputTextTokenCount" + // Anthropic Claude -> "/usage/input_tokens" + // Cohere Command -> "/prompt" + // Cohere Command R -> "/message" + // AI21 Jamba -> "/usage/prompt_tokens" + // Meta Llama -> "/prompt_token_count" + // Mistral AI -> "/prompt" + @Nullable + static String getInputTokens(Object target) { + BedrockJsonParser.LlmJson jsonBody = getJsonBody(target); + if (jsonBody == null) { + return null; + } + + // Try direct token counts first + String directCount = + findFirstMatchingPath( + jsonBody, + "/inputTextTokenCount", + "/prompt_token_count", + "/usage/input_tokens", + "/usage/prompt_tokens", + "/usage/inputTokens"); + + if (directCount != null && !directCount.equals("null")) { + return directCount; + } + + // Fall back to token approximation + return approximateTokenCount(jsonBody, "/prompt", "/message"); + } + + // Model -> Path Mapping: + // Amazon Nova -> "/usage/outputTokens" + // Amazon Titan -> "/results/0/tokenCount" + // Anthropic Claude -> "/usage/output_tokens" + // Cohere Command -> "/generations/0/text" + // Cohere Command R -> "/text" + // AI21 Jamba -> "/usage/completion_tokens" + // Meta Llama -> "/generation_token_count" + // Mistral AI -> "/outputs/0/text" + @Nullable + static String getOutputTokens(Object target) { + BedrockJsonParser.LlmJson jsonBody = getJsonBody(target); + if (jsonBody == null) { + return null; + } + + // Try direct token counts first + String directCount = + findFirstMatchingPath( + jsonBody, + "/generation_token_count", + "/results/0/tokenCount", + "/usage/output_tokens", + "/usage/completion_tokens", + "/usage/outputTokens"); + + if (directCount != null && !directCount.equals("null")) { + return directCount; + } + + // Fall back to token approximation + return approximateTokenCount(jsonBody, "/text", "/outputs/0/text"); + } + + // Model -> Path Mapping: + // Amazon Nova -> "/stopReason" + // Amazon Titan -> "/results/0/completionReason" + // Anthropic Claude -> "/stop_reason" + // Cohere Command -> "/generations/0/finish_reason" + // Cohere Command R -> "/finish_reason" + // AI21 Jamba -> "/choices/0/finish_reason" + // Meta Llama -> "/stop_reason" + // Mistral AI -> "/outputs/0/stop_reason" + @Nullable + static String getFinishReasons(Object target) { + BedrockJsonParser.LlmJson jsonBody = getJsonBody(target); + String finishReason = + findFirstMatchingPath( + jsonBody, + "/stopReason", + "/finish_reason", + "/stop_reason", + "/results/0/completionReason", + "/generations/0/finish_reason", + "/choices/0/finish_reason", + "/outputs/0/stop_reason"); + + return finishReason != null ? "[" + finishReason + "]" : null; + } + + @Nullable + static String getLambdaName(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getLambdaName, request); + } + + @Nullable + static String getLambdaArn(Object request) { + if (request == null) { + return null; + } + return findNestedAccessorOrNull(request, "getConfiguration", "getFunctionArn"); + } + + @Nullable + static String getLambdaResourceId(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getLambdaResourceId, request); + } + + @Nullable + static String getSecretArn(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getSecretArn, request); + } + + @Nullable + static String getSnsTopicArn(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getSnsTopicArn, request); + } + + @Nullable + static String getStepFunctionsActivityArn(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getStepFunctionsActivityArn, request); + } + + @Nullable + static String getStateMachineArn(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getStateMachineArn, request); + } + + @Nullable + static String getTableArn(Object request) { + if (request == null) { + return null; + } + return findNestedAccessorOrNull(request, "getTable", "getTableArn"); + } + + @Nullable + static String getStreamArn(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getStreamArn, request); + } + + @Nullable + static String getAgentId(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getAgentId, request); + } + + @Nullable + static String getKnowledgeBaseId(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getKnowledgeBaseId, request); + } + + @Nullable + static String getDataSourceId(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getDataSourceId, request); + } + + @Nullable + static String getGuardrailId(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getGuardrailId, request); + } + + @Nullable + static String getGuardrailArn(Object request) { + if (request == null) { + return null; + } + return findNestedAccessorOrNull(request, "getGuardrailArn"); + } + + @Nullable + static String getModelId(Object request) { + if (request == null) { + return null; + } + RequestAccess access = REQUEST_ACCESSORS.get(request.getClass()); + return invokeOrNull(access.getModelId, request); + } + + // End of Amazon addition + + @Nullable + private static String invokeOrNull(@Nullable MethodHandle method, Object obj) { + if (method == null) { + return null; + } + try { + return (String) method.invoke(obj); + } catch (Throwable t) { + return null; + } + } + + // 2025-07-22: Amazon addition + @Nullable + private static T invokeOrNullGeneric( + @Nullable MethodHandle method, Object obj, Class returnType) { + if (method == null) { + return null; + } + try { + return returnType.cast(method.invoke(obj)); + } catch (Throwable e) { + return null; + } + } + + @Nullable private final MethodHandle getStreamArn; + @Nullable private final MethodHandle getAgentId; + @Nullable private final MethodHandle getKnowledgeBaseId; + @Nullable private final MethodHandle getDataSourceId; + @Nullable private final MethodHandle getGuardrailId; + @Nullable private final MethodHandle getModelId; + @Nullable private final MethodHandle getBody; + @Nullable private final MethodHandle getStateMachineArn; + @Nullable private final MethodHandle getStepFunctionsActivityArn; + @Nullable private final MethodHandle getSnsTopicArn; + @Nullable private final MethodHandle getSecretArn; + @Nullable private final MethodHandle getLambdaName; + @Nullable private final MethodHandle getLambdaResourceId; + + private RequestAccess(Class clz) { + getStreamArn = findAccessorOrNull(clz, "getStreamARN", String.class); + getAgentId = findAccessorOrNull(clz, "getAgentId", String.class); + getKnowledgeBaseId = findAccessorOrNull(clz, "getKnowledgeBaseId", String.class); + getDataSourceId = findAccessorOrNull(clz, "getDataSourceId", String.class); + getGuardrailId = findAccessorOrNull(clz, "getGuardrailId", String.class); + getModelId = findAccessorOrNull(clz, "getModelId", String.class); + getBody = findAccessorOrNull(clz, "getBody", ByteBuffer.class); + getStateMachineArn = findAccessorOrNull(clz, "getStateMachineArn", String.class); + getStepFunctionsActivityArn = findAccessorOrNull(clz, "getActivityArn", String.class); + getSnsTopicArn = findAccessorOrNull(clz, "getTopicArn", String.class); + getSecretArn = findAccessorOrNull(clz, "getARN", String.class); + getLambdaName = findAccessorOrNull(clz, "getFunctionName", String.class); + getLambdaResourceId = findAccessorOrNull(clz, "getUUID", String.class); + } + + /** + * Uses Java reflection to find a getter method on a class and create a MethodHandle for it. + * + * @param clz The class to search for the method + * @param methodName The name of the getter method (e.g., "getStreamARN") + * @param returnType The expected return type of the method + * @return A MethodHandle for the method, or null if not found + *

Example: For class PutRecordRequest with method "getStreamARN": + * findAccessorOrNull(PutRecordRequest.class, "getStreamARN", String.class) Creates a method + * handle that can invoke getStreamARN() on PutRecordRequest instances + */ + @Nullable + private static MethodHandle findAccessorOrNull( + Class clz, String methodName, Class returnType) { + try { + // Uses MethodHandles.publicLookup() to get access to public methods + // findVirtual finds an instance method with the given name and type + // methodType creates a method type with no parameters and the specified return type + return MethodHandles.publicLookup() + .findVirtual(clz, methodName, MethodType.methodType(returnType)); + } catch (Throwable t) { + // Returns null if method doesn't exist or can't be accessed + return null; + } + } + + /** + * Uses reflection to navigate through nested method calls and extract a String value. Unlike + * using method handles, this supports chained method calls where each method might return a + * different type of object. + * + * @param obj The initial object to start method calls from + * @param methodNames Variable list of method names to call in sequence + * @return The final String value, or null if any method in the chain fails or returns null + *

Example: For Lambda ARN: findNestedAccessorOrNull(request, "getConfiguration", + * "getFunctionArn") - First calls request.getConfiguration() to get a Configuration object, + * then calls configuration.getFunctionArn() to get the ARN string + */ + @Nullable + private static String findNestedAccessorOrNull(Object obj, String... methodNames) { + Object current = obj; + for (String methodName : methodNames) { + if (current == null) { + return null; + } + try { + Method method = current.getClass().getMethod(methodName); + current = method.invoke(current); + } catch (Exception e) { + return null; + } + } + return (current instanceof String) ? (String) current : null; + } + // End of Amazon addition +} diff --git a/instrumentation/aws-sdk/src/test/groovy/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/BedrockJsonParserTest.groovy b/instrumentation/aws-sdk/src/test/groovy/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/BedrockJsonParserTest.groovy new file mode 100644 index 0000000000..d415e40e2f --- /dev/null +++ b/instrumentation/aws-sdk/src/test/groovy/software/amazon/opentelemetry/javaagent/instrumentation/awssdk_v1_11/BedrockJsonParserTest.groovy @@ -0,0 +1,117 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.instrumentation.awssdk_v1_11 + +import spock.lang.Specification + +class BedrockJsonParserTest extends Specification { + def "should parse simple JSON object"() { + given: + String json = '{"key":"value"}' + + when: + def parsedJson = BedrockJsonParser.parse(json) + + then: + parsedJson.getJsonBody() == [key: "value"] + } + + def "should parse nested JSON object"() { + given: + String json = '{"parent":{"child":"value"}}' + + when: + def parsedJson = BedrockJsonParser.parse(json) + + then: + def parent = parsedJson.getJsonBody().get("parent") + parent instanceof Map + parent["child"] == "value" + } + + def "should parse JSON array"() { + given: + String json = '{"array":[1, "two", 1.0]}' + + when: + def parsedJson = BedrockJsonParser.parse(json) + + then: + def array = parsedJson.getJsonBody().get("array") + array instanceof List + array == [1, "two", 1.0] + } + + def "should parse escape sequences"() { + given: + String json = '{"escaped":"Line1\\nLine2\\tTabbed\\\"Quoted\\\"\\bBackspace\\fFormfeed\\rCarriageReturn\\\\Backslash\\/Slash\\u0041"}' + + when: + def parsedJson = BedrockJsonParser.parse(json) + + then: + parsedJson.getJsonBody().get("escaped") == + "Line1\nLine2\tTabbed\"Quoted\"\bBackspace\fFormfeed\rCarriageReturn\\Backslash/SlashA" + } + + def "should throw exception for malformed JSON"() { + given: + String malformedJson = '{"key":value}' + + when: + BedrockJsonParser.parse(malformedJson) + + then: + def ex = thrown(IllegalArgumentException) + ex.message.contains("Unexpected character") + } + + def "should resolve path in JSON object"() { + given: + String json = '{"parent":{"child":{"key":"value"}}}' + + when: + def parsedJson = BedrockJsonParser.parse(json) + def resolvedValue = BedrockJsonParser.JsonPathResolver.resolvePath(parsedJson, "/parent/child/key") + + then: + resolvedValue == "value" + } + + def "should resolve path in JSON array"() { + given: + String json = '{"array":[{"key":"value1"}, {"key":"value2"}]}' + + when: + def parsedJson = BedrockJsonParser.parse(json) + def resolvedValue = BedrockJsonParser.JsonPathResolver.resolvePath(parsedJson, "/array/1/key") + + then: + resolvedValue == "value2" + } + + def "should return null for invalid path resolution"() { + given: + String json = '{"parent":{"child":{"key":"value"}}}' + + when: + def parsedJson = BedrockJsonParser.parse(json) + def resolvedValue = BedrockJsonParser.JsonPathResolver.resolvePath(parsedJson, "/invalid/path") + + then: + resolvedValue == null + } +}