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

Author: anuraaga

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

@@ -21,7 +21,7 @@ import {
   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 {
@@ -67,13 +67,20 @@ type V3PluginCommand = AwsV3Command<any, any, any, any, any> & {
 
 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();
+    }
+
     const v3MiddlewareStackFileOldVersions = new InstrumentationNodeModuleFile(
       '@aws-sdk/middleware-stack/dist/cjs/MiddlewareStack.js',
       ['>=3.1.0 <3.35.0'],
@@ -341,6 +348,7 @@ export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentatio
           self.getConfig(),
           self._diag
         );
+        const startTime = hrTime();
         const span = self._startAwsV3Span(normalizedRequest, requestMetadata);
         const activeContextWithSpan = trace.setSpan(context.active(), span);
 
@@ -404,7 +412,8 @@ export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentatio
                     normalizedResponse,
                     span,
                     self.tracer,
-                    self.getConfig()
+                    self.getConfig(),
+                    startTime
                   );
                   self._callUserResponseHook(span, normalizedResponse);
                   return response;
@@ -464,4 +473,11 @@ export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentatio
       return originalFunction();
     }
   }
+
+  override _updateMetricInstruments() {
+    if (!this.servicesExtensions) {
+      this.servicesExtensions = new ServicesExtensions();
+    }
+    this.servicesExtensions.updateMetricInstruments(this.meter);
+  }
 }

plugins/node/opentelemetry-instrumentation-aws-sdk/src/semconv.ts

@@ -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;

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/ServiceExtension.ts

@@ -15,6 +15,8 @@
  */
 import {
   DiagLogger,
+  HrTime,
+  Meter,
   Span,
   SpanAttributes,
   SpanKind,
@@ -49,6 +51,9 @@ export interface ServiceExtension {
     response: NormalizedResponse,
     span: Span,
     tracer: Tracer,
-    config: AwsSdkInstrumentationConfig
+    config: AwsSdkInstrumentationConfig,
+    startTime: HrTime
   ) => void;
+
+  updateMetricInstruments?: (meter: Meter) => void;
 }

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/ServicesExtensions.ts

@@ -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 {
@@ -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);
+    }
   }
 }

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.ts

@@ -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,
@@ -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,
@@ -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,
+        });
       }
     }
 

plugins/node/opentelemetry-instrumentation-aws-sdk/test/aws-sdk-v3-s3.test.ts

@@ -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,

plugins/node/opentelemetry-instrumentation-aws-sdk/test/aws-sdk-v3-sqs.test.ts

@@ -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';
 

plugins/node/opentelemetry-instrumentation-aws-sdk/test/bedrock-runtime.test.ts

@@ -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,
@@ -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);
     });
   });