feat!: Remove bedrock-specific tracker method (#1385)

Author: jsonbailey

packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts

@@ -167,112 +167,6 @@ it('tracks success', () => {
   );
 });
 
-it('tracks Bedrock conversation with successful response', () => {
-  const tracker = new LDAIConfigTrackerImpl(
-    mockLdClient,
-    testRunId,
-    configKey,
-    variationKey,
-    version,
-    modelName,
-    providerName,
-    testContext,
-  );
-
-  const TOTAL_TOKENS = 100;
-  const PROMPT_TOKENS = 49;
-  const COMPLETION_TOKENS = 51;
-
-  const response = {
-    $metadata: { httpStatusCode: 200 },
-    metrics: { latencyMs: 500 },
-    usage: {
-      inputTokens: PROMPT_TOKENS,
-      outputTokens: COMPLETION_TOKENS,
-      totalTokens: TOTAL_TOKENS,
-    },
-  };
-
-  tracker.trackBedrockConverseMetrics(response);
-
-  expect(mockTrack).toHaveBeenCalledWith(
-    '$ld:ai:generation:success',
-    testContext,
-    getExpectedTrackData(),
-    1,
-  );
-
-  expect(mockTrack).not.toHaveBeenCalledWith(
-    '$ld:ai:generation:error',
-    expect.anything(),
-    expect.anything(),
-    expect.anything(),
-  );
-
-  expect(mockTrack).toHaveBeenCalledWith(
-    '$ld:ai:duration:total',
-    testContext,
-    getExpectedTrackData(),
-    500,
-  );
-
-  expect(mockTrack).toHaveBeenCalledWith(
-    '$ld:ai:tokens:total',
-    testContext,
-    getExpectedTrackData(),
-    TOTAL_TOKENS,
-  );
-
-  expect(mockTrack).toHaveBeenCalledWith(
-    '$ld:ai:tokens:input',
-    testContext,
-    getExpectedTrackData(),
-    PROMPT_TOKENS,
-  );
-
-  expect(mockTrack).toHaveBeenCalledWith(
-    '$ld:ai:tokens:output',
-    testContext,
-    getExpectedTrackData(),
-    COMPLETION_TOKENS,
-  );
-});
-
-it('tracks Bedrock conversation with error response', () => {
-  const tracker = new LDAIConfigTrackerImpl(
-    mockLdClient,
-    testRunId,
-    configKey,
-    variationKey,
-    version,
-    modelName,
-    providerName,
-    testContext,
-  );
-
-  const response = {
-    $metadata: { httpStatusCode: 400 },
-  };
-
-  tracker.trackBedrockConverseMetrics(response);
-
-  expect(mockTrack).toHaveBeenCalledTimes(1);
-
-  expect(mockTrack).toHaveBeenCalledWith(
-    '$ld:ai:generation:error',
-    testContext,
-    getExpectedTrackData(),
-    1,
-  );
-
-  expect(mockTrack).not.toHaveBeenCalledWith(
-    expect.stringMatching(/^\$ld:ai:tokens:/),
-    expect.anything(),
-    expect.anything(),
-    expect.anything(),
-  );
-});
-
 it('tracks tokens', () => {
   const tracker = new LDAIConfigTrackerImpl(
     mockLdClient,

packages/sdk/server-ai/__tests__/ManagedAgent.test.ts

@@ -33,7 +33,6 @@ describe('ManagedAgent', () => {
       trackFeedback: jest.fn(),
       trackTimeToFirstToken: jest.fn(),
       trackDurationOf: jest.fn(),
-      trackBedrockConverseMetrics: jest.fn(),
       getSummary: jest.fn().mockReturnValue({ success: true, resumptionToken: 'agent-resumption-token' }),
     } as any;
 

packages/sdk/server-ai/__tests__/ManagedModel.test.ts

@@ -24,7 +24,6 @@ describe('ManagedModel', () => {
       trackFeedback: jest.fn(),
       trackTimeToFirstToken: jest.fn(),
       trackDurationOf: jest.fn(),
-      trackBedrockConverseMetrics: jest.fn(),
       getSummary: jest.fn().mockReturnValue({}),
       trackJudgeResult: jest.fn(),
       resumptionToken: 'resumption-token-123',

packages/sdk/server-ai/__tests__/ManagedModelRun.test.ts

@@ -33,7 +33,6 @@ describe('ManagedModel.run() evaluations', () => {
       trackFeedback: jest.fn(),
       trackTimeToFirstToken: jest.fn(),
       trackDurationOf: jest.fn(),
-      trackBedrockConverseMetrics: jest.fn(),
       getSummary: jest
         .fn()
         .mockReturnValue({ success: true, resumptionToken: 'test-resumption-token' }),

packages/sdk/server-ai/__tests__/TokenUsage.test.ts

@@ -1,11 +1,16 @@
 /* eslint-disable no-console */
 import 'dotenv/config';
 
-import { BedrockRuntimeClient, ConverseCommand, Message } from '@aws-sdk/client-bedrock-runtime';
+import {
+  BedrockRuntimeClient,
+  ConverseCommand,
+  type ConverseCommandOutput,
+  Message,
+} from '@aws-sdk/client-bedrock-runtime';
 
 import { init, type LDContext } from '@launchdarkly/node-server-sdk';
 import { Observability } from '@launchdarkly/observability-node';
-import { initAi } from '@launchdarkly/server-sdk-ai';
+import { initAi, type LDAIMetrics } from '@launchdarkly/server-sdk-ai';
 
 const awsClient = new BedrockRuntimeClient({
   region: process.env.AWS_DEFAULT_REGION ?? 'us-east-1',
@@ -34,6 +39,23 @@ const context: LDContext = {
   name: 'Sandy',
 };
 
+// Extract LD AI metrics from a Bedrock Converse response. Passed to
+// tracker.trackMetricsOf, which wraps the AI call to track duration and
+// success automatically. The dedicated @launchdarkly/server-sdk-ai-bedrock
+// provider package will supply an equivalent extractor out of the box.
+function extractBedrockConverseMetrics(res: ConverseCommandOutput): LDAIMetrics {
+  return {
+    success: (res.$metadata?.httpStatusCode ?? 0) === 200,
+    tokens: res.usage
+      ? {
+          total: res.usage.totalTokens ?? 0,
+          input: res.usage.inputTokens ?? 0,
+          output: res.usage.outputTokens ?? 0,
+        }
+      : undefined,
+  };
+}
+
 function mapPromptToConversation(
   prompt: { role: 'user' | 'assistant' | 'system'; content: string }[],
 ): Message[] {
@@ -93,8 +115,8 @@ async function main() {
     console.log(`\nSending sample question to ${aiConfig.model?.name}: "${sampleQuestion}"`);
     console.log('Waiting for response...');
 
-    const completion = tracker.trackBedrockConverseMetrics(
-      await awsClient.send(
+    const completion = await tracker.trackMetricsOf(extractBedrockConverseMetrics, () =>
+      awsClient.send(
         new ConverseCommand({
           modelId: aiConfig.model?.name ?? 'no-model',
           messages: chatMessages,

packages/sdk/server-ai/examples/getting-started/bedrock/converse/src/index.ts

@@ -3,12 +3,7 @@ import { LDContext } from '@launchdarkly/js-server-sdk-common';
 import { LDAIConfigTracker } from './api/config';
 import { LDAIMetricSummary } from './api/model/types';
 import { LDJudgeResult } from './api/judge/types';
-import {
-  createBedrockTokenUsage,
-  LDAIMetrics,
-  LDFeedbackKind,
-  LDTokenUsage,
-} from './api/metrics';
+import { LDAIMetrics, LDFeedbackKind, LDTokenUsage } from './api/metrics';
 import { LDClientMin } from './LDClientMin';
 
 export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
@@ -275,31 +270,6 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
     }
   }
 
-  trackBedrockConverseMetrics<
-    TRes extends {
-      $metadata: { httpStatusCode?: number };
-      metrics?: { latencyMs?: number };
-      usage?: {
-        inputTokens?: number;
-        outputTokens?: number;
-        totalTokens?: number;
-      };
-    },
-  >(res: TRes): TRes {
-    if (res.$metadata?.httpStatusCode === 200) {
-      this.trackSuccess();
-    } else if (res.$metadata?.httpStatusCode && res.$metadata.httpStatusCode >= 400) {
-      this.trackError();
-    }
-    if (res.metrics && res.metrics.latencyMs) {
-      this.trackDuration(res.metrics.latencyMs);
-    }
-    if (res.usage) {
-      this.trackTokens(createBedrockTokenUsage(res.usage));
-    }
-    return res;
-  }
-
   trackTokens(tokens: LDTokenUsage): void {
     if (this._trackedMetrics.tokens !== undefined) {
       this._ldClient.logger?.warn(

packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts

@@ -191,31 +191,6 @@ export interface LDAIConfigTracker {
     metricsExtractor: (stream: TStream) => Promise<LDAIMetrics>,
   ): TStream;
 
-  /**
-   * Track an AI run which uses Bedrock.
-   *
-   * This function will track the duration of the AI run, the token usage, and the success or error status.
-   *
-   * @param res The result of the Bedrock operation.
-   * @returns The input operation.
-   *
-   * @remarks Subsequent calls emit only metrics not already recorded on this
-   * Tracker. Call createTracker on the AI Config to start a new run.
-   */
-  trackBedrockConverseMetrics<
-    TRes extends {
-      $metadata: { httpStatusCode?: number };
-      metrics?: { latencyMs?: number };
-      usage?: {
-        inputTokens?: number;
-        outputTokens?: number;
-        totalTokens?: number;
-      };
-    },
-  >(
-    res: TRes,
-  ): TRes;
-
   /**
    * Get a summary of the tracked metrics.
    */

packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts

@@ -1,4 +1,3 @@
-export * from './BedrockTokenUsage';
 export * from './LDFeedbackKind';
 export * from './LDAIMetrics';
 export * from './LDTokenUsage';