feat: Add getSummary method to the tracker. (#698)

launchdarkly · Dec 4, 2024 · 4df902d · 4df902d
1 parent 7a93a45
commit 4df902d
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 0 deletions.
diff --git a/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts b/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts
@@ -268,3 +268,39 @@ it('only tracks non-zero token counts', () => {
     expect.anything(),
   );
 });
+
+it('returns empty summary when no metrics tracked', () => {
+  const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, versionKey, testContext);
+
+  const summary = tracker.getSummary();
+
+  expect(summary).toEqual({});
+});
+
+it('summarizes tracked metrics', () => {
+  const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, versionKey, testContext);
+
+  tracker.trackDuration(1000);
+  tracker.trackTokens({
+    total: 100,
+    input: 40,
+    output: 60,
+  });
+  tracker.trackFeedback({ kind: LDFeedbackKind.Positive });
+  tracker.trackSuccess();
+
+  const summary = tracker.getSummary();
+
+  expect(summary).toEqual({
+    durationMs: 1000,
+    tokens: {
+      total: 100,
+      input: 40,
+      output: 60,
+    },
+    feedback: {
+      kind: 'positive',
+    },
+    success: true,
+  });
+});
diff --git a/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts b/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts
@@ -1,11 +1,14 @@
 import { LDContext } from '@launchdarkly/js-server-sdk-common';
 
 import { LDAIConfigTracker } from './api/config';
+import { LDAIMetricSummary } from './api/config/LDAIConfigTracker';
 import { createBedrockTokenUsage, LDFeedbackKind, LDTokenUsage } from './api/metrics';
 import { createOpenAiUsage } from './api/metrics/OpenAiUsage';
 import { LDClientMin } from './LDClientMin';
 
 export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
+  private _trackedMetrics: LDAIMetricSummary = {};
+
   constructor(
     private _ldClient: LDClientMin,
     private _configKey: string,
@@ -21,6 +24,7 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
   }
 
   trackDuration(duration: number): void {
+    this._trackedMetrics.durationMs = duration;
     this._ldClient.track('$ld:ai:duration:total', this._context, this._getTrackData(), duration);
   }
 
@@ -34,6 +38,7 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
   }
 
   trackFeedback(feedback: { kind: LDFeedbackKind }): void {
+    this._trackedMetrics.feedback = feedback;
     if (feedback.kind === LDFeedbackKind.Positive) {
       this._ldClient.track('$ld:ai:feedback:user:positive', this._context, this._getTrackData(), 1);
     } else if (feedback.kind === LDFeedbackKind.Negative) {
@@ -42,6 +47,7 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
   }
 
   trackSuccess(): void {
+    this._trackedMetrics.success = true;
     this._ldClient.track('$ld:ai:generation', this._context, this._getTrackData(), 1);
   }
 
@@ -88,6 +94,7 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
   }
 
   trackTokens(tokens: LDTokenUsage): void {
+    this._trackedMetrics.tokens = tokens;
     const trackData = this._getTrackData();
     if (tokens.total > 0) {
       this._ldClient.track('$ld:ai:tokens:total', this._context, trackData, tokens.total);
@@ -99,4 +106,11 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
       this._ldClient.track('$ld:ai:tokens:output', this._context, trackData, tokens.output);
     }
   }
+
+  /**
+   * Get a summary of the tracked metrics.
+   */
+  getSummary(): LDAIMetricSummary {
+    return { ...this._trackedMetrics };
+  }
 }
diff --git a/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts b/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts
@@ -1,5 +1,30 @@
 import { LDFeedbackKind, LDTokenUsage } from '../metrics';
 
+/**
+ * Metrics which have been tracked.
+ */
+export interface LDAIMetricSummary {
+  /**
+   * The duration of generation.
+   */
+  durationMs?: number;
+
+  /**
+   * Information about token usage.
+   */
+  tokens?: LDTokenUsage;
+
+  /**
+   * Was generation successful.
+   */
+  success?: boolean;
+
+  /**
+   * Any sentiment about the generation.
+   */
+  feedback?: { kind: LDFeedbackKind };
+}
+
 /**
  * The LDAIConfigTracker is used to track various details about AI operations.
  */
@@ -76,4 +101,9 @@ export interface LDAIConfigTracker {
   >(
     res: TRes,
   ): TRes;
+
+  /**
+   * Get a summary of the tracked metrics.
+   */
+  getSummary(): LDAIMetricSummary;
 }