Skip to content

Commit e6ae12d

Browse files
cursoragentatornsii
andcommitted
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 <atornsii@users.noreply.github.com>
1 parent 06ea7bd commit e6ae12d

4 files changed

Lines changed: 173 additions & 11 deletions

File tree

lib/server/ai/ai_config_tracker.rb

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,8 @@ def initialize
5151
# tracker reconstructed in another process share the original runId.
5252
#
5353
class AIConfigTracker
54-
attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary, :model_name, :provider_name
54+
attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary, :model_name, :provider_name,
55+
:model_key, :model_version
5556

5657
#
5758
# Initialize a new AIConfigTracker instance.
@@ -62,15 +63,20 @@ class AIConfigTracker
6263
# @param version [Integer] The version number
6364
# @param model_name [String] The name of the AI model being used
6465
# @param provider_name [String] The name of the AI provider
66+
# @param model_key [String, nil] The stable, unique key of the model used
67+
# @param model_version [Integer] The pinned version of the model used
6568
# @param context [LDContext] The context used for the flag evaluation
6669
#
67-
def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, context:, model_name:, provider_name:)
70+
def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, context:, model_name:, provider_name:,
71+
model_key: nil, model_version: 1)
6872
@ld_client = ld_client
6973
@variation_key = variation_key
7074
@config_key = config_key
7175
@version = version
7276
@model_name = model_name
7377
@provider_name = provider_name
78+
@model_key = model_key
79+
@model_version = model_version
7480
@context = context
7581
@summary = MetricSummary.new
7682
@run_id = run_id
@@ -82,7 +88,7 @@ def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, conte
8288
# a tracker in a different process (e.g. for deferred feedback).
8389
#
8490
# The token contains: runId, configKey, variationKey, version.
85-
# modelName and providerName are NOT included.
91+
# modelName, providerName, modelKey, and modelVersion are NOT included.
8692
#
8793
# @return [String] the resumption token
8894
#
@@ -328,8 +334,10 @@ def track_bedrock_converse_metrics(&block)
328334
version: @version,
329335
modelName: @model_name,
330336
providerName: @provider_name,
337+
modelVersion: @model_version,
331338
}
332339
data[:variationKey] = @variation_key if @variation_key && !@variation_key.empty?
340+
data[:modelKey] = @model_key if @model_key && !@model_key.empty?
333341
data
334342
end
335343

lib/server/ai/client.rb

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,14 @@ def to_h
3838
# The ModelConfig class represents an AI model configuration.
3939
#
4040
class ModelConfig
41-
attr_reader :name
41+
attr_reader :name, :model_key, :model_version
4242

43-
def initialize(name:, parameters: {}, custom: {})
43+
def initialize(name:, parameters: {}, custom: {}, model_key: nil, model_version: nil)
4444
@name = name
4545
@parameters = parameters
4646
@custom = custom
47+
@model_key = model_key
48+
@model_version = model_version
4749
end
4850

4951
#
@@ -75,11 +77,14 @@ def custom(key)
7577
end
7678

7779
def to_h
78-
{
80+
result = {
7981
name: @name,
8082
parameters: @parameters,
8183
custom: @custom,
8284
}
85+
result[:modelKey] = @model_key if @model_key && !@model_key.empty?
86+
result[:modelVersion] = @model_version unless @model_version.nil?
87+
result
8388
end
8489
end
8590

@@ -280,20 +285,25 @@ def _completion_config(key:, context:, default:, variables: nil)
280285
provider_config = ProviderConfig.new(provider_config.fetch(:name, ''))
281286
end
282287

288+
tracked_model_version = 1
283289
if (model = variation[:model]) && model.is_a?(Hash)
284290
parameters = variation[:model][:parameters]
285291
custom = variation[:model][:custom]
292+
tracked_model_version = (variation[:model][:modelVersion] || 1).to_i
286293
model = ModelConfig.new(
287294
name: variation[:model][:name],
288295
parameters: parameters,
289-
custom: custom
296+
custom: custom,
297+
model_key: variation[:model][:modelKey],
298+
model_version: tracked_model_version
290299
)
291300
end
292301

293302
variation_key = variation.dig(:_ldMeta, :variationKey) || ''
294303
version = variation.dig(:_ldMeta, :version) || 1
295304
model_name = model&.name || ''
296305
provider_name = provider_config&.name || ''
306+
model_key = model&.model_key
297307

298308
tracker_factory = lambda {
299309
LaunchDarkly::Server::AI::AIConfigTracker.new(
@@ -304,6 +314,8 @@ def _completion_config(key:, context:, default:, variables: nil)
304314
version: version,
305315
model_name: model_name,
306316
provider_name: provider_name,
317+
model_key: model_key,
318+
model_version: tracked_model_version,
307319
context: context
308320
)
309321
}

spec/server/ai/client_spec.rb

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,21 @@
8888
)
8989
.variation_for_all(1))
9090

91+
data_source.update(data_source.flag('model-config-with-key-version')
92+
.variations(
93+
{
94+
model: {
95+
name: 'gpt-4',
96+
modelKey: 'my-model',
97+
modelVersion: 2,
98+
},
99+
provider: { name: 'openai' },
100+
messages: [],
101+
_ldMeta: { enabled: true, variationKey: 'v1', version: 1 },
102+
}
103+
)
104+
.variation_for_all(0))
105+
91106
data_source
92107
end
93108

@@ -113,6 +128,28 @@
113128
expect(model.custom('non-existent')).to be_nil
114129
expect(model.custom('name')).to be_nil
115130
end
131+
132+
it 'to_h omits modelKey and modelVersion when unset' do
133+
model = described_class.new(name: 'fakeModel', parameters: { temperature: 0.5 })
134+
expect(model.model_key).to be_nil
135+
expect(model.model_version).to be_nil
136+
result = model.to_h
137+
expect(result).not_to have_key(:modelKey)
138+
expect(result).not_to have_key(:modelVersion)
139+
end
140+
141+
it 'to_h includes modelKey and modelVersion when set' do
142+
model = described_class.new(name: 'fakeModel', model_key: 'my-model', model_version: 2)
143+
result = model.to_h
144+
expect(result[:modelKey]).to eq('my-model')
145+
expect(result[:modelVersion]).to eq(2)
146+
end
147+
148+
it 'to_h omits empty modelKey' do
149+
model = described_class.new(name: 'fakeModel', model_key: '')
150+
result = model.to_h
151+
expect(result).not_to have_key(:modelKey)
152+
end
116153
end
117154

118155
describe LaunchDarkly::Server::AI::Client do
@@ -408,6 +445,64 @@
408445
expect(flag_data[:version]).to eq(1)
409446
expect(flag_data[:modelName]).to eq('fakeModel')
410447
expect(flag_data[:providerName]).to eq('fakeProvider')
448+
expect(flag_data[:modelVersion]).to eq(1)
449+
expect(flag_data).not_to have_key(:modelKey)
450+
end
451+
452+
it 'reads modelKey and modelVersion from flag payload' do
453+
context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
454+
config = ai_client.completion_config(key: 'model-config-with-key-version', context:)
455+
456+
expect(config.model).not_to be_nil
457+
expect(config.model.model_key).to eq('my-model')
458+
expect(config.model.model_version).to eq(2)
459+
end
460+
461+
it 'stamps modelKey and modelVersion on track data' do
462+
context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
463+
config = ai_client.completion_config(key: 'model-config-with-key-version', context:)
464+
465+
tracker = config.create_tracker
466+
expect(ld_client).to receive(:track).with(
467+
'$ld:ai:generation:success',
468+
context,
469+
hash_including(modelKey: 'my-model', modelVersion: 2),
470+
1
471+
)
472+
tracker.track_success
473+
end
474+
475+
it 'defaults modelVersion to 1 and omits modelKey when absent from payload' do
476+
context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
477+
config = ai_client.completion_config(key: 'model-config', context:, variables: { 'name' => 'World' })
478+
479+
tracker = config.create_tracker
480+
flag_data = tracker.send(:flag_data)
481+
482+
expect(flag_data[:modelVersion]).to eq(1)
483+
expect(flag_data).not_to have_key(:modelKey)
484+
end
485+
486+
it 'uses explicit modelVersion from payload when modelKey is absent' do
487+
td.update(td.flag('model-config-version-only')
488+
.variations(
489+
{
490+
model: { name: 'gpt-4', modelVersion: 3 },
491+
provider: { name: 'openai' },
492+
messages: [],
493+
_ldMeta: { enabled: true, variationKey: 'v1', version: 1 },
494+
}
495+
)
496+
.variation_for_all(0))
497+
498+
context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' })
499+
config = ai_client.completion_config(key: 'model-config-version-only', context:)
500+
501+
tracker = config.create_tracker
502+
flag_data = tracker.send(:flag_data)
503+
504+
expect(flag_data[:modelVersion]).to eq(3)
505+
expect(flag_data).not_to have_key(:modelKey)
411506
end
412507

413508
it 'create_tracker returns a tracker even for disabled configs from evaluation' do
@@ -434,6 +529,8 @@
434529
expect(restored.send(:flag_data)[:configKey]).to eq('model-config')
435530
expect(restored.send(:flag_data)[:modelName]).to eq('')
436531
expect(restored.send(:flag_data)[:providerName]).to eq('')
532+
expect(restored.send(:flag_data)[:modelVersion]).to eq(1)
533+
expect(restored.send(:flag_data)).not_to have_key(:modelKey)
437534
end
438535

439536
it 'each tracker has independent at-most-once tracking' do

spec/server/ai/config_tracker_spec.rb

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@
3030
end
3131

3232
let(:context) { LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) }
33-
let(:tracker_flag_data) { { runId: kind_of(String), variationKey: 'test-variation', configKey: 'test-config', version: 1, modelName: 'fakeModel', providerName: 'fakeProvider' } }
33+
let(:tracker_flag_data) {
34+
{ runId: kind_of(String), variationKey: 'test-variation', configKey: 'test-config', version: 1, modelName: 'fakeModel', providerName: 'fakeProvider', modelVersion: 1 } }
3435
let(:tracker) do
3536
described_class.new(
3637
ld_client: ld_client,
@@ -461,10 +462,49 @@
461462
expect(flag_data).to include(
462463
runId: kind_of(String),
463464
modelName: 'fakeModel',
464-
providerName: 'fakeProvider'
465+
providerName: 'fakeProvider',
466+
modelVersion: 1
465467
)
466468
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)
467469
end
470+
471+
it 'includes modelKey when set' do
472+
tracker_with_key = described_class.new(
473+
ld_client: ld_client,
474+
run_id: SecureRandom.uuid,
475+
config_key: 'test-config',
476+
context: context,
477+
variation_key: 'test-variation',
478+
version: 1,
479+
model_name: 'fakeModel',
480+
provider_name: 'fakeProvider',
481+
model_key: 'my-model',
482+
model_version: 2
483+
)
484+
485+
flag_data = tracker_with_key.send(:flag_data)
486+
expect(flag_data[:modelKey]).to eq('my-model')
487+
expect(flag_data[:modelVersion]).to eq(2)
488+
end
489+
490+
it 'omits modelKey when empty' do
491+
tracker_with_empty_key = described_class.new(
492+
ld_client: ld_client,
493+
run_id: SecureRandom.uuid,
494+
config_key: 'test-config',
495+
context: context,
496+
variation_key: 'test-variation',
497+
version: 1,
498+
model_name: 'fakeModel',
499+
provider_name: 'fakeProvider',
500+
model_key: '',
501+
model_version: 3
502+
)
503+
504+
flag_data = tracker_with_empty_key.send(:flag_data)
505+
expect(flag_data).not_to have_key(:modelKey)
506+
expect(flag_data[:modelVersion]).to eq(3)
507+
end
468508
end
469509

470510
describe '#resumption_token' do
@@ -478,12 +518,14 @@
478518
expect(decoded['version']).to eq(1)
479519
end
480520

481-
it 'does not include modelName or providerName' do
521+
it 'does not include modelName, providerName, modelKey, or modelVersion' do
482522
token = tracker.resumption_token
483523
decoded = JSON.parse(Base64.urlsafe_decode64(token))
484524

485525
expect(decoded).not_to have_key('modelName')
486526
expect(decoded).not_to have_key('providerName')
527+
expect(decoded).not_to have_key('modelKey')
528+
expect(decoded).not_to have_key('modelVersion')
487529
end
488530

489531
it 'contains the same runId as the tracker flag data' do
@@ -512,7 +554,7 @@
512554
expect(restored.send(:flag_data)[:version]).to eq(1)
513555
end
514556

515-
it 'sets modelName and providerName to empty strings' do
557+
it 'sets modelName and providerName to empty strings and modelVersion to 1' do
516558
token = tracker.resumption_token
517559

518560
restored = described_class.from_resumption_token(
@@ -523,6 +565,8 @@
523565

524566
expect(restored.send(:flag_data)[:modelName]).to eq('')
525567
expect(restored.send(:flag_data)[:providerName]).to eq('')
568+
expect(restored.send(:flag_data)[:modelVersion]).to eq(1)
569+
expect(restored.send(:flag_data)).not_to have_key(:modelKey)
526570
end
527571

528572
it 'can track events with the restored tracker' do
@@ -542,6 +586,7 @@
542586
version: 1,
543587
modelName: '',
544588
providerName: '',
589+
modelVersion: 1,
545590
}
546591

547592
expect(ld_client).to receive(:track).with(

0 commit comments

Comments
 (0)