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
88 changes: 85 additions & 3 deletions packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,39 @@ it('tracks OpenAI usage', async () => {
);
});

it('tracks error when OpenAI metrics function throws', async () => {
const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, variationKey, testContext);
jest.spyOn(global.Date, 'now').mockReturnValueOnce(1000).mockReturnValueOnce(2000);

const error = new Error('OpenAI API error');
await expect(
tracker.trackOpenAIMetrics(async () => {
throw error;
}),
).rejects.toThrow(error);

expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:duration:total',
testContext,
{ configKey, variationKey },
1000,
);

expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:generation',
testContext,
{ configKey, variationKey },
1,
);

expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:generation:error',
testContext,
{ configKey, variationKey },
1,
);
});

it('tracks Bedrock conversation with successful response', () => {
const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, variationKey, testContext);

Expand Down Expand Up @@ -196,11 +229,22 @@ it('tracks Bedrock conversation with error response', () => {
$metadata: { httpStatusCode: 400 },
};

// TODO: We may want a track failure.

tracker.trackBedrockConverseMetrics(response);

expect(mockTrack).not.toHaveBeenCalled();
expect(mockTrack).toHaveBeenCalledTimes(2);
expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:generation',
testContext,
{ configKey, variationKey },
1,
);

expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:generation:error',
testContext,
{ configKey, variationKey },
1,
);
});

it('tracks tokens', () => {
Expand Down Expand Up @@ -304,3 +348,41 @@ it('summarizes tracked metrics', () => {
success: true,
});
});

it('tracks duration when async function throws', async () => {
const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, variationKey, testContext);
jest.spyOn(global.Date, 'now').mockReturnValueOnce(1000).mockReturnValueOnce(2000);

const error = new Error('test error');
await expect(
tracker.trackDurationOf(async () => {
throw error;
}),
).rejects.toThrow(error);

expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:duration:total',
testContext,
{ configKey, variationKey },
1000,
);
});

it('tracks error', () => {
const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, variationKey, testContext);
tracker.trackError();

expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:generation',
testContext,
{ configKey, variationKey },
1,
);

expect(mockTrack).toHaveBeenCalledWith(
'$ld:ai:generation:error',
testContext,
{ configKey, variationKey },
1,
);
});
38 changes: 27 additions & 11 deletions packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,15 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {

async trackDurationOf<TRes>(func: () => Promise<TRes>): Promise<TRes> {
const startTime = Date.now();
const result = await func();
const endTime = Date.now();
const duration = endTime - startTime; // duration in milliseconds
this.trackDuration(duration);
return result;
try {
// Be sure to await here so that we can track the duration of the function and also handle errors.
const result = await func();
return result;
} finally {
const endTime = Date.now();
const duration = endTime - startTime; // duration in milliseconds
this.trackDuration(duration);
}
}

trackFeedback(feedback: { kind: LDFeedbackKind }): void {
Expand All @@ -49,6 +53,13 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
trackSuccess(): void {
this._trackedMetrics.success = true;
this._ldClient.track('$ld:ai:generation', this._context, this._getTrackData(), 1);
this._ldClient.track('$ld:ai:generation:success', this._context, this._getTrackData(), 1);
}

trackError(): void {
this._trackedMetrics.success = false;
this._ldClient.track('$ld:ai:generation', this._context, this._getTrackData(), 1);
this._ldClient.track('$ld:ai:generation:error', this._context, this._getTrackData(), 1);
}

async trackOpenAIMetrics<
Expand All @@ -60,12 +71,17 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
};
},
>(func: () => Promise<TRes>): Promise<TRes> {
const result = await this.trackDurationOf(func);
this.trackSuccess();
if (result.usage) {
this.trackTokens(createOpenAiUsage(result.usage));
try {
const result = await this.trackDurationOf(func);
this.trackSuccess();
if (result.usage) {
this.trackTokens(createOpenAiUsage(result.usage));
}
return result;
} catch (err) {
this.trackError();
throw err;
}
return result;
}

trackBedrockConverseMetrics<
Expand All @@ -82,7 +98,7 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
if (res.$metadata?.httpStatusCode === 200) {
this.trackSuccess();
} else if (res.$metadata?.httpStatusCode && res.$metadata.httpStatusCode >= 400) {
// Potentially add error tracking in the future.
this.trackError();
}
if (res.metrics && res.metrics.latencyMs) {
this.trackDuration(res.metrics.latencyMs);
Expand Down
19 changes: 19 additions & 0 deletions packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ export interface LDAIConfigTracker {
*/
trackSuccess(): void;

/**
* An error was encountered during generation.
*/
trackError(): void;

/**
* Track sentiment about the generation.
*
Expand All @@ -59,6 +64,12 @@ export interface LDAIConfigTracker {

/**
* Track the duration of execution of the provided function.
*
* If the provided function throws, then this method will also throw.
* In the case the provided function throws, this function will still record the duration.
*
* This function does not automatically record an error when the function throws.
*
* @param func The function to track the duration of.
* @returns The result of the function.
*/
Expand All @@ -67,6 +78,12 @@ export interface LDAIConfigTracker {
/**
* Track an OpenAI operation.
*
* This function will track the duration of the operation, the token usage, and the success or error status.
*
* If the provided function throws, then this method will also throw.
* In the case the provided function throws, this function will record the duration and an error.
* A failed operation will not have any token usage data.
*
* @param func Function which executes the operation.
* @returns The result of the operation.
*/
Expand All @@ -85,6 +102,8 @@ export interface LDAIConfigTracker {
/**
* Track an operation which uses Bedrock.
*
* This function will track the duration of the operation, the token usage, and the success or error status.
*
* @param res The result of the Bedrock operation.
* @returns The input operation.
*/
Expand Down
Loading