Class: LlmMemory::Hippocampus

Inherits:
Object
  • Object
show all
Defined in:
lib/llm_memory/hippocampus.rb

Instance Method Summary collapse

Constructor Details

#initialize(embedding_name: :openai, chunk_size: 1024, chunk_overlap: 50, store: :redis, index_name: "llm_memory") ⇒ Hippocampus

Returns a new instance of Hippocampus.



9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
 
# File 'lib/llm_memory/hippocampus.rb', line 9

def initialize(
  embedding_name: :openai,
  chunk_size: 1024,
  chunk_overlap: 50,
  store: :redis,
  index_name: "llm_memory"
)
  LlmMemory.configure

  embedding_class = EmbeddingManager.embeddings[embedding_name]
  raise "Embedding '#{embedding_name}' not found." unless embedding_class
  @embedding_instance = embedding_class.new

  store_class = StoreManager.stores[store]
  raise "Store '#{store}' not found." unless store_class
  @store = store_class.new(index_name: index_name)

  # char count, not word count
  @chunk_size = chunk_size
  @chunk_overlap = chunk_overlap
end

Instance Method Details

#add_vectors(docs) ⇒ Object 



72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
 
# File 'lib/llm_memory/hippocampus.rb', line 72

def add_vectors(docs)
  # embed documents and add vector
  result = []
  docs.each do |doc|
    content = doc[:content]
    metadata = doc[:metadata]
    vector = @embedding_instance.embed_document(content)
    result.push({
      content: content,
      metadata: metadata,
      vector: vector
    })
  end
  result
end

#forget(key) ⇒ Object 



60
61
62
 
# File 'lib/llm_memory/hippocampus.rb', line 60

def forget(key)
  @store.delete(key)
end

#forget_all ⇒ Object 



56
57
58
 
# File 'lib/llm_memory/hippocampus.rb', line 56

def forget_all
  @store.drop_index if @store.index_exists?
end

#get(key) ⇒ Object 



68
69
70
 
# File 'lib/llm_memory/hippocampus.rb', line 68

def get(key)
  @store.get(key)
end

#list(*args) ⇒ Object 



64
65
66
 
# File 'lib/llm_memory/hippocampus.rb', line 64

def list(*args)
  @store.list(*args)
end

#make_chunks(docs) ⇒ Object 



88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
 
# File 'lib/llm_memory/hippocampus.rb', line 88

def make_chunks(docs)
  result = []
  docs.each do |item|
    content = item[:content]
    metadata = item[:metadata]
    if content.length > @chunk_size
      start_index = 0
      while start_index < content.length
        end_index = [start_index + @chunk_size, content.length].min
        chunk = content[start_index...end_index]
        result << {content: chunk, metadata: metadata}
        break if end_index == content.length
        start_index += @chunk_size - @chunk_overlap
      end
    else
      result << {content: content, metadata: metadata}
    end
  end
  result
end

#memorize(docs) ⇒ Object 



43
44
45
46
47
48
49
 
# File 'lib/llm_memory/hippocampus.rb', line 43

def memorize(docs)
  validate_documents(docs)
  docs = make_chunks(docs)
  docs = add_vectors(docs)
  @store.create_index unless @store.index_exists?
  @store.add(data: docs)
end

#query(query_str, limit: 3) ⇒ Object 



51
52
53
54
 
# File 'lib/llm_memory/hippocampus.rb', line 51

def query(query_str, limit: 3)
  vector = @embedding_instance.embed_document(query_str)
  @store.search(query: vector, k: limit)
end

#validate_documents(documents) ⇒ Object

validate the document format



32
33
34
35
36
37
38
39
40
41
 
# File 'lib/llm_memory/hippocampus.rb', line 32

def validate_documents(documents)
  is_valid = documents.all? do |hash|
    hash.is_a?(Hash) &&
      hash.key?(:content) && hash[:content].is_a?(String) &&
      hash.key?(:metadata) && hash[:metadata].is_a?(Hash)
  end
  unless is_valid
    raise "Your documents need to have an array of hashes (content: string and metadata: hash)"
  end
end