From e6ae12ddb2e3b1b7a20f88987460e196e3abab63 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 03:18:09 +0000 Subject: [PATCH] feat(server-ai): stamp modelKey and modelVersion on AI usage events (AIC-2857) Read modelKey and modelVersion from the AI Config variation payload and expose them on ModelConfig. Stamp both fields on all AIConfigTracker metric event payloads alongside existing modelName/providerName fields. Default modelVersion to 1 when absent. Exclude both fields from the resumption token. Additive and backward compatible with older payloads. Co-authored-by: Anthony Torns II --- lib/server/ai/ai_config_tracker.rb | 14 +++- lib/server/ai/client.rb | 20 ++++-- spec/server/ai/client_spec.rb | 97 +++++++++++++++++++++++++++ spec/server/ai/config_tracker_spec.rb | 53 +++++++++++++-- 4 files changed, 173 insertions(+), 11 deletions(-) diff --git a/lib/server/ai/ai_config_tracker.rb b/lib/server/ai/ai_config_tracker.rb index e2277b7..8213148 100644 --- a/lib/server/ai/ai_config_tracker.rb +++ b/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 diff --git a/lib/server/ai/client.rb b/lib/server/ai/client.rb index 63eb6ba..431be97 100644 --- a/lib/server/ai/client.rb +++ b/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,13 +285,17 @@ 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 @@ -294,6 +303,7 @@ def _completion_config(key:, context:, default:, variables: nil) 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 ) } diff --git a/spec/server/ai/client_spec.rb b/spec/server/ai/client_spec.rb index 0a5e6e3..a399953 100644 --- a/spec/server/ai/client_spec.rb +++ b/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 diff --git a/spec/server/ai/config_tracker_spec.rb b/spec/server/ai/config_tracker_spec.rb index a11a365..c21c4e0 100644 --- a/spec/server/ai/config_tracker_spec.rb +++ b/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(