sqlite_search 0.0.1

data/lib/sqlite_search/fts5/definition.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require "sqlite_search/fts5"
+
+module SqliteSearch
+  module Fts5
+    # Immutable per-scope configuration.
+    class Definition
+      attr_reader :model, :name, :columns, :weights, :tokenizer, :table_name
+
+      def initialize(model:, name:, against:, tokenizer:)
+        @model = model
+        @name = name.to_sym
+        @columns = SqliteSearch::Fts5.columns_for(against)
+        @weights = SqliteSearch::Fts5.weights_for(against)
+        @tokenizer = tokenizer
+        @table_name = "#{model.table_name}_#{@name}_fts"
+      end
+
+      def column_names = columns.map(&:to_s)
+
+      # SQLite bm25() returns lower (more negative) = better. We ORDER BY it
+      # ascending, and expose -bm25 as the rank so higher = better.
+      def bm25_expression(connection)
+        args = [connection.quote_table_name(table_name)]
+        args.concat(weights.map { |w| format("%g", w) }) if weights
+        "bm25(#{args.join(", ")})"
+      end
+
+      def rank_column = "#{name}_rank"
+    end
+  end
+end

data/lib/sqlite_search/fts5.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module SqliteSearch
+  module Fts5
+    module_function
+
+    # against: :body | [:title, :body] | { title: 2.0, body: 1.0 }
+    def columns_for(against)
+      case against
+      when Hash then against.keys.map(&:to_sym)
+      when Array then against.map(&:to_sym)
+      else [against.to_sym]
+      end
+    end
+
+    def weights_for(against)
+      against.is_a?(Hash) ? against.values.map(&:to_f) : nil
+    end
+
+    # Extended onto the .none relation returned for blank queries so that
+    # .order_by_rank chains safely (returns the same empty relation).
+    module NullRank
+      def order_by_rank = self
+    end
+  end
+end

data/lib/sqlite_search/hybrid.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module SqliteSearch
+  module Hybrid
+    module_function
+
+    # Reciprocal Rank Fusion of two id lists (each already in rank order).
+    # Returns [[id, score], ...] sorted by fused score desc.
+    def rrf(ids_a, ids_b, k: 60)
+      require "neighbor"
+      a = ids_a.map { |id| {id: id} }
+      b = ids_b.map { |id| {id: id} }
+      Neighbor::Reranking.rrf(a, b, k: k).map { |row| [row[:result][:id], row[:score]] }
+    end
+
+    # Best-effort rerank. `fused` is [[id, score], ...]; returns the same shape
+    # reordered by the reranker (RRF scores preserved), or the input unchanged
+    # on any reranker failure.
+    def rerank(query, fused, model:, scope_name:, reranker:)
+      return fused unless reranker
+      ids = fused.map(&:first)
+      by_id = model.where(model.primary_key => ids).index_by { |r| r.public_send(model.primary_key) }
+      records = ids.filter_map { |id| by_id[id] }
+      scores = fused.to_h
+      begin
+        reordered = reranker.call(query, records, model: model, scope: scope_name)
+        reordered.filter_map do |rec|
+          id = rec.public_send(model.primary_key)
+          [id, scores[id]] if scores.key?(id)
+        end
+      rescue => e
+        if defined?(Rails) && Rails.respond_to?(:logger) && Rails.logger
+          Rails.logger.warn { "sqlite_search: rerank failed, using fused order (#{e.class}: #{e.message})" }
+        end
+        fused
+      end
+    end
+  end
+end

data/lib/sqlite_search/migration.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require "sqlite_search/fts5"
+require "sqlite_search/vec"
+
+module SqliteSearch
+  # Mixed into ActiveRecord::Migration. Wraps create_virtual_table so the FTS5
+  # table is representable in schema.rb (unlike triggers), and optionally seeds
+  # it from the source table.
+  module Migration
+    def create_fts5_index(table, name, against:, tokenizer: "porter unicode61", primary_key: "id", backfill: false)
+      columns = SqliteSearch::Fts5.columns_for(against)
+      fts_table = "#{table}_#{name}_fts"
+
+      options = columns.map(&:to_s)
+      options << "tokenize = '#{tokenizer}'"
+      connection.create_virtual_table(fts_table, :fts5, options)
+
+      backfill_fts5_index(table, fts_table, columns, primary_key) if backfill
+    end
+
+    # vec0 indexes are keyed by an integer `id` column holding the source row's
+    # primary-key value. distance: sets the vec0 distance_metric (:cosine,
+    # :euclidean, or :taxicab) and must match the vec_scope's distance.
+    def create_vec_index(table, name, dimensions:, distance: :cosine)
+      vec_table = "#{table}_#{name}_vec"
+      connection.create_virtual_table(vec_table, :vec0, [
+        "id integer primary key",
+        "embedding float[#{dimensions}] distance_metric=#{SqliteSearch::Vec.vec0_metric(distance)}"
+      ])
+    end
+
+    private
+
+    def backfill_fts5_index(table, fts_table, columns, primary_key)
+      col_list = columns.map { |c| connection.quote_column_name(c) }.join(", ")
+      non_blank = columns.map { |c| "COALESCE(#{connection.quote_column_name(c)}, '')" }.join(" || ")
+      connection.execute(<<~SQL.squish)
+        INSERT INTO #{connection.quote_table_name(fts_table)} (rowid, #{col_list})
+        SELECT #{connection.quote_column_name(primary_key)}, #{col_list}
+        FROM #{connection.quote_table_name(table)}
+        WHERE (#{non_blank}) <> ''
+      SQL
+    end
+  end
+end

data/lib/sqlite_search/model.rb

@@ -0,0 +1,187 @@
+# frozen_string_literal: true
+
+require "active_support/concern"
+require "sqlite_search/fts5/definition"
+require "sqlite_search/vec/definition"
+
+module SqliteSearch
+  module Model
+    extend ActiveSupport::Concern
+
+    class_methods do
+      def sqlite_search_fts5_definitions
+        own = (@sqlite_search_fts5_definitions ||= {})
+        return own unless superclass.respond_to?(:sqlite_search_fts5_definitions)
+        superclass.sqlite_search_fts5_definitions.merge(own)
+      end
+
+      def fts5_scope(name, against:, tokenizer: "porter unicode61")
+        definition = Fts5::Definition.new(model: self, name: name, against: against, tokenizer: tokenizer)
+        (@sqlite_search_fts5_definitions ||= {})[definition.name] = definition
+
+        scope name, ->(query = nil, prefix: false, raw: nil) do
+          match = raw || SqliteSearch::Query.build(query, prefix: prefix)
+          next none.extending(SqliteSearch::Fts5::NullRank) if match.nil? || match.to_s.empty?
+
+          fts = connection.quote_table_name(definition.table_name)
+          pk = "#{connection.quote_table_name(table_name)}.#{connection.quote_column_name(primary_key)}"
+
+          # bm25() is only valid in a query that MATCHes the fts table, so
+          # order_by_rank joins the fts table and re-applies MATCH here.
+          rank_module = Module.new do
+            define_method(:order_by_rank) do
+              bm25 = definition.bm25_expression(connection)
+              joins("JOIN #{fts} ON #{fts}.rowid = #{pk}")
+                .where("#{fts} MATCH ?", match)
+                .select("#{connection.quote_table_name(table_name)}.*, -#{bm25} AS #{definition.rank_column}")
+                .order(Arel.sql(bm25))
+            end
+          end
+
+          where("#{pk} IN (SELECT rowid FROM #{fts} WHERE #{fts} MATCH ?)", match).extending(rank_module)
+        end
+
+        cols = definition.column_names
+        # FTS5 sync runs inside the transaction (after_save/after_destroy), not
+        # after_commit: an FTS5 write is a cheap local write, so keeping it in the
+        # transaction makes the index atomic with the row. A failed write rolls
+        # both back rather than leaving a committed row with a stale index.
+        after_save do
+          if (saved_changes.keys & cols).any?
+            SqliteSearch::Fts5::Backend.new(definition).sync(self)
+          end
+        end
+        after_destroy do
+          SqliteSearch::Fts5::Backend.new(definition).remove(self)
+        end
+      end
+
+      def reindex(name = nil)
+        definitions = name ? [sqlite_search_fts5_definitions.fetch(name.to_sym)] : sqlite_search_fts5_definitions.values
+        definitions.each { |definition| SqliteSearch::Fts5::Backend.new(definition).rebuild(self) }
+      end
+
+      def sqlite_search_vec_definitions
+        own = (@sqlite_search_vec_definitions ||= {})
+        return own unless superclass.respond_to?(:sqlite_search_vec_definitions)
+        superclass.sqlite_search_vec_definitions.merge(own)
+      end
+
+      def vec_scope(name, against:, dimensions:, distance: :cosine, embedder: nil, sync: :async)
+        SqliteSearch::Vec.load!
+        definition = SqliteSearch::Vec::Definition.new(
+          model: self, name: name, against: against, dimensions: dimensions, distance: distance, embedder: embedder
+        )
+        (@sqlite_search_vec_definitions ||= {})[definition.name] = definition
+
+        scope name, ->(query = nil, k: 20, threshold: nil) do
+          next none if query.nil? || query.to_s.strip.empty?
+
+          vector = definition.embed(query.to_s, record_model: klass)
+
+          caller_conditions = all.only(:where, :joins)
+          src = connection.quote_table_name(table_name)
+          vec = connection.quote_table_name(definition.table_name)
+          pkc = connection.quote_column_name(primary_key)
+
+          knn = definition.neighbor_model
+            .joins("JOIN #{src} ON #{src}.#{pkc} = #{vec}.id")
+            .merge(caller_conditions)
+            .nearest_neighbors(:embedding, vector, distance: definition.distance)
+            .limit(k)
+          hits = knn.map { |r| [r.id, r.neighbor_distance] }
+          if threshold
+            hits = if definition.cosine?
+              hits.select { |(_, dist)| (1.0 - dist) >= threshold } # threshold = minimum cosine similarity
+            else
+              hits.select { |(_, dist)| dist <= threshold }         # threshold = maximum distance
+            end
+          end
+          next none if hits.empty?
+
+          ids = hits.map(&:first)
+          distances = hits.to_h
+          pk_sql = "#{quoted_table_name}.#{pkc}"
+          # Order and per-row scores are computed in Ruby (from the KNN), so carry
+          # them as selected CASE columns: <name>_distance is then a real attribute.
+          order = Arel.sql(SqliteSearch::Sql.id_case(pk_sql, ids.each_with_index.to_h, connection))
+          cols = ["#{quoted_table_name}.*", "#{SqliteSearch::Sql.id_case(pk_sql, distances, connection)} AS #{definition.distance_method}"]
+          if definition.cosine?
+            cols << "#{SqliteSearch::Sql.id_case(pk_sql, distances.transform_values { |d| 1.0 - d }, connection)} AS #{definition.similarity_method}"
+          end
+          where(primary_key => ids).order(order).select(cols.join(", "))
+        end
+
+        vec_cols = definition.column_names
+        vec_sync = sync
+        SqliteSearch.ensure_embed_job! unless vec_sync == :inline
+        after_save_commit do
+          if (saved_changes.keys & vec_cols).any?
+            backend = SqliteSearch::Vec::Backend.new(definition)
+            if vec_sync == :inline
+              backend.embed_and_store(self)
+            else
+              SqliteSearch::EmbedJob.perform_later(self.class.name, public_send(self.class.primary_key), definition.name.to_s)
+            end
+          end
+        end
+        after_destroy_commit do
+          SqliteSearch::Vec::Backend.new(definition).remove(self)
+        end
+      end
+
+      def reembed(name = nil)
+        definitions = name ? [sqlite_search_vec_definitions.fetch(name.to_sym)] : sqlite_search_vec_definitions.values
+        definitions.each { |definition| SqliteSearch::Vec::Backend.new(definition).reembed(self) }
+      end
+
+      def sqlite_search_hybrid_definitions
+        own = (@sqlite_search_hybrid_definitions ||= {})
+        return own unless superclass.respond_to?(:sqlite_search_hybrid_definitions)
+        superclass.sqlite_search_hybrid_definitions.merge(own)
+      end
+
+      def hybrid_scope(name, fts5:, vec:, k: 60)
+        fts5_name = fts5.to_sym
+        vec_name = vec.to_sym
+        if name.to_sym == fts5_name || name.to_sym == vec_name
+          raise SqliteSearch::Error, "hybrid_scope :#{name} must not reuse its own fts5:/vec: scope name (it would overwrite that scope)."
+        end
+        unless sqlite_search_fts5_definitions.key?(fts5_name)
+          raise SqliteSearch::Error, "hybrid_scope :#{name} references fts5: :#{fts5_name}, but no such fts5_scope is declared on #{self.name}."
+        end
+        unless sqlite_search_vec_definitions.key?(vec_name)
+          raise SqliteSearch::Error, "hybrid_scope :#{name} references vec: :#{vec_name}, but no such vec_scope is declared on #{self.name}."
+        end
+        (@sqlite_search_hybrid_definitions ||= {})[name.to_sym] = {fts5: fts5_name, vec: vec_name, k: k}
+        rrf_k = k
+        score_method = "#{name}_score"
+
+        scope name, ->(query = nil, limit: 20, rerank: true) do
+          next none if query.nil? || query.to_s.strip.empty?
+          # candidate pool per arm before fusion; capped so a large limit: can't over-fetch
+          pool = [limit * 3, 100].min
+
+          fts_ids = all.public_send(fts5_name, query).order_by_rank.limit(pool).pluck(primary_key)
+          vec_ids = all.public_send(vec_name, query, k: pool).pluck(primary_key)
+
+          fused = SqliteSearch::Hybrid.rrf(fts_ids, vec_ids, k: rrf_k)
+          next none if fused.empty?
+
+          reranker = rerank ? SqliteSearch.config.reranker : nil
+          fused = SqliteSearch::Hybrid.rerank(query, fused, model: klass, scope_name: name.to_sym, reranker: reranker)
+          next none if fused.empty?
+          fused = fused.first(limit)
+          ids = fused.map(&:first)
+          scores = fused.to_h
+          pk_sql = "#{quoted_table_name}.#{connection.quote_column_name(primary_key)}"
+          # Fused rank and score come from Ruby, so carry them as SQL: the score
+          # becomes a real <name>_score attribute (works with pluck, first, etc.).
+          order = Arel.sql(SqliteSearch::Sql.id_case(pk_sql, ids.each_with_index.to_h, connection))
+          score_col = "#{SqliteSearch::Sql.id_case(pk_sql, scores, connection)} AS #{score_method}"
+          where(primary_key => ids).order(order).select("#{quoted_table_name}.*, #{score_col}")
+        end
+      end
+    end
+  end
+end

data/lib/sqlite_search/query.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module SqliteSearch
+  # Converts free user text into a safe FTS5 MATCH string.
+  # Only quoted phrases and term characters (Unicode alphanumerics and
+  # underscore; this does NOT strip to ASCII, so accented/non-Latin words
+  # like "café" or "日本語" survive) make it through; everything else (FTS5
+  # operators, punctuation) is dropped, so untrusted input cannot inject
+  # MATCH syntax. Terms are AND-joined. Returns nil when nothing usable.
+  class Query
+    PHRASE = /"([^"]+)"/
+    TERM_CHARS = /[^[:alnum:]_]+/
+    RESERVED = /\A(?:AND|OR|NOT|NEAR)\z/
+
+    def self.build(raw, prefix: false)
+      new(raw, prefix: prefix).to_match
+    end
+
+    def initialize(raw, prefix: false)
+      @raw = raw.to_s.scrub
+      @prefix = prefix
+    end
+
+    def to_match
+      tokens = phrases + words
+      return nil if tokens.empty?
+
+      tokens[-1] = "#{tokens[-1]}*" if @prefix && !quoted?(tokens[-1])
+      tokens.join(" AND ")
+    end
+
+    private
+
+    def phrases
+      @raw.scan(PHRASE).map { |(inner)| %("#{inner.strip}") }.reject { |p| p == '""' }
+    end
+
+    def words
+      @raw.gsub(PHRASE, " ").split(TERM_CHARS).reject { |w| w.empty? || w.match?(RESERVED) }
+    end
+
+    def quoted?(token)
+      token.start_with?('"')
+    end
+  end
+end

data/lib/sqlite_search/railtie.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require "rails/railtie"
+
+module SqliteSearch
+  class Railtie < Rails::Railtie
+    initializer "sqlite_search.active_record" do
+      ActiveSupport.on_load(:active_record) do
+        require "sqlite_search/model"
+        include SqliteSearch::Model
+      end
+      ActiveSupport.on_load(:active_record) do
+        require "sqlite_search/migration"
+        ActiveRecord::Migration.include(SqliteSearch::Migration)
+      end
+    end
+
+    rake_tasks do
+      load File.expand_path("../tasks/sqlite_search.rake", __dir__)
+    end
+  end
+end

data/lib/sqlite_search/sql.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module SqliteSearch
+  module Sql
+    module_function
+
+    # Build "CASE <pk_sql> WHEN <id> THEN <value> ... END" from an id => value
+    # map, quoting each id and value through the connection. Used to carry
+    # per-row values computed in Ruby (KNN order, distances, fused scores) into
+    # the SQL as an ORDER expression or a selected column.
+    def id_case(pk_sql, values, connection)
+      whens = values.map { |id, value| "WHEN #{connection.quote(id)} THEN #{connection.quote(value)}" }.join(" ")
+      "CASE #{pk_sql} #{whens} END"
+    end
+  end
+end

data/lib/sqlite_search/vec/backend.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module SqliteSearch
+  module Vec
+    class Backend
+      def initialize(definition)
+        @definition = definition
+      end
+
+      def embed_and_store(record)
+        id = record.public_send(record.class.primary_key)
+        unless id.is_a?(Integer)
+          raise SqliteSearch::Error,
+            "vec indexing requires an integer primary key (used as the vec rowid), " \
+            "but #{record.class.name}##{record.class.primary_key} is #{id.class}."
+        end
+
+        model = @definition.neighbor_model
+        text = SqliteSearch::Vec.text_for(record, @definition.columns)
+        if text.empty?
+          model.where(id: id).delete_all
+          return
+        end
+
+        vector = @definition.embed(text, record_model: record.class) # compute first; if this raises, old row is untouched
+        model.transaction do
+          model.where(id: id).delete_all
+          model.create!(id: id, embedding: vector)
+        end
+      end
+
+      def remove(record)
+        id = record.public_send(record.class.primary_key)
+        @definition.neighbor_model.where(id: id).delete_all
+      end
+
+      def reembed(model)
+        @definition.neighbor_model.delete_all
+        model.find_each { |record| embed_and_store(record) }
+      end
+    end
+  end
+end

data/lib/sqlite_search/vec/definition.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require "sqlite_search/vec"
+
+module SqliteSearch
+  module Vec
+    class Definition
+      attr_reader :model, :name, :columns, :dimensions, :distance, :table_name
+
+      def initialize(model:, name:, against:, dimensions:, distance: :cosine, embedder: nil)
+        @model = model
+        @name = name.to_sym
+        @columns = SqliteSearch::Vec.columns_for(against)
+        @dimensions = dimensions
+        @distance = distance.to_sym
+        @embedder = embedder
+        @table_name = "#{model.table_name}_#{@name}_vec"
+
+        SqliteSearch::Vec.vec0_metric(@distance) # validates the distance is supported
+      end
+
+      def column_names = columns.map(&:to_s)
+      def cosine? = distance == :cosine
+      def distance_method = "#{name}_distance"
+      def similarity_method = "#{name}_similarity"
+
+      def embed(text, record_model: model)
+        callable = @embedder || SqliteSearch.config.embedder
+        raise SqliteSearch::Error, "No embedder registered. Call SqliteSearch.embedder { |text, model:, scope:| ... }." unless callable
+        vector = callable.call(text, model: record_model, scope: name)
+        if vector.nil? || vector.length != dimensions
+          raise SqliteSearch::Error, "Embedder returned #{vector&.length.inspect}-dim vector for #{model.name}##{name}, expected #{dimensions}."
+        end
+        vector
+      end
+
+      def neighbor_model
+        @neighbor_model ||= begin
+          tbl = table_name
+          dims = dimensions
+          Class.new(ActiveRecord::Base) do
+            self.table_name = tbl
+            self.primary_key = "id"
+            has_neighbors :embedding, dimensions: dims
+          end
+        end
+      end
+    end
+  end
+end

data/lib/sqlite_search/vec.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module SqliteSearch
+  module Vec
+    # Supported distance metrics, mapped to the vec0 table's distance_metric name.
+    # neighbor computes the distance with the matching vec_distance_* function at
+    # query time, so the scope's distance is what ranks results; the table metric
+    # is kept in step to keep the schema honest. inner_product is queryable by
+    # neighbor but is not a valid vec0 distance_metric, so it is not offered.
+    DISTANCE_METRICS = {cosine: "cosine", euclidean: "l2", taxicab: "l1"}.freeze
+
+    @loaded = false
+
+    module_function
+
+    def vec0_metric(distance)
+      DISTANCE_METRICS.fetch(distance.to_sym) do
+        raise SqliteSearch::Error,
+          "Unsupported distance #{distance.inspect}. Use one of: #{DISTANCE_METRICS.keys.join(", ")}."
+      end
+    end
+
+    # Lazily load the optional vector dependencies. Raises a clear error when
+    # they are missing so a declared vec_scope fails at boot, not mid-query.
+    def load!
+      return if @loaded
+      begin
+        require "neighbor"
+        require "sqlite_vec"
+      rescue LoadError => e
+        raise SqliteSearch::Error,
+          "Vector search needs the `neighbor` and `sqlite-vec` gems. " \
+          "Add them to your Gemfile. (#{e.message})"
+      end
+      Neighbor::SQLite.initialize!
+      @loaded = true
+    end
+
+    def columns_for(against)
+      case against
+      when Array then against.map(&:to_sym)
+      else [against.to_sym]
+      end
+    end
+
+    # Text handed to the embedder: the against columns joined, blanks dropped.
+    def text_for(record, columns)
+      columns.map { |c| record.public_send(c) }.reject { |v| v.nil? || v.to_s.strip.empty? }.join("\n")
+    end
+  end
+end

data/lib/sqlite_search/version.rb

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

data/lib/sqlite_search.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require "active_record"
+require "sqlite_search/version"
+require "sqlite_search/errors"
+require "sqlite_search/sql"
+require "sqlite_search/config"
+require "sqlite_search/vec"
+require "sqlite_search/vec/definition"
+require "sqlite_search/vec/backend"
+require "sqlite_search/hybrid"
+require "sqlite_search/embed_job"
+require "sqlite_search/query"
+require "sqlite_search/fts5"
+require "sqlite_search/migration"
+require "sqlite_search/fts5/definition"
+require "sqlite_search/fts5/backend"
+require "sqlite_search/model"
+
+module SqliteSearch
+end
+
+require "sqlite_search/railtie" if defined?(Rails::Railtie)

data/lib/tasks/sqlite_search.rake

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+namespace :sqlite_search do
+  desc "Rebuild FTS5 index. Usage: rake sqlite_search:reindex[Post,by_body] (scope optional)"
+  task :reindex, [:model, :scope] => :environment do |_t, args|
+    raise ArgumentError, "model is required" unless args[:model]
+    klass = args[:model].constantize
+    klass.reindex(args[:scope])
+    puts "Reindexed #{klass}#{" (#{args[:scope]})" if args[:scope]}."
+  end
+
+  desc "Re-embed a vec index. Usage: rake sqlite_search:reembed[Post,semantic]"
+  task :reembed, [:model, :scope] => :environment do |_t, args|
+    raise ArgumentError, "model is required" unless args[:model]
+    klass = args[:model].constantize
+    klass.reembed(args[:scope])
+    puts "Re-embedded #{klass}#{" (#{args[:scope]})" if args[:scope]}."
+  end
+end