maglev-rb 0.1.1 → 0.2.1

data/lib/maglev/index_diagnostics.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+require_relative "index_state"
+
+module Maglev
+  class IndexDiagnostics
+    State = Data.define(:owner_type, :owner_id, :status, :active_index_version, :chunk_count,
+      :last_success_at, :latest_failure, :rebuild_required)
+
+    class << self
+      def record_started(owner_type:, owner_id:, index_version:)
+        return unless available?
+
+        previous = IndexState.find_by(owner_type: owner_type.to_s, owner_id: owner_id)
+        persist(owner_type, owner_id, status: previous&.last_success_at ? "rebuilding" : "indexing",
+          active_index_version: index_version, rebuild_required: true)
+      end
+
+      def record_success(owner_type:, owner_id:, index_version:, chunk_count:)
+        return unless available?
+
+        persist(owner_type, owner_id, status: "ready", active_index_version: index_version,
+          chunk_count: chunk_count, last_success_at: Time.now.utc, latest_failure_class: nil,
+          latest_failure_at: nil, rebuild_required: false)
+      end
+
+      def record_failure(owner_type:, owner_id:, index_version:, error:)
+        return unless available?
+
+        persist(owner_type, owner_id, status: "failed", active_index_version: index_version,
+          latest_failure_class: error.class.name, latest_failure_at: Time.now.utc, rebuild_required: true)
+      end
+
+      def record_unindexed(owner_type:, owner_id:)
+        return unless available?
+
+        persist(owner_type, owner_id, status: "not_indexed", active_index_version: nil,
+          chunk_count: 0, rebuild_required: false)
+      end
+
+      def status(owner_type:, owner_id:)
+        return empty_state(owner_type, owner_id) unless available?
+
+        row = IndexState.find_by(owner_type: owner_type.to_s, owner_id: owner_id)
+        return empty_state(owner_type, owner_id) unless row
+
+        State.new(owner_type: row.owner_type, owner_id: row.owner_id, status: row.status.to_sym,
+          active_index_version: row.active_index_version, chunk_count: row.chunk_count,
+          last_success_at: row.last_success_at, latest_failure: failure_for(row),
+          rebuild_required: row.rebuild_required)
+      end
+
+      private
+
+      def available?
+        IndexState.table_exists?
+      rescue ActiveRecord::ConnectionNotDefined, ActiveRecord::StatementInvalid
+        false
+      end
+
+      def persist(owner_type, owner_id, **attributes)
+        row = IndexState.create_or_find_by!(owner_type: owner_type.to_s, owner_id: owner_id) do |state|
+          state.assign_attributes(attributes)
+        end
+        row.with_lock do
+          row.update!(attributes)
+        end
+      end
+
+      def empty_state(owner_type, owner_id)
+        State.new(owner_type: owner_type.to_s, owner_id: owner_id, status: :not_indexed,
+          active_index_version: nil, chunk_count: 0, last_success_at: nil,
+          latest_failure: nil, rebuild_required: false)
+      end
+
+      def failure_for(row)
+        return unless row.latest_failure_class
+
+        {error_class: row.latest_failure_class, occurred_at: row.latest_failure_at}.freeze
+      end
+    end
+  end
+end

data/lib/maglev/index_identity.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+require "digest"
+require "json"
+
+require_relative "chunker"
+require_relative "errors"
+
+module Maglev
+  class IndexIdentity
+    PAYLOAD_NAMESPACE = "maglev-index"
+    FORMAT_VERSION = 1
+
+    def initialize(configuration:, adapter:, chunk_size:)
+      @configuration = configuration
+      @adapter = adapter
+      @chunk_size = chunk_size
+    end
+
+    def to_s
+      Digest::SHA256.hexdigest(JSON.generate([
+        PAYLOAD_NAMESPACE,
+        FORMAT_VERSION,
+        required_string(@configuration.embedding_model, "embedding model"),
+        positive_integer(@configuration.embedding_dimensions, "embedding dimensions"),
+        adapter_id,
+        adapter_version,
+        required_string(Chunker::ALGORITHM_VERSION, "chunking algorithm version"),
+        positive_integer(@chunk_size, "chunk size"),
+        required_string(@configuration.application_index_version, "application index version")
+      ]))
+    end
+
+    private
+
+    def adapter_id
+      value = @configuration.embedding_adapter_id
+      if value.nil?
+        unless @adapter.respond_to?(:maglev_adapter_id)
+          raise ConfigurationError, "embedding adapter ID must be configured or implemented by the adapter"
+        end
+        value = @adapter.maglev_adapter_id
+      end
+      required_string(value, "embedding adapter ID")
+    end
+
+    def adapter_version
+      value = @configuration.embedding_adapter_version
+      if value.nil?
+        unless @adapter.respond_to?(:maglev_adapter_version)
+          raise ConfigurationError, "embedding adapter version must be configured or implemented by the adapter"
+        end
+        value = @adapter.maglev_adapter_version
+      end
+      required_string(value, "embedding adapter version")
+    end
+
+    def required_string(value, name)
+      return value if value.is_a?(String) && !value.empty?
+
+      raise ConfigurationError, "#{name} must be a non-empty string"
+    end
+
+    def positive_integer(value, name)
+      return value if value.is_a?(Integer) && value.positive?
+
+      raise ConfigurationError, "#{name} must be a positive integer"
+    end
+  end
+end

data/lib/maglev/index_state.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require "active_record"
+
+module Maglev
+  class IndexState < ActiveRecord::Base
+    self.table_name = "maglev_index_states"
+  end
+end

data/lib/maglev/indexer.rb

@@ -2,50 +2,115 @@
 
 require "digest"
 
-require_relative "adapters/ruby_llm_embedding"
+require_relative "adapters/faraday_embedding"
 require_relative "chunk"
 require_relative "chunker"
 require_relative "errors"
+require_relative "index_identity"
 require_relative "provider_call"
 require_relative "vector_stores/document"
+require_relative "source_extractor"
+require_relative "index_diagnostics"
+require_relative "vector_stores/document_id"
 
 module Maglev
   class Indexer
     SOURCE = "snapshot"
-
-    def initialize(record, chunk_model: Chunk, embedding_adapter: Maglev.configuration.embedding_adapter, embedding_dimensions: Maglev.configuration.embedding_dimensions, chunk_size: Maglev.configuration.chunk_size, vector_store: Maglev.configuration.vector_store)
+    IdentityConfiguration = Struct.new(
+      :embedding_model,
+      :embedding_dimensions,
+      :embedding_adapter_id,
+      :embedding_adapter_version,
+      :application_index_version
+    )
+    private_constant :IdentityConfiguration
+
+    def initialize(record, chunk_model: Chunk, embedding_adapter: Maglev.configuration.embedding_adapter, embedding_dimensions: Maglev.configuration.embedding_dimensions, chunk_size: Maglev.configuration.chunk_size, vector_store: Maglev.configuration.vector_store, index_identity: nil, provider_call: ProviderCall.new)
       @record = record
       @chunk_model = chunk_model
-      @embedding_adapter = embedding_adapter || Adapters::RubyLLMEmbedding.new
+      @embedding_adapter = embedding_adapter || Adapters::FaradayEmbedding.new
       @embedding_dimensions = embedding_dimensions
       @chunk_size = chunk_size
       @vector_store = vector_store
+      @index_identity = index_identity
+      @provider_call = provider_call
     end
 
     def index
       validate_owner_id!
       validate_storage_dimensions!
+      prepare_index_identity
+      IndexDiagnostics.record_started(owner_type: @record.class.name, owner_id: @record.id, index_version: @index_version)
       ActiveSupport::Notifications.instrument("maglev.index.start", owner_type: @record.class.name, owner_id: @record.id)
+      @embedding_count = 0
 
+      index_metadata = nil
       chunk_count = loop do
-        snapshot = @record.maglev_snapshot
-        chunks = Chunker.new(max_characters: @chunk_size).call(snapshot)
+        snapshot_result = if @record.respond_to?(:maglev_snapshot_result) && @record.class.respond_to?(:maglev_config) && @record.class.maglev_config
+          @record.maglev_snapshot_result
+        end
+        snapshot = snapshot_result ? snapshot_result.to_s : @record.maglev_snapshot
+        source_chunks = chunks_for(snapshot, snapshot_result)
+        all_chunks = source_chunks
+        chunks = all_chunks.first(Maglev.configuration.snapshot_max_chunks)
+        budget_metadata = snapshot_result&.metadata || {truncated: false, sources: []}
+        if all_chunks.length > chunks.length && budget_metadata.fetch(:sources, []).none? { |source| source[:kind] == :chunks }
+          budget_metadata = budget_metadata.merge(
+            truncated: true,
+            sources: budget_metadata.fetch(:sources, []) + [{
+              kind: :chunks,
+              path: "snapshot.chunks",
+              original_chunks: all_chunks.length,
+              retained_chunks: chunks.length
+            }]
+          )
+        end
+        index_metadata = budget_metadata.merge(
+          original_chunk_count: all_chunks.length,
+          retained_chunk_count: chunks.length,
+          chunk_limit: Maglev.configuration.snapshot_max_chunks
+        )
+        index_metadata = deep_freeze(index_metadata)
         prepared = @vector_store ? prepare_documents(chunks) : prepare_chunks(chunks)
         break chunks.length if persist_if_current(snapshot, prepared)
       end
 
-      ActiveSupport::Notifications.instrument("maglev.index.success", owner_type: @record.class.name, owner_id: @record.id, chunk_count: chunk_count)
+      ActiveSupport::Notifications.instrument("maglev.index.success", owner_type: @record.class.name, owner_id: @record.id, chunk_count: chunk_count, budget: index_metadata)
+      IndexDiagnostics.record_success(owner_type: @record.class.name, owner_id: @record.id,
+        index_version: @index_version, chunk_count: chunk_count)
     rescue => error
-      ActiveSupport::Notifications.instrument("maglev.index.failure", owner_type: @record.class.name, owner_id: @record.id, error_class: error.class.name)
+      ActiveSupport::Notifications.instrument("maglev.index.failure", owner_type: @record.class.name, owner_id: @record.id, error_class: error.class.name, budget: index_metadata)
+      IndexDiagnostics.record_failure(owner_type: @record.class.name, owner_id: @record.id,
+        index_version: @index_version, error: error)
       raise
     end
 
     def unindex
-      identity_scope.delete_all
+      if @vector_store
+        @vector_store.delete_by_owner(owner_type: @record.class.name, owner_id: @record.id)
+      else
+        identity_scope.delete_all
+      end
+
+      IndexDiagnostics.record_unindexed(owner_type: @record.class.name, owner_id: @record.id)
     end
 
     private
 
+    def chunks_for(snapshot, snapshot_result)
+      unless snapshot_result
+        return Chunker.new(max_characters: @chunk_size, max_chunks: nil).call(snapshot).each_with_index.map do |content, index|
+          {content: content, source_identity: SOURCE, source_type: :snapshot, chunk_index: index}
+        end
+      end
+
+      SourceExtractor.new.call(snapshot).flat_map do |source|
+        Chunker.new(max_characters: @chunk_size, max_chunks: nil).call(source.content).each_with_index.map do |content, index|
+          {content: content, source_identity: source.identity, source_type: source.type, chunk_index: index}
+        end
+      end
+    end
+
     def persist_if_current(snapshot, prepared)
       with_owner_lock do
         next false unless @record.maglev_snapshot == snapshot
@@ -55,11 +120,11 @@ module Maglev
     end
 
     def prepare_chunks(chunks)
-      chunks.each_with_index.map do |content, chunk_index|
-        checksum = Digest::SHA256.hexdigest(content)
-        existing = identity_scope.find_by(chunk_index: chunk_index, content_checksum: checksum)
-        embedding = embed(content) unless existing
-        {content: content, chunk_index: chunk_index, checksum: checksum, embedding: embedding}
+      chunks.map do |chunk|
+        checksum = Digest::SHA256.hexdigest(chunk[:content])
+        existing = identity_scope_for(chunk).find_by(chunk_index: chunk[:chunk_index], content_checksum: checksum, index_version: @index_version)
+        embedding = embed(chunk[:content]) unless existing
+        chunk.merge(checksum: checksum, embedding: embedding)
       end
     end
 
@@ -68,48 +133,82 @@ module Maglev
 
       @chunk_model.transaction do
         chunks.each { |chunk| persist_chunk(chunk) }
-        delete_obsolete_chunks(chunks.length)
+        delete_obsolete_chunks(chunks)
       end
       true
     end
 
     def matching_chunk?(chunk)
-      identity_scope.find_by(chunk_index: chunk[:chunk_index], content_checksum: chunk[:checksum])
+      identity_scope_for(chunk).find_by(chunk_index: chunk[:chunk_index], content_checksum: chunk[:checksum], index_version: @index_version)
     end
 
     def persist_chunk(chunk)
       return if matching_chunk?(chunk)
 
-      identity_scope.where(chunk_index: chunk[:chunk_index]).delete_all
+      identity_scope_for(chunk).where(chunk_index: chunk[:chunk_index]).delete_all
 
-      @chunk_model.create!(
+      attributes = {
         **identity,
         owner: @record,
-        source: SOURCE,
+        source: chunk[:source_identity],
         chunk_index: chunk[:chunk_index],
         content: chunk[:content],
         content_checksum: chunk[:checksum],
-        embedding_model: Maglev.configuration.embedding_model,
+        embedding_model: @identity_configuration.embedding_model,
+        index_version: @index_version,
         embedding: chunk[:embedding]
-      )
+      }
+      if source_metadata_columns?
+        attributes[:source_identity] = chunk[:source_identity]
+        attributes[:source_type] = chunk[:source_type]
+        attributes[:tenant_id] = Maglev.configuration.tenant_id(record: @record)
+      end
+      @chunk_model.create!(attributes)
     end
 
     def prepare_documents(chunks)
-      chunks.each_with_index.map { |content, chunk_index| document_for(content, chunk_index) }
+      prepared = chunks.map do |chunk|
+        {
+          id: stable_document_id(chunk[:source_identity], chunk[:chunk_index]),
+          content: chunk[:content],
+          source_identity: chunk[:source_identity], source_type: chunk[:source_type],
+          chunk_index: chunk[:chunk_index],
+          checksum: Digest::SHA256.hexdigest(chunk[:content])
+        }
+      end
+      existing_by_id = @vector_store.fetch(ids: prepared.map { |chunk| chunk[:id] }).to_h { |document| [document.id, document] }
+
+      prepared.map do |chunk|
+        existing = existing_by_id[chunk[:id]]
+        reusable = existing && existing.content_checksum == chunk[:checksum] && existing.index_version == @index_version
+        document_for(chunk, embedding: reusable ? existing.embedding : embed(chunk[:content]))
+      end
     end
 
     def persist_documents(documents)
-      @vector_store.delete_by_owner(owner_type: @record.class.name, owner_id: @record.id)
-      @vector_store.upsert(documents: documents)
+      @vector_store.replace_owner(owner_type: @record.class.name, owner_id: @record.id, documents: documents)
       true
     end
 
-    def delete_obsolete_chunks(chunk_count)
-      identity_scope.where.not(chunk_index: (0...chunk_count).to_a).delete_all
+    def delete_obsolete_chunks(chunks)
+      retained = chunks.group_by { |chunk| chunk[:source_identity] }
+      if retained.keys == [SOURCE]
+        identity_scope.where(source: SOURCE).where.not(chunk_index: retained.fetch(SOURCE).map { |chunk| chunk[:chunk_index] }).delete_all
+        return
+      end
+
+      identity_scope.where.not(source: retained.keys).delete_all
+      retained.each do |source, source_chunks|
+        identity_scope.where(source: source).where.not(chunk_index: source_chunks.map { |chunk| chunk[:chunk_index] }).delete_all
+      end
     end
 
     def identity_scope
-      @chunk_model.where(identity.merge(source: SOURCE))
+      @chunk_model.where(identity)
+    end
+
+    def identity_scope_for(chunk)
+      @chunk_model.where(identity.merge(source: chunk[:source_identity]))
     end
 
     def identity
@@ -145,30 +244,81 @@ module Maglev
         "#{@chunk_model.table_name}.embedding vector(#{database_dimensions})"
     end
 
+    def source_metadata_columns?
+      !@chunk_model.respond_to?(:columns_hash) || @chunk_model.columns_hash.key?("source_identity")
+    end
+
     def embed(content)
-      embedding = ProviderCall.new.call(operation: "embed") { @embedding_adapter.embed(content) }
+      embedding = @provider_call.call(operation: "embed") { perform_embedding(content) }
       validate_embedding!(embedding)
       embedding
     end
 
+    def perform_embedding(content)
+      if @embedding_count >= Maglev.configuration.snapshot_max_chunks
+        raise RetryableProviderError, "snapshot changed while indexing and exhausted the per-execution embedding budget"
+      end
+
+      @embedding_count += 1
+      @embedding_adapter.embed(content)
+    end
+
+    def deep_freeze(value)
+      case value
+      when Hash
+        value.to_h { |key, item| [key, deep_freeze(item)] }.freeze
+      when Array
+        value.map { |item| deep_freeze(item) }.freeze
+      else
+        value.freeze
+      end
+    end
+
+    def prepare_index_identity
+      configuration = Maglev.configuration
+      @identity_configuration = IdentityConfiguration.new(
+        embedding_model: configuration.embedding_model,
+        embedding_dimensions: @embedding_dimensions,
+        embedding_adapter_id: configuration.embedding_adapter_id,
+        embedding_adapter_version: configuration.embedding_adapter_version,
+        application_index_version: configuration.application_index_version
+      )
+      identity = @index_identity || IndexIdentity.new(
+        configuration: @identity_configuration,
+        adapter: @embedding_adapter,
+        chunk_size: @chunk_size
+      )
+      @index_version = identity.to_s
+    end
+
     def with_owner_lock(&block)
       return yield unless @record.respond_to?(:with_lock)
 
       @record.with_lock(&block)
     end
 
-    def document_for(content, chunk_index)
+    def stable_document_id(source_identity, chunk_index)
+      VectorStores::DocumentId.build(owner_type: @record.class.name, owner_id: @record.id,
+        source_identity: source_identity, chunk_index: chunk_index)
+    end
+
+    def document_for(chunk, embedding:)
       VectorStores::Document.new(
+        id: chunk[:id],
         owner_type: @record.class.name,
         owner_id: @record.id,
         owner_model_name: @record.class.name,
         owner: @record,
-        source: SOURCE,
-        chunk_index: chunk_index,
-        content: content,
-        content_checksum: Digest::SHA256.hexdigest(content),
-        embedding_model: Maglev.configuration.embedding_model,
-        embedding: embed(content)
+        source: chunk[:source_identity],
+        source_identity: chunk[:source_identity],
+        source_type: chunk[:source_type],
+        tenant_id: Maglev.configuration.tenant_id(record: @record),
+        chunk_index: chunk[:chunk_index],
+        content: chunk[:content],
+        content_checksum: chunk[:checksum],
+        embedding_model: @identity_configuration.embedding_model,
+        index_version: @index_version,
+        embedding: embedding
       )
     end
   end

data/lib/maglev/knowledge_config.rb

@@ -4,9 +4,9 @@ require_relative "errors"
 
 module Maglev
   class KnowledgeConfig
-    Relation = Struct.new(:name, :depth, :limit, :inverse) do
-      def initialize(name:, depth:, limit:, inverse: nil)
-        super(name.to_s, depth, limit, inverse&.to_s)
+    Relation = Struct.new(:name, :depth, :limit, :inverse, :order) do
+      def initialize(name:, depth:, limit:, inverse: nil, order: nil)
+        super(name.to_s, depth, limit, inverse&.to_s, order&.dup&.freeze)
         freeze
       end
     end
@@ -97,8 +97,9 @@ module Maglev
         @tags.concat(tags)
       end
 
-      def include_related(association, depth:, limit:, inverse: nil)
-        @relations << Relation.new(name: association, depth: depth, limit: limit, inverse: inverse)
+      def include_related(association, depth:, limit:, inverse: nil, order: nil)
+        normalized_order = normalize_order(order)
+        @relations << Relation.new(name: association, depth: depth, limit: limit, inverse: inverse, order: normalized_order)
       end
 
       def expose_attached(*names)
@@ -131,6 +132,27 @@ module Maglev
       def normalize(values)
         values.map(&:to_s).uniq
       end
+
+      VALID_DIRECTIONS = %i[asc desc].freeze
+
+      def normalize_order(order)
+        return nil if order.nil?
+
+        case order
+        when Symbol, String
+          {order.to_sym => :asc}
+        when Hash
+          raise ConfigurationError, "Maglev relation order Hash cannot be empty" if order.empty?
+
+          order.each do |column, direction|
+            raise ConfigurationError, "Maglev relation order direction must be :asc or :desc, got: #{direction.inspect}" unless VALID_DIRECTIONS.include?(direction)
+            raise ConfigurationError, "Maglev relation order column must be a Symbol or String, got: #{column.inspect}" unless column.is_a?(Symbol) || column.is_a?(String)
+          end
+          order.to_h { |col, dir| [col.to_sym, dir.to_sym] }
+        else
+          raise ConfigurationError, "Maglev relation order must be a Symbol, String, or Hash, got: #{order.class}"
+        end
+      end
     end
   end
 end

data/lib/maglev/knowledge_registry.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/string/inflections"
+
+module Maglev
+  class KnowledgeRegistry
+    class << self
+      def register(model_name)
+        mutex.synchronize do
+          @model_names ||= []
+          @model_names << model_name.to_s unless @model_names.include?(model_name.to_s)
+        end
+      end
+
+      def model_names
+        mutex.synchronize { (@model_names || []).dup.freeze }
+      end
+
+      def rebuild!
+        model_names.each do |model_name|
+          model = model_name.safe_constantize
+          model.rebuild_maglev_registration if model&.respond_to?(:rebuild_maglev_registration)
+        end
+      end
+
+      private
+
+      def mutex
+        @mutex ||= Mutex.new
+      end
+    end
+  end
+end