kuda-pilot conversation topic embedding backfill script

backfill_conversation_topic_embeddings.rb

# frozen_string_literal: true

# Backfill Voyage conversation topic embeddings via API2 for kuda-pilot.
# Prefer this over Maintenance::EnqueueConversationTopicEmbeddingsTask —
# that task's JobIteration keyset queries hit statement timeouts on prod.
#
# Calls API2 /internal/conversation_topic_embeddings/enqueue directly (60s
# timeout). API2 schedules the Turbopuffer upsert asynchronously and returns
# 202 quickly; this script tracks enqueue success, not embed completion.
#
# Iterates day-by-day and plucks ids (avoids find_each id-cursor timeouts on
# large org+started_at scopes). Raises the session statement_timeout for counts/plucks.
#
# Usage (Rails console / ECS Rails shell):
#
#   load "scripts/backfill_conversation_topic_embeddings.rb"
#
# Smoke test first (recommended — full window can be 50k+ rows):
#
#   LIMIT = 50
#   load "scripts/backfill_conversation_topic_embeddings.rb"
#
# Resume after interrupt (set START_AFTER_ID to last_id from logs):
#
#   START_AFTER_ID = 12345
#   load "scripts/backfill_conversation_topic_embeddings.rb"

ORGANIZATION_SLUG = 'kuda-pilot'
LOOKBACK_DAYS = 7
LOG_EVERY = 25
TIMEOUT_SECONDS = 60
OPEN_TIMEOUT = 5
STATEMENT_TIMEOUT = '5min'
DISCARD_GOOD_JOBS = true
LOOKBACK_HOURS = 6
DRY_RUN = false
LIMIT = nil # e.g. 50 for a smoke test; nil = all

# Optional resume cursor. Assign in console before load:
#   START_AFTER_ID = 12345
start_after_id = Object.const_defined?(:START_AFTER_ID) ? Object.const_get(:START_AFTER_ID) : nil

organization = Organization.find_by!(slug: ORGANIZATION_SLUG)
url = "#{RulebaseConfig.api2_base_url}/internal/conversation_topic_embeddings/enqueue"
key = RulebaseConfig.internal_api_key
since = LOOKBACK_DAYS.days.ago
now = Time.current

format_duration = lambda do |seconds|
  seconds = seconds.to_i
  hours = seconds / 3600
  minutes = (seconds % 3600) / 60
  secs = seconds % 60
  format('%02d:%02d:%02d', hours, minutes, secs)
end

topic_embed_text_present = lambda do |conversation|
  conversation.subject.present? ||
    conversation.summary.present? ||
    conversation.transcript.present?
end

log_progress = lambda do |stats, process_total, last_id, started_monotonic|
  elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_monotonic
  rate = stats[:processed].positive? && elapsed.positive? ? stats[:processed] / elapsed : 0.0
  remaining = process_total ? (process_total - stats[:processed]) : nil
  eta_seconds = remaining && rate.positive? ? (remaining / rate) : nil

  payload = {
    processed: stats[:processed],
    ok: stats[:ok],
    skipped: stats[:skipped],
    err: stats[:err],
    last_id: last_id,
    elapsed: format_duration.call(elapsed),
    rate_per_sec: format('%.2f', rate),
    resume_hint: last_id ? "START_AFTER_ID = #{last_id}" : nil
  }
  if process_total&.positive?
    payload[:total] = process_total
    payload[:pct] = format('%.1f%%', (stats[:processed] * 100.0) / process_total)
    payload[:eta] = eta_seconds ? format_duration.call(eta_seconds) : 'n/a'
  end
  puts(payload.inspect)
end

ActsAsTenant.with_tenant(organization) do
  previous_timeout = ActiveRecord::Base.connection.select_value('SHOW statement_timeout')
  ActiveRecord::Base.connection.execute("SET statement_timeout = '#{STATEMENT_TIMEOUT}'")

  begin
    base_scope = organization.conversations.where(spam: false).where(started_at: since..)
    base_scope = base_scope.where(id: (start_after_id + 1)..) if start_after_id

    candidate_total = base_scope.count
    process_total = LIMIT ? [candidate_total, LIMIT].min : candidate_total

    puts({
      organization_slug: organization.slug,
      organization_id: organization.id,
      lookback_days: LOOKBACK_DAYS,
      since: since.iso8601,
      candidate_total: candidate_total,
      process_total: process_total,
      start_after_id: start_after_id,
      limit: LIMIT,
      log_every: LOG_EVERY,
      timeout_seconds: TIMEOUT_SECONDS,
      statement_timeout: STATEMENT_TIMEOUT,
      discard_good_jobs: DISCARD_GOOD_JOBS,
      dry_run: DRY_RUN,
      enqueue_url: url,
      note: 'Iterates day-by-day; ~3–4 req/s ⇒ full 50k window can take several hours'
    }.inspect)

    if DISCARD_GOOD_JOBS
      jobs = GoodJob::Job
             .where(finished_at: nil, job_class: 'EnqueueConversationTopicEmbeddingJob')
             .where(created_at: LOOKBACK_HOURS.hours.ago..)
      job_count = jobs.count
      puts "Discarding unfinished EnqueueConversationTopicEmbeddingJob rows: #{job_count}"
      unless DRY_RUN
        discarded = 0
        jobs.find_each do |job|
          job.discard_job('Discarded during kuda-pilot conversation topic backfill')
          discarded += 1
        end
        puts "Discarded: #{discarded}"
      end
    end

    if DRY_RUN
      puts 'Dry run complete — no embeddings requested.'
    elsif process_total.zero?
      puts 'Nothing to process.'
    else
      stats = { ok: 0, skipped: 0, err: 0, processed: 0 }
      errors = []
      started_monotonic = Process.clock_gettime(Process::CLOCK_MONOTONIC)
      last_id = start_after_id

      LOOKBACK_DAYS.downto(0) do |days_ago|
        break if LIMIT && stats[:processed] >= LIMIT

        day_start = days_ago.days.ago.beginning_of_day
        day_end = days_ago.zero? ? now : (days_ago - 1).days.ago.beginning_of_day
        day_start = [day_start, since].max

        next if day_start >= day_end

        day_scope = organization.conversations
                                .where(spam: false)
                                .where(started_at: day_start...day_end)
        day_scope = day_scope.where(id: (start_after_id + 1)..) if start_after_id

        # Pluck ids via (organization_id, started_at) — avoid find_each id-cursor timeouts.
        day_ids = day_scope.order(:started_at, :id).pluck(:id)
        puts({
          day: day_start.to_date.iso8601,
          day_ids: day_ids.size,
          processed_so_far: stats[:processed]
        }.inspect)

        day_ids.each do |conversation_id|
          break if LIMIT && stats[:processed] >= LIMIT

          conversation = organization.conversations.find_by(id: conversation_id)
          unless conversation
            stats[:skipped] += 1
            stats[:processed] += 1
            next
          end

          begin
            unless topic_embed_text_present.call(conversation)
              stats[:skipped] += 1
            else
              response = Faraday.post(url) do |req|
                req.headers['Authorization'] = "Bearer #{key}"
                req.headers['Content-Type'] = 'application/json'
                req.options.open_timeout = OPEN_TIMEOUT
                req.options.timeout = TIMEOUT_SECONDS
                req.body = {
                  conversation_id: conversation.id,
                  organization_id: organization.id
                }.to_json
              end

              if response.success?
                stats[:ok] += 1
              else
                stats[:err] += 1
                errors << {
                  conversation_id: conversation.id,
                  http_status: response.status,
                  body: response.body.to_s.truncate(200)
                }
              end
            end
          rescue StandardError => e
            stats[:err] += 1
            errors << {
              conversation_id: conversation.id,
              error_class: e.class.name,
              message: e.message.truncate(200)
            }
          end

          stats[:processed] += 1
          last_id = conversation.id

          if (stats[:processed] % LOG_EVERY).zero? ||
             (process_total && stats[:processed] == process_total)
            log_progress.call(stats, process_total, last_id, started_monotonic)
          end
        end
      end

      log_progress.call(stats, process_total, last_id, started_monotonic) unless (stats[:processed] % LOG_EVERY).zero?

      elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_monotonic
      puts({
        done: true,
        stats: stats,
        elapsed: format_duration.call(elapsed),
        last_id: last_id,
        error_sample: errors.first(10),
        error_count: errors.size,
        resume_hint: last_id ? "START_AFTER_ID = #{last_id}" : nil,
        note: 'ok = API2 accepted enqueue (202); Turbopuffer upsert runs async in API2 workers'
      }.inspect)
    end
  ensure
    ActiveRecord::Base.connection.execute(
      "SET statement_timeout = #{ActiveRecord::Base.connection.quote(previous_timeout)}"
    )
  end
end