feat: Add ability to track time to first token for LDAIConfigTracker (#37)

Author: moribellamy

ldai/testing/test_tracker.py

@@ -75,6 +75,21 @@ def test_tracks_duration_of(client: LDClient):
     assert calls[0].args[3] == pytest.approx(10, rel=10)
 
 
+def test_tracks_time_to_first_token(client: LDClient):
+    context = Context.create('user-key')
+    tracker = LDAIConfigTracker(client, "variation-key", "config-key", context)
+    tracker.track_time_to_first_token(100)
+
+    client.track.assert_called_with(  # type: ignore
+        '$ld:ai:tokens:ttf',
+        context,
+        {'variationKey': 'variation-key', 'configKey': 'config-key'},
+        100
+    )
+
+    assert tracker.get_summary().time_to_first_token == 100
+
+
 def test_tracks_duration_of_with_exception(client: LDClient):
     context = Context.create('user-key')
     tracker = LDAIConfigTracker(client, "variation-key", "config-key", context)

ldai/tracker.py

@@ -40,6 +40,7 @@ def __init__(self):
         self._success = None
         self._feedback = None
         self._usage = None
+        self._time_to_first_token = None
 
     @property
     def duration(self) -> Optional[int]:
@@ -57,6 +58,10 @@ def feedback(self) -> Optional[Dict[str, FeedbackKind]]:
     def usage(self) -> Optional[TokenUsage]:
         return self._usage
 
+    @property
+    def time_to_first_token(self) -> Optional[int]:
+        return self._time_to_first_token
+
 
 class LDAIConfigTracker:
     """
@@ -102,6 +107,17 @@ def track_duration(self, duration: int) -> None:
             '$ld:ai:duration:total', self._context, self.__get_track_data(), duration
         )
 
+    def track_time_to_first_token(self, time_to_first_token: int) -> None:
+        """
+        Manually track the time to first token of an AI operation.
+
+        :param time_to_first_token: Time to first token in milliseconds.
+        """
+        self._summary._time_to_first_token = time_to_first_token
+        self._ld_client.track(
+            '$ld:ai:tokens:ttf', self._context, self.__get_track_data(), time_to_first_token
+        )
+
     def track_duration_of(self, func):
         """
         Automatically track the duration of an AI operation.