diff --git a/ldai/client.py b/ldai/client.py index 6f488f3..e8b3c1f 100644 --- a/ldai/client.py +++ b/ldai/client.py @@ -187,6 +187,7 @@ def config( self._client, variation.get('_ldMeta', {}).get('variationKey', ''), key, + int(variation.get('_ldMeta', {}).get('version', 1)), context, ) diff --git a/ldai/testing/test_model_config.py b/ldai/testing/test_model_config.py index 48fd00c..6e2c40c 100644 --- a/ldai/testing/test_model_config.py +++ b/ldai/testing/test_model_config.py @@ -15,7 +15,7 @@ def td() -> TestData: 'model': {'name': 'fakeModel', 'parameters': {'temperature': 0.5, 'maxTokens': 4096}, 'custom': {'extra-attribute': 'value'}}, 'provider': {'name': 'fakeProvider'}, 'messages': [{'role': 'system', 'content': 'Hello, {{name}}!'}], - '_ldMeta': {'enabled': True, 'variationKey': 'abcd'}, + '_ldMeta': {'enabled': True, 'variationKey': 'abcd', 'version': 1}, }, "green", ) @@ -31,7 +31,7 @@ def td() -> TestData: {'role': 'system', 'content': 'Hello, {{name}}!'}, {'role': 'user', 'content': 'The day is, {{day}}!'}, ], - '_ldMeta': {'enabled': True, 'variationKey': 'abcd'}, + '_ldMeta': {'enabled': True, 'variationKey': 'abcd', 'version': 1}, }, "green", ) @@ -44,7 +44,7 @@ def td() -> TestData: { 'model': {'name': 'fakeModel', 'parameters': {'extra-attribute': 'I can be anything I set my mind/type to'}}, 'messages': [{'role': 'system', 'content': 'Hello, {{ldctx.name}}! Is your last name {{ldctx.last}}?'}], - '_ldMeta': {'enabled': True, 'variationKey': 'abcd'}, + '_ldMeta': {'enabled': True, 'variationKey': 'abcd', 'version': 1}, } ) .variation_for_all(0) @@ -56,7 +56,7 @@ def td() -> TestData: { 'model': {'name': 'fakeModel', 'parameters': {'extra-attribute': 'I can be anything I set my mind/type to'}}, 'messages': [{'role': 'system', 'content': 'Hello, {{ldctx.user.name}}! Do you work for {{ldctx.org.shortname}}?'}], - '_ldMeta': {'enabled': True, 'variationKey': 'abcd'}, + '_ldMeta': {'enabled': True, 'variationKey': 'abcd', 'version': 1}, } ) .variation_for_all(0) @@ -68,7 +68,7 @@ def td() -> TestData: { 'model': {'name': 'fakeModel', 'parameters': {'temperature': 0.1}}, 'messages': [{'role': 'system', 'content': 'Hello, {{name}}!'}], - '_ldMeta': {'enabled': False, 'variationKey': 'abcd'}, + '_ldMeta': {'enabled': False, 'variationKey': 'abcd', 'version': 1}, } ) .variation_for_all(0) diff --git a/ldai/testing/test_tracker.py b/ldai/testing/test_tracker.py index 14215a9..3c584b7 100644 --- a/ldai/testing/test_tracker.py +++ b/ldai/testing/test_tracker.py @@ -18,7 +18,7 @@ def td() -> TestData: 'model': {'name': 'fakeModel', 'parameters': {'temperature': 0.5, 'maxTokens': 4096}, 'custom': {'extra-attribute': 'value'}}, 'provider': {'name': 'fakeProvider'}, 'messages': [{'role': 'system', 'content': 'Hello, {{name}}!'}], - '_ldMeta': {'enabled': True, 'variationKey': 'abcd'}, + '_ldMeta': {'enabled': True, 'variationKey': 'abcd', 'version': 1}, }, "green", ) @@ -38,7 +38,7 @@ def client(td: TestData) -> LDClient: def test_summary_starts_empty(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 1, context) assert tracker.get_summary().duration is None assert tracker.get_summary().feedback is None @@ -48,13 +48,13 @@ def test_summary_starts_empty(client: LDClient): def test_tracks_duration(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) tracker.track_duration(100) client.track.assert_called_with( # type: ignore '$ld:ai:duration:total', context, - {'variationKey': 'variation-key', 'configKey': 'config-key'}, + {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 100 ) @@ -63,7 +63,7 @@ def test_tracks_duration(client: LDClient): def test_tracks_duration_of(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) tracker.track_duration_of(lambda: sleep(0.01)) calls = client.track.mock_calls # type: ignore @@ -71,19 +71,19 @@ def test_tracks_duration_of(client: LDClient): 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[2] == {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3} 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 = LDAIConfigTracker(client, "variation-key", "config-key", 3, 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'}, + {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 100 ) @@ -92,7 +92,7 @@ def test_tracks_time_to_first_token(client: LDClient): def test_tracks_duration_of_with_exception(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) def sleep_and_throw(): sleep(0.01) @@ -109,21 +109,21 @@ def sleep_and_throw(): 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[2] == {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3} assert calls[0].args[3] == pytest.approx(10, rel=10) def test_tracks_token_usage(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) tokens = TokenUsage(300, 200, 100) tracker.track_tokens(tokens) calls = [ - call('$ld:ai:tokens:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 300), - call('$ld:ai:tokens:input', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 200), - call('$ld:ai:tokens:output', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 100), + call('$ld:ai:tokens:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 300), + call('$ld:ai:tokens:input', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 200), + call('$ld:ai:tokens:output', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 100), ] client.track.assert_has_calls(calls) # type: ignore @@ -133,7 +133,7 @@ def test_tracks_token_usage(client: LDClient): def test_tracks_bedrock_metrics(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) bedrock_result = { '$metadata': {'httpStatusCode': 200}, @@ -149,12 +149,12 @@ def test_tracks_bedrock_metrics(client: LDClient): tracker.track_bedrock_converse_metrics(bedrock_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:duration:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 50), - 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), + call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:success', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:duration:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 50), + call('$ld:ai:tokens:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 330), + call('$ld:ai:tokens:input', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 220), + call('$ld:ai:tokens:output', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 110), ] client.track.assert_has_calls(calls) # type: ignore @@ -166,7 +166,7 @@ def test_tracks_bedrock_metrics(client: LDClient): def test_tracks_bedrock_metrics_with_error(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) bedrock_result = { '$metadata': {'httpStatusCode': 500}, @@ -182,12 +182,12 @@ def test_tracks_bedrock_metrics_with_error(client: LDClient): tracker.track_bedrock_converse_metrics(bedrock_result) 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), - call('$ld:ai:duration:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 50), - 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), + call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:error', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:duration:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 50), + call('$ld:ai:tokens:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 330), + call('$ld:ai:tokens:input', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 220), + call('$ld:ai:tokens:output', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 110), ] client.track.assert_has_calls(calls) # type: ignore @@ -199,7 +199,7 @@ def test_tracks_bedrock_metrics_with_error(client: LDClient): def test_tracks_openai_metrics(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) class Result: def __init__(self): @@ -216,11 +216,11 @@ 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), + call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:success', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:tokens:total', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 330), + call('$ld:ai:tokens:input', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 220), + call('$ld:ai:tokens:output', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 110), ] client.track.assert_has_calls(calls, any_order=False) # type: ignore @@ -230,7 +230,7 @@ def to_dict(self): def test_tracks_openai_metrics_with_exception(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) def raise_exception(): raise ValueError("Something went wrong") @@ -242,8 +242,8 @@ def raise_exception(): 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), + call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:error', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), ] client.track.assert_has_calls(calls, any_order=False) # type: ignore @@ -260,14 +260,14 @@ def raise_exception(): ) def test_tracks_feedback(client: LDClient, kind: FeedbackKind, label: str): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) tracker.track_feedback({'kind': kind}) client.track.assert_called_with( # type: ignore f'$ld:ai:feedback:user:{label}', context, - {'variationKey': 'variation-key', 'configKey': 'config-key'}, + {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1 ) assert tracker.get_summary().feedback == {'kind': kind} @@ -275,12 +275,12 @@ def test_tracks_feedback(client: LDClient, kind: FeedbackKind, label: str): def test_tracks_success(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) tracker.track_success() 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:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:success', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), ] client.track.assert_has_calls(calls) # type: ignore @@ -290,12 +290,12 @@ def test_tracks_success(client: LDClient): def test_tracks_error(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) tracker.track_error() 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), + call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:error', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), ] client.track.assert_has_calls(calls) # type: ignore @@ -305,15 +305,15 @@ def test_tracks_error(client: LDClient): def test_error_overwrites_success(client: LDClient): context = Context.create('user-key') - tracker = LDAIConfigTracker(client, "variation-key", "config-key", context) + tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, context) tracker.track_success() tracker.track_error() 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:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 1), - call('$ld:ai:generation:error', context, {'variationKey': 'variation-key', 'configKey': 'config-key'}, 1), + call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:success', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), + call('$ld:ai:generation:error', context, {'variationKey': 'variation-key', 'configKey': 'config-key', 'version': 3}, 1), ] client.track.assert_has_calls(calls) # type: ignore diff --git a/ldai/tracker.py b/ldai/tracker.py index b29ec6f..ddaa75a 100644 --- a/ldai/tracker.py +++ b/ldai/tracker.py @@ -69,7 +69,7 @@ class LDAIConfigTracker: """ def __init__( - self, ld_client: LDClient, variation_key: str, config_key: str, context: Context + self, ld_client: LDClient, variation_key: str, config_key: str, version: int, context: Context ): """ Initialize an AI configuration tracker. @@ -77,11 +77,13 @@ def __init__( :param ld_client: LaunchDarkly client instance. :param variation_key: Variation key for tracking. :param config_key: Configuration key for tracking. + :param version: Version of the variation. :param context: Context for evaluation. """ self._ld_client = ld_client self._variation_key = variation_key self._config_key = config_key + self._version = version self._context = context self._summary = LDAIMetricSummary() @@ -94,6 +96,7 @@ def __get_track_data(self): return { 'variationKey': self._variation_key, 'configKey': self._config_key, + 'version': self._version, } def track_duration(self, duration: int) -> None: