maglev-rb 0.1.1

data/lib/maglev/retriever.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require_relative "adapters/ruby_llm_embedding"
+require_relative "authorization"
+require_relative "chunk"
+require_relative "provider_call"
+require_relative "search_result"
+require_relative "vector_stores/pgvector"
+
+module Maglev
+  class Retriever
+    def initialize(model_class, chunk_model: Chunk, embedding_adapter: Maglev.configuration.embedding_adapter, authorization: Authorization.new, vector_store: Maglev.configuration.vector_store)
+      @model_class = model_class
+      @chunk_model = chunk_model
+      @embedding_adapter = embedding_adapter || Adapters::RubyLLMEmbedding.new
+      @authorization = authorization
+      @vector_store = vector_store
+    end
+
+    def search(query, limit:, owner: nil, user: nil)
+      embedding = ProviderCall.new.call(operation: "embed") { @embedding_adapter.embed(query) }
+      return search_vector_store(embedding, limit: limit, owner: owner, user: user) if @vector_store
+
+      scope = @chunk_model.where(owner_model_name: @model_class.name)
+      scope = scope.where(owner: owner) if owner
+      scope = apply_authorization_scope(scope, user) unless owner
+      rows = scope
+        .nearest_neighbors(:embedding, embedding, distance: "cosine")
+
+      unique_owner_results(rows).first(limit)
+    end
+
+    private
+
+    def search_vector_store(embedding, limit:, owner:, user:)
+      documents = @vector_store.search(vector: embedding, filters: filters_for(owner), limit: limit)
+      unique_owner_results(documents).first(limit)
+    end
+
+    def filters_for(owner)
+      filters = {owner_model_name: @model_class.name}
+      if owner
+        filters[:owner_type] = owner.class.name
+        filters[:owner_id] = owner.id
+      end
+      filters
+    end
+
+    def apply_authorization_scope(scope, user)
+      return scope unless @authorization.configured? && user
+
+      authorized_scope = @authorization.scope(model: @model_class, user: user)
+      if authorized_scope.respond_to?(:select)
+        scope.where(owner_id: authorized_scope.select(:id))
+      else
+        scope
+      end
+    end
+
+    def unique_owner_results(rows)
+      seen_owners = {}
+      results = []
+
+      rows.each do |row|
+        owner = owner_for(row)
+        next if seen_owners[owner]
+
+        seen_owners[owner] = true
+        results << SearchResult.new(
+          owner: owner,
+          content: row.content,
+          source: row.source,
+          distance: row.respond_to?(:neighbor_distance) ? row.neighbor_distance : row.distance,
+          chunk_index: row.respond_to?(:chunk_index) ? row.chunk_index : nil
+        )
+      end
+
+      results
+    end
+
+    def owner_for(row)
+      if row.owner
+        row.owner
+      elsif row.respond_to?(:owner_type) && row.respond_to?(:owner_id)
+        row.owner_type.constantize.find(row.owner_id)
+      end
+    end
+  end
+end

data/lib/maglev/schema_compiler.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require_relative "configuration"
+require_relative "errors"
+
+module Maglev
+  class SchemaCompiler
+    SUPPORTED_MACROS = %i[belongs_to has_one has_many].freeze
+
+    CompiledSchema = Struct.new(:model_class, :relations)
+    CompiledRelation = Struct.new(:name, :depth, :limit, :inverse, :macro, :related_class)
+
+    def initialize(config, max_depth: Maglev.configuration.max_relation_depth)
+      @config = config
+      @max_depth = max_depth
+    end
+
+    def compile
+      CompiledSchema.new(@config.model_class, @config.relations.map { |relation| compile_relation(relation) }.freeze)
+    end
+
+    private
+
+    def compile_relation(relation)
+      reflection = @config.model_class.reflect_on_association(relation.name.to_sym)
+      unless reflection
+        raise ConfigurationError, "Unknown Maglev association #{@config.model_class.name}.#{relation.name}"
+      end
+
+      unless SUPPORTED_MACROS.include?(reflection.macro)
+        raise ConfigurationError, "Unsupported Maglev association #{@config.model_class.name}.#{relation.name}: #{reflection.macro}"
+      end
+
+      if relation.depth > @max_depth
+        raise ConfigurationError, "Maglev association #{@config.model_class.name}.#{relation.name} exceeds maximum depth #{@max_depth}"
+      end
+
+      related_class = reflection.klass
+      unless related_class.respond_to?(:maglev_config) && related_class.maglev_config
+        raise ConfigurationError, "Related Maglev model #{related_class.name} must declare has_knowledge"
+      end
+
+      inverse = relation.inverse || reflection.inverse_of&.name&.to_s
+      unless inverse && related_class.reflect_on_association(inverse.to_sym)
+        raise ConfigurationError, "Maglev association #{@config.model_class.name}.#{relation.name} requires an inverse for invalidation"
+      end
+
+      CompiledRelation.new(relation.name, relation.depth, relation.limit, inverse, reflection.macro, related_class)
+    end
+  end
+end

data/lib/maglev/search_result.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Maglev
+  class SearchResult
+    attr_reader :owner, :content, :source, :distance, :chunk_index
+
+    def initialize(owner:, content:, source:, distance:, chunk_index: nil)
+      @owner = owner
+      @content = content
+      @source = source
+      @distance = distance
+      @chunk_index = chunk_index
+      freeze
+    end
+
+    def similarity
+      return nil if distance.nil?
+
+      1.0 - distance.to_f
+    end
+  end
+end

data/lib/maglev/snapshot.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Maglev
+  class Snapshot
+    def initialize(lines)
+      @text = lines.join("\n").freeze
+      freeze
+    end
+
+    def to_s
+      @text
+    end
+  end
+end

data/lib/maglev/snapshot_builder.rb

@@ -0,0 +1,204 @@
+# frozen_string_literal: true
+
+require "cgi"
+
+require_relative "attachment_extractor"
+require_relative "snapshot"
+
+module Maglev
+  class SnapshotBuilder
+    def initialize(record, config, path: nil, visited: nil, attachment_extractor: nil, remaining_depth: Maglev.configuration.max_relation_depth)
+      @record = record
+      @config = config
+      @path = path
+      @visited = visited || {}
+      @attachment_extractor = attachment_extractor || Maglev.configuration.attachment_extractor || AttachmentExtractor.new
+      @remaining_depth = remaining_depth
+    end
+
+    def build
+      return Snapshot.new([]) if visited?
+
+      mark_visited
+      Snapshot.new(lines)
+    end
+
+    private
+
+    def lines
+      [
+        header,
+        *attribute_lines,
+        *tag_lines,
+        *attachment_lines,
+        *rich_text_lines,
+        *relation_lines
+      ]
+    end
+
+    def header
+      label = "#{@record.class.name}##{record_id}"
+      @path ? "#{@path} #{label}" : label
+    end
+
+    def record_id
+      id = @record.respond_to?(:id) ? @record.id : nil
+      id.nil? ? "new_record" : id
+    end
+
+    def attribute_lines
+      @config.exposed_attributes.filter_map do |attribute|
+        value = @record.public_send(attribute)
+        "#{attribute_path(attribute)}: #{value}" unless value.nil?
+      end
+    end
+
+    def tag_lines
+      return [] if @config.tags.empty?
+
+      ["tags: #{@config.tags.join(", ")}"]
+    end
+
+    def relation_lines
+      return [] unless @remaining_depth.positive?
+
+      @config.relations.flat_map do |relation|
+        remaining_depth = [@remaining_depth, relation.depth].min - 1
+        relation_records, collection = relation_records(relation)
+        relation_records.each_with_index.flat_map do |related_record, index|
+          relation_path = path_for(relation, index, collection: collection)
+          SnapshotBuilder.new(
+            related_record,
+            related_record.class.maglev_config,
+            path: relation_path,
+            visited: @visited,
+            attachment_extractor: @attachment_extractor,
+            remaining_depth: remaining_depth
+          ).build.to_s.split("\n")
+        end
+      end
+    end
+
+    def attachment_lines
+      @config.attached_sources.flat_map do |source|
+        attachment_blobs(source.name).flat_map do |blob|
+          document = extract_attachment(source.name, blob)
+          if document.extracted?
+            ["#{document.source_identifier}.text: #{document.text}"]
+          else
+            ["#{document.source_identifier}.skipped: #{document.metadata[:reason]}"]
+          end
+        end
+      end
+    end
+
+    def attachment_blobs(source_name)
+      value = @record.public_send(source_name)
+      attachments = if value.respond_to?(:attachments)
+        value.attachments.to_a
+      elsif value.respond_to?(:attachment)
+        [value.attachment].compact
+      elsif value.respond_to?(:blob)
+        [value]
+      elsif value.respond_to?(:to_ary)
+        value.to_ary
+      else
+        [value].compact
+      end
+
+      attachments.map { |attachment| attachment.respond_to?(:blob) ? attachment.blob : attachment }
+    end
+
+    def extract_attachment(source_name, blob)
+      @attachment_extractor.extract(blob, source_name: source_name)
+    rescue
+      ExtractedDocument.skipped(
+        source_identifier: "#{source_name}[blob:#{blob_identifier(blob)}]",
+        reason: "extraction_failed",
+        metadata: {
+          filename: blob.respond_to?(:filename) ? blob.filename.to_s : "unknown",
+          content_type: blob.respond_to?(:content_type) ? blob.content_type.to_s : nil
+        }
+      )
+    end
+
+    def rich_text_lines
+      @config.rich_text_sources.filter_map do |source|
+        text = rich_text_value(source.name)
+        "#{rich_text_identifier(source.name)}.text: #{text}" unless text.nil? || text.empty?
+      end
+    end
+
+    def rich_text_value(source_name)
+      value = @record.public_send(source_name)
+      text = if value.respond_to?(:body) && value.body.respond_to?(:to_plain_text)
+        value.body.to_plain_text
+      elsif value.respond_to?(:to_plain_text)
+        value.to_plain_text
+      elsif value.respond_to?(:body) && value.body.respond_to?(:to_html)
+        sanitize_html(value.body.to_html.to_s)
+      else
+        sanitize_html(value.to_s)
+      end
+      text.squeeze(" ").strip
+    end
+
+    def rich_text_identifier(source_name)
+      "rich_text.#{source_name}"
+    end
+
+    def blob_identifier(blob)
+      if blob.respond_to?(:id) && blob.id
+        blob.id
+      elsif blob.respond_to?(:key) && blob.key
+        blob.key
+      elsif blob.respond_to?(:filename)
+        blob.filename.to_s
+      else
+        "unknown"
+      end
+    end
+
+    def sanitize_html(html)
+      without_scripts = html.gsub(%r{<script\b[^>]*>.*?</script>}im, " ")
+        .gsub(%r{<style\b[^>]*>.*?</style>}im, " ")
+      CGI.unescapeHTML(without_scripts.gsub(/<[^>]+>/, " ").squeeze(" ").strip)
+    end
+
+    def relation_records(relation)
+      value = @record.public_send(relation.name)
+      records, collection = if value.nil?
+        [[], false]
+      elsif value.respond_to?(:to_ary)
+        [value.to_ary, true]
+      elsif value.respond_to?(:limit)
+        [value.limit(relation.limit).to_a, true]
+      else
+        [[value], false]
+      end
+
+      [records.first(relation.limit), collection]
+    end
+
+    def path_for(relation, index, collection:)
+      segment = collection ? "#{relation.name}[#{index}]" : relation.name
+      @path ? "#{@path}.#{segment}" : segment
+    end
+
+    def attribute_path(attribute)
+      @path ? "#{@path}.#{attribute}" : attribute
+    end
+
+    def visited?
+      @visited.key?(record_key)
+    end
+
+    def mark_visited
+      @visited[record_key] = true
+    end
+
+    def record_key
+      [@record.class.name, record_id]
+    end
+  end
+end

data/lib/maglev/vector_stores/base.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Maglev
+  module VectorStores
+    class Base
+      def upsert(documents:)
+        raise NotImplementedError, "#{self.class.name} must implement #upsert"
+      end
+
+      def search(vector:, filters:, limit:)
+        raise NotImplementedError, "#{self.class.name} must implement #search"
+      end
+
+      def delete(ids:)
+        raise NotImplementedError, "#{self.class.name} must implement #delete"
+      end
+
+      def delete_by_owner(owner_type:, owner_id:)
+        raise NotImplementedError, "#{self.class.name} must implement #delete_by_owner"
+      end
+
+      def healthcheck
+        raise NotImplementedError, "#{self.class.name} must implement #healthcheck"
+      end
+
+      def capabilities
+        {}
+      end
+    end
+  end
+end

data/lib/maglev/vector_stores/document.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Maglev
+  module VectorStores
+    class Document
+      attr_reader :id, :owner_type, :owner_id, :owner_model_name, :source,
+        :chunk_index, :content, :content_checksum, :embedding_model,
+        :embedding, :owner, :distance
+
+      def initialize(owner_type:, owner_id:, owner_model_name:, source:, chunk_index:,
+        content:, content_checksum:, embedding_model:, embedding:, id: nil,
+        owner: nil, distance: nil)
+        @owner_type = owner_type
+        @owner_id = owner_id
+        @owner_model_name = owner_model_name
+        @source = source
+        @chunk_index = chunk_index
+        @content = content
+        @content_checksum = content_checksum
+        @embedding_model = embedding_model
+        @embedding = embedding
+        @id = id || "#{owner_type}:#{owner_id}:#{source}:#{chunk_index}"
+        @owner = owner
+        @distance = distance
+        freeze
+      end
+
+      def metadata
+        {
+          owner_type: owner_type,
+          owner_id: owner_id,
+          owner_model_name: owner_model_name,
+          source: source,
+          chunk_index: chunk_index
+        }
+      end
+    end
+  end
+end

data/lib/maglev/vector_stores/memory.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+require_relative "base"
+require_relative "document"
+
+module Maglev
+  module VectorStores
+    class Memory < Base
+      def initialize
+        @documents = {}
+      end
+
+      def upsert(documents:)
+        documents.each { |document| @documents[document.id] = document }
+      end
+
+      def search(vector:, filters:, limit:)
+        @documents.values
+          .select { |document| matches_filters?(document, filters) }
+          .map { |document| with_distance(document, cosine_distance(vector, document.embedding)) }
+          .sort_by(&:distance)
+          .first(limit)
+      end
+
+      def delete(ids:)
+        ids.each { |id| @documents.delete(id) }
+      end
+
+      def delete_by_owner(owner_type:, owner_id:)
+        @documents.delete_if do |_id, document|
+          document.owner_type == owner_type && document.owner_id == owner_id
+        end
+      end
+
+      def healthcheck
+        :ok
+      end
+
+      def capabilities
+        {metadata_filtering: true, in_memory: true}
+      end
+
+      private
+
+      def matches_filters?(document, filters)
+        filters.all? { |key, value| document.metadata.fetch(key) == value }
+      end
+
+      def cosine_distance(left, right)
+        dot = left.zip(right).sum { |a, b| a.to_f * b.to_f }
+        left_norm = Math.sqrt(left.sum { |value| value.to_f**2 })
+        right_norm = Math.sqrt(right.sum { |value| value.to_f**2 })
+        return 1.0 if left_norm.zero? || right_norm.zero?
+
+        1.0 - (dot / (left_norm * right_norm))
+      end
+
+      def with_distance(document, distance)
+        Document.new(
+          id: document.id,
+          owner_type: document.owner_type,
+          owner_id: document.owner_id,
+          owner_model_name: document.owner_model_name,
+          source: document.source,
+          chunk_index: document.chunk_index,
+          content: document.content,
+          content_checksum: document.content_checksum,
+          embedding_model: document.embedding_model,
+          embedding: document.embedding,
+          owner: document.owner,
+          distance: distance
+        )
+      end
+    end
+  end
+end

data/lib/maglev/vector_stores/pgvector.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+require_relative "../chunk"
+require_relative "base"
+require_relative "document"
+
+module Maglev
+  module VectorStores
+    class Pgvector < Base
+      def initialize(chunk_model: Chunk)
+        @chunk_model = chunk_model
+      end
+
+      def upsert(documents:)
+        @chunk_model.transaction do
+          documents.each do |document|
+            scope = identity_scope(document)
+            existing = scope.find_by(chunk_index: document.chunk_index, content_checksum: document.content_checksum)
+            next if existing
+
+            scope.where(chunk_index: document.chunk_index).delete_all
+            @chunk_model.create!(
+              owner_type: document.owner_type,
+              owner_id: document.owner_id,
+              owner_model_name: document.owner_model_name,
+              owner: document.owner,
+              source: document.source,
+              chunk_index: document.chunk_index,
+              content: document.content,
+              content_checksum: document.content_checksum,
+              embedding_model: document.embedding_model,
+              embedding: document.embedding
+            )
+          end
+        end
+      end
+
+      def search(vector:, filters:, limit:)
+        scope = filters.reduce(@chunk_model.all) do |current_scope, (key, value)|
+          current_scope.where(key => value)
+        end
+        scope.nearest_neighbors(:embedding, vector, distance: "cosine")
+          .first(limit)
+          .map { |row| document_for(row) }
+      end
+
+      def delete(ids:)
+        ids.each do |id|
+          owner_type, owner_id, source, chunk_index = id.split(":", 4)
+          @chunk_model.where(owner_type: owner_type, owner_id: owner_id, source: source, chunk_index: chunk_index).delete_all
+        end
+      end
+
+      def delete_by_owner(owner_type:, owner_id:)
+        @chunk_model.where(owner_type: owner_type, owner_id: owner_id).delete_all
+      end
+
+      def healthcheck
+        @chunk_model.connection.active? ? :ok : :unavailable
+      end
+
+      def capabilities
+        {metadata_filtering: true, pgvector: true}
+      end
+
+      private
+
+      def identity_scope(document)
+        @chunk_model.where(
+          owner_type: document.owner_type,
+          owner_id: document.owner_id,
+          owner_model_name: document.owner_model_name,
+          source: document.source
+        )
+      end
+
+      def document_for(row)
+        Document.new(
+          owner_type: row.owner_type,
+          owner_id: row.owner_id,
+          owner_model_name: row.owner_model_name,
+          source: row.source,
+          chunk_index: row.chunk_index,
+          content: row.content,
+          content_checksum: row.content_checksum,
+          embedding_model: row.embedding_model,
+          embedding: row.embedding,
+          owner: row.owner,
+          distance: row.respond_to?(:neighbor_distance) ? row.neighbor_distance : row.distance
+        )
+      end
+    end
+  end
+end

data/lib/maglev/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Maglev
+  VERSION = "0.1.1"
+end

data/lib/maglev.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require_relative "maglev/version"
+require_relative "maglev/configuration"
+require_relative "maglev/provider_configuration"
+require_relative "maglev/errors"
+require_relative "maglev/authorization"
+require_relative "maglev/provider_call"
+require_relative "maglev/knowledge_config"
+require_relative "maglev/snapshot"
+require_relative "maglev/snapshot_builder"
+require_relative "maglev/chunker"
+require_relative "maglev/embedding_adapter"
+require_relative "maglev/adapters/ruby_llm_embedding"
+require_relative "maglev/generation_adapter"
+require_relative "maglev/adapters/ruby_llm_generation"
+require_relative "maglev/extracted_document"
+require_relative "maglev/attachment_extractor"
+require_relative "maglev/adapters/ruby_llm_attachment_extractor"
+require_relative "maglev/search_result"
+require_relative "maglev/vector_stores/base"
+require_relative "maglev/vector_stores/document"
+require_relative "maglev/vector_stores/memory"
+require_relative "maglev/vector_stores/pgvector"
+require_relative "maglev/response"
+require_relative "maglev/prompt_builder"
+require_relative "maglev/context_assembler"
+require_relative "maglev/context_preview"
+require_relative "maglev/indexer"
+require_relative "maglev/reindex_job"
+require_relative "maglev/retriever"
+require_relative "maglev/schema_compiler"
+require_relative "maglev/dependency_graph"
+require_relative "maglev/content_source_graph"
+require_relative "maglev/answerer"
+
+require_relative "maglev/railtie" if defined?(Rails::Railtie)