@aws/aws-distro-opentelemetry-node-autoinstrumentation
Version:
This package provides Amazon Web Services distribution of the OpenTelemetry Node Instrumentation, which allows for auto-instrumentation of NodeJS applications.
409 lines • 24.6 kB
JavaScript
;
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
Object.defineProperty(exports, "__esModule", { value: true });
exports.BedrockRuntimeServiceExtension = exports.BedrockServiceExtension = exports.BedrockAgentRuntimeServiceExtension = exports.BedrockAgentServiceExtension = void 0;
const api_1 = require("@opentelemetry/api");
const aws_attribute_keys_1 = require("../../../aws-attribute-keys");
const aws_span_processing_util_1 = require("../../../aws-span-processing-util");
const AGENT_ID = 'agentId';
const KNOWLEDGE_BASE_ID = 'knowledgeBaseId';
const DATA_SOURCE_ID = 'dataSourceId';
const GUARDRAIL_ID = 'guardrailId';
const GUARDRAIL_ARN = 'guardrailArn';
const MODEL_ID = 'modelId';
const AWS_BEDROCK_SYSTEM = 'aws.bedrock';
const AGENT_OPERATIONS = [
'CreateAgentActionGroup',
'CreateAgentAlias',
'DeleteAgentActionGroup',
'DeleteAgentAlias',
'DeleteAgent',
'DeleteAgentVersion',
'GetAgentActionGroup',
'GetAgentAlias',
'GetAgent',
'GetAgentVersion',
'ListAgentActionGroups',
'ListAgentAliases',
'ListAgentKnowledgeBases',
'ListAgentVersions',
'PrepareAgent',
'UpdateAgentActionGroup',
'UpdateAgentAlias',
'UpdateAgent',
];
const KNOWLEDGE_BASE_OPERATIONS = [
'AssociateAgentKnowledgeBase',
'CreateDataSource',
'DeleteKnowledgeBase',
'DisassociateAgentKnowledgeBase',
'GetAgentKnowledgeBase',
'GetKnowledgeBase',
'ListDataSources',
'UpdateAgentKnowledgeBase',
];
const DATA_SOURCE_OPERATIONS = ['DeleteDataSource', 'GetDataSource', 'UpdateDataSource'];
// The following constants map the way we present the data in telemetry to how they appear in request/responses
// e.g. we put `aws.bedrock.knowledge_base.id` into trace data by finding `knowledgeBaseId`
const agentOperationAttributeKeyMapping = { [aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]: AGENT_ID };
const knowledgeBaseOperationAttributeKeyMapping = {
[aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]: KNOWLEDGE_BASE_ID,
};
const dataSourceOperationAttributeKeyMapping = {
[aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]: DATA_SOURCE_ID,
[aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]: KNOWLEDGE_BASE_ID,
};
// This map allows us to get all relevant attribute key mappings for a given operation
const operationToBedrockAgentAttributesMap = {};
for (const operation of AGENT_OPERATIONS) {
operationToBedrockAgentAttributesMap[operation] = agentOperationAttributeKeyMapping;
}
for (const operation of KNOWLEDGE_BASE_OPERATIONS) {
operationToBedrockAgentAttributesMap[operation] = knowledgeBaseOperationAttributeKeyMapping;
}
for (const operation of DATA_SOURCE_OPERATIONS) {
operationToBedrockAgentAttributesMap[operation] = dataSourceOperationAttributeKeyMapping;
}
// This class is an extension for <a
// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock.html">
// Agents for Amazon Bedrock</a>.
// This class primarily identify three types of resource based operations: AGENT_OPERATIONS,
// KNOWLEDGE_BASE_OPERATIONS, and DATA_SOURCE_OPERATIONS. We only support operations that are related to
// the resource and where the context contains the resource ID.
class BedrockAgentServiceExtension {
requestPreSpanHook(request, config, diag) {
var _a;
const spanAttributes = {};
const isIncoming = false;
const spanKind = api_1.SpanKind.CLIENT;
let spanName;
const operation = request.commandName;
if (operation && operationToBedrockAgentAttributesMap[operation]) {
const bedrockAgentServiceInfo = operationToBedrockAgentAttributesMap[operation];
for (const serviceInfo of Object.entries(bedrockAgentServiceInfo)) {
const [attributeKey, requestParamKey] = serviceInfo;
const requestParamValue = (_a = request.commandInput) === null || _a === void 0 ? void 0 : _a[requestParamKey];
if (requestParamValue) {
spanAttributes[attributeKey] = requestParamValue;
}
}
}
return {
isIncoming,
spanAttributes,
spanKind,
spanName,
};
}
responseHook(response, span, tracer, config) {
const operation = response.request.commandName;
if (operation && operationToBedrockAgentAttributesMap[operation]) {
const bedrockAgentServiceInfo = operationToBedrockAgentAttributesMap[operation];
for (const serviceInfo of Object.entries(bedrockAgentServiceInfo)) {
const [attributeKey, responseParamKey] = serviceInfo;
const responseParamValue = response.data[responseParamKey];
if (responseParamValue) {
span.setAttribute(attributeKey, responseParamValue);
}
}
}
}
}
exports.BedrockAgentServiceExtension = BedrockAgentServiceExtension;
// This class is an extension for <a
// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock_Runtime.html">
// Agents for Amazon Bedrock Runtime</a>.
class BedrockAgentRuntimeServiceExtension {
requestPreSpanHook(request, config, diag) {
var _a, _b;
const spanAttributes = {};
const isIncoming = false;
const spanKind = api_1.SpanKind.CLIENT;
let spanName;
const agentId = (_a = request.commandInput) === null || _a === void 0 ? void 0 : _a[AGENT_ID];
const knowledgeBaseId = (_b = request.commandInput) === null || _b === void 0 ? void 0 : _b[KNOWLEDGE_BASE_ID];
if (agentId) {
spanAttributes[aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID] = agentId;
}
if (knowledgeBaseId) {
spanAttributes[aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID] = knowledgeBaseId;
}
return {
isIncoming,
spanAttributes,
spanKind,
spanName,
};
}
}
exports.BedrockAgentRuntimeServiceExtension = BedrockAgentRuntimeServiceExtension;
// This class is an extension for <a
// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock.html">Bedrock</a>.
class BedrockServiceExtension {
// Must be implemented, returning empty metadata
requestPreSpanHook(request, config, diag) {
const spanAttributes = {};
const isIncoming = false;
const spanKind = api_1.SpanKind.CLIENT;
let spanName;
return {
isIncoming,
spanAttributes,
spanKind,
spanName,
};
}
responseHook(response, span, tracer, config) {
const guardrailId = response.data[GUARDRAIL_ID];
const guardrailArn = response.data[GUARDRAIL_ARN];
if (guardrailId) {
span.setAttribute(aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, guardrailId);
}
if (guardrailArn) {
span.setAttribute(aws_attribute_keys_1.AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN, guardrailArn);
}
}
}
exports.BedrockServiceExtension = BedrockServiceExtension;
// This class is an extension for <a
// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock_Runtime.html">
// Amazon Bedrock Runtime</a>.
class BedrockRuntimeServiceExtension {
requestPreSpanHook(request, config, diag) {
var _a, _b, _c, _d, _e, _f, _g, _h;
const spanAttributes = {};
const isIncoming = false;
const spanKind = api_1.SpanKind.CLIENT;
let spanName;
const modelId = (_a = request.commandInput) === null || _a === void 0 ? void 0 : _a[MODEL_ID];
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_SYSTEM] = AWS_BEDROCK_SYSTEM;
if (modelId) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL] = modelId;
}
if ((_b = request.commandInput) === null || _b === void 0 ? void 0 : _b.body) {
const requestBody = JSON.parse(request.commandInput.body);
if (modelId.includes('amazon.titan')) {
if (((_c = requestBody.textGenerationConfig) === null || _c === void 0 ? void 0 : _c.temperature) !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] =
requestBody.textGenerationConfig.temperature;
}
if (((_d = requestBody.textGenerationConfig) === null || _d === void 0 ? void 0 : _d.topP) !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.textGenerationConfig.topP;
}
if (((_e = requestBody.textGenerationConfig) === null || _e === void 0 ? void 0 : _e.maxTokenCount) !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] =
requestBody.textGenerationConfig.maxTokenCount;
}
}
else if (modelId.includes('amazon.nova')) {
if (((_f = requestBody.inferenceConfig) === null || _f === void 0 ? void 0 : _f.temperature) !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.inferenceConfig.temperature;
}
if (((_g = requestBody.inferenceConfig) === null || _g === void 0 ? void 0 : _g.top_p) !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.inferenceConfig.top_p;
}
if (((_h = requestBody.inferenceConfig) === null || _h === void 0 ? void 0 : _h.max_new_tokens) !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.inferenceConfig.max_new_tokens;
}
}
else if (modelId.includes('anthropic.claude')) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
}
else if (modelId.includes('meta.llama')) {
if (requestBody.max_gen_len !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_gen_len;
}
if (requestBody.temperature !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
}
else if (modelId.includes('cohere.command-r')) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.p !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.p;
}
if (requestBody.message !== undefined) {
// NOTE: We approximate the token count since this value is not directly available in the body
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
// https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.message.length / 6);
}
}
else if (modelId.includes('cohere.command')) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.p !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.p;
}
if (requestBody.prompt !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.prompt.length / 6);
}
}
else if (modelId.includes('ai21.jamba')) {
if (requestBody.max_tokens !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
}
else if (modelId.includes('mistral')) {
if (requestBody.prompt !== undefined) {
// NOTE: We approximate the token count since this value is not directly available in the body
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
// https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.prompt.length / 6);
}
if (requestBody.max_tokens !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens;
}
if (requestBody.temperature !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature;
}
if (requestBody.top_p !== undefined) {
spanAttributes[aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_REQUEST_TOP_P] = requestBody.top_p;
}
}
}
return {
isIncoming,
spanAttributes,
spanKind,
spanName,
};
}
responseHook(response, span, tracer, config) {
var _a, _b, _c, _d, _e, _f, _g, _h, _j, _k, _l, _m, _o, _p, _q, _r, _s, _t, _u, _v;
const currentModelId = (_a = response.request.commandInput) === null || _a === void 0 ? void 0 : _a.modelId;
if ((_b = response.data) === null || _b === void 0 ? void 0 : _b.body) {
const decodedResponseBody = new TextDecoder().decode(response.data.body);
const responseBody = JSON.parse(decodedResponseBody);
if (currentModelId.includes('amazon.titan')) {
if (responseBody.inputTextTokenCount !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.inputTextTokenCount);
}
if (((_d = (_c = responseBody.results) === null || _c === void 0 ? void 0 : _c[0]) === null || _d === void 0 ? void 0 : _d.tokenCount) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.results[0].tokenCount);
}
if (((_f = (_e = responseBody.results) === null || _e === void 0 ? void 0 : _e[0]) === null || _f === void 0 ? void 0 : _f.completionReason) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.results[0].completionReason,
]);
}
}
else if (currentModelId.includes('amazon.nova')) {
if (responseBody.usage !== undefined) {
if (responseBody.usage.inputTokens !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.inputTokens);
}
if (responseBody.usage.outputTokens !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.outputTokens);
}
}
if (responseBody.stopReason !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stopReason]);
}
}
else if (currentModelId.includes('anthropic.claude')) {
if (((_g = responseBody.usage) === null || _g === void 0 ? void 0 : _g.input_tokens) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.input_tokens);
}
if (((_h = responseBody.usage) === null || _h === void 0 ? void 0 : _h.output_tokens) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.output_tokens);
}
if (responseBody.stop_reason !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]);
}
}
else if (currentModelId.includes('meta.llama')) {
if (responseBody.prompt_token_count !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.prompt_token_count);
}
if (responseBody.generation_token_count !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.generation_token_count);
}
if (responseBody.stop_reason !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]);
}
}
else if (currentModelId.includes('cohere.command-r')) {
if (responseBody.text !== undefined) {
// NOTE: We approximate the token count since this value is not directly available in the body
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
// https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.text.length / 6));
}
if (responseBody.finish_reason !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.finish_reason]);
}
}
else if (currentModelId.includes('cohere.command')) {
if (((_k = (_j = responseBody.generations) === null || _j === void 0 ? void 0 : _j[0]) === null || _k === void 0 ? void 0 : _k.text) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS,
// NOTE: We approximate the token count since this value is not directly available in the body
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
// https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
Math.ceil(responseBody.generations[0].text.length / 6));
}
if (((_m = (_l = responseBody.generations) === null || _l === void 0 ? void 0 : _l[0]) === null || _m === void 0 ? void 0 : _m.finish_reason) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.generations[0].finish_reason,
]);
}
}
else if (currentModelId.includes('ai21.jamba')) {
if (((_o = responseBody.usage) === null || _o === void 0 ? void 0 : _o.prompt_tokens) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.prompt_tokens);
}
if (((_p = responseBody.usage) === null || _p === void 0 ? void 0 : _p.completion_tokens) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.completion_tokens);
}
if (((_r = (_q = responseBody.choices) === null || _q === void 0 ? void 0 : _q[0]) === null || _r === void 0 ? void 0 : _r.finish_reason) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.choices[0].finish_reason,
]);
}
}
else if (currentModelId.includes('mistral')) {
if (((_t = (_s = responseBody.outputs) === null || _s === void 0 ? void 0 : _s[0]) === null || _t === void 0 ? void 0 : _t.text) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_USAGE_OUTPUT_TOKENS,
// NOTE: We approximate the token count since this value is not directly available in the body
// According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing.
// https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-prepare.html
Math.ceil(responseBody.outputs[0].text.length / 6));
}
if (((_v = (_u = responseBody.outputs) === null || _u === void 0 ? void 0 : _u[0]) === null || _v === void 0 ? void 0 : _v.stop_reason) !== undefined) {
span.setAttribute(aws_span_processing_util_1.AwsSpanProcessingUtil.GEN_AI_RESPONSE_FINISH_REASONS, [
responseBody.outputs[0].stop_reason,
]);
}
}
}
}
}
exports.BedrockRuntimeServiceExtension = BedrockRuntimeServiceExtension;
//# sourceMappingURL=bedrock.js.map