diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 2d5ba42..01182ea 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -930,21 +930,28 @@ def __evaluate( provider_config = ProviderConfig(provider.get('name', '')) model = None + 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) + 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, 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 +963,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 38e4fdc..27aeb1d 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 1cc5ee3..2288c03 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: Optional[int] = None, 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, when present in the payload. :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, @@ -219,6 +226,10 @@ def __get_track_data(self) -> dict: } 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 6d0f014..11f6dc2 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,87 @@ 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, None, id='omits_model_version_when_absent'), + 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 + 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 ff6201f..9fce018 100644 --- a/packages/sdk/server-ai/tests/test_tracker.py +++ b/packages/sdk/server-ai/tests/test_tracker.py @@ -270,6 +270,36 @@ def _base_td() -> dict: } +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 +804,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 +819,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 +986,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 "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