truffler 0.1.0

data/lib/truffler/jobs/resume_job.rb

@@ -0,0 +1,61 @@
+module Truffler
+  module Jobs
+    # 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
+    # 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.
+    class ResumeJob < ActiveJob::Base
+      queue_as { Truffler.config.queue_name }
+
+      class_attribute :embedding_sweep_limit, default: 1_000
+      class_attribute :embedding_sweep_interval, default: 1.hour
+
+      def perform(record_type = nil)
+        models = record_type ? [ record_type.safe_constantize ].compact : Truffler.registry.models
+        models.select { |model| model.respond_to?(:truffler_definition) && model.truffler_definition }.each { |model| sweep(model) }
+      end
+
+      private
+
+      def sweep(model)
+        states = Records::RecordState.for_model(model)
+        cutoff = Time.current - Truffler.config.resume_pending_after
+        requeued = requeue(states.where(status: "failed").or(states.where(status: "labeling").where(claimed_at: ...cutoff)))
+        waiting = states.where(status: "pending").where(updated_at: ...cutoff)
+
+        live = (requeued.select { |_, priority| priority == "live" }.map(&:first) +
+          waiting.where(priority: "live").distinct.pluck(:tenant_key)).uniq
+        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
+        sweep_embeddings(model)
+      end
+
+      def sweep_embeddings(model)
+        return unless Embeddings.managed?(model.truffler_definition)
+
+        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)
+      end
+
+      # Returns the distinct [tenant_key, priority] pairs it moved to pending.
+      def requeue(scope)
+        pairs = scope.distinct.pluck(:tenant_key, :priority)
+        scope.update_all(status: "pending", attempts: 0, claimed_at: nil, updated_at: Time.current) if pairs.any?
+        pairs
+      end
+    end
+  end
+end

data/lib/truffler/jobs/smart_search_job.rb

@@ -0,0 +1,14 @@
+module Truffler
+  module Jobs
+    # Plans one Smart run: budget, encoding wait, filters, chunk fan-out.
+    # The only argument is the run id; the query waits in the cache
+    # (encrypted on encrypted models).
+    class SmartSearchJob < ActiveJob::Base
+      queue_as { Truffler.config.queue_name }
+
+      def perform(run_id)
+        SmartSearch::Dispatcher.new.call(SmartSearch::Run.load(run_id))
+      end
+    end
+  end
+end

data/lib/truffler/label_definition.rb

@@ -0,0 +1,153 @@
+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
+  # key, which makes the vocabulary per-tenant.
+  #
+  # A label with `from:` is supplied by the host: its answer is read from the
+  # record in the shape Jev answers normalize to and Jev is never asked. Its
+  # question is optional, `watch:` names extra columns that refresh it, and
+  # `version:` forces a refresh when its logic changes.
+  class LabelDefinition
+    TYPES = %i[noul choice score].freeze
+    KEY = /\A[a-z][a-z0-9_]*\z/
+
+    attr_reader :key, :type, :instructions, :filter_at, :boost, :watch, :version
+    # Intent weight a filter decision adds to the KTD20 query vector (default 0).
+    attr_reader :filter_weight
+
+    def initialize(key, type, question: nil, criteria: nil, options: nil, legend: nil, filter_at: nil, boost: nil, filter_weight: 0.0,
+      description: nil, from: nil, watch: nil, version: nil)
+      @key = key.to_s
+      @type = type.to_sym
+      @instructions = question
+      @criteria = criteria
+      @options = options
+      @legend = legend
+      @filter_at = filter_at&.to_f
+      @boost = boost&.to_f
+      @filter_weight = Float(filter_weight)
+      @description = description
+      @from = from
+      @watch = Array(watch).map(&:to_s)
+      @version = version
+      validate!
+    end
+
+    def supplied?
+      !@from.nil?
+    end
+
+    # The label's wording for query encoding: the description, else the question, else the key.
+    def description
+      @description.presence || instructions.presence || key
+    end
+
+    def per_tenant?
+      @options.respond_to?(:call)
+    end
+
+    def options(tenant_key = nil)
+      options = per_tenant? ? @options.call(tenant_key) : @options
+      options = Array(options).to_h { |option| [ option, nil ] } unless options.is_a?(Hash)
+      raise DefinitionError, "#{key}: choice options for tenant #{tenant_key.inspect} are empty" if options.empty?
+
+      options.to_h { |option, description| [ option.to_s, description ] }
+    end
+
+    def question(tenant_key = nil)
+      wording = instructions.presence || description
+      Questions.build do |questions|
+        case type
+        when :noul then questions.noul(key, instructions: wording, criteria: @criteria)
+        when :choice then questions.choice(key, instructions: wording, criteria: options(tenant_key))
+        when :score then questions.score(key, instructions: wording, criteria: @legend)
+        end
+      end.fetch(key)
+    end
+
+    # Stored supplied values change with their shape and version, never with
+    # the Jev model. The description and option wording are digested too
+    # because query encoding asks about them.
+    def supplied_fingerprint(tenant_key = nil)
+      Canonical.digest(supplied: true, type: type, options: (options(tenant_key) if type == :choice),
+        levels: (levels if type == :score), description: description, version: version)
+    end
+
+    # {storage_key => value} read from the record, or nil when the host has
+    # no answer. Raises InvalidSuppliedAnswer for a value out of shape.
+    def supplied_values(record, tenant_key = nil)
+      value = @from.call(record)
+      return if value.nil?
+
+      case type
+      when :noul then { key => probability(value) }
+      when :score then { key => level(value) }
+      when :choice then choice_values(value, options(tenant_key).keys)
+      end
+    end
+
+    def storage_keys(tenant_key = nil)
+      type == :choice ? options(tenant_key).keys.map { |option| "#{key}:#{option}" } : [ key ]
+    end
+
+    # The label part of its Jev question id ("<tag>__<question_key>").
+    def question_key
+      key
+    end
+
+    private
+
+    def levels
+      Array(@legend).size
+    end
+
+    def probability(value)
+      return value ? 1.0 : 0.0 if [ true, false ].include?(value)
+      return value.to_f if value.is_a?(Numeric) && value.between?(0, 1)
+
+      raise InvalidSuppliedAnswer, "#{key}: expected a probability from 0 to 1 or true/false, got #{value.class}"
+    end
+
+    def level(value)
+      return value.to_f / (levels - 1) if value.is_a?(Integer) && value.between?(0, levels - 1)
+
+      raise InvalidSuppliedAnswer, "#{key}: expected a level index from 0 to #{levels - 1}"
+    end
+
+    def choice_values(value, options)
+      probabilities = value.is_a?(Hash) ? value.to_h { |option, share| [ option.to_s, share ] } : { value.to_s => 1.0 }
+      unknown = probabilities.keys - options
+      raise InvalidSuppliedAnswer, "#{key}: #{unknown.size} answer option(s) are not declared options" if unknown.any?
+
+      options.to_h { |option| [ "#{key}:#{option}", probabilities.key?(option) ? probability(probabilities[option]) : 0.0 ] }
+    end
+
+    def validate!
+      raise DefinitionError, "label #{key.inspect} must match #{KEY.source} without a double underscore" unless valid_key?
+      raise DefinitionError, "label #{key.inspect} is reserved for lens dimensions" if key == Lenses::KEY_PREFIX
+      raise DefinitionError, "#{key}: type must be one of #{TYPES.join(', ')}" unless TYPES.include?(type)
+      validate_supplied!
+      raise DefinitionError, "#{key}: choice labels need options:" if type == :choice && @options.blank?
+      raise DefinitionError, "#{key}: score labels need a legend: of at least two levels" if type == :score && levels < 2
+
+      question unless per_tenant?
+    rescue ArgumentError => error
+      raise DefinitionError, error.message
+    end
+
+    def validate_supplied!
+      unless supplied?
+        raise DefinitionError, "#{key}: a question is required" if instructions.blank?
+        raise DefinitionError, "#{key}: watch: and version: need from:" if @watch.any? || !@version.nil?
+
+        return
+      end
+      raise DefinitionError, "#{key}: from: must be callable with the record" unless @from.respond_to?(:call)
+    end
+
+    def valid_key?
+      KEY.match?(key) && !key.include?(Questions::SEPARATOR)
+    end
+  end
+end

data/lib/truffler/labeling/backfill.rb

@@ -0,0 +1,189 @@
+module Truffler
+  module Labeling
+    # Relabels one model's records whose labels are missing, stale (labeled
+    # under another vocabulary version), failed, or demoted to backfill
+    # priority. It walks newest-first below an id cursor, splits each page by
+    # tenant, packs batch_size records per request at backfill priority, and
+    # asks only the stale questions. Live pending rows belong to the flush job
+    # and are never touched.
+    #
+    # Resumable: the cursor moves past a page only once the whole page is
+    # done, and records already current are skipped, so a rerun never asks
+    # Jev about them again. A spend cap stops the run before a request would
+    # exceed it; host-supplied labels cost nothing, so they are still written
+    # once the cap is reached. A Jev error releases the claimed rows and ends
+    # the run with `:client_error`, so the caller keeps the spend metered so far.
+    class Backfill
+      Result = Data.define(:status, :labeled, :requests, :cost, :cursor)
+
+      STATES = Records::RecordState.table_name
+
+      class SpendCapReached < StandardError; end
+
+      # Wraps the client to meter spend per request and refuse a request whose
+      # estimated cost would push spend past the cap.
+      class SpendMeter
+        attr_reader :spent, :requests
+
+        def initialize(client, cap:, spent:, config: Truffler.config)
+          @client = client
+          @cap = cap
+          @spent = spent.to_f
+          @requests = 0
+          @config = config
+        end
+
+        def ask(state:, questions:, **options)
+          raise SpendCapReached if @cap && @spent + estimate(state, questions) > @cap
+
+          answers = @client.ask(state: state, questions: questions, **options)
+          @requests += 1
+          @spent += answers.usage&.cost.to_f
+          answers
+        end
+
+        private
+
+        def estimate(state, questions)
+          @config.cost_for(Tokens.estimate({ state: state, questions: questions }))
+        end
+      end
+
+      def self.status(model)
+        new(model).status
+      end
+
+      attr_reader :model, :batch_size, :page_size
+
+      def initialize(model, spend_cap: Truffler.config.backfill_spend_cap, batch_size: Truffler.config.batch_size,
+        page_size: nil, cursor: nil, spent: 0.0, client: Truffler.config.client, budget: Budget.new)
+        @model = model
+        @batch_size = batch_size
+        @page_size = page_size || batch_size * 5
+        @cursor = cursor
+        @meter = SpendMeter.new(client, cap: spend_cap, spent: spent)
+        @budget = budget
+        @versions = {}
+      end
+
+      def run(max_pages: nil)
+        @labeled = 0
+        @started_spent = @meter.spent
+        cursor = @cursor
+        pages = 0
+
+        loop do
+          scanned, rows = page(cursor)
+          return result(:complete, nil) if scanned.empty?
+          return result(:paused, cursor) if max_pages && pages >= max_pages
+
+          rows.group_by(&:last).each do |tenant_key, tenant_rows|
+            tenant_rows.map(&:first).each_slice(batch_size) do |ids|
+              stop = label(ids, tenant_key)
+              return result(stop, cursor) if stop
+            end
+          end
+
+          cursor = scanned.last
+          pages += 1
+          return result(:complete, nil) if scanned.size < page_size
+        end
+      end
+
+      def status
+        counts = states.group(:status).count
+        labeled = states.where(status: "labeled").group(:tenant_key, :vocabulary_version).count
+        stale = labeled.sum { |(tenant_key, version), count| version == version_for(tenant_key) ? 0 : count }
+        { total: model.count, missing: model.joins(state_join).where("#{STATES}.id IS NULL").count,
+          pending: counts["pending"].to_i, labeling: counts["labeling"].to_i, labeled: counts["labeled"].to_i,
+          failed: counts["failed"].to_i, stale: stale, current: labeled.values.sum - stale }
+      end
+
+      private
+
+      def definition
+        model.truffler_definition
+      end
+
+      def record_type
+        model.polymorphic_name
+      end
+
+      def states
+        Records::RecordState.for_model(model)
+      end
+
+      def queue
+        @queue ||= Queue.new(model)
+      end
+
+      def version_for(tenant_key)
+        @versions[tenant_key] ||= definition.vocabulary.version(tenant_key: tenant_key, all_users: true)
+      end
+
+      def result(status, cursor)
+        Result.new(status: status, labeled: @labeled, requests: @meter.requests, cost: @meter.spent - @started_spent,
+          cursor: cursor)
+      end
+
+      # Returns the scanned ids (for the cursor) and the [id, tenant_key] rows
+      # that still need labels. Per-tenant vocabularies (per-tenant choices or
+      # lenses) compare versions here because each tenant has its own.
+      def page(cursor)
+        pk = model.primary_key
+        scope = model.joins(state_join).where(needs_labeling_sql)
+        scope = scope.where(model.arel_table[pk].lt(cursor)) if cursor
+        tenant = definition.scoped? ? model.arel_table[definition.tenant_column] : Arel.sql("NULL")
+        plucked = scope.reorder(pk => :desc).limit(page_size)
+          .pluck(model.arel_table[pk], tenant, Arel.sql("#{STATES}.status"), Arel.sql("#{STATES}.vocabulary_version"))
+
+        rows = plucked.filter_map do |id, tenant_key, status, version|
+          tenant_key = tenant_key&.to_s
+          [ id, tenant_key ] unless status == "labeled" && version == version_for(tenant_key)
+        end
+        [ plucked.map(&:first), rows ]
+      end
+
+      def state_join
+        pk = "#{model.quoted_table_name}.#{model.connection.quote_column_name(model.primary_key)}"
+        ActiveRecord::Base.sanitize_sql_array([
+          "LEFT OUTER JOIN #{STATES} ON #{STATES}.record_type = ? AND #{STATES}.record_id = #{pk}", record_type
+        ])
+      end
+
+      def needs_labeling_sql
+        stale = if definition.per_tenant_vocabulary? || Lenses::Lens.active.for_model(model).exists?
+          "#{STATES}.status = 'labeled'"
+        else
+          ActiveRecord::Base.sanitize_sql_array([
+            "(#{STATES}.status = 'labeled' AND (#{STATES}.vocabulary_version IS NULL OR #{STATES}.vocabulary_version <> ?))",
+            version_for(nil)
+          ])
+        end
+        "#{STATES}.id IS NULL OR #{STATES}.status = 'failed' OR " \
+          "(#{STATES}.status = 'pending' AND #{STATES}.priority = 'backfill') OR #{stale}"
+      end
+
+      # Labels one tenant chunk. Returns nil when done, or the status that
+      # stops the run; claimed rows that were not labeled go back to pending
+      # at backfill priority.
+      def label(ids, tenant_key)
+        claimed = queue.claim_backfill(ids, tenant_key)
+        return if claimed.empty?
+
+        begin
+          Labeler.new(model, client: @meter, budget: @budget).label(claimed, priority: :backfill)
+        rescue BudgetExhausted, SpendCapReached => error
+          queue.demote(claimed)
+          return error.is_a?(SpendCapReached) ? :spend_cap_reached : :budget_denied
+        rescue ClientError, IncompleteAnswers => error
+          queue.release(claimed, error)
+          return :client_error
+        ensure
+          @labeled += Records::RecordState.where(id: claimed.map(&:id), status: "labeled").count
+        end
+        nil
+      end
+    end
+  end
+end

data/lib/truffler/labeling/labeler.rb

@@ -0,0 +1,161 @@
+module Truffler
+  module Labeling
+    # Labels claimed record states of one tenant. It asks Jev only the labels
+    # that are missing or whose fingerprint is stale, packs records into as
+    # few requests as the budget allows, stores each answer as numeric label
+    # rows, and marks states labeled as each request lands, so a failure
+    # partway leaves finished records finished. Stale host-supplied labels
+    # are written first, before any budget slot or Jev call, so a Jev outage
+    # or denial never holds them back.
+    class Labeler
+      Result = Data.define(:labeled, :requests, :cost, :demoted)
+
+      attr_reader :model, :client, :budget
+
+      def initialize(model, client: Truffler.config.client, budget: Budget.new)
+        @model = model
+        @client = client
+        @budget = budget
+      end
+
+      def label(states, priority:)
+        return Result.new(labeled: 0, requests: 0, cost: 0.0, demoted: false) if states.empty?
+
+        tenant_key = tenant_key_of(states)
+        records = load_records(states)
+        vocabulary = definition.vocabulary
+        @labels = vocabulary.labels_for(tenant_key: tenant_key, all_users: true)
+        fingerprints = vocabulary.fingerprints(tenant_key: tenant_key, all_users: true)
+        version = vocabulary.version(tenant_key: tenant_key, all_users: true)
+        states_by_id = states.index_by { |state| state.record_id.to_s }
+
+        stored = stored_fingerprints(records)
+        askable = askable_labels
+        stale = records.map { |record| [ record, stale_keys(stored[record.id.to_s].to_h, fingerprints, tenant_key, askable) ] }
+        supplied = askable.select(&:supplied?).map(&:key)
+        supplier = Supplied.new(model)
+        written = supplier.write(stale.map { |record, keys| [ record, keys & supplied ] }, tenant_key: tenant_key)
+        current, pending = stale.map { |record, keys| [ record, keys - supplied ] }.partition { |_, keys| keys.empty? }
+        settled = current.reject { |record, _| supplier.failed_ids.include?(record.id) }
+        Records::RecordState.mark_labeled(settled.map { |record, _| states_by_id[record.id.to_s].id }, version: version)
+        Embeddings::LabelVector.new(model).write(current.map { |record, _| record.id } - written, tenant_key: tenant_key)
+
+        requests = RequestBuilder.new(definition, tenant_key: tenant_key, labels: @labels).build(pending)
+        cost = 0.0
+        requests.each_with_index do |request, index|
+          decision = budget.acquire(priority: priority, tenant_key: (tenant_key if index.zero?),
+            records: index.zero? ? pending.size : 1)
+          if decision.demoted?
+            retry_failed_supplied(supplier.failed_ids, states_by_id)
+            return Result.new(labeled: current.size, requests: index, cost: cost, demoted: true)
+          end
+          raise BudgetExhausted, "no Jev budget for #{priority} labeling" if decision.denied?
+
+          answers = client.ask(state: request.state, questions: request.questions, priority: decision.priority)
+          cost += answers.usage&.cost.to_f
+          charge_lenses(request, answers.usage&.cost.to_f)
+          store(request, answers, fingerprints, tenant_key, version, states_by_id)
+        end
+
+        retry_failed_supplied(supplier.failed_ids, states_by_id)
+        Result.new(labeled: records.size, requests: requests.size, cost: cost, demoted: false)
+      end
+
+      private
+
+      # A failed `from` must not leave its record looking current, or backfill
+      # would never ask again. Supplied labels cost nothing, so retrying at
+      # backfill priority is free.
+      def retry_failed_supplied(record_ids, states_by_id)
+        ids = record_ids.filter_map { |id| states_by_id[id.to_s]&.id }
+        Records::RecordState.mark_supplied_failed(ids, error_class: SuppliedLabelFailed.name) if ids.any?
+      end
+
+      def definition
+        model.truffler_definition
+      end
+
+      def record_type
+        model.polymorphic_name
+      end
+
+      def tenant_key_of(states)
+        tenants = states.map(&:tenant_key).uniq
+        raise TenantMismatch, "labeling holds records from exactly one tenant" if tenants.size > 1
+
+        tenants.first
+      end
+
+      def load_records(states)
+        records = model.where(model.primary_key => states.map(&:record_id)).to_a
+        found = records.map { |record| record.id.to_s }
+        gone = states.reject { |state| found.include?(state.record_id.to_s) }
+        Records::RecordState.where(id: gone.map(&:id)).delete_all if gone.any?
+        records
+      end
+
+      def stale_keys(stored, fingerprints, tenant_key, askable)
+        askable.reject do |label|
+          label.storage_keys(tenant_key).all? { |key| stored[key] == fingerprints[label.key] }
+        end.map(&:key)
+      end
+
+      # Every label except those of lenses at their spend cap (R43).
+      def askable_labels
+        lens_labels = @labels.values.grep(Lenses::LensLabel)
+        @lenses = nil
+        return @labels.values if lens_labels.empty?
+
+        @lenses = Lenses::Lens.where(id: lens_labels.map(&:lens_id).uniq).index_by(&:id)
+        @labels.values.reject { |label| label.is_a?(Lenses::LensLabel) && @lenses[label.lens_id]&.spend_cap_reached? }
+      end
+
+      # Splits a request's cost over its questions and adds each lens's share
+      # to that lens's spend.
+      def charge_lenses(request, cost)
+        return if @lenses.blank? || cost.zero?
+
+        keys = request.entries.values.flat_map { |_, keys| keys }
+        keys.filter_map { |key| @labels[key] }.grep(Lenses::LensLabel).group_by(&:lens_id).each do |lens_id, asked|
+          @lenses[lens_id]&.record_spend!(cost * asked.size / keys.size)
+        end
+      end
+
+      def stored_fingerprints(records)
+        rows = Records::Label.where(record_type: record_type, record_id: records.map(&:id)).pluck(:record_id, :label_key, :fingerprint)
+        rows.each_with_object({}) { |(id, key, print), map| (map[id.to_s] ||= {})[key] = print }
+      end
+
+      def store(request, answers, fingerprints, tenant_key, version, states_by_id)
+        now = Time.current
+        rows = request.entries.flat_map do |tag, (record, keys)|
+          keys.flat_map do |key|
+            label = @labels.fetch(key)
+            label_rows(label, Questions.tagged_id(tag, label.question_key), answers, tenant_key).map do |label_key, value|
+              { record_type: record_type, record_id: record.id, tenant_key: tenant_key, label_key: label_key,
+                value: value, fingerprint: fingerprints[key], labeled_at: now }
+            end
+          end
+        end
+
+        Records::Label.transaction do
+          request.entries.each_value do |record, keys|
+            keys.each { |key| Records::Label.where(record_type: record_type, record_id: record.id).for_label(key).delete_all }
+          end
+          Records::Label.insert_all!(rows) if rows.any?
+          Embeddings::LabelVector.new(model).write(request.entries.values.map { |record, _| record.id }, tenant_key: tenant_key)
+          ids = request.entries.values.map { |record, _| states_by_id[record.id.to_s].id }
+          Records::RecordState.mark_labeled(ids, version: version)
+        end
+      end
+
+      def label_rows(label, id, answers, tenant_key)
+        case label.type
+        when :noul then [ [ label.key, answers.noul(id) ] ]
+        when :score then [ [ label.key, answers.score(id) ] ]
+        when :choice then label.options(tenant_key).keys.map { |option| [ "#{label.key}:#{option}", answers.probability(id, option) ] }
+        end
+      end
+    end
+  end
+end