feat: Add support for tracking errors. (#715)

Author: kinyoklion

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,
+  );
+});

packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts

@@ -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 {
@@ -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<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<
@@ -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);

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.
    */