Class: Llmemory::VectorStore::MemoryStore

Inherits:
Base
  • Object
show all
Defined in:
lib/llmemory/vector_store/memory_store.rb

Instance Method Summary collapse

Constructor Details

#initialize(embedding_provider: nil) ⇒ MemoryStore

Returns a new instance of MemoryStore.



8
9
10
11
# File 'lib/llmemory/vector_store/memory_store.rb', line 8

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

Instance Method Details

#embed(text) ⇒ Object



13
14
15
16
# File 'lib/llmemory/vector_store/memory_store.rb', line 13

def embed(text)
  return Array.new(1536, 0.0) unless @embedding_provider&.respond_to?(:embed)
  @embedding_provider.embed(text)
end

#search(query_embedding, top_k: 10, user_id: nil) ⇒ Object



24
25
26
27
28
29
30
31
32
33
# File 'lib/llmemory/vector_store/memory_store.rb', line 24

def search(query_embedding, top_k: 10, user_id: nil)
  query = query_embedding.to_a.map(&:to_f)
  return [] if query.empty?
  entries = user_id ? @entries.select { |k, _| k.to_s.start_with?("#{user_id}:") } : @entries
  scores = entries.map do |id, data|
    sim = cosine_similarity(query, data[:embedding])
    { id: id, score: sim, metadata: data[:metadata] }
  end
  scores.sort_by { |s| -s[:score] }.first(top_k)
end

#search_by_text(query_text, top_k: 10, user_id: nil) ⇒ Object



35
36
37
38
39
# File 'lib/llmemory/vector_store/memory_store.rb', line 35

def search_by_text(query_text, top_k: 10, user_id: nil)
  return [] unless @embedding_provider&.respond_to?(:embed)
  query_embedding = @embedding_provider.embed(query_text)
  search(query_embedding, top_k: top_k, user_id: user_id)
end

#store(id:, embedding:, metadata: {}, user_id: nil) ⇒ Object



18
19
20
21
22
# File 'lib/llmemory/vector_store/memory_store.rb', line 18

def store(id:, embedding:, metadata: {}, user_id: nil)
  key = user_id ? "#{user_id}:#{id}" : id.to_s
  @entries[key] = { embedding: embedding.to_a.map(&:to_f), metadata: ( || {}).merge("user_id" => user_id) }
  id
end