From a04361b211c17ee20d0aaa7a7862f9fe1d8df59e Mon Sep 17 00:00:00 2001 From: Anthony Torns II Date: Tue, 7 Jul 2026 19:56:15 -0700 Subject: [PATCH 1/2] feat(server-ai): stamp modelKey and modelVersion on AI usage events Read pinned model key/version from the AI Config variation payload and include them on tracker metric events to support version-aware cost attribution. Co-authored-by: Cursor --- packages/sdk/server-ai/src/ldai/client.py | 7 + packages/sdk/server-ai/src/ldai/models.py | 28 +++- packages/sdk/server-ai/src/ldai/tracker.py | 12 +- .../sdk/server-ai/tests/test_model_config.py | 124 ++++++++++++++++++ packages/sdk/server-ai/tests/test_tracker.py | 57 +++++++- 5 files changed, 219 insertions(+), 9 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 2d5ba42c..4f8c912e 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -930,21 +930,26 @@ def __evaluate( provider_config = ProviderConfig(provider.get('name', '')) model = None + tracked_model_version = 1 if 'model' in variation and isinstance(variation['model'], dict): parameters = variation['model'].get('parameters', None) custom = variation['model'].get('custom', None) region = variation['model'].get('region', None) + tracked_model_version = int(variation['model'].get('modelVersion', 1)) model = ModelConfig( name=variation['model']['name'], parameters=parameters, custom=custom, region=region, + model_key=variation['model'].get('modelKey'), + model_version=tracked_model_version, ) variation_key = variation.get('_ldMeta', {}).get('variationKey', '') version = int(variation.get('_ldMeta', {}).get('version', 1)) model_name = model.name if model else '' provider_name = provider_config.name if provider_config else '' + model_key = model.model_key if model else None def tracker_factory() -> LDAIConfigTracker: return LDAIConfigTracker( @@ -956,6 +961,8 @@ def tracker_factory() -> LDAIConfigTracker: context=context, model_name=model_name, provider_name=provider_name, + model_key=model_key, + model_version=tracked_model_version, graph_key=graph_key, ) diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index 38e4fdc9..27aeb1d6 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -58,17 +58,23 @@ def __init__( parameters: Optional[Dict[str, Any]] = None, custom: Optional[Dict[str, Any]] = None, region: Optional[str] = None, + model_key: Optional[str] = None, + model_version: Optional[int] = None, ): """ :param name: The name of the model. :param parameters: Additional model-specific parameters. :param custom: Additional customer provided data. :param region: The region the model is deployed in. + :param model_key: The stable, unique key of the model. + :param model_version: The pinned version of the model. """ self._name = name self._parameters = parameters self._custom = custom self._region = region + self._model_key = model_key + self._model_version = model_version @property def name(self) -> str: @@ -108,16 +114,36 @@ def region(self) -> Optional[str]: """ return self._region + @property + def model_key(self) -> Optional[str]: + """ + The stable, unique key of the model (used for direct lookup; distinct from + ``name``, which is not guaranteed unique). + """ + return self._model_key + + @property + def model_version(self) -> Optional[int]: + """ + The pinned version of the model that this config variation references. + """ + return self._model_version + def to_dict(self) -> dict: """ Render the given model config as a dictionary object. """ - return { + result: Dict[str, Any] = { 'name': self._name, 'parameters': self._parameters, 'custom': self._custom, 'region': self._region, } + if self._model_key: + result['modelKey'] = self._model_key + if self._model_version is not None: + result['modelVersion'] = self._model_version + return result class ProviderConfig: diff --git a/packages/sdk/server-ai/src/ldai/tracker.py b/packages/sdk/server-ai/src/ldai/tracker.py index 1cc5ee3e..382d4560 100644 --- a/packages/sdk/server-ai/src/ldai/tracker.py +++ b/packages/sdk/server-ai/src/ldai/tracker.py @@ -110,6 +110,8 @@ def __init__( context: Context, model_name: str, provider_name: str, + model_key: Optional[str] = None, + model_version: int = 1, graph_key: Optional[str] = None, ): """ @@ -123,6 +125,8 @@ def __init__( :param context: Context for evaluation. :param model_name: Name of the model used. :param provider_name: Name of the provider used. + :param model_key: Stable, unique key of the model used. + :param model_version: Pinned version of the model used. :param graph_key: When set, include ``graphKey`` in all event payloads (e.g. config-level metrics inside a graph). """ @@ -132,6 +136,8 @@ def __init__( self._version = version self._model_name = model_name self._provider_name = provider_name + self._model_key = model_key + self._model_version = model_version self._context = context self._graph_key = graph_key self._run_id = run_id @@ -147,7 +153,8 @@ def resumption_token(self) -> str: The token contains ``runId``, ``configKey``, ``version``, and optionally ``variationKey`` and ``graphKey`` (omitted when empty). - ``modelName`` and ``providerName`` are **not** included. + ``modelName``, ``providerName``, ``modelKey``, and ``modelVersion`` are + **not** included. """ data: dict = { "runId": self._run_id, @@ -216,9 +223,12 @@ def __get_track_data(self) -> dict: "version": self._version, "modelName": self._model_name, "providerName": self._provider_name, + "modelVersion": self._model_version, } if self._variation_key: data["variationKey"] = self._variation_key + if self._model_key: + data["modelKey"] = self._model_key if self._graph_key: data['graphKey'] = self._graph_key return data diff --git a/packages/sdk/server-ai/tests/test_model_config.py b/packages/sdk/server-ai/tests/test_model_config.py index 6d0f0147..b563e995 100644 --- a/packages/sdk/server-ai/tests/test_model_config.py +++ b/packages/sdk/server-ai/tests/test_model_config.py @@ -132,6 +132,23 @@ def td() -> TestData: .variation_for_all(1) ) + td.update( + td.flag('model-config-with-key-version') + .variations( + { + 'model': { + 'name': 'gpt-4', + 'modelKey': 'my-model', + 'modelVersion': 2, + }, + 'provider': {'name': 'openai'}, + 'messages': [], + '_ldMeta': {'enabled': True, 'variationKey': 'v1', 'version': 1}, + }, + ) + .variation_for_all(0) + ) + return td @@ -163,6 +180,32 @@ def test_model_config_handles_custom(): assert model.get_custom('name') is None +def test_model_config_to_dict_omits_model_key_and_version_when_unset(): + model = ModelConfig('fakeModel', parameters={'temperature': 0.5}) + assert model.model_key is None + assert model.model_version is None + result = model.to_dict() + assert 'modelKey' not in result + assert 'modelVersion' not in result + + +def test_model_config_to_dict_includes_model_key_and_version_when_set(): + model = ModelConfig( + 'fakeModel', + model_key='my-model', + model_version=2, + ) + result = model.to_dict() + assert result['modelKey'] == 'my-model' + assert result['modelVersion'] == 2 + + +def test_model_config_to_dict_omits_empty_model_key(): + model = ModelConfig('fakeModel', model_key='') + result = model.to_dict() + assert 'modelKey' not in result + + def test_uses_default_on_invalid_flag(ldai_client: LDAIClient): context = Context.create('user-key') default = AICompletionConfigDefault( @@ -560,3 +603,84 @@ def test_create_tracker_each_call_has_different_run_id(): run_id_1 = success_calls[0].args[2]['runId'] run_id_2 = success_calls[1].args[2]['runId'] assert run_id_1 != run_id_2 + + +def test_model_config_reads_model_key_and_version_from_flag(ldai_client: LDAIClient): + context = Context.create('user-key') + result = ldai_client.completion_config('model-config-with-key-version', context) + + assert result.model is not None + assert result.model.model_key == 'my-model' + assert result.model.model_version == 2 + + +def test_create_tracker_stamps_model_key_and_version_on_track_data(): + from unittest.mock import Mock + + mock_client = Mock() + mock_client.variation.return_value = { + '_ldMeta': {'enabled': True, 'variationKey': 'var-abc', 'version': 7}, + 'model': { + 'name': 'gpt-4', + 'modelKey': 'my-model', + 'modelVersion': 2, + }, + 'provider': {'name': 'openai'}, + 'messages': [] + } + + client = LDAIClient(mock_client) + context = Context.create('user-key') + + config = client.completion_config('my-config-key', context) + tracker = config.create_tracker() + tracker.track_success() + + success_calls = [ + c for c in mock_client.track.call_args_list + if c.args[0] == '$ld:ai:generation:success' + ] + assert len(success_calls) == 1 + track_data = success_calls[0].args[2] + assert track_data['modelKey'] == 'my-model' + assert track_data['modelVersion'] == 2 + + +@pytest.mark.parametrize( + 'model_payload,expected_model_key,expected_model_version', + [ + pytest.param({'name': 'gpt-4'}, None, 1, id='no_model_key_defaults_version_to_one'), + pytest.param({'name': 'gpt-4', 'modelVersion': 3}, None, 3, id='omits_model_key_when_absent'), + ], +) +def test_create_tracker_model_key_and_version_defaults( + model_payload, expected_model_key, expected_model_version, +): + from unittest.mock import Mock + + mock_client = Mock() + mock_client.variation.return_value = { + '_ldMeta': {'enabled': True, 'variationKey': 'var-abc', 'version': 7}, + 'model': model_payload, + 'provider': {'name': 'openai'}, + 'messages': [] + } + + client = LDAIClient(mock_client) + context = Context.create('user-key') + + config = client.completion_config('my-config-key', context) + tracker = config.create_tracker() + tracker.track_success() + + success_calls = [ + c for c in mock_client.track.call_args_list + if c.args[0] == '$ld:ai:generation:success' + ] + assert len(success_calls) == 1 + track_data = success_calls[0].args[2] + if expected_model_key is None: + assert 'modelKey' not in track_data + else: + assert track_data['modelKey'] == expected_model_key + assert track_data['modelVersion'] == expected_model_version diff --git a/packages/sdk/server-ai/tests/test_tracker.py b/packages/sdk/server-ai/tests/test_tracker.py index ff6201f4..8497ec5f 100644 --- a/packages/sdk/server-ai/tests/test_tracker.py +++ b/packages/sdk/server-ai/tests/test_tracker.py @@ -68,7 +68,8 @@ def test_tracks_duration(client: LDClient): "$ld:ai:duration:total", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1}, 100, ) @@ -111,7 +112,8 @@ def test_tracks_time_to_first_token(client: LDClient): "$ld:ai:tokens:ttf", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1}, 100, ) @@ -162,7 +164,8 @@ def test_tracks_token_usage(client: LDClient): tracker.track_tokens(tokens) _td = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"} + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1} calls = [ call("$ld:ai:tokens:total", context, _td, 300), call("$ld:ai:tokens:input", context, _td, 200), @@ -195,7 +198,8 @@ def test_tracks_feedback(client: LDClient, kind: FeedbackKind, label: str): f"$ld:ai:feedback:user:{label}", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1}, 1, ) assert tracker.get_summary().feedback == {"kind": kind} @@ -211,7 +215,8 @@ def test_tracks_success(client: LDClient): tracker.track_success() _std = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"} + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1} calls = [call("$ld:ai:generation:success", context, _std, 1)] client.track.assert_has_calls(calls) # type: ignore @@ -229,7 +234,8 @@ def test_tracks_error(client: LDClient): tracker.track_error() _etd2 = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"} + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1} calls = [call("$ld:ai:generation:error", context, _etd2, 1)] client.track.assert_has_calls(calls) # type: ignore @@ -252,7 +258,8 @@ def test_error_after_success_is_blocked(client: LDClient): "$ld:ai:generation:success", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1}, 1, ) @@ -267,9 +274,40 @@ def _base_td() -> dict: "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", + "modelVersion": 1, } +def test_track_data_includes_model_key_when_set(client: LDClient): + context = Context.create("user-key") + tracker = LDAIConfigTracker( + ld_client=client, run_id="test-run-id", config_key="config-key", + variation_key="variation-key", version=3, model_name="fakeModel", + provider_name="fakeProvider", context=context, + model_key="my-model", model_version=2, + ) + tracker.track_success() + + track_data = client.track.call_args[0][2] # type: ignore + assert track_data["modelKey"] == "my-model" + assert track_data["modelVersion"] == 2 + + +def test_track_data_omits_model_key_when_empty(client: LDClient): + context = Context.create("user-key") + tracker = LDAIConfigTracker( + ld_client=client, run_id="test-run-id", config_key="config-key", + variation_key="variation-key", version=3, model_name="fakeModel", + provider_name="fakeProvider", context=context, + model_key="", model_version=3, + ) + tracker.track_success() + + track_data = client.track.call_args[0][2] # type: ignore + assert "modelKey" not in track_data + assert track_data["modelVersion"] == 3 + + def test_config_tracker_includes_graph_key_when_provided(client: LDClient): context = Context.create("user-key") tracker = LDAIConfigTracker( @@ -774,6 +812,7 @@ def test_resumption_token_round_trip(client: LDClient): ld_client=client, run_id="test-run-id", config_key="cfg-key", variation_key="var-key", version=5, model_name="gpt-4", provider_name="openai", context=context, + model_key="my-model", model_version=2, ) token = tracker.resumption_token @@ -788,6 +827,8 @@ def test_resumption_token_round_trip(client: LDClient): # modelName and providerName should NOT be in the token assert "modelName" not in decoded assert "providerName" not in decoded + assert "modelKey" not in decoded + assert "modelVersion" not in decoded def test_resumption_token_has_no_padding(client: LDClient): @@ -953,6 +994,8 @@ def test_client_create_tracker_from_resumption_token(): # modelName and providerName are empty when reconstructed from token assert track_data["modelName"] == "" assert track_data["providerName"] == "" + assert track_data["modelVersion"] == 1 + assert "modelKey" not in track_data # Context should be the new one, not the original assert feedback_calls[0].args[1] == context From 372e99e80d2ff4e79f0a21a76d966f5c1d9ebb58 Mon Sep 17 00:00:00 2001 From: Anthony Torns II Date: Thu, 9 Jul 2026 14:43:52 -0700 Subject: [PATCH 2/2] fix(server-ai): omit modelVersion from events when payload lacks it Older variations are not backfilled with modelVersion, so defaulting to 1 mis-stamped events. Only include the field when the payload provides it and let the backend apply its version-1 fallback on read. Co-authored-by: Cursor --- packages/sdk/server-ai/src/ldai/client.py | 6 +++-- packages/sdk/server-ai/src/ldai/tracker.py | 7 +++--- .../sdk/server-ai/tests/test_model_config.py | 7 ++++-- packages/sdk/server-ai/tests/test_tracker.py | 24 +++++++------------ 4 files changed, 21 insertions(+), 23 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 4f8c912e..01182eaf 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -930,12 +930,14 @@ def __evaluate( provider_config = ProviderConfig(provider.get('name', '')) model = None - tracked_model_version = 1 + tracked_model_version = None if 'model' in variation and isinstance(variation['model'], dict): parameters = variation['model'].get('parameters', None) custom = variation['model'].get('custom', None) region = variation['model'].get('region', None) - tracked_model_version = int(variation['model'].get('modelVersion', 1)) + raw_model_version = variation['model'].get('modelVersion') + if raw_model_version is not None: + tracked_model_version = int(raw_model_version) model = ModelConfig( name=variation['model']['name'], parameters=parameters, diff --git a/packages/sdk/server-ai/src/ldai/tracker.py b/packages/sdk/server-ai/src/ldai/tracker.py index 382d4560..2288c03b 100644 --- a/packages/sdk/server-ai/src/ldai/tracker.py +++ b/packages/sdk/server-ai/src/ldai/tracker.py @@ -111,7 +111,7 @@ def __init__( model_name: str, provider_name: str, model_key: Optional[str] = None, - model_version: int = 1, + model_version: Optional[int] = None, graph_key: Optional[str] = None, ): """ @@ -126,7 +126,7 @@ def __init__( :param model_name: Name of the model used. :param provider_name: Name of the provider used. :param model_key: Stable, unique key of the model used. - :param model_version: Pinned version of the model used. + :param model_version: Pinned version of the model used, when present in the payload. :param graph_key: When set, include ``graphKey`` in all event payloads (e.g. config-level metrics inside a graph). """ @@ -223,12 +223,13 @@ def __get_track_data(self) -> dict: "version": self._version, "modelName": self._model_name, "providerName": self._provider_name, - "modelVersion": self._model_version, } if self._variation_key: data["variationKey"] = self._variation_key if self._model_key: data["modelKey"] = self._model_key + if self._model_version is not None: + data["modelVersion"] = self._model_version if self._graph_key: data['graphKey'] = self._graph_key return data diff --git a/packages/sdk/server-ai/tests/test_model_config.py b/packages/sdk/server-ai/tests/test_model_config.py index b563e995..11f6dc28 100644 --- a/packages/sdk/server-ai/tests/test_model_config.py +++ b/packages/sdk/server-ai/tests/test_model_config.py @@ -649,7 +649,7 @@ def test_create_tracker_stamps_model_key_and_version_on_track_data(): @pytest.mark.parametrize( 'model_payload,expected_model_key,expected_model_version', [ - pytest.param({'name': 'gpt-4'}, None, 1, id='no_model_key_defaults_version_to_one'), + pytest.param({'name': 'gpt-4'}, None, None, id='omits_model_version_when_absent'), pytest.param({'name': 'gpt-4', 'modelVersion': 3}, None, 3, id='omits_model_key_when_absent'), ], ) @@ -683,4 +683,7 @@ def test_create_tracker_model_key_and_version_defaults( assert 'modelKey' not in track_data else: assert track_data['modelKey'] == expected_model_key - assert track_data['modelVersion'] == expected_model_version + if expected_model_version is None: + assert 'modelVersion' not in track_data + else: + assert track_data['modelVersion'] == expected_model_version diff --git a/packages/sdk/server-ai/tests/test_tracker.py b/packages/sdk/server-ai/tests/test_tracker.py index 8497ec5f..9fce0185 100644 --- a/packages/sdk/server-ai/tests/test_tracker.py +++ b/packages/sdk/server-ai/tests/test_tracker.py @@ -68,8 +68,7 @@ def test_tracks_duration(client: LDClient): "$ld:ai:duration:total", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, 100, ) @@ -112,8 +111,7 @@ def test_tracks_time_to_first_token(client: LDClient): "$ld:ai:tokens:ttf", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, 100, ) @@ -164,8 +162,7 @@ def test_tracks_token_usage(client: LDClient): tracker.track_tokens(tokens) _td = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1} + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"} calls = [ call("$ld:ai:tokens:total", context, _td, 300), call("$ld:ai:tokens:input", context, _td, 200), @@ -198,8 +195,7 @@ def test_tracks_feedback(client: LDClient, kind: FeedbackKind, label: str): f"$ld:ai:feedback:user:{label}", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, 1, ) assert tracker.get_summary().feedback == {"kind": kind} @@ -215,8 +211,7 @@ def test_tracks_success(client: LDClient): tracker.track_success() _std = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1} + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"} calls = [call("$ld:ai:generation:success", context, _std, 1)] client.track.assert_has_calls(calls) # type: ignore @@ -234,8 +229,7 @@ def test_tracks_error(client: LDClient): tracker.track_error() _etd2 = {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1} + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"} calls = [call("$ld:ai:generation:error", context, _etd2, 1)] client.track.assert_has_calls(calls) # type: ignore @@ -258,8 +252,7 @@ def test_error_after_success_is_blocked(client: LDClient): "$ld:ai:generation:success", context, {"runId": ANY, "variationKey": "variation-key", "configKey": "config-key", - "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1}, + "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider"}, 1, ) @@ -274,7 +267,6 @@ def _base_td() -> dict: "version": 3, "modelName": "fakeModel", "providerName": "fakeProvider", - "modelVersion": 1, } @@ -994,7 +986,7 @@ def test_client_create_tracker_from_resumption_token(): # modelName and providerName are empty when reconstructed from token assert track_data["modelName"] == "" assert track_data["providerName"] == "" - assert track_data["modelVersion"] == 1 + assert "modelVersion" not in track_data assert "modelKey" not in track_data # Context should be the new one, not the original assert feedback_calls[0].args[1] == context