Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 0 additions & 106 deletions packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 0 additions & 1 deletion packages/sdk/server-ai/__tests__/ManagedAgent.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
1 change: 0 additions & 1 deletion packages/sdk/server-ai/__tests__/ManagedModel.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
1 change: 0 additions & 1 deletion packages/sdk/server-ai/__tests__/ManagedModelRun.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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' }),
Expand Down
39 changes: 0 additions & 39 deletions packages/sdk/server-ai/__tests__/TokenUsage.test.ts

This file was deleted.

Original file line number Diff line number Diff line change
@@ -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',
Expand Down Expand Up @@ -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[] {
Expand Down Expand Up @@ -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,
Expand Down
32 changes: 1 addition & 31 deletions packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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(
Expand Down
25 changes: 0 additions & 25 deletions packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
Expand Down
13 changes: 0 additions & 13 deletions packages/sdk/server-ai/src/api/metrics/BedrockTokenUsage.ts

This file was deleted.

1 change: 0 additions & 1 deletion packages/sdk/server-ai/src/api/metrics/index.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
export * from './BedrockTokenUsage';
export * from './LDFeedbackKind';
export * from './LDAIMetrics';
export * from './LDTokenUsage';
Loading