From 02f1d3daa711319a620a55b50481083980ab18f7 Mon Sep 17 00:00:00 2001 From: Ryan Lamb <4955475+kinyoklion@users.noreply.github.com> Date: Tue, 17 Dec 2024 15:08:19 -0800 Subject: [PATCH] feat: Add support for tracking errors. (#715) --- .../__tests__/LDAIConfigTrackerImpl.test.ts | 88 ++++++++++++++++++- .../server-ai/src/LDAIConfigTrackerImpl.ts | 38 +++++--- .../src/api/config/LDAIConfigTracker.ts | 19 ++++ 3 files changed, 131 insertions(+), 14 deletions(-) diff --git a/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts b/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts index baddd6dcd..73ccd8e2c 100644 --- a/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts +++ b/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts @@ -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); @@ -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', () => { @@ -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, + ); +}); diff --git a/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts b/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts index 17e49387b..c7906d271 100644 --- a/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts +++ b/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts @@ -30,11 +30,15 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker { async trackDurationOf(func: () => Promise): Promise { 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 { @@ -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< @@ -60,12 +71,17 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker { }; }, >(func: () => Promise): Promise { - 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< @@ -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); diff --git a/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts b/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts index 3348a13d1..9cfc55c86 100644 --- a/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts +++ b/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts @@ -50,6 +50,11 @@ export interface LDAIConfigTracker { */ trackSuccess(): void; + /** + * An error was encountered during generation. + */ + trackError(): void; + /** * Track sentiment about the generation. * @@ -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. */ @@ -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. */ @@ -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. */