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

Author: anuraaga

package-lock.json

@@ -47,7 +47,7 @@
     "@opentelemetry/core": "^2.0.0",
     "@opentelemetry/instrumentation": "^0.200.0",
     "@opentelemetry/propagation-utils": "^0.31.0",
-    "@opentelemetry/semantic-conventions": "^1.27.0"
+    "@opentelemetry/semantic-conventions": "^1.31.0"
   },
   "devDependencies": {
     "@aws-sdk/client-bedrock-runtime": "^3.587.0",

plugins/node/opentelemetry-instrumentation-aws-sdk/package.json

@@ -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,7 +67,8 @@ type V3PluginCommand = AwsV3Command<any, any, any, any, any> & {
 
 export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentationConfig> {
   static readonly component = 'aws-sdk';
-  private servicesExtensions: ServicesExtensions = new ServicesExtensions();
+  // need declare since initialized in callbacks from super constructor
+  private declare servicesExtensions: ServicesExtensions;
 
   constructor(config: AwsSdkInstrumentationConfig = {}) {
     super(PACKAGE_NAME, PACKAGE_VERSION, config);
@@ -341,6 +342,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 +406,8 @@ export class AwsInstrumentation extends InstrumentationBase<AwsSdkInstrumentatio
                     normalizedResponse,
                     span,
                     self.tracer,
-                    self.getConfig()
+                    self.getConfig(),
+                    startTime
                   );
                   self._callUserResponseHook(span, normalizedResponse);
                   return response;
@@ -464,4 +467,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/aws-sdk.ts

@@ -109,6 +109,16 @@ export const ATTR_GEN_AI_RESPONSE_FINISH_REASONS =
  */
 export const ATTR_GEN_AI_SYSTEM = 'gen_ai.system' 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;
+
 /**
  * The number of tokens used in the GenAI input (prompt).
  *
@@ -138,3 +148,13 @@ 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;
+
+/**
+ * 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_OUTPUT = 'output' as const;

plugins/node/opentelemetry-instrumentation-aws-sdk/src/semconv.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/ServiceExtension.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/ServicesExtensions.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,60 @@ 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_OUTPUT,
 } 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) {
+    // https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/#metric-gen_aiclienttokenusage
+    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,
+        ],
+      },
+    });
+
+    // https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/#metric-gen_aiclientoperationduration
+    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,
@@ -262,15 +312,22 @@ 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
+        );
       case 'InvokeModel':
         return this.responseHookInvokeModel(response, span, tracer, config);
     }
@@ -280,16 +337,38 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
     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_OUTPUT,
+        });
       }
     }
 

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.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-s3.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';