diff --git a/examples/chatbot/aws-bedrock/hello_bedrock.rb b/examples/chatbot/aws-bedrock/hello_bedrock.rb index ecc5699..e88e911 100644 --- a/examples/chatbot/aws-bedrock/hello_bedrock.rb +++ b/examples/chatbot/aws-bedrock/hello_bedrock.rb @@ -29,10 +29,13 @@ def initialize(ai_config, bedrock_client) @bedrock_client = bedrock_client end + # Returns [response_content, resumption_token]. The resumption token can be + # used to reconstruct a tracker for deferred operations like feedback. def ask_agent(question) @messages << LaunchDarkly::Server::AI::Message.new('user', question) + tracker = @ai_config.create_tracker begin - response = ai_config.tracker.track_bedrock_converse_metrics do + response = tracker.track_bedrock_converse_metrics do @bedrock_client.converse( map_converse_arguments( ai_config.model.name, @@ -41,15 +44,17 @@ def ask_agent(question) ) end @messages << LaunchDarkly::Server::AI::Message.new('assistant', response.output.message.content[0].text) - response.output.message.content[0].text + [response.output.message.content[0].text, tracker.resumption_token] rescue StandardError => e - "An error occured: #{e.message}" + ["An error occured: #{e.message}", nil] end end - def agent_was_helpful(helpful) + def agent_was_helpful(helpful, tracker) + return if tracker.nil? + kind = helpful ? :positive : :negative - ai_config.tracker.track_feedback(kind: kind) + tracker.track_feedback(kind: kind) end def map_converse_arguments(model_id, messages) @@ -89,7 +94,7 @@ def map_converse_arguments(model_id, messages) # Pass a default for improved resiliency when the flag is unavailable or LaunchDarkly is unreachable; omit for a disabled default. # Example: -# default = LaunchDarkly::Server::AI::AIConfig.new( +# default = LaunchDarkly::Server::AI::AIConfigDefault.new( # enabled: true, # model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'my-model'), # provider: LaunchDarkly::Server::AI::ProviderConfig.new(name: 'my-provider'), @@ -105,18 +110,27 @@ def map_converse_arguments(model_id, messages) chatbot = BedrockChatbot.new(ai_config, bedrock_client) +last_resumption_token = nil + loop do print "Ask a question (or type 'exit'): " question = gets&.chomp break if question.nil? || question.strip.downcase == 'exit' - response = chatbot.ask_agent(question) + response, last_resumption_token = chatbot.ask_agent(question) puts "AI Response: #{response}" end print "Was the chat helpful? [yes/no]: " feedback = gets&.chomp -chatbot.agent_was_helpful(feedback == 'yes') unless feedback.nil? +unless feedback.nil? || last_resumption_token.nil? + tracker = LaunchDarkly::Server::AI::AIConfigTracker.from_resumption_token( + token: last_resumption_token, + ld_client:, + context: + ) + chatbot.agent_was_helpful(feedback == 'yes', tracker) +end ld_client.close diff --git a/examples/chatbot/openai/hello_openai.rb b/examples/chatbot/openai/hello_openai.rb index 39ca137..bf7bfe2 100644 --- a/examples/chatbot/openai/hello_openai.rb +++ b/examples/chatbot/openai/hello_openai.rb @@ -35,10 +35,13 @@ def initialize(ai_config, openai_client) @openai_client = openai_client end + # Returns [response_content, resumption_token]. The resumption token can be + # used to reconstruct a tracker for deferred operations like feedback. def ask_agent(question) @messages << LaunchDarkly::Server::AI::Message.new('user', question) + tracker = @ai_config.create_tracker begin - completion = ai_config.tracker.track_openai_metrics do + completion = tracker.track_openai_metrics do @openai_client.chat.completions.create( model: ai_config.model.name, messages: @messages.map(&:to_h) @@ -46,15 +49,17 @@ def ask_agent(question) end response_content = completion[:choices][0][:message][:content] @messages << LaunchDarkly::Server::AI::Message.new('assistant', response_content) - response_content + [response_content, tracker.resumption_token] rescue StandardError => e - "An error occurred: #{e.message}" + ["An error occurred: #{e.message}", nil] end end - def agent_was_helpful(helpful) + def agent_was_helpful(helpful, tracker) + return if tracker.nil? + kind = helpful ? :positive : :negative - ai_config.tracker.track_feedback(kind: kind) + tracker.track_feedback(kind: kind) end end @@ -78,7 +83,7 @@ def agent_was_helpful(helpful) # Pass a default for improved resiliency when the flag is unavailable or LaunchDarkly is unreachable; omit for a disabled default. # Example: -# default = LaunchDarkly::Server::AI::AIConfig.new( +# default = LaunchDarkly::Server::AI::AIConfigDefault.new( # enabled: true, # model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'my-model'), # provider: LaunchDarkly::Server::AI::ProviderConfig.new(name: 'my-provider'), @@ -94,18 +99,27 @@ def agent_was_helpful(helpful) chatbot = Chatbot.new(ai_config, OpenAI::Client.new(api_key: openai_api_key)) +last_resumption_token = nil + loop do print "Ask a question (or type 'exit'): " question = gets&.chomp break if question.nil? || question.strip.downcase == 'exit' - response = chatbot.ask_agent(question) + response, last_resumption_token = chatbot.ask_agent(question) puts "AI Response: #{response}" end print "Was the chat helpful? [yes/no]: " feedback = gets&.chomp -chatbot.agent_was_helpful(feedback == 'yes') unless feedback.nil? +unless feedback.nil? || last_resumption_token.nil? + tracker = LaunchDarkly::Server::AI::AIConfigTracker.from_resumption_token( + token: last_resumption_token, + ld_client:, + context: + ) + chatbot.agent_was_helpful(feedback == 'yes', tracker) +end ld_client.close diff --git a/launchdarkly-server-sdk-ai.gemspec b/launchdarkly-server-sdk-ai.gemspec index dd290f2..dc7d749 100644 --- a/launchdarkly-server-sdk-ai.gemspec +++ b/launchdarkly-server-sdk-ai.gemspec @@ -19,6 +19,7 @@ Gem::Specification.new do |spec| spec.require_paths = ['lib'] spec.required_ruby_version = '>= 3.1.0' + spec.add_dependency 'base64' spec.add_dependency 'launchdarkly-server-sdk', '~> 8.5' spec.add_dependency 'logger' spec.add_dependency 'mustache', '~> 1.1' diff --git a/lib/server/ai/ai_config_tracker.rb b/lib/server/ai/ai_config_tracker.rb index c603e46..e2277b7 100644 --- a/lib/server/ai/ai_config_tracker.rb +++ b/lib/server/ai/ai_config_tracker.rb @@ -1,5 +1,7 @@ # frozen_string_literal: true +require 'base64' +require 'json' require 'ldclient-rb' module LaunchDarkly @@ -39,7 +41,14 @@ def initialize end # - # The AIConfigTracker class is used to track AI configuration usage. + # The AIConfigTracker records metrics for a single AI run. Unless + # otherwise noted, the tracker's methods are not safe for concurrent use. + # + # All events a tracker emits share a runId (a UUIDv4) so LaunchDarkly can + # correlate them in metrics views. See individual track methods for their + # specific semantics. Call create_tracker on the AI Config to start a new + # run. A resumption token preserves the runId, so events emitted by a + # 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 @@ -55,7 +64,7 @@ class AIConfigTracker # @param provider_name [String] The name of the AI provider # @param context [LDContext] The context used for the flag evaluation # - def initialize(ld_client:, variation_key:, config_key:, version:, model_name:, provider_name:, context:) + def initialize(ld_client:, run_id:, config_key:, variation_key:, version:, context:, model_name:, provider_name:) @ld_client = ld_client @variation_key = variation_key @config_key = config_key @@ -64,14 +73,62 @@ def initialize(ld_client:, variation_key:, config_key:, version:, model_name:, p @provider_name = provider_name @context = context @summary = MetricSummary.new + @run_id = run_id + @logger = LaunchDarkly::Server::AI.default_logger + end + + # + # Returns a URL-safe Base64-encoded JSON token that can be used to reconstruct + # a tracker in a different process (e.g. for deferred feedback). + # + # The token contains: runId, configKey, variationKey, version. + # modelName and providerName are NOT included. + # + # @return [String] the resumption token + # + def resumption_token + payload = { runId: @run_id, configKey: @config_key } + payload[:variationKey] = @variation_key if @variation_key && !@variation_key.empty? + payload[:version] = @version + Base64.urlsafe_encode64(JSON.generate(payload), padding: false) + end + + # + # Reconstructs a tracker from a resumption token. + # + # @param token [String] A URL-safe Base64-encoded JSON resumption token + # @param ld_client [LDClient] The LaunchDarkly client instance + # @param context [LDContext] The context for track events + # @return [AIConfigTracker] A new tracker instance + # + def self.from_resumption_token(token:, ld_client:, context:) + json = Base64.urlsafe_decode64(token) + payload = JSON.parse(json) + + new( + ld_client: ld_client, + run_id: payload['runId'], + config_key: payload['configKey'], + variation_key: payload.fetch('variationKey', ''), + version: payload['version'], + context: context, + model_name: '', + provider_name: '' + ) end # - # Track the duration of an AI operation + # Track the duration of an AI run. + # + # Records at most once per Tracker; further calls are ignored. # # @param duration [Integer] The duration in milliseconds # def track_duration(duration) + unless @summary.duration.nil? + @logger&.warn("Skipping track_duration: duration already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}") + return + end @summary.duration = duration @ld_client.track( '$ld:ai:duration:total', @@ -96,11 +153,18 @@ def track_duration_of(&block) end # - # Track time to first token + # Track time to first token. + # + # Records at most once per Tracker; further calls are ignored. # # @param duration [Integer] The duration in milliseconds # def track_time_to_first_token(time_to_first_token) + unless @summary.time_to_first_token.nil? + @logger&.warn("Skipping track_time_to_first_token: time-to-first-token already recorded on this tracker. " \ + "Call create_tracker on the AI Config for a new run. #{flag_data}") + return + end @summary.time_to_first_token = time_to_first_token @ld_client.track( '$ld:ai:tokens:ttf', @@ -111,11 +175,17 @@ def track_time_to_first_token(time_to_first_token) end # - # Track user feedback + # Track user feedback. + # + # Records at most once per Tracker; further calls are ignored. # # @param kind [Symbol] The kind of feedback (:positive or :negative) # def track_feedback(kind:) + unless @summary.feedback.nil? + @logger&.warn("Skipping track_feedback: feedback already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}") + return + end @summary.feedback = kind event_name = kind == :positive ? '$ld:ai:feedback:user:positive' : '$ld:ai:feedback:user:negative' @ld_client.track( @@ -127,9 +197,17 @@ def track_feedback(kind:) end # - # Track a successful AI generation + # Track a successful AI generation. + # + # Records at most once per Tracker. track_success and track_error share + # state; only one of the two can record per Tracker, and subsequent + # calls are ignored. # def track_success + unless @summary.success.nil? + @logger&.warn("Skipping track_success: success/error already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}") + return + end @summary.success = true @ld_client.track( '$ld:ai:generation:success', @@ -140,9 +218,17 @@ def track_success end # - # Track an error in AI generation + # Track an error in AI generation. + # + # Records at most once per Tracker. track_success and track_error share + # state; only one of the two can record per Tracker, and subsequent + # calls are ignored. # def track_error + unless @summary.success.nil? + @logger&.warn("Skipping track_error: success/error already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}") + return + end @summary.success = false @ld_client.track( '$ld:ai:generation:error', @@ -153,11 +239,17 @@ def track_error end # - # Track token usage + # Track token usage. + # + # Records at most once per Tracker; further calls are ignored. # # @param token_usage [TokenUsage] An object containing token usage details # def track_tokens(token_usage) + unless @summary.usage.nil? + @logger&.warn("Skipping track_tokens: token usage already recorded on this tracker. Call create_tracker on the AI Config for a new run. #{flag_data}") + return + end @summary.usage = token_usage if token_usage.total.positive? @ld_client.track( @@ -191,6 +283,10 @@ def track_tokens(token_usage) # If the provided block raises, this method will also raise. # A failed operation will not have any token usage data. # + # Subsequent calls re-run the inner block but emit only metrics not + # already recorded on this Tracker. Call create_tracker on the AI + # Config to start a new run. + # # @yield The block to track. # @return The result of the tracked block. # @@ -208,6 +304,10 @@ def track_openai_metrics(&block) # Track AWS Bedrock conversation operations. # This method tracks the duration, token usage, and success/error status. # + # Subsequent calls re-run the inner block but emit only metrics not + # already recorded on this Tracker. Call create_tracker on the AI + # Config to start a new run. + # # @yield The block to track. # @return [Hash] The original response hash. # @@ -222,13 +322,15 @@ def track_bedrock_converse_metrics(&block) end private def flag_data - { - variationKey: @variation_key, + data = { + runId: @run_id, configKey: @config_key, version: @version, modelName: @model_name, providerName: @provider_name, } + data[:variationKey] = @variation_key if @variation_key && !@variation_key.empty? + data end private def openai_to_token_usage(usage) diff --git a/lib/server/ai/client.rb b/lib/server/ai/client.rb index 87d49e7..63eb6ba 100644 --- a/lib/server/ai/client.rb +++ b/lib/server/ai/client.rb @@ -2,6 +2,7 @@ require 'ldclient-rb' require 'mustache' +require 'securerandom' require_relative 'ai_config_tracker' require_relative 'sdk_info' @@ -100,26 +101,65 @@ def to_h end # - # The AIConfig class represents an AI configuration. + # The AIConfigDefault class represents a user-provided fallback AI + # configuration. # - class AIConfig - attr_reader :enabled, :messages, :tracker, :model, :provider + # Pass an instance of this class as the +default:+ parameter to + # {Client#completion_config} to control the fallback values when a flag + # is not found or cannot be evaluated. + # + # This is an input-only type: it is what an application supplies to the + # SDK, never what the SDK returns. The SDK always returns an + # {AIConfig}, which carries a tracker factory; AIConfigDefault does not. + # + class AIConfigDefault + attr_reader :enabled, :messages, :model, :provider - def initialize(enabled: nil, model: nil, messages: nil, tracker: nil, provider: nil) + # + # Returns a new AIConfigDefault with enabled: false and no model, + # messages, or provider. + # + # @return [AIConfigDefault] a new disabled fallback config + # + def self.disabled + new(enabled: false) + end + + def initialize(enabled: false, model: nil, messages: nil, provider: nil) @enabled = enabled @messages = messages - @tracker = tracker @model = model @provider = provider end - # - # Returns a new disabled AIConfig instance. - # - # @return [AIConfig] a new disabled config - # - def self.disabled - new(enabled: false) + def to_h + { + _ldMeta: { + enabled: @enabled || false, + }, + messages: @messages.is_a?(Array) ? @messages.map { |msg| msg&.to_h } : nil, + model: @model&.to_h, + provider: @provider&.to_h, + } + end + end + + # + # The AIConfig class represents an AI configuration returned by the SDK. + # + # Instances are created by {Client#completion_config} and always carry a + # tracker factory; see {#create_tracker}. Do not instantiate directly. + # For application-supplied fallback values, use {AIConfigDefault}. + # + class AIConfig + attr_reader :enabled, :messages, :model, :provider + + def initialize(tracker_factory:, enabled: nil, model: nil, messages: nil, provider: nil) + @enabled = enabled + @messages = messages + @model = model + @provider = provider + @tracker_factory = tracker_factory end def to_h @@ -132,6 +172,18 @@ def to_h provider: @provider&.to_h, } end + + # + # Creates a new tracker for a fresh AI run. Each call mints a new + # runId (a UUIDv4) that LaunchDarkly uses to correlate the tracker's + # events in metrics views. Call this once per AI run; metrics from + # different runIds cannot be combined. + # + # @return [AIConfigTracker] a new tracker instance + # + def create_tracker + @tracker_factory.call + end end # @@ -172,14 +224,26 @@ def initialize(ld_client) # # @param key [String] The key of the configuration flag # @param context [LDContext] The context used when evaluating the flag - # @param default [AIConfig] The default value to use if the flag is not found + # @param default [AIConfigDefault] The default value to use if the flag is not found # @param variables [Hash] Optional variables for rendering messages # @return [AIConfig] An AIConfig instance containing the configuration data # def completion_config(key:, context:, default: nil, variables: nil) @ld_client.track(TRACK_USAGE_COMPLETION_CONFIG, context, key, 1) - _completion_config(key:, context:, default: default || AIConfig.disabled, variables:) + _completion_config(key:, context:, default: default || AIConfigDefault.disabled, variables:) + end + + # + # Reconstructs a tracker from a resumption token, allowing deferred tracking + # (e.g. feedback from a different process). + # + # @param token [String] A resumption token obtained from AIConfigTracker#resumption_token + # @param context [LDContext] The context for track events + # @return [AIConfigTracker] A new tracker instance + # + def create_tracker(token:, context:) + AIConfigTracker.from_resumption_token(token: token, ld_client: @ld_client, context: context) end # @deprecated Use {#completion_config} instead. @@ -226,20 +290,28 @@ def _completion_config(key:, context:, default:, variables: nil) ) end - tracker = LaunchDarkly::Server::AI::AIConfigTracker.new( - ld_client: @ld_client, - variation_key: variation.dig(:_ldMeta, :variationKey) || '', - config_key: key, - version: variation.dig(:_ldMeta, :version) || 1, - model_name: model&.name || '', - provider_name: provider_config&.name || '', - context: - ) + variation_key = variation.dig(:_ldMeta, :variationKey) || '' + version = variation.dig(:_ldMeta, :version) || 1 + model_name = model&.name || '' + provider_name = provider_config&.name || '' + + tracker_factory = lambda { + LaunchDarkly::Server::AI::AIConfigTracker.new( + ld_client: @ld_client, + run_id: SecureRandom.uuid, + variation_key: variation_key, + config_key: key, + version: version, + model_name: model_name, + provider_name: provider_name, + context: context + ) + } AIConfig.new( enabled: variation.dig(:_ldMeta, :enabled) || false, messages:, - tracker:, + tracker_factory:, model:, provider: provider_config ) diff --git a/spec/server/ai/client_spec.rb b/spec/server/ai/client_spec.rb index 43fdb35..0a5e6e3 100644 --- a/spec/server/ai/client_spec.rb +++ b/spec/server/ai/client_spec.rb @@ -153,7 +153,7 @@ temperature: 0.5, maxTokens: 4096 }) messages = [LaunchDarkly::Server::AI::Message.new('system', 'Hello, {{name}}!')] - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: model, messages: messages @@ -174,7 +174,7 @@ it 'interpolates variables in model config messages' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fakeModel'), messages: [LaunchDarkly::Server::AI::Message.new('system', 'Hello, {{name}}!')] @@ -195,7 +195,7 @@ it 'returns config with messages interpolated as empty when no variables are provided' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fakeModel'), messages: [] @@ -216,7 +216,7 @@ it 'handles provider config correctly' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user', name: 'Sandy' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fake-model'), messages: [] @@ -227,8 +227,9 @@ expect(config.provider).not_to be_nil expect(config.provider.name).to eq('fakeProvider') - expect(config.tracker).not_to be_nil - expect(config.tracker.send(:flag_data)).to include( + tracker = config.create_tracker + expect(tracker).not_to be_nil + expect(tracker.send(:flag_data)).to include( modelName: 'fakeModel', providerName: 'fakeProvider' ) @@ -236,7 +237,7 @@ it 'interpolates context variables in messages using ldctx' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user', name: 'Sandy', last: 'Beaches' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fake-model'), messages: [] @@ -262,7 +263,7 @@ org_context = LaunchDarkly::LDContext.create({ key: 'org-key', kind: 'org', name: 'LaunchDarkly', shortname: 'LD' }) context = LaunchDarkly::LDContext.create_multi([user_context, org_context]) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fake-model'), messages: [] @@ -285,7 +286,7 @@ it 'handles multiple messages and variable interpolation' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fake-model'), messages: [] @@ -308,7 +309,7 @@ it 'returns disabled config when flag is off' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fake-model'), messages: [] @@ -325,7 +326,7 @@ it 'returns disabled config with nil model/messages/provider when initial config is disabled' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: true, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fake-model'), messages: [] @@ -341,7 +342,7 @@ it 'returns enabled config with nil model/messages/provider when initial config is enabled' do context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) - default = LaunchDarkly::Server::AI::AIConfig.new( + default = LaunchDarkly::Server::AI::AIConfigDefault.new( enabled: false, model: LaunchDarkly::Server::AI::ModelConfig.new(name: 'fake-model'), messages: [] @@ -353,8 +354,9 @@ expect(config.model).to be_nil expect(config.messages).to be_nil expect(config.provider).to be_nil - expect(config.tracker).not_to be_nil - expect(config.tracker.send(:flag_data)).to include( + tracker = config.create_tracker + expect(tracker).not_to be_nil + expect(tracker.send(:flag_data)).to include( modelName: '', providerName: '' ) @@ -375,15 +377,111 @@ expect(config.enabled).to be false end + + it 'create_tracker returns a new tracker with a fresh runId each time' do + context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) + config = ai_client.completion_config(key: 'model-config', context:, variables: { 'name' => 'World' }) + + tracker1 = config.create_tracker + tracker2 = config.create_tracker + + expect(tracker1).to be_a(LaunchDarkly::Server::AI::AIConfigTracker) + expect(tracker2).to be_a(LaunchDarkly::Server::AI::AIConfigTracker) + expect(tracker1).not_to equal(tracker2) + + run_id1 = tracker1.send(:flag_data)[:runId] + run_id2 = tracker2.send(:flag_data)[:runId] + expect(run_id1).not_to eq(run_id2) + expect(run_id1).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) + expect(run_id2).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 'create_tracker preserves flag metadata across calls' 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[:configKey]).to eq('model-config') + expect(flag_data[:variationKey]).to eq('abcd') + expect(flag_data[:version]).to eq(1) + expect(flag_data[:modelName]).to eq('fakeModel') + expect(flag_data[:providerName]).to eq('fakeProvider') + end + + it 'create_tracker returns a tracker even for disabled configs from evaluation' do + context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) + config = ai_client.completion_config(key: 'off-config', context:) + + expect(config.enabled).to be false + tracker = config.create_tracker + expect(tracker).to be_a(LaunchDarkly::Server::AI::AIConfigTracker) + end + + it 'round-trips a tracker through a resumption token via client create_tracker' 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 + token = tracker.resumption_token + original_run_id = tracker.send(:flag_data)[:runId] + + restored = ai_client.create_tracker(token: token, context: context) + + expect(restored).to be_a(LaunchDarkly::Server::AI::AIConfigTracker) + expect(restored.send(:flag_data)[:runId]).to eq(original_run_id) + 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('') + end + + it 'each tracker has independent at-most-once tracking' do + context = LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) + config = ai_client.completion_config(key: 'model-config', context:, variables: { 'name' => 'World' }) + + tracker1 = config.create_tracker + tracker2 = config.create_tracker + + tracker1.track_duration(100) + tracker2.track_duration(200) + + expect(tracker1.summary.duration).to eq(100) + expect(tracker2.summary.duration).to eq(200) + end end describe LaunchDarkly::Server::AI::AIConfig do - it 'disabled class method returns a disabled AIConfig' do + it 'does not expose a public disabled class method' do + expect(described_class).not_to respond_to(:disabled) + end + + it 'requires a tracker_factory to construct' do + expect { described_class.new(enabled: false) }.to raise_error(ArgumentError, /tracker_factory/) + end + + it 'is not a subclass of AIConfigDefault' do + expect(described_class.ancestors).not_to include(LaunchDarkly::Server::AI::AIConfigDefault) + end + end + + describe LaunchDarkly::Server::AI::AIConfigDefault do + it 'defaults to a disabled configuration' do + config = described_class.new + expect(config).to be_a(described_class) + expect(config.enabled).to be false + expect(config.messages).to be_nil + expect(config.model).to be_nil + expect(config.provider).to be_nil + end + + it 'disabled class method returns a disabled AIConfigDefault' do config = described_class.disabled expect(config).to be_a(described_class) expect(config.enabled).to be false expect(config.messages).to be_nil expect(config.model).to be_nil + expect(config.provider).to be_nil end it 'disabled class method returns a new instance each call' do @@ -391,6 +489,18 @@ second = described_class.disabled expect(first).not_to be(second) end + + it 'serializes to a hash matching the variation format' do + model = LaunchDarkly::Server::AI::ModelConfig.new(name: 'test-model') + messages = [LaunchDarkly::Server::AI::Message.new('system', 'Hello')] + config = described_class.new(enabled: true, model: model, messages: messages) + hash = config.to_h + + expect(hash[:_ldMeta][:enabled]).to be true + expect(hash[:model][:name]).to eq('test-model') + expect(hash[:messages].length).to eq(1) + expect(hash[:messages][0][:content]).to eq('Hello') + end end end end diff --git a/spec/server/ai/config_tracker_spec.rb b/spec/server/ai/config_tracker_spec.rb index ef470e4..a11a365 100644 --- a/spec/server/ai/config_tracker_spec.rb +++ b/spec/server/ai/config_tracker_spec.rb @@ -1,5 +1,8 @@ # frozen_string_literal: true +require 'base64' +require 'json' +require 'securerandom' require 'launchdarkly-server-sdk' require 'launchdarkly-server-sdk-ai' @@ -27,10 +30,11 @@ end let(:context) { LaunchDarkly::LDContext.create({ key: 'user-key', kind: 'user' }) } - let(:tracker_flag_data) { { 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' } } let(:tracker) do described_class.new( ld_client: ld_client, + run_id: SecureRandom.uuid, config_key: tracker_flag_data[:configKey], context: context, variation_key: tracker_flag_data[:variationKey], @@ -322,25 +326,107 @@ expect(tracker.summary.success).to be false end - it 'overwrites success with error if both are tracked' do + it 'does not track error if success has already been tracked' do expect(ld_client).to receive(:track).with( '$ld:ai:generation:success', context, tracker_flag_data, 1 ) + + tracker.track_success + expect(tracker.summary.success).to be true + tracker.track_error + expect(tracker.summary.success).to be true + end + end + + describe 'at-most-once tracking' do + it 'only tracks duration once' do expect(ld_client).to receive(:track).with( - '$ld:ai:generation:error', + '$ld:ai:duration:total', context, tracker_flag_data, - 1 - ) + 100 + ).once + tracker.track_duration(100) + tracker.track_duration(200) + expect(tracker.summary.duration).to eq(100) + end + it 'only tracks time to first token once' do + expect(ld_client).to receive(:track).with( + '$ld:ai:tokens:ttf', + context, + tracker_flag_data, + 100 + ).once + tracker.track_time_to_first_token(100) + tracker.track_time_to_first_token(200) + expect(tracker.summary.time_to_first_token).to eq(100) + end + + it 'only tracks tokens once' do + tokens1 = LaunchDarkly::Server::AI::TokenUsage.new(total: 300, input: 200, output: 100) + tokens2 = LaunchDarkly::Server::AI::TokenUsage.new(total: 600, input: 400, output: 200) + expect(ld_client).to receive(:track).with( + '$ld:ai:tokens:total', + context, + tracker_flag_data, + 300 + ).once + expect(ld_client).to receive(:track).with( + '$ld:ai:tokens:input', + context, + tracker_flag_data, + 200 + ).once + expect(ld_client).to receive(:track).with( + '$ld:ai:tokens:output', + context, + tracker_flag_data, + 100 + ).once + tracker.track_tokens(tokens1) + tracker.track_tokens(tokens2) + expect(tracker.summary.usage).to eq(tokens1) + end + + it 'only tracks success once' do + expect(ld_client).to receive(:track).with( + '$ld:ai:generation:success', + context, + tracker_flag_data, + 1 + ).once + tracker.track_success tracker.track_success expect(tracker.summary.success).to be true + end + + it 'only tracks error once' do + expect(ld_client).to receive(:track).with( + '$ld:ai:generation:error', + context, + tracker_flag_data, + 1 + ).once + tracker.track_error tracker.track_error expect(tracker.summary.success).to be false end + + it 'only tracks feedback once' do + expect(ld_client).to receive(:track).with( + '$ld:ai:feedback:user:positive', + context, + tracker_flag_data, + 1 + ).once + tracker.track_feedback(kind: :positive) + tracker.track_feedback(kind: :negative) + expect(tracker.summary.feedback).to eq(:positive) + end end describe '#summary' do @@ -370,11 +456,123 @@ end describe '#flag_data' do - it 'includes model_name and provider_name in flag data' do - expect(tracker.send(:flag_data)).to include( + it 'includes runId, model_name, and provider_name in flag data' do + flag_data = tracker.send(:flag_data) + expect(flag_data).to include( + runId: kind_of(String), modelName: 'fakeModel', providerName: 'fakeProvider' ) + 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 + end + + describe '#resumption_token' do + it 'returns a URL-safe base64-encoded JSON string' do + token = tracker.resumption_token + decoded = JSON.parse(Base64.urlsafe_decode64(token)) + + expect(decoded['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) + expect(decoded['configKey']).to eq('test-config') + expect(decoded['variationKey']).to eq('test-variation') + expect(decoded['version']).to eq(1) + end + + it 'does not include modelName or providerName' 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') + end + + it 'contains the same runId as the tracker flag data' do + flag_data = tracker.send(:flag_data) + token = tracker.resumption_token + decoded = JSON.parse(Base64.urlsafe_decode64(token)) + + expect(decoded['runId']).to eq(flag_data[:runId]) + end + end + + describe '.from_resumption_token' do + it 'reconstructs a tracker that uses the same runId' do + original_token = tracker.resumption_token + original_run_id = tracker.send(:flag_data)[:runId] + + restored = described_class.from_resumption_token( + token: original_token, + ld_client: ld_client, + context: context + ) + + expect(restored.send(:flag_data)[:runId]).to eq(original_run_id) + expect(restored.send(:flag_data)[:configKey]).to eq('test-config') + expect(restored.send(:flag_data)[:variationKey]).to eq('test-variation') + expect(restored.send(:flag_data)[:version]).to eq(1) + end + + it 'sets modelName and providerName to empty strings' do + token = tracker.resumption_token + + restored = described_class.from_resumption_token( + token: token, + ld_client: ld_client, + context: context + ) + + expect(restored.send(:flag_data)[:modelName]).to eq('') + expect(restored.send(:flag_data)[:providerName]).to eq('') + end + + it 'can track events with the restored tracker' do + token = tracker.resumption_token + original_run_id = tracker.send(:flag_data)[:runId] + + restored = described_class.from_resumption_token( + token: token, + ld_client: ld_client, + context: context + ) + + expected_data = { + runId: original_run_id, + variationKey: 'test-variation', + configKey: 'test-config', + version: 1, + modelName: '', + providerName: '', + } + + expect(ld_client).to receive(:track).with( + '$ld:ai:feedback:user:positive', + context, + expected_data, + 1 + ) + + restored.track_feedback(kind: :positive) + end + + it 'round-trips an empty variation_key as an empty string' do + empty_tracker = described_class.new( + ld_client: ld_client, + run_id: SecureRandom.uuid, + config_key: 'test-config', + context: context, + variation_key: '', + version: 1, + model_name: 'fakeModel', + provider_name: 'fakeProvider' + ) + + restored = described_class.from_resumption_token( + token: empty_tracker.resumption_token, + ld_client: ld_client, + context: context + ) + + expect(restored.variation_key).to eq('') end end