From 1fc9e769a420a22eb8a1b64a1ec9dc88eb165339 Mon Sep 17 00:00:00 2001
From: Mori Bellamy <mbellamy@launchdarkly.com>
Date: Wed, 22 Jan 2025 15:40:45 -0800
Subject: [PATCH] feat: track timeToFirstToken in LDAIConfigTracker

---
 .../__tests__/LDAIConfigTrackerImpl.test.ts          | 12 ++++++++++++
 packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts  | 10 ++++++++++
 .../server-ai/src/api/config/LDAIConfigTracker.ts    | 12 ++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts b/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts
index 73ccd8e2c..0c26f0029 100644
--- a/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts
+++ b/packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts
@@ -46,6 +46,18 @@ it('tracks duration of async function', async () => {
   );
 });
 
+it('tracks time to first token', () => {
+  const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, variationKey, testContext);
+  tracker.trackTimeToFirstToken(1000);
+
+  expect(mockTrack).toHaveBeenCalledWith(
+    '$ld:ai:tokens:ttf',
+    testContext,
+    { configKey, variationKey },
+    1000,
+  );
+});
+
 it('tracks positive feedback', () => {
   const tracker = new LDAIConfigTrackerImpl(mockLdClient, configKey, variationKey, testContext);
   tracker.trackFeedback({ kind: LDFeedbackKind.Positive });
diff --git a/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts b/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts
index c7906d271..b7884ff70 100644
--- a/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts
+++ b/packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts
@@ -41,6 +41,16 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
     }
   }
 
+  trackTimeToFirstToken(timeToFirstTokenMs: number) {
+    this._trackedMetrics.timeToFirstTokenMs = timeToFirstTokenMs;
+    this._ldClient.track(
+      '$ld:ai:tokens:ttf',
+      this._context,
+      this._getTrackData(),
+      timeToFirstTokenMs,
+    );
+  }
+
   trackFeedback(feedback: { kind: LDFeedbackKind }): void {
     this._trackedMetrics.feedback = feedback;
     if (feedback.kind === LDFeedbackKind.Positive) {
diff --git a/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts b/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts
index 9cfc55c86..2f92aa938 100644
--- a/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts
+++ b/packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts
@@ -23,6 +23,11 @@ export interface LDAIMetricSummary {
    * Any sentiment about the generation.
    */
   feedback?: { kind: LDFeedbackKind };
+
+  /**
+   * Time to first token for this generation.
+   */
+  timeToFirstTokenMs?: number;
 }
 
 /**
@@ -62,6 +67,13 @@ export interface LDAIConfigTracker {
    */
   trackFeedback(feedback: { kind: LDFeedbackKind }): void;
 
+  /**
+   * Track the time to first token for this generation.
+   *
+   * @param timeToFirstTokenMs The duration in milliseconds.
+   */
+  trackTimeToFirstToken(timeToFirstTokenMs: number): void;
+
   /**
    * Track the duration of execution of the provided function.
    *