truffler 0.1.4 → 0.1.6

data/lib/truffler/configuration.rb

@@ -19,6 +19,14 @@ module Truffler
       :smart_run_ttl, :smart_candidate_pool, :broadcaster
     # Generic nouns that are never required keywords on their own (Search::Filler).
     attr_accessor :filler_words
+    # ->(model, tenant_key) { true/false } for tenant-scoped models; nil indexes every tenant.
+    attr_accessor :tenant_enabled
+    # :tenant (a spend ledger per tenant for scoped models) or :app (one per model).
+    attr_accessor :backfill_spend_cap_scope
+    # Choice options below this probability store no row and read as 0.0 (nil stores every option).
+    attr_accessor :choice_min_probability
+    # Query encoding omits choice options the tenant has no label row for (QueryEncoding::PresentOptions).
+    attr_accessor :skip_empty_options
 
     def initialize(env: ENV)
       @model = "jev-latest"
@@ -53,6 +61,10 @@ module Truffler
       @smart_candidate_pool = 200
       @broadcaster = nil
       @filler_words = Search::Filler::DEFAULT_WORDS.dup
+      @tenant_enabled = nil
+      @backfill_spend_cap_scope = :tenant
+      @choice_min_probability = 0.05
+      @skip_empty_options = false
     end
 
     def client

data/lib/truffler/current.rb

@@ -0,0 +1,30 @@
+module Truffler
+  # State scoped to one unit of work: a keystroke search, a query encoding,
+  # a labeler batch, a Smart run step. Inside `scope`, a per-tenant choice
+  # `options:` callable is resolved once per (label, tenant) instead of at
+  # every vocabulary, fingerprint, and wording read. Nested scopes share the
+  # outermost one, which clears everything when it ends, so nothing carries
+  # over to the next search or job. Outside a scope nothing is memoized.
+  module Current
+    KEY = :truffler_current_options
+
+    def self.scope
+      return yield if ActiveSupport::IsolatedExecutionState.key?(KEY)
+
+      ActiveSupport::IsolatedExecutionState[KEY] = {}
+      begin
+        yield
+      ensure
+        ActiveSupport::IsolatedExecutionState.delete(KEY)
+      end
+    end
+
+    def self.options(label, tenant_key)
+      memo = ActiveSupport::IsolatedExecutionState[KEY]
+      return yield unless memo
+
+      key = [ label, tenant_key&.to_s ]
+      memo.fetch(key) { memo[key] = yield }
+    end
+  end
+end

data/lib/truffler/definition.rb

@@ -116,6 +116,42 @@ module Truffler
       @weak_below || DEFAULT_WEAK_BELOW
     end
 
+    attr_accessor :index_if, :index_scope
+
+    # config.tenant_enabled is asked only for tenant-scoped models; an
+    # unscoped model is always enabled.
+    def tenant_enabled?(tenant_key)
+      check = Truffler.config.tenant_enabled
+      return true if check.nil? || !scoped?
+
+      check.call(model, tenant_key) ? true : false
+    end
+
+    # Whether the after-commit hooks, Queue, and EmbedJob handle this record:
+    # its tenant is enabled and `index_if` (when declared) accepts it.
+    def indexable?(record)
+      tenant_enabled?(tenant_key_for(record)) && (index_if.nil? || index_if.call(record) ? true : false)
+    end
+
+    # The relation the batch paths (backfills, sweeps) page over.
+    def index_relation(relation = model.all)
+      index_scope ? index_scope.call(relation) : relation
+    end
+
+    # The tenant a backfill spend ledger row belongs to: the tenant for
+    # scoped models under backfill_spend_cap_scope :tenant, else nil (app-wide).
+    def ledger_tenant(tenant_key)
+      tenant_key if scoped? && Truffler.config.backfill_spend_cap_scope.to_sym == :tenant
+    end
+
+    attr_writer :invite_on_pending_encoding
+
+    # Whether a model with a `keyword` source shows the Smart search row while
+    # the query's encoding is in flight (AE10). Default true.
+    def invite_on_pending_encoding
+      @invite_on_pending_encoding.nil? || @invite_on_pending_encoding
+    end
+
     private
 
     def table_available?
@@ -204,6 +240,22 @@ module Truffler
       def weak_below(count)
         @definition.weak_below = Integer(count)
       end
+
+      def index_if(callable)
+        raise DefinitionError, "index_if must be callable with the record" unless callable.respond_to?(:call)
+
+        @definition.index_if = callable
+      end
+
+      def index_scope(callable)
+        raise DefinitionError, "index_scope must be callable with a relation" unless callable.respond_to?(:call)
+
+        @definition.index_scope = callable
+      end
+
+      def invite_on_pending_encoding(enabled)
+        @definition.invite_on_pending_encoding = enabled ? true : false
+      end
     end
   end
 end

data/lib/truffler/embeddings/backfill.rb

@@ -3,11 +3,15 @@ module Truffler
     # Finds records whose embedding is missing or was made under another
     # fingerprint (model, width, or fields changed) and enqueues EmbedJob for
     # them, newest first, `batch_size` jobs at a time behind an id cursor.
-    # ResumeJob runs a bounded pass on every sweep; hosts call `enqueue` with
-    # no limit after enabling embeddings or changing the model, width, or
-    # fields.
+    # It pages over the definition's index_scope, one tenant at a time when
+    # given `tenant_key:`, skipping disabled tenants, with a NOT EXISTS
+    # anti-join against truffler_record_states so each page stops at its
+    # limit instead of materializing every current id. ResumeJob runs a
+    # bounded pass on every sweep; hosts call `enqueue` with no limit after
+    # enabling embeddings or changing the model, width, or fields.
     class Backfill
       BATCH_SIZE = 1_000
+      STATES = Records::RecordState.table_name
 
       attr_reader :model
 
@@ -15,25 +19,30 @@ module Truffler
         @model = model
       end
 
-      def stale_ids(limit: nil, before: nil)
+      def stale_ids(limit: nil, before: nil, tenant_key: nil)
         definition = model.truffler_definition
         return [] unless Embeddings.managed?(definition)
+        return [] if tenant_key && !definition.tenant_enabled?(tenant_key)
 
-        current = Records::RecordState.for_model(model).where(embedding_fingerprint: Embeddings.fingerprint(definition))
-          .where.not(embedded_at: nil).select(:record_id)
-        scope = model.where.not(model.primary_key => current)
-        scope = scope.where(model.primary_key => ...before) if before
-        scope.order(model.primary_key => :desc).limit(limit).pluck(model.primary_key)
+        pk = model.arel_table[model.primary_key]
+        scope = definition.index_relation(model.all).where(current_embedding_missing_sql(definition))
+        if tenant_key && definition.scoped?
+          scope = scope.where(definition.tenant_column => tenant_key)
+        elsif definition.scoped? && Truffler.config.tenant_enabled
+          scope = scope.where(definition.tenant_column => tenant_keys)
+        end
+        scope = scope.where(pk.lt(before)) if before
+        scope.reorder(pk.desc).limit(limit).pluck(pk)
       end
 
-      def enqueue(limit: nil, batch_size: BATCH_SIZE)
+      def enqueue(limit: nil, batch_size: BATCH_SIZE, tenant_key: nil)
         count = 0
         cursor = nil
         loop do
           take = limit ? [ batch_size, limit - count ].min : batch_size
           break unless take.positive?
 
-          ids = stale_ids(limit: take, before: cursor)
+          ids = stale_ids(limit: take, before: cursor, tenant_key: tenant_key)
           break if ids.empty?
 
           ActiveJob.perform_all_later(ids.map { |id| Jobs::EmbedJob.new(model.polymorphic_name, id) })
@@ -43,6 +52,26 @@ module Truffler
         end
         count
       end
+
+      # The enabled tenants with records in the index scope, for per-tenant sweeps.
+      def tenant_keys
+        definition = model.truffler_definition
+        return [ nil ] unless definition.scoped?
+
+        definition.index_relation(model.all).reorder(nil).distinct.pluck(definition.tenant_column)
+          .map(&:to_s).select { |key| definition.tenant_enabled?(key) }.sort
+      end
+
+      private
+
+      def current_embedding_missing_sql(definition)
+        pk = "#{model.quoted_table_name}.#{model.connection.quote_column_name(model.primary_key)}"
+        ActiveRecord::Base.sanitize_sql_array([
+          "NOT EXISTS (SELECT 1 FROM #{STATES} WHERE #{STATES}.record_type = ? AND #{STATES}.record_id = #{pk} " \
+          "AND #{STATES}.embedding_fingerprint = ? AND #{STATES}.embedded_at IS NOT NULL)",
+          model.polymorphic_name, Embeddings.fingerprint(definition)
+        ])
+      end
     end
   end
 end

data/lib/truffler/embeddings/label_vector.rb

@@ -33,7 +33,7 @@ module Truffler
             label_vocabulary_version: version, created_at: now, updated_at: now }
         end
         Records::Embedding.upsert_all(rows, unique_by: %i[record_type record_id],
-          update_only: %i[tenant_key label_vector label_vocabulary_version updated_at])
+          update_only: %i[tenant_key label_vector label_vocabulary_version])
         rows.size
       end
 

data/lib/truffler/embeddings/neighbor_store.rb

@@ -2,9 +2,13 @@ module Truffler
   module Embeddings
     # Cosine distance computed by the database: pgvector's `<=>` on a vector
     # column, or sqlite-vec's `vec_distance_cosine` on float32 blobs (the
-    # host loads the extension). Because similarity is plain SQL, search can
-    # score text inline in its one query as an exact scan of the tenant's
-    # rows; an approximate (HNSW) index is an opt-in host migration.
+    # host loads the extension).
+    #
+    # On Postgres, search reads text similarity from the tenant's top `k`
+    # neighbors (`neighbors_sql`, one `ORDER BY embedding <=> q LIMIT k`
+    # that an HNSW index can serve), so records outside the top K score no
+    # text similarity. Pass `top_k: false` for an exact inline similarity per
+    # row, the sqlite-vec default; `top_k: true` uses the join there too.
     class NeighborStore < VectorStore
       TABLE = "truffler_embeddings".freeze
 
@@ -35,16 +39,31 @@ module Truffler
         end
       end
 
-      def initialize(dialect: nil)
+      attr_reader :k
+
+      def initialize(dialect: nil, k: DEFAULT_K, top_k: nil)
         @dialect = dialect
+        @k = Integer(k)
+        @top_k = top_k
       end
 
-      def nearest(model, tenant_key:, vector:, k: DEFAULT_K)
+      def nearest(model, tenant_key:, vector:, k: self.k)
         distance = distance_sql(model, "#{TABLE}.embedding", vector)
         embeddings(model, tenant_key, vector.size).order(Arel.sql(distance)).limit(k)
           .pluck(:record_id, Arel.sql("1 - #{distance}")).map { |id, similarity| [ id, similarity.to_f ] }
       end
 
+      # The tenant's `k` nearest vectors as `(record_id, similarity)` rows,
+      # for search to LEFT JOIN on record_id; nil when this store scores
+      # inline instead.
+      def neighbors_sql(model, tenant_key:, vector:, k: self.k)
+        return unless top_k?(model)
+
+        distance = distance_sql(model, "#{TABLE}.embedding", vector)
+        embeddings(model, tenant_key, vector.size).reorder(Arel.sql(distance)).limit(k)
+          .select(Arel.sql("#{TABLE}.record_id AS record_id"), Arel.sql("1 - #{distance} AS similarity")).to_sql
+      end
+
       def similarity_sql(model, tenant_key:, vector:, k: nil)
         similarity = "1 - #{distance_sql(model, "#{TABLE}.embedding", vector)}"
         subquery = embeddings(model, tenant_key, vector.size).where("#{TABLE}.record_id = #{primary_key_sql(model)}")
@@ -55,11 +74,18 @@ module Truffler
         true
       end
 
+      def top_k?(model)
+        @top_k.nil? ? dialect(model) == :postgres : @top_k
+      end
+
       private
 
-      def distance_sql(model, column_sql, vector)
+      def dialect(model)
         @dialect ||= self.class.dialect(model.connection)
-        self.class.distance_sql(@dialect, model.connection, column_sql, vector)
+      end
+
+      def distance_sql(model, column_sql, vector)
+        self.class.distance_sql(dialect(model), model.connection, column_sql, vector)
       end
     end
   end

data/lib/truffler/embeddings/vector_store.rb

@@ -5,7 +5,12 @@ module Truffler
     # `similarity_sql` returns a scalar SQL expression over the model's table
     # that search can put in its SELECT or ORDER BY: stores that compute
     # similarity in the database (`inline_sql?`) scan every row exactly, and
-    # the rest fall back to a CASE over the top-K neighbors.
+    # the rest fall back to a CASE over the top-K neighbors. A store may also
+    # answer `neighbors_sql` with a `(record_id, similarity)` subquery, which
+    # search LEFT JOINs instead (see NeighborStore).
+    #
+    # `config.vector_store` is one of ADAPTERS or a store instance, such as
+    # `NeighborStore.new(k: 500)` or a subclass of this class.
     class VectorStore
       ADAPTERS = %i[auto ruby neighbor].freeze
       DEFAULT_K = 200
@@ -14,12 +19,13 @@ module Truffler
         embeddings = model.truffler_definition.embeddings
         return unless embeddings
         return ColumnStore.new(embeddings[:column]) if embeddings.key?(:column)
+        return config.vector_store if config.vector_store.respond_to?(:similarity_sql)
 
         case config.vector_store&.to_sym
         when :ruby then RubyStore.new
         when :neighbor then NeighborStore.new
         when :auto then NeighborStore.available?(model.connection) ? NeighborStore.new : RubyStore.new
-        else raise Error, "config.vector_store must be one of #{ADAPTERS.join(', ')}"
+        else raise Error, "config.vector_store must be one of #{ADAPTERS.join(', ')} or a store instance"
         end
       end
 
@@ -58,7 +64,7 @@ module Truffler
             fingerprint: fingerprint, embedding: Records::Embedding.encode(vector), dimensions: vector.size,
             created_at: now, updated_at: now },
           unique_by: %i[record_type record_id],
-          update_only: %i[tenant_key fingerprint embedding dimensions updated_at]
+          update_only: %i[tenant_key fingerprint embedding dimensions]
         )
       end
 

data/lib/truffler/jobs/backfill_job.rb

@@ -1,7 +1,9 @@
 module Truffler
   module Jobs
     # Backfills one model's stale, missing, failed, and demoted labels at
-    # backfill priority. Arguments are the record type, the id cursor, the
+    # backfill priority, over one tenant when given `tenant_key:` (the
+    # tenant's own spend ledger) or else the whole model's indexed records.
+    # Arguments are the record type, the tenant key, the id cursor, the
     # spend so far, the cap, the retry attempt, and the count of budget
     # denials in a row, never record text. A budget denial reschedules the job
     # from its cursor after the same backoff a waiting Labeling::Backfill uses
@@ -15,16 +17,23 @@ module Truffler
 
       queue_as { Truffler.config.queue_name }
 
-      def perform(record_type, cursor: nil, spent: 0.0, spend_cap: Truffler.config.backfill_spend_cap, max_pages: MAX_PAGES,
+      # The job's tenant keyword: the tenant for scoped models, none otherwise.
+      def self.tenant_argument(model, tenant_key)
+        model.truffler_definition.scoped? && !tenant_key.nil? ? { tenant_key: tenant_key.to_s } : {}
+      end
+
+      def perform(record_type, tenant_key: nil, cursor: nil, spent: 0.0, spend_cap: Truffler.config.backfill_spend_cap, max_pages: MAX_PAGES,
         attempt: 0, denials: 0)
         model = record_type.safe_constantize
         return unless model.respond_to?(:truffler_definition) && model.truffler_definition
 
-        result = Labeling::Backfill.new(model, cursor: cursor, spent: spent, spend_cap: spend_cap).run(max_pages: max_pages)
-        Instrumentation.instrument(:backfill, record_type: record_type, outcome: result.status,
+        result = Labeling::Backfill.new(model, tenant_key: tenant_key, cursor: cursor, spent: spent, spend_cap: spend_cap)
+          .run(max_pages: max_pages)
+        Instrumentation.instrument(:backfill, record_type: record_type, tenant_key: tenant_key, outcome: result.status,
           labeled_count: result.labeled, request_count: result.requests, cost: result.cost)
 
-        follow_up = { cursor: result.cursor, spent: spent + result.cost, spend_cap: spend_cap, max_pages: max_pages }
+        follow_up = { **self.class.tenant_argument(model, tenant_key), cursor: result.cursor, spent: spent + result.cost,
+          spend_cap: spend_cap, max_pages: max_pages }
         case result.status
         when :budget_denied then retry_after_denial(record_type, follow_up, result, denials)
         when :paused then self.class.perform_later(record_type, **follow_up)

data/lib/truffler/jobs/embed_job.rb

@@ -3,6 +3,7 @@ module Truffler
     # Embeds one record's declared fields and stores the vector. Arguments are
     # the record type and id only. On an embedder failure nothing is written,
     # `embedded_at` stays as it was, and the job retries; labels are untouched.
+    # A record the definition no longer indexes is skipped.
     class EmbedJob < ActiveJob::Base
       queue_as { Truffler.config.queue_name }
 
@@ -17,6 +18,7 @@ module Truffler
 
         record = model.find_by(model.primary_key => record_id)
         return Records::Embedding.where(record_type: record_type, record_id: record_id).delete_all unless record
+        return unless definition.indexable?(record)
 
         settings = definition.embeddings
         fingerprint = Embeddings.fingerprint(definition)

data/lib/truffler/jobs/label_flush_job.rb

@@ -4,7 +4,7 @@ module Truffler
     # and tenant key only. On a Jev or budget failure the claimed rows go back
     # to pending (or failed after max_attempts) and the job retries. Rows over
     # the tenant's live cap drop to backfill priority and a delayed backfill
-    # is scheduled for them.
+    # of that tenant is scheduled for them. A disabled tenant is left alone.
     class LabelFlushJob < ActiveJob::Base
       RETRYABLE = [ ClientError, BudgetExhausted, IncompleteAnswers ].freeze
 
@@ -18,6 +18,8 @@ module Truffler
 
         queue = Labeling::Queue.new(model)
         queue.clear_marker(tenant_key)
+        return unless model.truffler_definition.tenant_enabled?(tenant_key)
+
         states = queue.claim(tenant_key, priority: :live, limit: Truffler.config.batch_size)
         return if states.empty?
 
@@ -30,7 +32,7 @@ module Truffler
 
         if result.demoted
           queue.demote(states)
-          queue.schedule_backfill
+          queue.schedule_backfill(tenant_key)
         end
         queue.schedule(tenant_key) if queue.pending?(tenant_key, priority: :live)
       end

data/lib/truffler/jobs/resume_job.rb

@@ -3,12 +3,13 @@ module Truffler
     # A periodic sweep hosts schedule (every few minutes) so labeling resumes
     # after a Jev outage or a crashed worker. It returns failed rows and rows
     # stuck in labeling to pending, reschedules a flush for tenants whose live
-    # rows have waited past `resume_pending_after`, and starts a backfill for
-    # rows waiting at backfill priority. For models with gem-managed
+    # rows have waited past `resume_pending_after`, and starts a BackfillJob
+    # per tenant with rows waiting at backfill priority. Disabled tenants
+    # (config.tenant_enabled) get neither. For models with gem-managed
     # embeddings it also enqueues up to `embedding_sweep_limit` missing or
-    # stale embeddings, at most once per `embedding_sweep_interval`, so an
-    # embedder outage cannot pile duplicate jobs onto the queue. Pass a record
-    # type to sweep one model.
+    # stale embeddings, tenant by tenant for scoped models, at most once per
+    # `embedding_sweep_interval`, so an embedder outage cannot pile duplicate
+    # jobs onto the queue. Pass a record type to sweep one model.
     class ResumeJob < ActiveJob::Base
       queue_as { Truffler.config.queue_name }
 
@@ -28,16 +29,20 @@ module Truffler
         requeued = requeue(states.where(status: "failed").or(states.where(status: "labeling").where(claimed_at: ...cutoff)))
         waiting = states.where(status: "pending").where(updated_at: ...cutoff)
 
+        definition = model.truffler_definition
         live = (requeued.select { |_, priority| priority == "live" }.map(&:first) +
-          waiting.where(priority: "live").distinct.pluck(:tenant_key)).uniq
+          waiting.where(priority: "live").distinct.pluck(:tenant_key)).uniq.select { |key| definition.tenant_enabled?(key) }
         queue = Labeling::Queue.new(model)
         live.each do |tenant_key|
           queue.clear_marker(tenant_key)
           queue.schedule(tenant_key)
         end
 
-        backfill = requeued.any? { |_, priority| priority == "backfill" } || waiting.exists?(priority: "backfill")
-        BackfillJob.perform_later(model.polymorphic_name) if backfill
+        backfill = (requeued.select { |_, priority| priority == "backfill" }.map(&:first) +
+          waiting.where(priority: "backfill").distinct.pluck(:tenant_key)).uniq
+        backfill.select { |key| definition.tenant_enabled?(key) }.each do |tenant_key|
+          BackfillJob.perform_later(model.polymorphic_name, **BackfillJob.tenant_argument(model, tenant_key))
+        end
         sweep_embeddings(model)
       end
 
@@ -47,7 +52,15 @@ module Truffler
         marker = "truffler:embedding_sweep:#{model.polymorphic_name}"
         return unless Truffler.config.cache_store.write(marker, true, unless_exist: true, expires_in: embedding_sweep_interval)
 
-        Embeddings::Backfill.new(model).enqueue(limit: embedding_sweep_limit)
+        backfill = Embeddings::Backfill.new(model)
+        return backfill.enqueue(limit: embedding_sweep_limit) unless model.truffler_definition.scoped?
+
+        remaining = embedding_sweep_limit
+        backfill.tenant_keys.each do |tenant_key|
+          break unless remaining.positive?
+
+          remaining -= backfill.enqueue(limit: remaining, tenant_key: tenant_key)
+        end
       end
 
       # Returns the distinct [tenant_key, priority] pairs it moved to pending.

data/lib/truffler/label_definition.rb

@@ -1,7 +1,9 @@
 module Truffler
   # One typed label question. Nouls store their probability, scores their
   # normalized position, and choices one row per option ("label:option") with
-  # that option's probability. Choice options may be a callable of the tenant
+  # that option's probability, for options at or above
+  # config.choice_min_probability plus the most likely option; a missing
+  # option row reads as 0.0. Choice options may be a callable of the tenant
   # key, which makes the vocabulary per-tenant; a tenant it gives no options
   # (`{}` or nil) simply does not have the label. An option's value is its
   # description, or `{ description:, search: }` to give query encoding a
@@ -104,10 +106,19 @@ module Truffler
       case type
       when :noul then { key => probability(value) }
       when :score then { key => level(value) }
-      when :choice then choice_values(value, options(tenant_key).keys)
+      when :choice then self.class.sparse_choice(choice_values(value, options(tenant_key).keys))
       end
     end
 
+    # Drops choice option rows ({"label:option" => probability}) below
+    # config.choice_min_probability, always keeping the most likely option.
+    def self.sparse_choice(values, min: Truffler.config.choice_min_probability)
+      return values if min.nil? || values.empty?
+
+      top = values.max_by { |_, probability| probability.to_f }.first
+      values.select { |key, probability| key == top || probability.to_f >= min }
+    end
+
     def storage_keys(tenant_key = nil)
       type == :choice ? options(tenant_key).keys.map { |option| "#{key}:#{option}" } : [ key ]
     end
@@ -120,7 +131,10 @@ module Truffler
     private
 
     def option_entries(tenant_key)
-      options = per_tenant? ? @options.call(tenant_key) : @options
+      per_tenant? ? Current.options(self, tenant_key) { build_option_entries(@options.call(tenant_key)).freeze } : build_option_entries(@options)
+    end
+
+    def build_option_entries(options)
       options = Array(options).to_h { |option| [ option, nil ] } unless options.is_a?(Hash)
       options = options.to_h { |option, value| [ option.to_s, option_entry(option, value) ] }
       raise DefinitionError, "#{key}: the option name #{NO_OPTION} is reserved" if options.key?(NO_OPTION)