|
| 1 | +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +import { Attributes, DiagLogger, Span, SpanKind, Tracer } from '@opentelemetry/api'; |
| 5 | +import { |
| 6 | + AwsSdkInstrumentationConfig, |
| 7 | + NormalizedRequest, |
| 8 | + NormalizedResponse, |
| 9 | +} from '@opentelemetry/instrumentation-aws-sdk'; |
| 10 | +import { AWS_ATTRIBUTE_KEYS } from '../../../aws-attribute-keys'; |
| 11 | +import { RequestMetadata, ServiceExtension } from '../../../third-party/otel/aws/services/ServiceExtension'; |
| 12 | +import { AwsSpanProcessingUtil } from '../../../aws-span-processing-util'; |
| 13 | + |
| 14 | +const AGENT_ID: string = 'agentId'; |
| 15 | +const KNOWLEDGE_BASE_ID: string = 'knowledgeBaseId'; |
| 16 | +const DATA_SOURCE_ID: string = 'dataSourceId'; |
| 17 | +const GUARDRAIL_ID: string = 'guardrailId'; |
| 18 | +const MODEL_ID: string = 'modelId'; |
| 19 | +const AWS_BEDROCK_SYSTEM: string = 'aws_bedrock'; |
| 20 | + |
| 21 | +const AGENT_OPERATIONS = [ |
| 22 | + 'CreateAgentActionGroup', |
| 23 | + 'CreateAgentAlias', |
| 24 | + 'DeleteAgentActionGroup', |
| 25 | + 'DeleteAgentAlias', |
| 26 | + 'DeleteAgent', |
| 27 | + 'DeleteAgentVersion', |
| 28 | + 'GetAgentActionGroup', |
| 29 | + 'GetAgentAlias', |
| 30 | + 'GetAgent', |
| 31 | + 'GetAgentVersion', |
| 32 | + 'ListAgentActionGroups', |
| 33 | + 'ListAgentAliases', |
| 34 | + 'ListAgentKnowledgeBases', |
| 35 | + 'ListAgentVersions', |
| 36 | + 'PrepareAgent', |
| 37 | + 'UpdateAgentActionGroup', |
| 38 | + 'UpdateAgentAlias', |
| 39 | + 'UpdateAgent', |
| 40 | +]; |
| 41 | + |
| 42 | +const KNOWLEDGE_BASE_OPERATIONS = [ |
| 43 | + 'AssociateAgentKnowledgeBase', |
| 44 | + 'CreateDataSource', |
| 45 | + 'DeleteKnowledgeBase', |
| 46 | + 'DisassociateAgentKnowledgeBase', |
| 47 | + 'GetAgentKnowledgeBase', |
| 48 | + 'GetKnowledgeBase', |
| 49 | + 'ListDataSources', |
| 50 | + 'UpdateAgentKnowledgeBase', |
| 51 | +]; |
| 52 | + |
| 53 | +const DATA_SOURCE_OPERATIONS = ['DeleteDataSource', 'GetDataSource', 'UpdateDataSource']; |
| 54 | + |
| 55 | +// The following constants map the way we present the data in telemetry to how they appear in request/responses |
| 56 | +// e.g. we put `aws.bedrock.knowledge_base.id` into trace data by finding `knowledgeBaseId` |
| 57 | +const agentOperationAttributeKeyMapping = { [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]: AGENT_ID }; |
| 58 | +const knowledgeBaseOperationAttributeKeyMapping = { |
| 59 | + [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]: KNOWLEDGE_BASE_ID, |
| 60 | +}; |
| 61 | +const dataSourceOperationAttributeKeyMapping = { |
| 62 | + [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]: DATA_SOURCE_ID, |
| 63 | +}; |
| 64 | + |
| 65 | +// This map allows us to get all relevant attribute key mappings for a given operation |
| 66 | +const operationToBedrockAgentAttributesMap: { [key: string]: { [key: string]: string } } = {}; |
| 67 | +for (const operation of AGENT_OPERATIONS) { |
| 68 | + operationToBedrockAgentAttributesMap[operation] = agentOperationAttributeKeyMapping; |
| 69 | +} |
| 70 | +for (const operation of KNOWLEDGE_BASE_OPERATIONS) { |
| 71 | + operationToBedrockAgentAttributesMap[operation] = knowledgeBaseOperationAttributeKeyMapping; |
| 72 | +} |
| 73 | +for (const operation of DATA_SOURCE_OPERATIONS) { |
| 74 | + operationToBedrockAgentAttributesMap[operation] = dataSourceOperationAttributeKeyMapping; |
| 75 | +} |
| 76 | + |
| 77 | +// This class is an extension for <a |
| 78 | +// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock.html"> |
| 79 | +// Agents for Amazon Bedrock</a>. |
| 80 | +// This class primarily identify three types of resource based operations: AGENT_OPERATIONS, |
| 81 | +// KNOWLEDGE_BASE_OPERATIONS, and DATA_SOURCE_OPERATIONS. We only support operations that are related to |
| 82 | +// the resource and where the context contains the resource ID. |
| 83 | +export class BedrockAgentServiceExtension implements ServiceExtension { |
| 84 | + requestPreSpanHook( |
| 85 | + request: NormalizedRequest, |
| 86 | + config: AwsSdkInstrumentationConfig, |
| 87 | + diag: DiagLogger |
| 88 | + ): RequestMetadata { |
| 89 | + const spanAttributes: Attributes = {}; |
| 90 | + const isIncoming = false; |
| 91 | + const spanKind: SpanKind = SpanKind.CLIENT; |
| 92 | + let spanName: string | undefined; |
| 93 | + |
| 94 | + const operation: string = request.commandName; |
| 95 | + if (operation && operationToBedrockAgentAttributesMap[operation]) { |
| 96 | + const bedrockAgentServiceInfo = operationToBedrockAgentAttributesMap[operation]; |
| 97 | + for (const serviceInfo of Object.entries(bedrockAgentServiceInfo)) { |
| 98 | + const [attributeKey, requestParamKey] = serviceInfo; |
| 99 | + const requestParamValue = request.commandInput?.[requestParamKey]; |
| 100 | + |
| 101 | + if (requestParamValue) { |
| 102 | + spanAttributes[attributeKey] = requestParamValue; |
| 103 | + } |
| 104 | + } |
| 105 | + } |
| 106 | + |
| 107 | + return { |
| 108 | + isIncoming, |
| 109 | + spanAttributes, |
| 110 | + spanKind, |
| 111 | + spanName, |
| 112 | + }; |
| 113 | + } |
| 114 | + |
| 115 | + responseHook(response: NormalizedResponse, span: Span, tracer: Tracer, config: AwsSdkInstrumentationConfig): void { |
| 116 | + const operation: string = response.request.commandName; |
| 117 | + if (operation && operationToBedrockAgentAttributesMap[operation]) { |
| 118 | + const bedrockAgentServiceInfo = operationToBedrockAgentAttributesMap[operation]; |
| 119 | + for (const serviceInfo of Object.entries(bedrockAgentServiceInfo)) { |
| 120 | + const [attributeKey, responseParamKey] = serviceInfo; |
| 121 | + const responseParamValue = response.data[responseParamKey]; |
| 122 | + |
| 123 | + if (responseParamValue) { |
| 124 | + span.setAttribute(attributeKey, responseParamValue); |
| 125 | + } |
| 126 | + } |
| 127 | + } |
| 128 | + } |
| 129 | +} |
| 130 | + |
| 131 | +// This class is an extension for <a |
| 132 | +// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock_Runtime.html"> |
| 133 | +// Agents for Amazon Bedrock Runtime</a>. |
| 134 | +export class BedrockAgentRuntimeServiceExtension implements ServiceExtension { |
| 135 | + requestPreSpanHook( |
| 136 | + request: NormalizedRequest, |
| 137 | + config: AwsSdkInstrumentationConfig, |
| 138 | + diag: DiagLogger |
| 139 | + ): RequestMetadata { |
| 140 | + const spanAttributes: Attributes = {}; |
| 141 | + const isIncoming = false; |
| 142 | + const spanKind: SpanKind = SpanKind.CLIENT; |
| 143 | + let spanName: string | undefined; |
| 144 | + |
| 145 | + const agentId = request.commandInput?.[AGENT_ID]; |
| 146 | + const knowledgeBaseId = request.commandInput?.[KNOWLEDGE_BASE_ID]; |
| 147 | + |
| 148 | + if (agentId) { |
| 149 | + spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID] = agentId; |
| 150 | + } |
| 151 | + if (knowledgeBaseId) { |
| 152 | + spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID] = knowledgeBaseId; |
| 153 | + } |
| 154 | + |
| 155 | + return { |
| 156 | + isIncoming, |
| 157 | + spanAttributes, |
| 158 | + spanKind, |
| 159 | + spanName, |
| 160 | + }; |
| 161 | + } |
| 162 | +} |
| 163 | + |
| 164 | +// This class is an extension for <a |
| 165 | +// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock.html">Bedrock</a>. |
| 166 | +export class BedrockServiceExtension implements ServiceExtension { |
| 167 | + // Must be implemented, returning empty metadata |
| 168 | + requestPreSpanHook( |
| 169 | + request: NormalizedRequest, |
| 170 | + config: AwsSdkInstrumentationConfig, |
| 171 | + diag: DiagLogger |
| 172 | + ): RequestMetadata { |
| 173 | + const spanAttributes: Attributes = {}; |
| 174 | + const isIncoming = false; |
| 175 | + const spanKind: SpanKind = SpanKind.CLIENT; |
| 176 | + let spanName: string | undefined; |
| 177 | + return { |
| 178 | + isIncoming, |
| 179 | + spanAttributes, |
| 180 | + spanKind, |
| 181 | + spanName, |
| 182 | + }; |
| 183 | + } |
| 184 | + responseHook(response: NormalizedResponse, span: Span, tracer: Tracer, config: AwsSdkInstrumentationConfig): void { |
| 185 | + const guardrail_id = response.data[GUARDRAIL_ID]; |
| 186 | + |
| 187 | + if (guardrail_id) { |
| 188 | + span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, guardrail_id); |
| 189 | + } |
| 190 | + } |
| 191 | +} |
| 192 | + |
| 193 | +// This class is an extension for <a |
| 194 | +// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock_Runtime.html"> |
| 195 | +// Amazon Bedrock Runtime</a>. |
| 196 | +export class BedrockRuntimeServiceExtension implements ServiceExtension { |
| 197 | + requestPreSpanHook( |
| 198 | + request: NormalizedRequest, |
| 199 | + config: AwsSdkInstrumentationConfig, |
| 200 | + diag: DiagLogger |
| 201 | + ): RequestMetadata { |
| 202 | + const spanAttributes: Attributes = {}; |
| 203 | + const isIncoming = false; |
| 204 | + const spanKind: SpanKind = SpanKind.CLIENT; |
| 205 | + let spanName: string | undefined; |
| 206 | + |
| 207 | + const modelId = request.commandInput?.[MODEL_ID]; |
| 208 | + |
| 209 | + spanAttributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM] = AWS_BEDROCK_SYSTEM; |
| 210 | + if (modelId) { |
| 211 | + spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL] = modelId; |
| 212 | + } |
| 213 | + |
| 214 | + return { |
| 215 | + isIncoming, |
| 216 | + spanAttributes, |
| 217 | + spanKind, |
| 218 | + spanName, |
| 219 | + }; |
| 220 | + } |
| 221 | +} |
0 commit comments