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

lib/server/ai/ai_config_tracker.rb

@@ -51,7 +51,8 @@ def initialize
       # tracker reconstructed in another process share the original runId.
       #
       class AIConfigTracker
-        attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary, :model_name, :provider_name
+        attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary, :model_name, :provider_name,
+                    :model_key, :model_version
 
         #
         # Initialize a new AIConfigTracker instance.
@@ -62,15 +63,20 @@ class AIConfigTracker
         # @param version [Integer] The version number
         # @param model_name [String] The name of the AI model being used
         # @param provider_name [String] The name of the AI provider
+        # @param model_key [String, nil] The stable, unique key of the model used
+        # @param model_version [Integer] The pinned version of the model used
         # @param context [LDContext] The context used for the flag evaluation
         #
-        def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, context:, model_name:, provider_name:)
+        def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, context:, model_name:, provider_name:,
+                       model_key: nil, model_version: 1)
           @ld_client = ld_client
           @variation_key = variation_key
           @config_key = config_key
           @version = version
           @model_name = model_name
           @provider_name = provider_name
+          @model_key = model_key
+          @model_version = model_version
           @context = context
           @summary = MetricSummary.new
           @run_id = run_id
@@ -82,7 +88,7 @@ def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, conte
         # a tracker in a different process (e.g. for deferred feedback).
         #
         # The token contains: runId, configKey, variationKey, version.
-        # modelName and providerName are NOT included.
+        # modelName, providerName, modelKey, and modelVersion are NOT included.
         #
         # @return [String] the resumption token
         #
@@ -328,8 +334,10 @@ def track_bedrock_converse_metrics(&block)
             version: @version,
             modelName: @model_name,
             providerName: @provider_name,
+            modelVersion: @model_version,
           }
           data[:variationKey] = @variation_key if @variation_key && !@variation_key.empty?
+          data[:modelKey] = @model_key if @model_key && !@model_key.empty?
           data
         end
 

lib/server/ai/client.rb

@@ -38,12 +38,14 @@ def to_h
       # The ModelConfig class represents an AI model configuration.
       #
       class ModelConfig
-        attr_reader :name
+        attr_reader :name, :model_key, :model_version
 
-        def initialize(name:, parameters: {}, custom: {})
+        def initialize(name:, parameters: {}, custom: {}, model_key: nil, model_version: nil)
           @name = name
           @parameters = parameters
           @custom = custom
+          @model_key = model_key
+          @model_version = model_version
         end
 
         #
@@ -75,11 +77,14 @@ def custom(key)
         end
 
         def to_h
-          {
+          result = {
             name: @name,
             parameters: @parameters,
             custom: @custom,
           }
+          result[:modelKey] = @model_key if @model_key && !@model_key.empty?
+          result[:modelVersion] = @model_version unless @model_version.nil?
+          result
         end
       end
 
@@ -280,20 +285,25 @@ def _completion_config(key:, context:, default:, variables: nil)
             provider_config = ProviderConfig.new(provider_config.fetch(:name, ''))
           end
 
+          tracked_model_version = 1
           if (model = variation[:model]) && model.is_a?(Hash)
             parameters = variation[:model][:parameters]
             custom = variation[:model][:custom]
+            tracked_model_version = (variation[:model][:modelVersion] || 1).to_i
             model = ModelConfig.new(
               name: variation[:model][:name],
               parameters: parameters,
-              custom: custom
+              custom: custom,
+              model_key: variation[:model][:modelKey],
+              model_version: tracked_model_version
             )
           end
 
           variation_key = variation.dig(:_ldMeta, :variationKey) || ''
           version = variation.dig(:_ldMeta, :version) || 1
           model_name = model&.name || ''
           provider_name = provider_config&.name || ''
+          model_key = model&.model_key
 
           tracker_factory = lambda {
             LaunchDarkly::Server::AI::AIConfigTracker.new(
@@ -304,6 +314,8 @@ def _completion_config(key:, context:, default:, variables: nil)
               version: version,
               model_name: model_name,
               provider_name: provider_name,
+              model_key: model_key,
+              model_version: tracked_model_version,
               context: context
             )
           }

spec/server/ai/client_spec.rb

@@ -88,6 +88,21 @@
         )
         .variation_for_all(1))
 
+    data_source.update(data_source.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))
+
     data_source
   end
 
@@ -113,6 +128,28 @@
       expect(model.custom('non-existent')).to be_nil
       expect(model.custom('name')).to be_nil
     end
+
+    it 'to_h omits modelKey and modelVersion when unset' do
+      model = described_class.new(name: 'fakeModel', parameters: { temperature: 0.5 })
+      expect(model.model_key).to be_nil
+      expect(model.model_version).to be_nil
+      result = model.to_h
+      expect(result).not_to have_key(:modelKey)
+      expect(result).not_to have_key(:modelVersion)
+    end
+
+    it 'to_h includes modelKey and modelVersion when set' do
+      model = described_class.new(name: 'fakeModel', model_key: 'my-model', model_version: 2)
+      result = model.to_h
+      expect(result[:modelKey]).to eq('my-model')
+      expect(result[:modelVersion]).to eq(2)
+    end
+
+    it 'to_h omits empty modelKey' do
+      model = described_class.new(name: 'fakeModel', model_key: '')
+      result = model.to_h
+      expect(result).not_to have_key(:modelKey)
+    end
   end
 
   describe LaunchDarkly::Server::AI::Client do
@@ -408,6 +445,64 @@
         expect(flag_data[:version]).to eq(1)
         expect(flag_data[:modelName]).to eq('fakeModel')
         expect(flag_data[:providerName]).to eq('fakeProvider')
+        expect(flag_data[:modelVersion]).to eq(1)
+        expect(flag_data).not_to have_key(:modelKey)
+      end
+
+      it 'reads modelKey and modelVersion from flag payload' do
+        context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
+        config = ai_client.completion_config(key: 'model-config-with-key-version', context:)
+
+        expect(config.model).not_to be_nil
+        expect(config.model.model_key).to eq('my-model')
+        expect(config.model.model_version).to eq(2)
+      end
+
+      it 'stamps modelKey and modelVersion on track data' do
+        context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
+        config = ai_client.completion_config(key: 'model-config-with-key-version', context:)
+
+        tracker = config.create_tracker
+        expect(ld_client).to receive(:track).with(
+          '$ld:ai:generation:success',
+          context,
+          hash_including(modelKey: 'my-model', modelVersion: 2),
+          1
+        )
+        tracker.track_success
+      end
+
+      it 'defaults modelVersion to 1 and omits modelKey when absent from payload' do
+        context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
+        config = ai_client.completion_config(key: 'model-config', context:, variables: { 'name' => 'World' })
+
+        tracker = config.create_tracker
+        flag_data = tracker.send(:flag_data)
+
+        expect(flag_data[:modelVersion]).to eq(1)
+        expect(flag_data).not_to have_key(:modelKey)
+      end
+
+      it 'uses explicit modelVersion from payload when modelKey is absent' do
+        td.update(td.flag('model-config-version-only')
+          .variations(
+            {
+              model: { name: 'gpt-4', modelVersion: 3 },
+              provider: { name: 'openai' },
+              messages: [],
+              _ldMeta: { enabled: true, variationKey: 'v1', version: 1 },
+            }
+          )
+          .variation_for_all(0))
+
+        context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
+        config = ai_client.completion_config(key: 'model-config-version-only', context:)
+
+        tracker = config.create_tracker
+        flag_data = tracker.send(:flag_data)
+
+        expect(flag_data[:modelVersion]).to eq(3)
+        expect(flag_data).not_to have_key(:modelKey)
       end
 
       it 'create_tracker returns a tracker even for disabled configs from evaluation' do
@@ -434,6 +529,8 @@
         expect(restored.send(:flag_data)[:configKey]).to eq('model-config')
         expect(restored.send(:flag_data)[:modelName]).to eq('')
         expect(restored.send(:flag_data)[:providerName]).to eq('')
+        expect(restored.send(:flag_data)[:modelVersion]).to eq(1)
+        expect(restored.send(:flag_data)).not_to have_key(:modelKey)
       end
 
       it 'each tracker has independent at-most-once tracking' do

spec/server/ai/config_tracker_spec.rb

@@ -30,7 +30,8 @@
   end
 
   let(:context) { LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) }
-  let(:tracker_flag_data) { { runId: kind_of(String), variationKey: 'test-variation', configKey: 'test-config', version: 1, modelName: 'fakeModel', providerName: 'fakeProvider' } }
+  let(:tracker_flag_data) {
+ { runId: kind_of(String), variationKey: 'test-variation', configKey: 'test-config', version: 1, modelName: 'fakeModel', providerName: 'fakeProvider', modelVersion: 1 } }
   let(:tracker) do
     described_class.new(
       ld_client: ld_client,
@@ -461,10 +462,49 @@
       expect(flag_data).to include(
         runId: kind_of(String),
         modelName: 'fakeModel',
-        providerName: 'fakeProvider'
+        providerName: 'fakeProvider',
+        modelVersion: 1
       )
       expect(flag_data[:runId]).to match(/\A[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\z/i)
     end
+
+    it 'includes modelKey when set' do
+      tracker_with_key = described_class.new(
+        ld_client: ld_client,
+        run_id: SecureRandom.uuid,
+        config_key: 'test-config',
+        context: context,
+        variation_key: 'test-variation',
+        version: 1,
+        model_name: 'fakeModel',
+        provider_name: 'fakeProvider',
+        model_key: 'my-model',
+        model_version: 2
+      )
+
+      flag_data = tracker_with_key.send(:flag_data)
+      expect(flag_data[:modelKey]).to eq('my-model')
+      expect(flag_data[:modelVersion]).to eq(2)
+    end
+
+    it 'omits modelKey when empty' do
+      tracker_with_empty_key = described_class.new(
+        ld_client: ld_client,
+        run_id: SecureRandom.uuid,
+        config_key: 'test-config',
+        context: context,
+        variation_key: 'test-variation',
+        version: 1,
+        model_name: 'fakeModel',
+        provider_name: 'fakeProvider',
+        model_key: '',
+        model_version: 3
+      )
+
+      flag_data = tracker_with_empty_key.send(:flag_data)
+      expect(flag_data).not_to have_key(:modelKey)
+      expect(flag_data[:modelVersion]).to eq(3)
+    end
   end
 
   describe '#resumption_token' do
@@ -478,12 +518,14 @@
       expect(decoded['version']).to eq(1)
     end
 
-    it 'does not include modelName or providerName' do
+    it 'does not include modelName, providerName, modelKey, or modelVersion' do
       token = tracker.resumption_token
       decoded = JSON.parse(Base64.urlsafe_decode64(token))
 
       expect(decoded).not_to have_key('modelName')
       expect(decoded).not_to have_key('providerName')
+      expect(decoded).not_to have_key('modelKey')
+      expect(decoded).not_to have_key('modelVersion')
     end
 
     it 'contains the same runId as the tracker flag data' do
@@ -512,7 +554,7 @@
       expect(restored.send(:flag_data)[:version]).to eq(1)
     end
 
-    it 'sets modelName and providerName to empty strings' do
+    it 'sets modelName and providerName to empty strings and modelVersion to 1' do
       token = tracker.resumption_token
 
       restored = described_class.from_resumption_token(
@@ -523,6 +565,8 @@
 
       expect(restored.send(:flag_data)[:modelName]).to eq('')
       expect(restored.send(:flag_data)[:providerName]).to eq('')
+      expect(restored.send(:flag_data)[:modelVersion]).to eq(1)
+      expect(restored.send(:flag_data)).not_to have_key(:modelKey)
     end
 
     it 'can track events with the restored tracker' do
@@ -542,6 +586,7 @@
         version: 1,
         modelName: '',
         providerName: '',
+        modelVersion: 1,
       }
 
       expect(ld_client).to receive(:track).with(