Class: Llmemory::Memory

Inherits:
Object
  • Object
show all
Defined in:
lib/llmemory/memory.rb

Constant Summary collapse

DEFAULT_SESSION_ID = 
"default"
STATE_KEY_MESSAGES = 
:messages

Instance Method Summary collapse

Constructor Details

#initialize(user_id:, session_id: DEFAULT_SESSION_ID, checkpoint: nil, long_term: nil, long_term_type: nil, retrieval_engine: nil, api_key: nil) ⇒ Memory

Returns a new instance of Memory.



13
14
15
16
17
18
19
20
21
 
# File 'lib/llmemory/memory.rb', line 13

def initialize(user_id:, session_id: DEFAULT_SESSION_ID, checkpoint: nil, long_term: nil, long_term_type: nil, retrieval_engine: nil, api_key: nil)
  @user_id = user_id
  @session_id = session_id
  @checkpoint = checkpoint || ShortTerm::Checkpoint.new(user_id: user_id, session_id: session_id)
  @llm = api_key.to_s.empty? ? nil : Llmemory::LLM.client(api_key: api_key)
  type = long_term_type || Llmemory.configuration.long_term_type || :file_based
  @long_term = long_term || build_long_term(type)
  @retrieval_engine = retrieval_engine || Retrieval::Engine.new(@long_term, llm: @llm)
end

Instance Method Details

#add_message(role:, content:) ⇒ Object 



23
24
25
26
27
28
 
# File 'lib/llmemory/memory.rb', line 23

def add_message(role:, content:)
  msgs = messages
  msgs << { role: role.to_sym, content: content.to_s }
  save_state(messages: msgs, **preserved_flush_state)
  true
end

#check_context_window! ⇒ Object 



151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
 
# File 'lib/llmemory/memory.rb', line 151

def check_context_window!
  return false if messages.empty?

  flushed = false
  if should_auto_consolidate? && Llmemory.configuration.memory_flush_enabled
    consolidate!
    flushed = true
  end

  compacted = false
  if should_compact?
    compacted = compact!
  end

  flushed || compacted
end

#clear_session! ⇒ Object 



83
84
85
86
 
# File 'lib/llmemory/memory.rb', line 83

def clear_session!
  @checkpoint.clear_state
  true
end

#compact!(max_bytes: nil) ⇒ Object 



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

def compact!(max_bytes: nil)
  max = max_bytes || Llmemory.configuration.compact_max_bytes
  msgs = messages
  current_bytes = messages_byte_size(msgs)
  return false if current_bytes <= max

  flushed = flush_memory_before_compaction!(msgs)

  old_msgs, recent_msgs = split_messages_by_bytes(msgs, max)
  return false if old_msgs.empty?

  summary = summarize_messages(old_msgs)
  compacted = [{ role: :system, content: summary }] + recent_msgs
  state = restore_state_for_save
  flush_ts = flushed ? Time.now : (state[:last_flush_at] || state["last_flush_at"])
  save_state(messages: compacted, last_compact_at: Time.now, last_flush_at: flush_ts)
  true
end

#consolidate! ⇒ Object 



75
76
77
78
79
80
81
 
# File 'lib/llmemory/memory.rb', line 75

def consolidate!
  msgs = messages
  return true if msgs.empty?
  conversation_text = msgs.map { |m| format_message(m) }.join("\n")
  @long_term.memorize(conversation_text)
  true
end

#context_tokens ⇒ Object 



116
117
118
 
# File 'lib/llmemory/memory.rb', line 116

def context_tokens
  estimated_tokens(messages)
end

#last_user_message ⇒ Object 



54
55
56
57
58
 
# File 'lib/llmemory/memory.rb', line 54

def last_user_message
  msgs = messages
  idx = msgs.rindex { |m| (m[:role] || m["role"]).to_s == "user" }
  idx ? (msgs[idx][:content] || msgs[idx]["content"]).to_s : ""
end

#maybe_flush_memory! ⇒ Object 



107
108
109
110
111
112
113
114
 
# File 'lib/llmemory/memory.rb', line 107

def maybe_flush_memory!
  return false unless Llmemory.configuration.memory_flush_enabled
  msgs = messages
  return false if msgs.empty?
  return false if estimated_tokens(msgs) < Llmemory.configuration.memory_flush_threshold_tokens

  consolidate!
end

#messages ⇒ Object 



30
31
32
33
34
35
36
 
# File 'lib/llmemory/memory.rb', line 30

def messages
  state = @checkpoint.restore_state
  return [] unless state.is_a?(Hash)
  list = state[STATE_KEY_MESSAGES] || state[STATE_KEY_MESSAGES.to_s]
  list = list.is_a?(Array) ? list.dup : []
  sanitize_messages(list)
end

#prune!(mode: nil) ⇒ Object 



60
61
62
63
64
65
66
67
68
69
70
71
72
73
 
# File 'lib/llmemory/memory.rb', line 60

def prune!(mode: nil)
  return false unless Llmemory.configuration.prune_tool_results_enabled

  msgs = messages
  return false if msgs.empty?

  mode ||= Llmemory.configuration.prune_tool_results_mode
  pruner = ShortTerm::Pruner.new(
    soft_trim_max_bytes: Llmemory.configuration.prune_tool_results_max_bytes
  )
  pruned = pruner.prune!(msgs, mode: mode)
  save_state(messages: pruned, **preserved_flush_state)
  true
end

#recall_for(query: nil, max_tokens: nil) ⇒ Object 



45
46
47
48
49
50
51
52
 
# File 'lib/llmemory/memory.rb', line 45

def recall_for(query: nil, max_tokens: nil)
  return "" unless Llmemory.configuration.auto_recall_enabled

  effective_query = query || last_user_message
  return "" if effective_query.to_s.strip.empty?

  retrieve(effective_query, max_tokens: max_tokens)
end

#retrieve(query, max_tokens: nil) ⇒ Object 



38
39
40
41
42
43
 
# File 'lib/llmemory/memory.rb', line 38

def retrieve(query, max_tokens: nil)
  msgs = pruned_messages
  short_context = format_short_term_context(msgs)
  long_context = @retrieval_engine.retrieve_for_inference(query, user_id: @user_id, max_tokens: max_tokens)
  combine_contexts(short_context, long_context)
end

#should_auto_consolidate? ⇒ Boolean

Returns:

  • (Boolean) 


120
121
122
123
124
 
# File 'lib/llmemory/memory.rb', line 120

def should_auto_consolidate?
  ctx = context_tokens
  threshold = Llmemory.configuration.context_window_tokens - Llmemory.configuration.reserve_tokens
  ctx >= threshold
end

#should_compact? ⇒ Boolean

Returns:

  • (Boolean) 


126
127
128
129
130
 
# File 'lib/llmemory/memory.rb', line 126

def should_compact?
  ctx = context_tokens
  threshold = Llmemory.configuration.context_window_tokens - Llmemory.configuration.reserve_tokens
  ctx >= threshold
end

#user_id ⇒ Object 



168
169
170
 
# File 'lib/llmemory/memory.rb', line 168

def user_id
  @user_id
end

#with_overflow_recovery(max_retries: 2, &block) ⇒ Object 



132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
 
# File 'lib/llmemory/memory.rb', line 132

def with_overflow_recovery(max_retries: 2, &block)
  return yield unless Llmemory.configuration.overflow_recovery_enabled
  return yield unless block_given?

  retries = 0
  begin
    yield
  rescue Llmemory::LLMError => e
    msg = e.message.to_s.downcase
    overflow = msg.include?("context") || msg.include?("token") || msg.include?("overflow") || msg.include?("limit")
    raise unless overflow && retries < max_retries

    prune! if Llmemory.configuration.prune_tool_results_enabled
    compact!
    retries += 1
    retry
  end
end