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