stretchy-model 0.6.6 → 0.7.0

data/lib/stretchy/machine_learning/connector.rb

@@ -0,0 +1,130 @@
+module Stretchy
+  module MachineLearning
+    class Connector
+
+      cattr_reader :client do 
+        Stretchy.configuration.client.connector
+      end
+
+      class Settings
+        def initialize(valid_keys = [])
+          @valid_keys = valid_keys
+          @settings = {}
+        end
+
+        def as_json(*)
+          @settings.as_json
+        end
+      
+        def method_missing(method, *args, &block)
+          if block_given?
+            @settings[method] = self.class.new.tap { |obj| obj.instance_eval(&block) }
+          elsif args.empty?
+            value = @settings[method]
+            value.is_a?(Hash) ? Elasticsearch::Model::HashWrapper[value] : value
+          elsif args.length == 1 && @valid_keys.empty? || @valid_keys.include?(method)
+            @settings[method] = args.first
+          else
+            super
+          end
+        end
+      
+        def respond_to_missing?(method, include_private = false)
+          @settings.key?(method) || super
+        end
+      end
+      
+      class << self
+        include Errors
+        
+        METHODS = [
+          :description,
+          :version,
+          :protocol,
+          :credentials,
+          :parameters,
+          :actions,
+          :name
+        ].freeze
+        
+        def settings
+          @settings ||= Settings.new(METHODS)
+        end
+    
+        delegate *METHODS, to: :settings
+
+        def name(name = nil)
+          settings.name(name) if name 
+          settings.name || to_s.split('::').last.underscore
+        end
+
+        def registry
+          @registery ||= Stretchy::MachineLearning::Registry.register(class_name: self.name, class_type: 'connector')
+        end
+
+        def id
+          @id || registry.model_id
+        end
+
+        concerning :API do
+          def create!
+            response = client.post(body: self.to_hash.as_json).with_indifferent_access
+            registry.update(model_id: response.dig(:connector_id))
+            @id = response.dig(:connector_id)
+          end
+
+          def delete!
+            begin
+              response = client.delete(connector_id: self.id).with_indifferent_access
+            rescue "#{Stretchy.configuration.search_backend_const}::Transport::Transport::Errors::NotFound".constantize => e
+              raise Stretchy::MachineLearning::Errors::ConnectorMissingError
+              
+            ensure
+              if response.dig(:result) == 'deleted'
+                registry.delete
+                @registry = nil
+                @id = nil
+              end
+              response
+            end
+          end
+
+          def update!
+            client.put(connector_id: self.id, body: self.to_hash)
+          end
+
+          def find(id = nil)
+            begin
+              id = self.id if id.nil?
+              client.get(connector_id: id)
+            rescue ArgumentError => e
+              raise Stretchy::MachineLearning::Errors::ConnectorMissingError
+            rescue "#{Stretchy.configuration.search_backend_const}::Transport::Transport::Errors::NotFound".constantize => e
+              # raise Stretchy::MachineLearning::Errors::ConnectorMissingError
+            end
+          end
+
+          def exists?
+            self.find(self.id).present?
+          end
+
+          def to_hash
+            acts = self.actions.as_json
+            acts[:request_body] = actions[:request_body]
+            {
+              name: self.name,
+              description: self.description,
+              version: self.version,
+              protocol: self.protocol,
+              credential: self.credentials,
+              parameters: self.parameters,
+              actions: [acts]
+            }.compact
+
+          end
+        end
+     
+      end
+    end
+  end
+end

data/lib/stretchy/machine_learning/errors.rb

@@ -0,0 +1,25 @@
+module Stretchy::MachineLearning::Errors
+  class ModelNotDeployedError < StandardError
+    def initialize(msg="Model is not deployed. Please run `rake stretchy:ml:deploy` to deploy it.")
+      super
+    end
+  end
+
+  class ModelNotRegisteredError < StandardError
+    def initialize(msg="Model is not registered. Please run `rake stretchy:ml:register` to register it.")
+      super
+    end
+  end
+
+  class ConnectorMissingError < StandardError
+    def initialize(msg="Connector is missing. Please run `rake stretchy:ml:connector:create` to create it.")
+      super
+    end
+  end
+
+  class AgentNotRegisteredError < StandardError
+    def initialize(msg="Agent is not registered. Please run `rake stretchy:ml:agent:create` to create it.")
+      super
+    end
+  end
+end

data/lib/stretchy/machine_learning/model.rb

@@ -1,6 +1,9 @@
+require 'stretchy/machine_learning/errors'
+
 module Stretchy 
   module MachineLearning
     class Model
+
       PRETRAINED_MODELS = {
         :neural_sparse => {
           :encoding => 'amazon/neural-sparse/opensearch-neural-sparse-encoding-v1',
@@ -31,7 +34,165 @@ module Stretchy
       end
 
       class << self
-        attr_accessor :model, :group_id 
+        include Errors
+        # delegate :find, :status, :deployed?, :registered?, :task_id, :deploy_id, :model_id, :register, :deploy, :undeploy, :delete, to: :model
+
+        METHODS = [
+                    :model,
+                    :group_id, 
+                    :version, 
+                    :description, 
+                    :model_format, 
+                    :enabled, 
+                    :connector_id, 
+                    :connector, 
+                    :function_name, 
+                    :model_config, 
+                    :model_content_hash_value, 
+                    :url  
+                  ]
+
+        def settings
+          @settings ||= {}
+        end
+
+        METHODS.each do |method|
+          define_method(method) do |args = nil|
+            return settings[method] unless args.present?
+            settings[method] = args
+
+            if method == :connector
+              connector_class = "#{args.to_s.camelize}".constantize
+              settings[:connector] = connector_class
+              # raise ConnectorMissingError if connector_class.id.nil?
+              settings[:connector_id] = connector_class.id
+            end
+
+            if method == :model
+              settings[:model] = model_lookup(args)
+            end
+          end
+        end
+
+        def model_id
+          @model_id || registry.model_id
+        end
+
+        def task_id
+          @task_id || registry.register_task_id
+        end
+
+        def deploy_id
+          @deploy_id || registry.deploy_task_id
+        end
+
+
+        def registry
+          @registry ||= Stretchy::MachineLearning::Registry.register(class_name: self.name, class_type: 'model')
+        end
+
+        def register
+          begin
+  
+            response = client.register(body: self.to_hash, deploy: true)
+  
+            @task_id = response['task_id']
+  
+            self.registry.update(register_task_id: @task_id)
+  
+            yield self if block_given? 
+  
+            registered?
+            self.registry.update(model_id: @model_id) 
+            @model_id
+          rescue => e
+            Stretchy.logger.error "Error registering model: #{e.message}"
+            false
+          end
+          true
+        end
+
+        def registered?
+          return false unless task_id
+          response = status
+          @model_id = response['model_id'] if response['model_id']
+          response['state'] == 'COMPLETED' && @model_id.present?
+        end
+  
+        def status
+          client.get_status(task_id: self.task_id)
+        end
+
+
+        def deploy
+          @deployed = nil
+
+          @deploy_id = client.deploy(id: self.model_id)['task_id']
+          self.registry.update(deploy_task_id: @deploy_id)
+          yield self if block_given? 
+          @deploy_id
+        end
+  
+        def undeploy
+          @deployed = nil
+          response = client.undeploy(id: self.model_id)
+          self.registry.update(deploy_task_id: nil)
+          yield self if block_given? 
+          response
+        end
+  
+        def deployed?
+          return @deployed if @deployed
+          response = client.get_model(id: self.model_id)
+          # raise "Model not deployed" if response['model_state'] == 'FAILED'
+          @deployed = response['model_state'] == 'DEPLOYED'
+        end
+  
+        def delete
+          self.registry.delete
+          client.delete_model(id: self.model_id)
+        end
+
+        def find 
+          begin
+            client.get_model(id: self.model_id)
+          rescue "#{Stretchy.configuration.search_backend_const}::Transport::Transport::Errors::InternalServerError".constantize => e
+            raise Stretchy::MachineLearning::Errors::ModelMissingError
+          end
+        end
+
+        def model_name(model_name = nil)
+          @model_name = model_name if model_name
+          @model_name || to_s.demodulize.underscore 
+        end
+  
+        def to_hash
+          {
+            name: self.model || self.model_name,
+            model_group_id: self.group_id,
+            version: self.version,
+            description: self.description,
+            model_format: self.model_format,
+            is_enabled: self.enabled?,
+            connector_id: self.connector.present? ? self.connector.id : nil,
+            function_name: self.function_name
+          }.compact
+        end
+  
+        def enabled?
+          self.enabled
+        end
+  
+        def wait_until_complete(max_attempts: 20, sleep_time: 4)
+          attempts = 0
+          loop do
+            result = yield
+            break if result
+            attempts += 1
+            break if attempts >= max_attempts
+            sleep(sleep_time)
+          end
+        end
 
         def all
           client.get_model
@@ -78,114 +239,6 @@ module Stretchy
         end
       end
 
-      attr_accessor :model,
-                    :group_id, 
-                    :version, 
-                    :description, 
-                    :model_format, 
-                    :enabled, 
-                    :connector_id, 
-                    :connector, 
-                    :function_name, 
-                    :model_config, 
-                    :model_content_hash_value, 
-                    :url
-
-      attr_reader :task_id, :model_id, :deploy_id
-
-      def initialize(args = {})
-        model_name = args.delete(:model)
-        args.each do |k,v|
-          self.send("#{k}=", v)
-        end
-        @model = self.class.model_lookup model_name
-      end
-
-      def register
-        begin
-          response = client.register(body: self.to_hash, deploy: true)
-
-          @task_id = response['task_id']
-
-          yield self if block_given? 
-
-          @model_id
-        rescue => e
-          Stretchy.logger.error "Error registering model: #{e.message}"
-          false
-        end
-        true
-      end
-
-      def registered?
-        response = status
-        @model_id = response['model_id'] if response['model_id']
-        response['state'] == 'COMPLETED' && @model_id.present?
-      end
-
-      def status
-        client.get_status(task_id: self.task_id)
-      end
-
-      def deploy
-        @deployed = nil
-        @deploy_id = client.deploy(id: self.model_id)['task_id']
-        yield self if block_given? 
-        @deploy_id
-      end
-
-      def undeploy
-        @deployed = nil
-        response = client.undeploy(id: self.model_id)
-        yield self if block_given? 
-        response
-      end
-
-      def deployed?
-        return @deployed if @deployed
-        response = client.get_model(id: self.model_id)
-        # raise "Model not deployed" if response['model_state'] == 'FAILED'
-        @deployed = response['model_state'] == 'DEPLOYED'
-      end
-
-      def delete
-        client.delete_model(id: self.model_id)
-      end
-
-      def client
-        @@client
-      end
-
-      def find 
-        client.get_model(id: self.model_id)
-      end
-
-      def to_hash
-        {
-          name: self.model,
-          model_group_id: self.group_id,
-          version: self.version,
-          description: self.description,
-          model_format: self.model_format,
-          is_enabled: self.enabled?
-        }.compact
-      end
-
-      def enabled?
-        self.enabled
-      end
-
-      def wait_until_complete(max_attempts: 20, sleep_time: 4)
-        attempts = 0
-        loop do
-          result = yield
-          break if result
-          attempts += 1
-          break if attempts >= max_attempts
-          sleep(sleep_time)
-        end
-      end
-
 
     end
   end

data/lib/stretchy/machine_learning/registry.rb

@@ -0,0 +1,19 @@
+module Stretchy::MachineLearning
+  class Registry < StretchyModel
+
+    index_name ".stretchy_ml_registry_#{Stretchy.env}"
+
+    attribute :model_id, :keyword
+    attribute :model_group_id, :keyword
+    attribute :deploy_task_id, :keyword
+    attribute :register_task_id, :keyword
+    attribute :class_name, :keyword
+    attribute :class_type, :keyword
+
+    def self.register(**args) 
+      self.create_index! unless index_exists?
+      where(class_name: args[:class_name]).first || create(**args)
+    end
+
+  end
+end

data/lib/stretchy/open_search_compatibility.rb

@@ -1,4 +1,5 @@
 require 'opensearch/api/namespace/machine_learning/model'
+require 'opensearch/api/namespace/connector'
 
 module Stretchy
     module OpenSearchCompatibility
@@ -82,6 +83,7 @@ module Stretchy
             ::Elasticsearch::Persistence::Repository.send(:include, patch)
             ::Elasticsearch::Persistence::Repository.send(:include, store)
             OpenSearch::API.send(:include, OpenSearch::API::MachineLearning::Models)
+            OpenSearch::API.send(:include, OpenSearch::API::Connector)
 
         end
 

data/lib/stretchy/pipelines/processor.rb

@@ -48,9 +48,11 @@ module Stretchy::Pipelines
       @type = type
       @description = opts[:description]
       @opts = opts
+      @model = opts.delete(:model)
     end
 
     def to_hash
+      @opts[:model_id] = @model.model_id if @model
       { type => @opts }
     end
   end

data/lib/stretchy/rails/railtie.rb

@@ -0,0 +1,11 @@
+module Stretchy
+  module Rails
+    class Railtie < ::Rails::Railtie
+
+      rake_tasks do
+        load "#{__dir__}/tasks/stretchy.rake"
+      end
+
+    end
+  end
+end

data/lib/stretchy/rails/tasks/connector/create.rake

@@ -0,0 +1,32 @@
+namespace :stretchy do
+  namespace :connector do
+    desc "Create connector"
+    task create: :environment do
+      klass = ENV['MODEL']
+      Rails.application.eager_load!
+      models = klass.nil? ? Stretchy::MachineLearning::Connector.descendants : [klass.constantize]
+        models.each do |model|
+          spinner = TTY::Spinner.new("Creating Connector #{model} ".ljust(JUSTIFICATION) + ":spinner", format: :dots)
+          spinner.auto_spin
+
+          # if model.exists?
+          #   spinner.stop(Rainbow("[EXISTS]").yellow)
+          #   next
+          # end
+
+          begin
+            response = model.create!
+          rescue Stretchy::MachineLearning::Errors::ConnectorMissingError => e
+            spinner.stop(Rainbow("[FAILED]").red)
+            next
+          end
+          status = if response
+            Rainbow("[SUCCESS]").green
+          else
+            Rainbow("[FAILED]").red
+          end
+          spinner.stop(status)
+        end
+      end
+  end
+end

data/lib/stretchy/rails/tasks/connector/delete.rake

@@ -0,0 +1,27 @@
+namespace :stretchy do
+  namespace :connector do
+    desc "Delete Connector"
+    task delete: :environment do
+      klass = ENV['MODEL']
+      Rails.application.eager_load!
+      models = klass.nil? ? Stretchy::MachineLearning::Connector.descendants : [klass.constantize]
+        models.each do |model|
+          spinner = TTY::Spinner.new("Deleting Connector #{model} ".ljust(JUSTIFICATION) + ":spinner", format: :dots)
+          spinner.auto_spin
+
+          begin
+            response = model.delete!
+          rescue Stretchy::MachineLearning::Errors::ConnectorMissingError => e
+            spinner.stop(Rainbow("[MISSING]").white)
+            next
+          end
+          status = if response['result'] == 'deleted'
+            Rainbow("[SUCCESS]").green
+          else
+            Rainbow("[#{response['result'].upcase}]").white
+          end
+          spinner.stop(status)
+        end
+      end
+  end
+end

data/lib/stretchy/rails/tasks/connector/status.rake

@@ -0,0 +1,31 @@
+namespace :stretchy do
+  namespace :connector do
+
+    desc "Check the status of all connectors"
+    task status: :environment do
+      klass = ENV['MODEL']
+
+      Rails.application.eager_load!
+      models = klass.nil? ? Stretchy::MachineLearning::Connector.descendants : [klass.constantize]
+
+      puts Rainbow("Connector").color :gray
+      models.each do |model|
+        
+        begin
+        response = model.find
+        rescue Stretchy::MachineLearning::Errors::ConnectorMissingError => e
+          puts "#{model}".ljust(JUSTIFICATION) + Rainbow("[MISSING]").white
+          next
+        end
+        status = if response
+          Rainbow("[CREATED]").green.bright
+        else
+          Rainbow("[MISSING]").white
+        end
+        puts "#{model}".ljust(JUSTIFICATION) + status
+      end
+      puts "\n\n"
+    end
+
+  end
+end

data/lib/stretchy/rails/tasks/connector/update.rake

@@ -0,0 +1,32 @@
+namespace :stretchy do
+  namespace :connector do
+    desc "Update connector"
+    task update: :environment do
+      klass = ENV['MODEL']
+      Rails.application.eager_load!
+      models = klass.nil? ? Stretchy::MachineLearning::Connector.descendants : [klass.constantize]
+        models.each do |model|
+          spinner = TTY::Spinner.new("Updating Connector #{model} ".ljust(JUSTIFICATION) + ":spinner", format: :dots)
+          spinner.auto_spin
+
+          # if model.exists?
+          #   spinner.stop(Rainbow("[EXISTS]").yellow)
+          #   next
+          # end
+
+          begin
+            response = model.update!
+          rescue Stretchy::MachineLearning::Errors::ConnectorMissingError => e
+            spinner.stop(Rainbow("[FAILED]").red)
+            next
+          end
+          status = if response
+            Rainbow("[SUCCESS]").green
+          else
+            Rainbow("[FAILED]").red
+          end
+          spinner.stop(status)
+        end
+      end
+  end
+end

data/lib/stretchy/rails/tasks/index/create.rake

@@ -0,0 +1,28 @@
+namespace :stretchy do
+  namespace :index do
+    desc "Create indices"
+    task create: :environment do
+      klass = ENV['MODEL']
+      Rails.application.eager_load!
+      models = klass.nil? ? StretchyModel.descendants : [klass.constantize]
+        models.each do |model|
+          next if model.name == "Stretchy::MachineLearning::Registry"
+          spinner = TTY::Spinner.new("Creating index #{model} (#{model.index_name}) ".ljust(JUSTIFICATION) + ":spinner", format: :dots)
+          spinner.auto_spin
+          if model.index_exists?
+            spinner.stop( Rainbow("[EXISTS]").orange)
+            next
+          end
+
+          response = model.create_index!
+          status = if response
+            Rainbow("[SUCCESS]").green
+          else
+            Rainbow("[FAILED]").red
+          end
+
+          spinner.stop(status)
+        end
+      end
+  end
+end

data/lib/stretchy/rails/tasks/index/delete.rake

@@ -0,0 +1,27 @@
+namespace :stretchy do
+  namespace :index do
+    desc "Delete indices"
+    task delete: :environment do
+      klass = ENV['MODEL']
+      models = klass.nil? ? StretchyModel.descendants : [klass.constantize]
+        models.each do |model|
+          next if model.name == "Stretchy::MachineLearning::Registry"
+          spinner = TTY::Spinner.new("Deleting index #{model} (#{model.index_name}) ".ljust(JUSTIFICATION) + ":spinner", format: :dots)
+          spinner.auto_spin
+          unless model.index_exists?
+            spinner.stop( Rainbow("[MISSING]").white)
+            next
+          end
+
+          response = model.delete_index!
+          status = if response
+            Rainbow("[SUCCESS]").green
+          else
+            Rainbow("[FAILED]").red
+          end
+
+          spinner.stop(status)
+        end
+      end
+  end
+end