Class: Engram::UseCases::RebuildEmbeddings

Inherits:
Object
  • Object
show all
Defined in:
lib/engram/use_cases/rebuild_embeddings.rb

Overview

Rebuild stored embeddings when the active embedder configuration changes. Returns counts useful for logging/observability:

- :processed: total rows examined in scope
- :updated: rows re-embedded and persisted
- :skipped: rows skipped because they were unchanged or incomplete
- :failed: count of rows that failed during rebuild
- :failed_ids: record IDs that failed during rebuild
- :failed_errors: error class/message keyed by failed record ID

Instance Method Summary collapse

Constructor Details

#initialize(store:, embedder:) ⇒ RebuildEmbeddings

Returns a new instance of RebuildEmbeddings.



14
15
16
17
 
# File 'lib/engram/use_cases/rebuild_embeddings.rb', line 14

def initialize(store:, embedder:)
  @store = store
  @embedder = embedder
end

Instance Method Details

#call(scope:, stale_only: true, batch_size: 100) ⇒ Object

Raises:

  • (ArgumentError) 


19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
 
# File 'lib/engram/use_cases/rebuild_embeddings.rb', line 19

def call(scope:, stale_only: true, batch_size: 100)
  batch_size = Integer(batch_size)
  raise ArgumentError, "batch_size must be greater than 0" unless batch_size.positive?

  counts = {
    processed: 0,
    updated: 0,
    skipped: 0,
    failed: 0,
    failed_ids: [],
    failed_errors: {}
  }

  after_id = nil
  legacy_index = 0
  legacy_records = nil
  processed_ids = Set.new
  processed_nil_records = Set.new
  loop do
    batch_result = fetch_batch(scope:, batch_size:, after_id:, legacy_records:, legacy_index:)
    legacy_records = batch_result[:legacy_records] if batch_result[:legacy_records]
    batch = batch_result[:records]
    if batch.empty?
      break if legacy_records

      # Reconcile with one stable snapshot: accepting keyset keywords does
      # not guarantee that an adapter returns complete, ordered pages.
      legacy_records = Array(@store.all(scope: scope))
      legacy_index = 0
      next
    end

    cursor_record = batch.reverse.find { |candidate| !candidate.id.nil? } unless legacy_records

    if !legacy_records && fallback_before_processing?(batch:, after_id:)
      legacy_records = Array(@store.all(scope: scope))
      legacy_index = 0
      batch = legacy_records.slice(legacy_index, batch_size) || []
      break if batch.empty?
    end

    batch.each do |record|
      identity = record.id.nil? ? record.object_id : record.id
      seen = record.id.nil? ? processed_nil_records : processed_ids
      next unless seen.add?(identity)

      counts[:processed] += 1

      if record.id.nil?
        counts[:skipped] += 1
        next
      end

      if stale_only && !stale?(record)
        counts[:skipped] += 1
        next
      end

      begin
        rebuilt = record.with(embedding: @embedder.embed(record.content))
        rebuilt = EmbeddingMetadata.attach(rebuilt, embedder: @embedder)
        updated = @store.update(scope: scope, id: record.id, record: rebuilt)
        raise Engram::Error, "memory update was not applied" unless updated && updated != 0

        counts[:updated] += 1
      rescue => error
        counts[:failed] += 1
        record_id = record.id
        counts[:failed_ids] << record_id
        counts[:failed_errors][record_id] = {
          class: error.class.name,
          message: error.message
        }
        warn "Rebuild failed for record ##{record_id}: #{error.class}: #{error.message}"
      end
    end

    if legacy_records
      legacy_index += batch.length
      next
    end

    if cursor_record.nil?
      legacy_records = Array(@store.all(scope: scope))
      legacy_index = 0
      next
    end

    if fallback_to_legacy_batching?(batch:, cursor_record:)
      legacy_records = Array(@store.all(scope: scope))
      legacy_index = 0
      next
    end

    after_id = cursor_record.id
  end

  {scope: scope, **counts}
end