feat(server-ai): stamp modelKey and modelVersion on AI usage events (AIC-2851)

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,
             )
 

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:

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

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