From 5d16dd44e0ee2f6a3a9453be1837d8e0f0d5e19a Mon Sep 17 00:00:00 2001 From: jsonbailey Date: Mon, 18 May 2026 15:12:42 -0700 Subject: [PATCH] feat!: Remove bedrock-specific tracker method Remove the trackBedrockConverseMetrics method from LDAIConfigTracker along with the createBedrockTokenUsage helper and the BedrockTokenUsage module. AWS Bedrock metric tracking will move to a dedicated @launchdarkly/server-sdk-ai-bedrock provider package, matching the pattern used by the OpenAI and LangChain providers. The getting-started/bedrock/converse example is kept and updated to use tracker.trackMetricsOf with an inline metrics extractor. This keeps the example functional and demonstrates the generic extractor pattern that the upcoming provider package will package as a helper. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../__tests__/LDAIConfigTrackerImpl.test.ts | 106 ------------------ .../server-ai/__tests__/ManagedAgent.test.ts | 1 - .../server-ai/__tests__/ManagedModel.test.ts | 1 - .../__tests__/ManagedModelRun.test.ts | 1 - .../server-ai/__tests__/TokenUsage.test.ts | 39 ------- .../bedrock/converse/src/index.ts | 30 ++++- .../server-ai/src/LDAIConfigTrackerImpl.ts | 32 +----- .../src/api/config/LDAIConfigTracker.ts | 25 ----- .../src/api/metrics/BedrockTokenUsage.ts | 13 --- .../sdk/server-ai/src/api/metrics/index.ts | 1 - 10 files changed, 27 insertions(+), 222 deletions(-) delete mode 100644 packages/sdk/server-ai/__tests__/TokenUsage.test.ts delete mode 100644 packages/sdk/server-ai/src/api/metrics/BedrockTokenUsage.ts diff --git a/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts b/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts index 7941fb9800..4b625f1bca 100644 --- a/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts +++ b/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, diff --git a/packages/sdk/server-ai/__tests__/ManagedAgent.test.ts b/packages/sdk/server-ai/__tests__/ManagedAgent.test.ts index c4bdc795e4..0b85b90cb5 100644 --- a/packages/sdk/server-ai/__tests__/ManagedAgent.test.ts +++ b/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; diff --git a/packages/sdk/server-ai/__tests__/ManagedModel.test.ts b/packages/sdk/server-ai/__tests__/ManagedModel.test.ts index bd947ee335..31a81a41d8 100644 --- a/packages/sdk/server-ai/__tests__/ManagedModel.test.ts +++ b/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', diff --git a/packages/sdk/server-ai/__tests__/ManagedModelRun.test.ts b/packages/sdk/server-ai/__tests__/ManagedModelRun.test.ts index a7033f2ff9..f2211758d2 100644 --- a/packages/sdk/server-ai/__tests__/ManagedModelRun.test.ts +++ b/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' }), diff --git a/packages/sdk/server-ai/__tests__/TokenUsage.test.ts b/packages/sdk/server-ai/__tests__/TokenUsage.test.ts deleted file mode 100644 index bde16c603f..0000000000 --- a/packages/sdk/server-ai/__tests__/TokenUsage.test.ts +++ /dev/null @@ -1,39 +0,0 @@ -import { createBedrockTokenUsage } from '../src/api/metrics'; - -it('createBedrockTokenUsage should create token usage with all values provided', () => { - const usage = createBedrockTokenUsage({ - totalTokens: 100, - inputTokens: 40, - outputTokens: 60, - }); - - expect(usage).toEqual({ - total: 100, - input: 40, - output: 60, - }); -}); - -it('createBedrockTokenUsage should default to 0 for missing values', () => { - const usage = createBedrockTokenUsage({}); - - expect(usage).toEqual({ - total: 0, - input: 0, - output: 0, - }); -}); - -it('createBedrockTokenUsage should handle explicitly undefined values', () => { - const usage = createBedrockTokenUsage({ - totalTokens: undefined, - inputTokens: 40, - outputTokens: undefined, - }); - - expect(usage).toEqual({ - total: 0, - input: 40, - output: 0, - }); -}); diff --git a/packages/sdk/server-ai/examples/getting-started/bedrock/converse/src/index.ts b/packages/sdk/server-ai/examples/getting-started/bedrock/converse/src/index.ts index 5ada50b8e0..5722a14ddc 100644 --- a/packages/sdk/server-ai/examples/getting-started/bedrock/converse/src/index.ts +++ b/packages/sdk/server-ai/examples/getting-started/bedrock/converse/src/index.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, diff --git a/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts b/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts index b7f7bcb8de..ab5a6fb918 100644 --- a/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts +++ b/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.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( diff --git a/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts b/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts index fdef019898..859cfd385f 100644 --- a/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts +++ b/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts @@ -191,31 +191,6 @@ export interface LDAIConfigTracker { metricsExtractor: (stream: TStream) => Promise, ): 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. */ diff --git a/packages/sdk/server-ai/src/api/metrics/BedrockTokenUsage.ts b/packages/sdk/server-ai/src/api/metrics/BedrockTokenUsage.ts deleted file mode 100644 index 0b7fc40e55..0000000000 --- a/packages/sdk/server-ai/src/api/metrics/BedrockTokenUsage.ts +++ /dev/null @@ -1,13 +0,0 @@ -import { LDTokenUsage } from './LDTokenUsage'; - -export function createBedrockTokenUsage(data: { - totalTokens?: number; - inputTokens?: number; - outputTokens?: number; -}): LDTokenUsage { - return { - total: data.totalTokens || 0, - input: data.inputTokens || 0, - output: data.outputTokens || 0, - }; -} diff --git a/packages/sdk/server-ai/src/api/metrics/index.ts b/packages/sdk/server-ai/src/api/metrics/index.ts index 5b34905c08..11fb6ffc6b 100644 --- a/packages/sdk/server-ai/src/api/metrics/index.ts +++ b/packages/sdk/server-ai/src/api/metrics/index.ts @@ -1,4 +1,3 @@ -export * from './BedrockTokenUsage'; export * from './LDFeedbackKind'; export * from './LDAIMetrics'; export * from './LDTokenUsage';