-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
This PR add AWS Bedrock and BedrockRuntime support for AWS Java SDK V1 with the following change: A. **Bedrock**: Extract `guardrailId` from API response, and add into `"aws.bedrock.guardrail.id"` span attribute. B. **Bedrock Agent**: Extract `agentId` from both API request and response, and add into `"aws.bedrock.agent.id"` span attribute. Extract `knowledgeBaseId` from API request, and add into `"aws.bedrock.knowledgebase.id"` span attribute. Extract `dataSourceId` from both API request and response,, and add into `"aws.bedrock.datasource.id"` span attribute. The instrumentation is on API operation level, we make sure only one attribute is extracted per API call, there will be no overlap/conflict to identify the resource. C. **Bedrock Agent Runtime**: Extract `agentId` from API request, and add into `"aws.bedrock.agent.id"` span attribute. Extract `knowledgeBaseId` from API request, and add into `"aws.bedrock.knowledgebase.id"` span attribute. D. **Bedrock Runtime**: Extract the following attributes and add into span according to [ Gen AI semantic-conventions](https://github.com/open-telemetry/semantic-conventions/blob/main/docs/attributes-registry/gen-ai.md): ``` gen_ai.request.model gen_ai.system ```
- Loading branch information
Showing
9 changed files
with
372 additions
and
19 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
133 changes: 133 additions & 0 deletions
133
...y/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsBedrockResourceType.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,133 @@ | ||
| /* | ||
| * Copyright The OpenTelemetry Authors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package io.opentelemetry.instrumentation.awssdk.v1_11; | ||
|
|
||
| import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_AGENT_ID; | ||
| import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_DATASOURCE_ID; | ||
| import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGEBASE_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_DATASOURCE_ID, RequestAccess::getDataSourceId), | ||
| KNOWLEDGE_BASE_TYPE(AWS_KNOWLEDGEBASE_ID, RequestAccess::getKnowledgeBaseId); | ||
|
|
||
| @SuppressWarnings("ImmutableEnumChecker") | ||
| private final AttributeKey<String> keyAttribute; | ||
|
|
||
| @SuppressWarnings("ImmutableEnumChecker") | ||
| private final Function<Object, String> attributeValueAccessor; | ||
|
|
||
| AwsBedrockResourceType( | ||
| AttributeKey<String> keyAttribute, Function<Object, String> attributeValueAccessor) { | ||
| this.keyAttribute = keyAttribute; | ||
| this.attributeValueAccessor = attributeValueAccessor; | ||
| } | ||
|
|
||
| public AttributeKey<String> getKeyAttribute() { | ||
| return keyAttribute; | ||
| } | ||
|
|
||
| public Function<Object, String> 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<String, AwsBedrockResourceType> BEDROCK_REQUEST_MAP = new HashMap<>(); | ||
| private static final Map<String, AwsBedrockResourceType> 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<String> 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<String> agentResponseClasses = | ||
| Arrays.asList( | ||
| "DeleteAgentAliasResult", | ||
| "DeleteAgentResult", | ||
| "DeleteAgentVersionResult", | ||
| "PrepareAgentResult"); | ||
| // DataSourceID | ||
| private static final List<String> dataSourceRequestClasses = | ||
| Arrays.asList("DeleteDataSourceRequest", "GetDataSourceRequest", "UpdateDataSourceRequest"); | ||
| private static final List<String> dataSourceResponseClasses = | ||
| Arrays.asList("DeleteDataSourceResult"); | ||
| // KnowledgeBaseID | ||
| private static final List<String> knowledgeBaseRequestClasses = | ||
| Arrays.asList( | ||
| "AssociateAgentKnowledgeBaseRequest", | ||
| "CreateDataSourceRequest", | ||
| "DeleteKnowledgeBaseRequest", | ||
| "DisassociateAgentKnowledgeBaseRequest", | ||
| "GetAgentKnowledgeBaseRequest", | ||
| "GetKnowledgeBaseRequest", | ||
| "ListDataSourcesRequest", | ||
| "UpdateAgentKnowledgeBaseRequest"); | ||
| private static final List<String> 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); | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.