Skip to content

Commit ce176e4

Browse files
authored
feat: Update AI tracker to include model & provider name for metrics generation (#11)
**Requirements** - [X] I have added test coverage for new or changed functionality - [ ] I have followed the repository's [pull request submission guidelines](../blob/main/CONTRIBUTING.md#submitting-pull-requests) - [ ] I have validated my changes against all supported platform versions **Related issues** Provide links to any issues in this repository or elsewhere relating to this pull request. **Describe the solution you've provided** Provide a clear and concise description of what you expect to happen. **Describe alternatives you've considered** Provide a clear and concise description of any alternative solutions or features you've considered. **Additional context** Add any other context about the pull request here.
2 parents dc13cfb + 6a95ccc commit ce176e4

File tree

4 files changed

+47
-5
lines changed

4 files changed

+47
-5
lines changed

lib/server/ai/ai_config_tracker.rb

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,13 +42,26 @@ def initialize
4242
# The AIConfigTracker class is used to track AI configuration usage.
4343
#
4444
class AIConfigTracker
45-
attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary
45+
attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary, :model_name, :provider_name
4646

47-
def initialize(ld_client:, variation_key:, config_key:, version:, context:)
47+
#
48+
# Initialize a new AIConfigTracker instance.
49+
#
50+
# @param ld_client [LDClient] The LaunchDarkly client instance
51+
# @param variation_key [String] The variation key from the flag evaluation
52+
# @param config_key [String] The configuration key
53+
# @param version [Integer] The version number
54+
# @param model_name [String] The name of the AI model being used
55+
# @param provider_name [String] The name of the AI provider
56+
# @param context [LDContext] The context used for the flag evaluation
57+
#
58+
def initialize(ld_client:, variation_key:, config_key:, version:, model_name:, provider_name:, context:)
4859
@ld_client = ld_client
4960
@variation_key = variation_key
5061
@config_key = config_key
5162
@version = version
63+
@model_name = model_name
64+
@provider_name = provider_name
5265
@context = context
5366
@summary = MetricSummary.new
5467
end
@@ -209,7 +222,13 @@ def track_bedrock_converse_metrics(&block)
209222
end
210223

211224
private def flag_data
212-
{ variationKey: @variation_key, configKey: @config_key, version: @version }
225+
{
226+
variationKey: @variation_key,
227+
configKey: @config_key,
228+
version: @version,
229+
modelName: @model_name,
230+
providerName: @provider_name,
231+
}
213232
end
214233

215234
private def openai_to_token_usage(usage)

lib/server/ai/client.rb

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,8 @@ def config(config_key, context, default_value = nil, variables = nil)
188188
variation_key: variation.dig(:_ldMeta, :variationKey) || '',
189189
config_key: config_key,
190190
version: variation.dig(:_ldMeta, :version) || 1,
191+
model_name: model&.name || '',
192+
provider_name: provider_config&.name || '',
191193
context: context
192194
)
193195

spec/server/ai/client_spec.rb

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,11 @@
213213

214214
expect(config.provider).not_to be_nil
215215
expect(config.provider.name).to eq('fakeProvider')
216+
expect(config.tracker).not_to be_nil
217+
expect(config.tracker.send(:flag_data)).to include(
218+
modelName: 'fakeModel',
219+
providerName: 'fakeProvider'
220+
)
216221
end
217222

218223
it 'interpolates context variables in messages using ldctx' do
@@ -334,6 +339,11 @@
334339
expect(config.model).to be_nil
335340
expect(config.messages).to be_nil
336341
expect(config.provider).to be_nil
342+
expect(config.tracker).not_to be_nil
343+
expect(config.tracker.send(:flag_data)).to include(
344+
modelName: '',
345+
providerName: ''
346+
)
337347
end
338348
end
339349
end

spec/server/ai/config_tracker_spec.rb

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,14 +27,16 @@
2727
end
2828

2929
let(:context) { LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) }
30-
let(:tracker_flag_data) { { variationKey: 'test-variation', configKey: 'test-config', version: 1 } }
30+
let(:tracker_flag_data) { { variationKey: 'test-variation', configKey: 'test-config', version: 1, modelName: 'fakeModel', providerName: 'fakeProvider' } }
3131
let(:tracker) do
3232
described_class.new(
3333
ld_client: ld_client,
3434
config_key: tracker_flag_data[:configKey],
3535
context: context,
3636
variation_key: tracker_flag_data[:variationKey],
37-
version: tracker_flag_data[:version]
37+
version: tracker_flag_data[:version],
38+
model_name: 'fakeModel',
39+
provider_name: 'fakeProvider'
3840
)
3941
end
4042

@@ -366,4 +368,13 @@
366368
expect(tracker.summary.time_to_first_token).to be_nil
367369
end
368370
end
371+
372+
describe '#flag_data' do
373+
it 'includes model_name and provider_name in flag data' do
374+
expect(tracker.send(:flag_data)).to include(
375+
modelName: 'fakeModel',
376+
providerName: 'fakeProvider'
377+
)
378+
end
379+
end
369380
end

0 commit comments

Comments
 (0)