Class: RagRuby::Stores::Memory

Inherits:
Base
  • Object
show all
Defined in:
lib/rag_ruby/stores/memory.rb

Defined Under Namespace

Classes: Entry

Instance Method Summary collapse

Constructor Details

#initialize(dimension: nil) ⇒ Memory

Returns a new instance of Memory.



8
9
10
11
 
# File 'lib/rag_ruby/stores/memory.rb', line 8

def initialize(dimension: nil)
  @dimension = dimension
  @entries = {}
end

Instance Method Details

#add(id, embedding:, metadata: {}, chunk: nil) ⇒ Object 



13
14
15
 
# File 'lib/rag_ruby/stores/memory.rb', line 13

def add(id, embedding:, metadata: {}, chunk: nil)
  @entries[id] = Entry.new(id: id, embedding: embedding, metadata: metadata, chunk: chunk)
end

#clear ⇒ Object 



46
47
48
 
# File 'lib/rag_ruby/stores/memory.rb', line 46

def clear
  @entries.clear
end

#count ⇒ Object 



42
43
44
 
# File 'lib/rag_ruby/stores/memory.rb', line 42

def count
  @entries.size
end

#delete(id) ⇒ Object 



38
39
40
 
# File 'lib/rag_ruby/stores/memory.rb', line 38

def delete(id)
  @entries.delete(id)
end

#search(embedding, top_k: 5, filter: nil, strategy: :similarity, lambda: 0.5, fetch_k: 20) ⇒ Object 



17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
 
# File 'lib/rag_ruby/stores/memory.rb', line 17

def search(embedding, top_k: 5, filter: nil, strategy: :similarity, lambda: 0.5, fetch_k: 20)
  results = @entries.values

  if filter
    results = results.select do |entry|
      filter.all? { |k, v| entry.metadata[k] == v }
    end
  end

  scored = results
    .map { |entry| [entry, cosine_similarity(embedding, entry.embedding)] }
    .sort_by { |_, score| -score }

  if strategy == :mmr
    mmr_select(scored, embedding, top_k: top_k, lambda: lambda, fetch_k: fetch_k)
  else
    scored.first(top_k)
      .map { |entry, score| { id: entry.id, score: score, metadata: entry.metadata, chunk: entry.chunk } }
  end
end