Skip to content

feat(instrumentation-aws-sdk): add gen ai metrics for bedrock #2771

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
May 12, 2025
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
diag,
SpanStatusCode,
} from '@opentelemetry/api';
import { suppressTracing } from '@opentelemetry/core';
import { hrTime, suppressTracing } from '@opentelemetry/core';
import { AttributeNames } from './enums';
import { ServicesExtensions } from './services';
import {
Expand Down Expand Up @@ -67,13 +67,20 @@

export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentationConfig> {
static readonly component = 'aws-sdk';
private servicesExtensions: ServicesExtensions = new ServicesExtensions();
// initialized in callbacks from super constructor for ordering reasons.
private declare servicesExtensions: ServicesExtensions;

constructor(config: AwsSdkInstrumentationConfig = {}) {
super(PACKAGE_NAME, PACKAGE_VERSION, config);
}

protected init(): InstrumentationModuleDefinition[] {
// Should always have been initialized in _updateMetricInstruments, but check again
// for safety.
if (!this.servicesExtensions) {
this.servicesExtensions = new ServicesExtensions();

Check warning on line 81 in plugins/node/opentelemetry-instrumentation-aws-sdk/src/aws-sdk.ts

View check run for this annotation

Codecov / codecov/patch

plugins/node/opentelemetry-instrumentation-aws-sdk/src/aws-sdk.ts#L81

Added line #L81 was not covered by tests
}

const v3MiddlewareStackFileOldVersions = new InstrumentationNodeModuleFile(
'@aws-sdk/middleware-stack/dist/cjs/MiddlewareStack.js',
['>=3.1.0 <3.35.0'],
Expand Down Expand Up @@ -341,6 +348,7 @@
self.getConfig(),
self._diag
);
const startTime = hrTime();
const span = self._startAwsV3Span(normalizedRequest, requestMetadata);
const activeContextWithSpan = trace.setSpan(context.active(), span);

Expand Down Expand Up @@ -404,7 +412,8 @@
normalizedResponse,
span,
self.tracer,
self.getConfig()
self.getConfig(),
startTime
);
self._callUserResponseHook(span, normalizedResponse);
return response;
Expand Down Expand Up @@ -464,4 +473,11 @@
return originalFunction();
}
}

override _updateMetricInstruments() {
if (!this.servicesExtensions) {
this.servicesExtensions = new ServicesExtensions();
}
this.servicesExtensions.updateMetricInstruments(this.meter);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -138,3 +138,23 @@ export const GEN_AI_OPERATION_NAME_VALUE_CHAT = 'chat' as const;
* Enum value "aws.bedrock" for attribute {@link ATTR_GEN_AI_SYSTEM}.
*/
export const GEN_AI_SYSTEM_VALUE_AWS_BEDROCK = 'aws.bedrock' as const;

/**
* The type of token being counted.
*
* @example input
* @example output
*
* @experimental This attribute is experimental and is subject to breaking changes in minor releases of `@opentelemetry/semantic-conventions`.
*/
export const ATTR_GEN_AI_TOKEN_TYPE = 'gen_ai.token.type' as const;

/**
* Enum value "input" for attribute {@link ATTR_GEN_AI_TOKEN_TYPE}.
*/
export const GEN_AI_TOKEN_TYPE_VALUE_INPUT = 'input' as const;

/**
* Enum value "output" for attribute {@link ATTR_GEN_AI_TOKEN_TYPE}.
*/
export const GEN_AI_TOKEN_TYPE_VALUE_COMPLETION = 'output' as const;
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
*/
import {
DiagLogger,
HrTime,
Meter,
Span,
SpanAttributes,
SpanKind,
Expand Down Expand Up @@ -49,6 +51,9 @@ export interface ServiceExtension {
response: NormalizedResponse,
span: Span,
tracer: Tracer,
config: AwsSdkInstrumentationConfig
config: AwsSdkInstrumentationConfig,
startTime: HrTime
) => void;

updateMetricInstruments?: (meter: Meter) => void;
}
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { Tracer, Span, DiagLogger } from '@opentelemetry/api';
import { Tracer, Span, DiagLogger, Meter, HrTime } from '@opentelemetry/api';
import { ServiceExtension, RequestMetadata } from './ServiceExtension';
import { SqsServiceExtension } from './sqs';
import {
Expand Down Expand Up @@ -64,9 +64,16 @@ export class ServicesExtensions implements ServiceExtension {
response: NormalizedResponse,
span: Span,
tracer: Tracer,
config: AwsSdkInstrumentationConfig
config: AwsSdkInstrumentationConfig,
startTime: HrTime
) {
const serviceExtension = this.services.get(response.request.serviceName);
serviceExtension?.responseHook?.(response, span, tracer, config);
serviceExtension?.responseHook?.(response, span, tracer, config, startTime);
}

updateMetricInstruments(meter: Meter) {
for (const serviceExtension of this.services.values()) {
serviceExtension.updateMetricInstruments?.(meter);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,16 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { Attributes, DiagLogger, Span, Tracer } from '@opentelemetry/api';
import {
Attributes,
DiagLogger,
Histogram,
HrTime,
Meter,
Span,
Tracer,
ValueType,
} from '@opentelemetry/api';
import { RequestMetadata, ServiceExtension } from './ServiceExtension';
import {
ATTR_GEN_AI_SYSTEM,
Expand All @@ -23,19 +32,57 @@ import {
ATTR_GEN_AI_REQUEST_TEMPERATURE,
ATTR_GEN_AI_REQUEST_TOP_P,
ATTR_GEN_AI_REQUEST_STOP_SEQUENCES,
ATTR_GEN_AI_TOKEN_TYPE,
ATTR_GEN_AI_USAGE_INPUT_TOKENS,
ATTR_GEN_AI_USAGE_OUTPUT_TOKENS,
ATTR_GEN_AI_RESPONSE_FINISH_REASONS,
GEN_AI_OPERATION_NAME_VALUE_CHAT,
GEN_AI_SYSTEM_VALUE_AWS_BEDROCK,
GEN_AI_TOKEN_TYPE_VALUE_INPUT,
GEN_AI_TOKEN_TYPE_VALUE_COMPLETION,
} from '../semconv';
import {
AwsSdkInstrumentationConfig,
NormalizedRequest,
NormalizedResponse,
} from '../types';
import {
hrTime,
hrTimeDuration,
hrTimeToMilliseconds,
} from '@opentelemetry/core';

export class BedrockRuntimeServiceExtension implements ServiceExtension {
private tokenUsage!: Histogram;
private operationDuration!: Histogram;

updateMetricInstruments(meter: Meter) {
this.tokenUsage = meter.createHistogram('gen_ai.client.token.usage', {
unit: '{token}',
description: 'Measures number of input and output tokens used',
valueType: ValueType.INT,
advice: {
explicitBucketBoundaries: [
1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304,
16777216, 67108864,
],
},
});
this.operationDuration = meter.createHistogram(
'gen_ai.client.operation.duration',
{
unit: 's',
description: 'GenAI operation duration',
advice: {
explicitBucketBoundaries: [
0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24,
20.48, 40.96, 81.92,
],
},
}
);
}

requestPreSpanHook(
request: NormalizedRequest,
config: AwsSdkInstrumentationConfig,
Expand Down Expand Up @@ -98,32 +145,61 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
response: NormalizedResponse,
span: Span,
tracer: Tracer,
config: AwsSdkInstrumentationConfig
config: AwsSdkInstrumentationConfig,
startTime: HrTime
) {
if (!span.isRecording()) {
return;
}

switch (response.request.commandName) {
case 'Converse':
return this.responseHookConverse(response, span, tracer, config);
return this.responseHookConverse(
response,
span,
tracer,
config,
startTime
);
}
}

private responseHookConverse(
response: NormalizedResponse,
span: Span,
tracer: Tracer,
config: AwsSdkInstrumentationConfig
config: AwsSdkInstrumentationConfig,
startTime: HrTime
) {
const { stopReason, usage } = response.data;

const sharedMetricAttrs: Attributes = {
[ATTR_GEN_AI_SYSTEM]: GEN_AI_SYSTEM_VALUE_AWS_BEDROCK,
[ATTR_GEN_AI_OPERATION_NAME]: GEN_AI_OPERATION_NAME_VALUE_CHAT,
[ATTR_GEN_AI_REQUEST_MODEL]: response.request.commandInput.modelId,
};

const durationSecs =
hrTimeToMilliseconds(hrTimeDuration(startTime, hrTime())) / 1000;
this.operationDuration.record(durationSecs, sharedMetricAttrs);

if (usage) {
const { inputTokens, outputTokens } = usage;
if (inputTokens !== undefined) {
span.setAttribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, inputTokens);

this.tokenUsage.record(inputTokens, {
...sharedMetricAttrs,
[ATTR_GEN_AI_TOKEN_TYPE]: GEN_AI_TOKEN_TYPE_VALUE_INPUT,
});
}
if (outputTokens !== undefined) {
span.setAttribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, outputTokens);

this.tokenUsage.record(outputTokens, {
...sharedMetricAttrs,
[ATTR_GEN_AI_TOKEN_TYPE]: GEN_AI_TOKEN_TYPE_VALUE_COMPLETION,
});
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,17 +15,12 @@
*/

import {
AwsInstrumentation,
AwsSdkRequestHookInformation,
AwsSdkResponseHookInformation,
} from '../src';
import {
getTestSpans,
registerInstrumentationTesting,
} from '@opentelemetry/contrib-test-utils';
const instrumentation = registerInstrumentationTesting(
new AwsInstrumentation()
);
import { getTestSpans } from '@opentelemetry/contrib-test-utils';
import { instrumentation } from './load-instrumentation';

import {
PutObjectCommand,
PutObjectCommandOutput,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,8 @@
// covered multiple `client-*` packages. Its tests could be merged into
// sqs.test.ts.

import { AwsInstrumentation } from '../src';
import {
getTestSpans,
registerInstrumentationTesting,
} from '@opentelemetry/contrib-test-utils';
registerInstrumentationTesting(new AwsInstrumentation());
import { getTestSpans } from '@opentelemetry/contrib-test-utils';
import './load-instrumentation';

import { SQS } from '@aws-sdk/client-sqs';

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,12 +27,8 @@
* keeping existing recordings, set NOCK_BACK_MODE to 'record'.
*/

import {
getTestSpans,
registerInstrumentationTesting,
} from '@opentelemetry/contrib-test-utils';
import { AwsInstrumentation } from '../src';
registerInstrumentationTesting(new AwsInstrumentation());
import { getTestSpans } from '@opentelemetry/contrib-test-utils';
import { metricReader } from './load-instrumentation';

import {
BedrockRuntimeClient,
Expand Down Expand Up @@ -151,6 +147,74 @@ describe('Bedrock', () => {
[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS]: 10,
[ATTR_GEN_AI_RESPONSE_FINISH_REASONS]: ['max_tokens'],
});

const { resourceMetrics } = await metricReader.collect();
expect(resourceMetrics.scopeMetrics.length).toBe(1);
const scopeMetrics = resourceMetrics.scopeMetrics[0];
const tokenUsage = scopeMetrics.metrics.filter(
m => m.descriptor.name === 'gen_ai.client.token.usage'
);
expect(tokenUsage.length).toBe(1);
expect(tokenUsage[0].descriptor).toMatchObject({
name: 'gen_ai.client.token.usage',
type: 'HISTOGRAM',
description: 'Measures number of input and output tokens used',
unit: '{token}',
});
expect(tokenUsage[0].dataPoints.length).toBe(2);
expect(tokenUsage[0].dataPoints).toEqual(
expect.arrayContaining([
expect.objectContaining({
value: expect.objectContaining({
sum: 8,
}),
attributes: {
'gen_ai.system': 'aws.bedrock',
'gen_ai.operation.name': 'chat',
'gen_ai.request.model': 'amazon.titan-text-lite-v1',
'gen_ai.token.type': 'input',
},
}),
expect.objectContaining({
value: expect.objectContaining({
sum: 10,
}),
attributes: {
'gen_ai.system': 'aws.bedrock',
'gen_ai.operation.name': 'chat',
'gen_ai.request.model': 'amazon.titan-text-lite-v1',
'gen_ai.token.type': 'output',
},
}),
])
);

const operationDuration = scopeMetrics.metrics.filter(
m => m.descriptor.name === 'gen_ai.client.operation.duration'
);
expect(operationDuration.length).toBe(1);
expect(operationDuration[0].descriptor).toMatchObject({
name: 'gen_ai.client.operation.duration',
type: 'HISTOGRAM',
description: 'GenAI operation duration',
unit: 's',
});
expect(operationDuration[0].dataPoints.length).toBe(1);
expect(operationDuration[0].dataPoints).toEqual([
expect.objectContaining({
value: expect.objectContaining({
sum: expect.any(Number),
}),
attributes: {
'gen_ai.system': 'aws.bedrock',
'gen_ai.operation.name': 'chat',
'gen_ai.request.model': 'amazon.titan-text-lite-v1',
},
}),
]);
expect(
(operationDuration[0].dataPoints[0].value as any).sum
).toBeGreaterThan(0);
});
});

Expand Down
Loading