feat: Implement the AIClient and AITracker classes

Gemfile

@@ -1,15 +1,6 @@
 # frozen_string_literal: true
 
-source "https://rubygems.org"
+source 'https://rubygems.org'
 
 # Specify your gem's dependencies in launchdarkly-server-sdk-ai.gemspec
 gemspec
-
-gem "rake", "~> 13.0"
-
-gem "rspec", "~> 3.0"
-
-gem "rubocop", "~> 1.21"
-gem "rubocop-performance", "~> 1.15"
-gem "rubocop-rake", "~> 0.6"
-gem "rubocop-rspec", "~> 2.27"

launchdarkly-server-sdk-ai.gemspec

@@ -28,5 +28,7 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency 'rake', '~> 13.0'
   spec.add_development_dependency 'rspec', '~> 3.0'
   spec.add_development_dependency 'rubocop', '~> 1.0'
+  spec.add_development_dependency 'rubocop-performance', '~> 1.15'
+  spec.add_development_dependency 'rubocop-rake', '~> 0.6'
   spec.add_development_dependency 'rubocop-rspec', '~> 2.0'
 end

lib/ldclient-ai.rb

@@ -4,8 +4,8 @@
 require 'mustache'
 
 require 'ldclient-ai/version'
-require 'ldclient-ai/ld_ai_client'
-require 'ldclient-ai/ld_ai_config_tracker'
+require 'ldclient-ai/client'
+require 'ldclient-ai/config_tracker'
 
 module LaunchDarkly
   #

lib/ldclient-ai/ld_ai_client.rb → lib/ldclient-ai/client.rb

@@ -2,18 +2,19 @@
 
 require 'ldclient-rb'
 require 'mustache'
-require_relative 'ld_ai_config_tracker'
+require_relative 'config_tracker'
 
 module LaunchDarkly
   #
   # Namespace for the LaunchDarkly AI SDK.
   #
   module AI
+    #
     # Holds AI role and content.
+    #
     class LDMessage
       attr_reader :role, :content
 
-      # TODO: Do we need to validate the role to only be 'system', 'user', or 'assistant'?
       def initialize(role, content)
         @role = role
         @content = content
@@ -27,36 +28,42 @@ def to_h
       end
     end
 
+    #
     # The ModelConfig class represents an AI model configuration.
+    #
     class ModelConfig
-      attr_reader :name, :parameters, :custom
+      attr_reader :name
 
       def initialize(name:, parameters: {}, custom: {})
         @name = name
         @parameters = parameters
         @custom = custom
       end
 
+      #
       # Retrieve model-specific parameters.
       #
       # Accessing a named, typed attribute (e.g. name) will result in the call
       # being delegated to the appropriate property.
       #
       # @param key [String] The parameter key to retrieve
       # @return [Object] The parameter value or nil if not found
-      def get_parameter(key)
+      #
+      def parameter(key)
         return @name if key == 'name'
-        return nil if @parameters.nil?
+        return nil unless @parameters.is_a?(Hash)
 
         @parameters[key]
       end
 
+      #
       # Retrieve customer provided data.
       #
       # @param key [String] The custom key to retrieve
       # @return [Object] The custom value or nil if not found
-      def get_custom(key)
-        return nil if @custom.nil?
+      #
+      def custom(key)
+        return nil unless @custom.is_a?(Hash)
 
         @custom[key]
       end
@@ -70,7 +77,9 @@ def to_h
       end
     end
 
+    #
     # Configuration related to the provider.
+    #
     class ProviderConfig
       attr_reader :name
 
@@ -85,7 +94,9 @@ def to_h
       end
     end
 
+    #
     # The AIConfig class represents an AI configuration.
+    #
     class AIConfig
       attr_reader :enabled, :messages, :variables, :tracker, :model, :provider
 
@@ -109,7 +120,9 @@ def to_h
       end
     end
 
+    #
     # The LDAIClient class is the main entry point for the LaunchDarkly AI SDK.
+    #
     class LDAIClient
       attr_reader :logger, :ld_client
 
@@ -120,35 +133,45 @@ def initialize(ld_client)
         @logger = LaunchDarkly::AI.default_logger
       end
 
+      #
       # Retrieves the AIConfig
+      #
       # @param config_key [String] The key of the configuration flag
       # @param context [LDContext] The context used when evaluating the flag
       # @param default_value [AIConfig] 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 config(config_key, context, default_value = nil, variables = nil)
         variation = @ld_client.variation(
           config_key,
           context,
           default_value.respond_to?(:to_h) ? default_value.to_h : nil
         )
 
-        variables ||= {}
-        variables[:ldctx] = context.to_h
-
-        messages = variation.fetch(:messages, nil)
-        if messages.is_a?(Array) && messages.all? { |msg| msg.is_a?(Hash) }
-          messages = messages.map do |message|
-            message[:content] = Mustache.render(message[:content], variables) if message[:content].is_a?(String)
-            message
+        all_variables = variables ? variables.dup : {}
+        all_variables[:ldctx] = context.to_h
+
+        #
+        # Process messages and provider configuration
+        #
+        messages = nil
+        if variation[:messages].is_a?(Array) && variation[:messages].all? { |msg| msg.is_a?(Hash) }
+          messages = variation[:messages].map do |message|
+            next unless message[:content].is_a?(String)
+
+            LDMessage.new(
+              message[:role],
+              Mustache.render(message[:content], variables)
+            )
           end
         end
 
-        if (provider_config = variation.fetch(:provider, nil)) && provider_config.is_a?(Hash)
+        if (provider_config = variation[:provider]) && provider_config.is_a?(Hash)
           provider_config = ProviderConfig.new(provider_config.fetch(:name, ''))
         end
 
-        if (model = variation.fetch(:model, nil)) && model.is_a?(Hash)
+        if (model = variation[:model]) && model.is_a?(Hash)
           parameters = variation[:model][:parameters]
           custom = variation[:model][:custom]
           model = ModelConfig.new(

lib/ldclient-ai/ld_ai_config_tracker.rb → lib/ldclient-ai/config_tracker.rb

@@ -4,21 +4,27 @@
 
 module LaunchDarkly
   module AI
+    #
     # Tracks token usage for AI operations.
+    #
     class TokenUsage
       attr_reader :total, :input, :output
 
+      #
       # @param total [Integer] Total number of tokens used.
       # @param input [Integer] Number of tokens in the prompt.
       # @param output [Integer] Number of tokens in the completion.
+      #
       def initialize(total: nil, input: nil, output: nil)
         @total = total
         @input = input
         @output = output
       end
     end
 
+    #
     # Summary of metrics which have been tracked.
+    #
     class LDAIMetricSummary
       attr_accessor :duration, :success, :feedback, :usage, :time_to_first_token
 
@@ -31,7 +37,9 @@ def initialize
       end
     end
 
+    #
     # The LDAIConfigTracker class is used to track AI configuration usage.
+    #
     class LDAIConfigTracker
       attr_reader :ld_client, :config_key, :context, :variation_key, :version, :summary
 
@@ -44,8 +52,11 @@ def initialize(ld_client:, variation_key:, config_key:, version:, context:)
         @summary = LDAIMetricSummary.new
       end
 
+      #
       # Track the duration of an AI operation
+      #
       # @param duration [Integer] The duration in milliseconds
+      #
       def track_duration(duration)
         @summary.duration = duration
         @ld_client.track(
@@ -56,9 +67,12 @@ def track_duration(duration)
         )
       end
 
+      #
       # Track the duration of a block of code
+      #
       # @yield The block to measure
       # @return The result of the block
+      #
       def track_duration_of
         start_time = Time.now
         yield
@@ -67,8 +81,11 @@ def track_duration_of
         track_duration(duration)
       end
 
+      #
       # Track time to first token
+      #
       # @param duration [Integer] The duration in milliseconds
+      #
       def track_time_to_first_token(time_to_first_token)
         @summary.time_to_first_token = time_to_first_token
         @ld_client.track(
@@ -79,8 +96,11 @@ def track_time_to_first_token(time_to_first_token)
         )
       end
 
+      #
       # Track user feedback
+      #
       # @param kind [Symbol] The kind of feedback (:positive or :negative)
+      #
       def track_feedback(kind:)
         @summary.feedback = kind
         event_name = kind == :positive ? '$ld:ai:feedback:user:positive' : '$ld:ai:feedback:user:negative'
@@ -92,7 +112,9 @@ def track_feedback(kind:)
         )
       end
 
+      #
       # Track a successful AI generation
+      #
       def track_success
         @summary.success = true
         @ld_client.track(
@@ -109,7 +131,9 @@ def track_success
         )
       end
 
+      #
       # Track an error in AI generation
+      #
       def track_error
         @summary.success = false
         @ld_client.track(
@@ -126,8 +150,11 @@ def track_error
         )
       end
 
+      #
       # Track token usage
+      #
       # @param token_usage [TokenUsage] An object containing token usage details
+      #
       def track_tokens(token_usage)
         @summary.usage = token_usage
         if token_usage.total.positive?
@@ -156,27 +183,32 @@ def track_tokens(token_usage)
         )
       end
 
+      #
       # Track OpenAI-specific operations.
       # This method tracks the duration, token usage, and success/error status.
       # If the provided block raises, this method will also raise.
       # A failed operation will not have any token usage data.
       #
       # @yield The block to track.
       # @return The result of the tracked block.
+      #
       def track_openai_metrics(&block)
         result = track_duration_of(&block)
         track_success
         track_tokens(openai_to_token_usage(result[:usage])) if result[:usage]
         result
-      rescue StandardError => e
+      rescue StandardError
         track_error
-        raise e
+        raise
       end
 
+      #
       # Track AWS Bedrock conversation operations.
       # This method tracks the duration, token usage, and success/error status.
+      #
       # @param res [Hash] Response hash from Bedrock.
       # @return [Hash] The original response hash.
+      #
       def track_bedrock_converse_metrics(res)
         status_code = res.dig(:'$metadata', :httpStatusCode) || 0
         if status_code == 200
@@ -191,21 +223,19 @@ def track_bedrock_converse_metrics(res)
         res
       end
 
-      private
-
-      def flag_data
+      private def flag_data
         { variationKey: @variation_key, configKey: @config_key, version: @version }
       end
 
-      def openai_to_token_usage(usage)
+      private def openai_to_token_usage(usage)
         TokenUsage.new(
           total: usage[:total_tokens] || usage['total_tokens'],
           input: usage[:prompt_tokens] || usage['prompt_tokens'],
           output: usage[:completion_tokens] || usage['completion_tokens']
         )
       end
 
-      def bedrock_to_token_usage(usage)
+      private def bedrock_to_token_usage(usage)
         TokenUsage.new(
           total: usage[:total_tokens] || usage['total_tokens'],
           input: usage[:input_tokens] || usage['input_tokens'],