Handle exceptions for track duration of and openai metrics

Author: keelerm84

ldai/testing/test_tracker.py

@@ -1,6 +1,7 @@
 from unittest.mock import MagicMock, call
 
 import pytest
+from time import sleep
 from ldclient import Config, Context, LDClient
 from ldclient.integrations.test_data import TestData
 
@@ -60,6 +61,43 @@ def test_tracks_duration(client: LDClient):
     assert tracker.get_summary().duration == 100
 
 
+def test_tracks_duration_of(client: LDClient):
+    context = Context.create('user-key')
+    tracker = LDAIConfigTracker(client, "variation-key", "config-key", context)
+    tracker.track_duration_of(lambda: sleep(0.01))
+
+    calls = client.track.mock_calls  # type: ignore
+
+    assert len(calls) == 1
+    assert calls[0].args[0] == '$ld:ai:duration:total'
+    assert calls[0].args[1] == context
+    assert calls[0].args[2] == {'variationKey': 'variation-key', 'configKey': 'config-key'}
+    assert calls[0].args[3] == pytest.approx(10)
+
+
+def test_tracks_duration_of_with_exception(client: LDClient):
+    context = Context.create('user-key')
+    tracker = LDAIConfigTracker(client, "variation-key", "config-key", context)
+
+    def sleep_and_throw():
+        sleep(0.01)
+        raise ValueError("Something went wrong")
+
+    try:
+        tracker.track_duration_of(sleep_and_throw)
+        assert False, "Should have thrown an exception"
+    except ValueError:
+        pass
+
+    calls = client.track.mock_calls  # type: ignore
+
+    assert len(calls) == 1
+    assert calls[0].args[0] == '$ld:ai:duration:total'
+    assert calls[0].args[1] == context
+    assert calls[0].args[2] == {'variationKey': 'variation-key', 'configKey': 'config-key'}
+    assert calls[0].args[3] == pytest.approx(10)
+
+
 def test_tracks_token_usage(client: LDClient):
     context = Context.create('user-key')
     tracker = LDAIConfigTracker(client, "variation-key", "config-key", context)
@@ -163,6 +201,8 @@ def to_dict(self):
     tracker.track_openai_metrics(lambda: Result())
 
     calls = [
+        call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 1),
+        call('$ld:ai:generation:success', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 1),
         call('$ld:ai:tokens:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 330),
         call('$ld:ai:tokens:input', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 220),
         call('$ld:ai:tokens:output', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 110),
@@ -173,6 +213,29 @@ def to_dict(self):
     assert tracker.get_summary().usage == TokenUsage(330, 220, 110)
 
 
+def test_tracks_openai_metrics_with_exception(client: LDClient):
+    context = Context.create('user-key')
+    tracker = LDAIConfigTracker(client, "variation-key", "config-key", context)
+
+    def raise_exception():
+        raise ValueError("Something went wrong")
+
+    try:
+        tracker.track_openai_metrics(raise_exception)
+        assert False, "Should have thrown an exception"
+    except ValueError:
+        pass
+
+    calls = [
+        call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 1),
+        call('$ld:ai:generation:error', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 1),
+    ]
+
+    client.track.assert_has_calls(calls, any_order=False)  # type: ignore
+
+    assert tracker.get_summary().usage is None
+
+
 @pytest.mark.parametrize(
     "kind,label",
     [

ldai/tracker.py

@@ -106,14 +106,20 @@ def track_duration_of(self, func):
         """
         Automatically track the duration of an AI operation.
 
+        An exception occurring during the execution of the function will still
+        track the duration. The exception will be re-thrown.
+
         :param func: Function to track.
         :return: Result of the tracked function.
         """
         start_time = time.time()
-        result = func()
-        end_time = time.time()
-        duration = int((end_time - start_time) * 1000)  # duration in milliseconds
-        self.track_duration(duration)
+        try:
+            result = func()
+        finally:
+            end_time = time.time()
+            duration = int((end_time - start_time) * 1000)  # duration in milliseconds
+            self.track_duration(duration)
+
         return result
 
     def track_feedback(self, feedback: Dict[str, FeedbackKind]) -> None:
@@ -169,9 +175,15 @@ def track_openai_metrics(self, func):
         :param func: Function to track.
         :return: Result of the tracked function.
         """
-        result = self.track_duration_of(func)
-        if hasattr(result, 'usage') and hasattr(result.usage, 'to_dict'):
-            self.track_tokens(_openai_to_token_usage(result.usage.to_dict()))
+        try:
+            result = self.track_duration_of(func)
+            self.track_success()
+            if hasattr(result, 'usage') and hasattr(result.usage, 'to_dict'):
+                self.track_tokens(_openai_to_token_usage(result.usage.to_dict()))
+        except Exception:
+            self.track_error()
+            raise
+
         return result
 
     def track_bedrock_converse_metrics(self, res: dict) -> dict: