diff --git a/ldai/testing/test_tracker.py b/ldai/testing/test_tracker.py index 3196bfb..14215a9 100644 --- a/ldai/testing/test_tracker.py +++ b/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) diff --git a/ldai/tracker.py b/ldai/tracker.py index 8f3c15c..b29ec6f 100644 --- a/ldai/tracker.py +++ b/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.