llmemory 0.2.6 → 0.2.7

data/lib/llmemory/memory.rb

@@ -34,7 +34,7 @@ module Llmemory
       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,
+        Retrieval::MultiSource.new(primary: @long_term, memory: self),
         llm: tracked_llm_client,
         feedback: Retrieval::FeedbackStore.new(store: @short_term_store)
       )
@@ -46,7 +46,7 @@ module Llmemory
       @working_memory ||= WorkingMemory.new(
         user_id: @user_id,
         session_id: @session_id,
-        store: build_short_term_store(@cipher)
+        store: @short_term_store
       )
     end
 
@@ -56,7 +56,8 @@ module Llmemory
       @episodic ||= LongTerm::Episodic::Memory.new(
         user_id: @user_id,
         storage: LongTerm::Episodic::Storages.build(cipher: @cipher),
-        cipher: @cipher
+        cipher: @cipher,
+        forget_log_store: @short_term_store
       )
     end
 
@@ -65,7 +66,8 @@ module Llmemory
       @procedural ||= LongTerm::Procedural::Memory.new(
         user_id: @user_id,
         storage: LongTerm::Procedural::Storages.build(cipher: @cipher),
-        cipher: @cipher
+        cipher: @cipher,
+        forget_log_store: @short_term_store
       )
     end
 
@@ -103,9 +105,14 @@ module Llmemory
     end
 
     def add_message(role:, content:)
-      msgs = messages
-      msgs << { role: role.to_sym, content: content.to_s }
-      save_state(messages: msgs, **preserved_flush_state)
+      @short_term_store.update(@user_id, @session_id) do |state|
+        state = normalize_state_hash(state)
+        list = state[STATE_KEY_MESSAGES]
+        list = list.is_a?(Array) ? list.dup : []
+        list << { role: role.to_sym, content: content.to_s }
+        list = sanitize_messages(list) if Llmemory.configuration.message_sanitizer_enabled
+        state.merge(STATE_KEY_MESSAGES => list, last_activity_at: Time.now, **preserved_flush_state_from(state))
+      end
       true
     end
 
@@ -164,6 +171,7 @@ module Llmemory
 
     def clear_session!
       @checkpoint.clear_state
+      working_memory.clear!
       true
     end
 
@@ -315,13 +323,23 @@ module Llmemory
     end
 
     def preserved_flush_state
-      state = restore_state_for_save
+      preserved_flush_state_from(restore_state_for_save)
+    end
+
+    def preserved_flush_state_from(state)
+      state = normalize_state_hash(state)
       {}.tap do |h|
-        h[:last_flush_at] = state[:last_flush_at] || state["last_flush_at"] if state[:last_flush_at] || state["last_flush_at"]
-        h[:last_compact_at] = state[:last_compact_at] || state["last_compact_at"] if state[:last_compact_at] || state["last_compact_at"]
+        h[:last_flush_at] = state[:last_flush_at] if state[:last_flush_at]
+        h[:last_compact_at] = state[:last_compact_at] if state[:last_compact_at]
       end
     end
 
+    def normalize_state_hash(state)
+      return {} unless state.is_a?(Hash)
+
+      state.transform_keys(&:to_sym)
+    end
+
     def estimated_tokens(msgs)
       (messages_byte_size(msgs) / 4.0).ceil
     end
@@ -358,7 +376,7 @@ module Llmemory
     end
 
     def build_long_term(long_term_type)
-      llm_opts = { llm: tracked_llm_client }
+      llm_opts = { llm: tracked_llm_client, forget_log_store: @short_term_store }
       case long_term_type.to_s.to_sym
       when :graph_based
         LongTerm::GraphBased::Memory.new(

data/lib/llmemory/memory_module.rb

@@ -54,9 +54,13 @@ module Llmemory
     end
 
     # Shared audit trail for #forget. Lazily built; override or inject by setting
-    # @forget_log if a specific backend/store is required.
+    # @forget_log, or pass @forget_log_store when constructing the memory.
     def forget_log
-      @forget_log ||= Llmemory::ForgetLog.new
+      @forget_log ||= Llmemory::ForgetLog.new(store: forget_log_store)
+    end
+
+    def forget_log_store
+      @forget_log_store
     end
   end
 end

data/lib/llmemory/provenance.rb

@@ -20,7 +20,7 @@ module Llmemory
       {
         sources: Array(sources).filter_map { |s| normalize_source(s) },
         method: method&.to_s,
-        confidence: confidence.nil? ? nil : confidence.to_f,
+        confidence: normalize_confidence(confidence),
         created_at: normalize_time(created_at)
       }
     end
@@ -60,5 +60,27 @@ module Llmemory
       value ||= Time.now
       value.respond_to?(:iso8601) ? value.iso8601 : value.to_s
     end
+
+    CONFIDENCE_LABELS = {
+      "high" => 0.9,
+      "medium" => 0.6,
+      "low" => 0.3
+    }.freeze
+
+    def normalize_confidence(value)
+      return nil if value.nil?
+
+      if value.is_a?(Numeric)
+        return [[value.to_f, 0.0].max, 1.0].min
+      end
+
+      label = value.to_s.downcase.strip
+      return CONFIDENCE_LABELS[label] if CONFIDENCE_LABELS.key?(label)
+
+      numeric = label.to_f
+      return nil if numeric.zero? && !label.match?(/\A0(\.0+)?\z/)
+
+      [[numeric, 0.0].max, 1.0].min
+    end
   end
 end

data/lib/llmemory/retrieval/feedback_store.rb

@@ -22,11 +22,14 @@ module Llmemory
 
       def record(user_id, item_id, delta)
         return if user_id.nil? || item_id.nil?
-        state = load(user_id)
+
         key = item_id.to_s
-        state[key] = (state[key] || 0) + delta.to_i
-        @store.save(user_id, SESSION_KEY, state)
-        state[key]
+        state = @store.update(user_id, SESSION_KEY) do |current|
+          current = load_from_state(current)
+          current[key] = (current[key] || 0) + delta.to_i
+          current
+        end
+        state ? state[key] : nil
       end
 
       def net(user_id, item_id)
@@ -41,7 +44,10 @@ module Llmemory
       private
 
       def load(user_id)
-        state = @store.load(user_id, SESSION_KEY)
+        load_from_state(@store.load(user_id, SESSION_KEY))
+      end
+
+      def load_from_state(state)
         return {} unless state.is_a?(Hash)
         state.each_with_object({}) { |(k, v), acc| acc[k.to_s] = v.to_i }
       end

data/lib/llmemory/retrieval/multi_source.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module Retrieval
+    # Combines semantic, episodic and procedural memory for unified retrieval.
+    class MultiSource
+      def initialize(primary:, memory:)
+        @primary = primary
+        @memory = memory
+      end
+
+      def user_id
+        @memory.user_id
+      end
+
+      def search_candidates(query, user_id: nil, top_k: 20)
+        uid = user_id || @memory.user_id
+        combined = []
+        combined.concat(@primary.search_candidates(query, user_id: uid, top_k: top_k))
+        combined.concat(@memory.episodic.search_candidates(query, user_id: uid, top_k: top_k))
+        combined.concat(@memory.procedural.search_candidates(query, user_id: uid, top_k: top_k))
+        combined
+      end
+    end
+  end
+end

data/lib/llmemory/retrieval/temporal_ranker.rb

@@ -9,14 +9,16 @@ module Llmemory
       end
 
       def rank(candidates, now: Time.now)
-        lambda_val = Math.log(2) / @half_life_days.to_f
+        half_life = @half_life_days.to_f
+        half_life = 30.0 if half_life <= 0
+        lambda_val = Math.log(2) / half_life
         weight = [@importance_weight.to_f, 0.0].max
 
         candidates.map do |c|
           score = (c[:score] || c["score"] || 1.0).to_f
           timestamp = c[:timestamp] || c["timestamp"]
           timestamp = Time.parse(timestamp.to_s) if timestamp.is_a?(String)
-          age_days = timestamp ? (now - timestamp).to_i / 86400 : 0
+          age_days = timestamp ? [(now - timestamp) / 86_400.0, 0.0].max : 0.0
 
           time_decay = if c[:evergreen] || c["evergreen"]
             1.0

data/lib/llmemory/retrieval.rb

@@ -2,6 +2,7 @@
 
 require_relative "retrieval/temporal_ranker"
 require_relative "retrieval/context_assembler"
+require_relative "retrieval/multi_source"
 require_relative "retrieval/engine"
 
 module Llmemory

data/lib/llmemory/short_term/pruner.rb

@@ -52,14 +52,26 @@ module Llmemory
       def soft_trim(content)
         return content if content.bytesize <= @soft_trim_max_bytes
 
-        head_chars = (@soft_trim_max_bytes * @head_ratio).to_i
-        tail_chars = (@soft_trim_max_bytes * @tail_ratio).to_i
+        head_bytes = (@soft_trim_max_bytes * @head_ratio).to_i
+        tail_bytes = (@soft_trim_max_bytes * @tail_ratio).to_i
 
-        head = content.byteslice(0, head_chars)
-        tail = content.bytesize > (head_chars + tail_chars) ? content.byteslice(-tail_chars, tail_chars) : ""
+        head = utf8_safe_byteslice(content, 0, head_bytes)
+        tail = if content.bytesize > (head_bytes + tail_bytes)
+                 utf8_safe_byteslice(content, -tail_bytes, tail_bytes)
+               else
+                 ""
+               end
 
         "#{head}\n...\n#{tail}"
       end
+
+      def utf8_safe_byteslice(str, start, length = nil)
+        slice = length.nil? ? str.byteslice(start) : str.byteslice(start, length)
+        return "" if slice.nil?
+
+        slice.force_encoding(str.encoding)
+        slice.scrub
+      end
     end
   end
 end

data/lib/llmemory/short_term/session_lifecycle.rb

@@ -30,10 +30,9 @@ module Llmemory
           state = @store.load(user_id, session_id)
           next unless state.is_a?(Hash)
 
-          last_activity = state[:last_activity_at] || state["last_activity_at"]
-          next if last_activity.nil?
+          last_time = parse_activity_time(state)
+          next if last_time.nil?
 
-          last_time = last_activity.is_a?(Time) ? last_activity : Time.parse(last_activity.to_s)
           if last_time < cutoff
             @store.delete(user_id, session_id)
             deleted += 1
@@ -52,10 +51,9 @@ module Llmemory
           state = @store.load(user_id, session_id)
           next unless state.is_a?(Hash)
 
-          last_activity = state[:last_activity_at] || state["last_activity_at"]
-          next if last_activity.nil?
+          last_time = parse_activity_time(state)
+          next if last_time.nil?
 
-          last_time = last_activity.is_a?(Time) ? last_activity : Time.parse(last_activity.to_s)
           if last_time < cutoff
             @store.delete(user_id, session_id)
             deleted += 1
@@ -70,10 +68,9 @@ module Llmemory
         sessions = user_sessions(user_id)
         return 0 if sessions.size <= max_entries
 
-        session_ages = sessions.map do |session_id|
+        session_ages = sessions.filter_map do |session_id|
           state = @store.load(user_id, session_id)
-          last_activity = state&.dig(:last_activity_at) || state&.dig("last_activity_at")
-          last_time = last_activity.is_a?(Time) ? last_activity : (last_activity ? Time.parse(last_activity.to_s) : Time.at(0))
+          last_time = parse_activity_time(state) || Time.at(0)
           [session_id, last_time]
         end
 
@@ -89,6 +86,18 @@ module Llmemory
         @store.list_sessions(user_id: user_id).reject { |s| self.class.pseudo_session?(s) }
       end
 
+      def parse_activity_time(state)
+        return nil unless state.is_a?(Hash)
+
+        last_activity = state[:last_activity_at] || state["last_activity_at"]
+        return nil if last_activity.nil?
+        return last_activity if last_activity.is_a?(Time)
+
+        Time.parse(last_activity.to_s)
+      rescue ArgumentError, TypeError
+        nil
+      end
+
       def build_store
         Stores.build
       end

data/lib/llmemory/short_term/stores/active_record_store.rb

@@ -1,13 +1,14 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "../../crypto/field_helpers"
+require_relative "../../active_record_helpers"
 
 module Llmemory
   module ShortTerm
     module Stores
       class ActiveRecordStore < Base
         include Llmemory::Crypto::FieldHelpers
+        include Llmemory::ActiveRecordHelpers
 
         def initialize(cipher: nil)
           @cipher = cipher || Llmemory.build_cipher
@@ -22,13 +23,15 @@ module Llmemory
         end
 
         def save(user_id, session_id, state)
-          record = Llmemory::ShortTerm::Stores::ActiveRecordCheckpoint.find_or_initialize_by(
-            user_id: user_id,
-            session_id: session_id
-          )
-          record.state = cipher.enabled? ? serialize_state(state) : state
-          record.updated_at = Time.current
-          record.save!
+          with_unique_retry do
+            record = Llmemory::ShortTerm::Stores::ActiveRecordCheckpoint.find_or_initialize_by(
+              user_id: user_id,
+              session_id: session_id
+            )
+            record.state = cipher.enabled? ? serialize_state(state) : state
+            record.updated_at = Time.current
+            record.save!
+          end
           true
         end
 
@@ -55,6 +58,32 @@ module Llmemory
           true
         end
 
+        def update(user_id, session_id, &block)
+          result = nil
+          Llmemory::ShortTerm::Stores::ActiveRecordCheckpoint.transaction do
+            record = Llmemory::ShortTerm::Stores::ActiveRecordCheckpoint.lock.find_by(
+              user_id: user_id,
+              session_id: session_id
+            )
+            current = if record
+                        raw = record.state
+                        raw.is_a?(Hash) ? raw.transform_keys(&:to_sym) : deserialize_state(raw)
+                      end
+            new_state = yield(current)
+            next if new_state.nil?
+
+            record ||= Llmemory::ShortTerm::Stores::ActiveRecordCheckpoint.new(
+              user_id: user_id,
+              session_id: session_id
+            )
+            record.state = cipher.enabled? ? serialize_state(new_state) : new_state
+            record.updated_at = Time.current
+            with_unique_retry { record.save! }
+            result = new_state
+          end
+          result
+        end
+
         def list_users
           Llmemory::ShortTerm::Stores::ActiveRecordCheckpoint.distinct.pluck(:user_id)
         end

data/lib/llmemory/short_term/stores/base.rb

@@ -16,6 +16,15 @@ module Llmemory
           raise NotImplementedError, "#{self.class}#delete must be implemented"
         end
 
+        # Atomic read-modify-write. Yields the current state (may be nil) and
+        # persists the block's return value when non-nil.
+        def update(user_id, session_id)
+          current = load(user_id, session_id)
+          new_state = yield(current)
+          save(user_id, session_id, new_state) unless new_state.nil?
+          new_state
+        end
+
         def list_users
           raise NotImplementedError, "#{self.class}#list_users must be implemented"
         end

data/lib/llmemory/short_term/stores/key_codec.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module ShortTerm
+    module Stores
+      # Escapes user/session identifiers so composite keys remain unambiguous
+      # when ids contain the separator character `:`.
+      module KeyCodec
+        ESCAPE = "%3A"
+        SEPARATOR = ":"
+
+        module_function
+
+        def encode(value)
+          value.to_s.gsub("%", "%25").gsub(SEPARATOR, ESCAPE)
+        end
+
+        def decode(value)
+          value.to_s.gsub(ESCAPE, SEPARATOR).gsub("%25", "%")
+        end
+
+        def composite_key(*parts)
+          parts.map { |part| encode(part) }.join(SEPARATOR)
+        end
+
+        def split_composite_key(key, parts: 2)
+          encoded = key.to_s.split(SEPARATOR, parts)
+          encoded.map { |part| decode(part) }
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/short_term/stores/memory_store.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
+require "json"
 require_relative "base"
+require_relative "key_codec"
 
 module Llmemory
   module ShortTerm
@@ -8,37 +10,66 @@ module Llmemory
       class MemoryStore < Base
         def initialize(cipher: nil)
           @store = {}
+          @mutexes = Hash.new { |h, k| h[k] = Mutex.new }
         end
 
         def save(user_id, session_id, state)
           key = key_for(user_id, session_id)
-          @store[key] = { state: state, updated_at: Time.now }
+          @mutexes[key].synchronize do
+            @store[key] = { state: deep_copy(state), updated_at: Time.now }
+          end
           true
         end
 
         def load(user_id, session_id)
           key = key_for(user_id, session_id)
-          @store.dig(key, :state)
+          @mutexes[key].synchronize do
+            deep_copy(@store.dig(key, :state))
+          end
         end
 
         def delete(user_id, session_id)
-          @store.delete(key_for(user_id, session_id))
+          key = key_for(user_id, session_id)
+          @mutexes[key].synchronize do
+            @store.delete(key)
+          end
           true
         end
 
+        def update(user_id, session_id, &block)
+          key = key_for(user_id, session_id)
+          @mutexes[key].synchronize do
+            current = deep_copy(@store.dig(key, :state))
+            new_state = yield(current)
+            unless new_state.nil?
+              @store[key] = { state: deep_copy(new_state), updated_at: Time.now }
+            end
+            deep_copy(new_state)
+          end
+        end
+
         def list_users
-          @store.keys.map { |k| k.split(":", 2).first }.uniq
+          @store.keys.map { |k| KeyCodec.split_composite_key(k, parts: 2).first }.uniq
         end
 
         def list_sessions(user_id:)
-          prefix = "#{user_id}:"
-          @store.keys.select { |k| k.start_with?(prefix) }.map { |k| k.split(":", 2).last }
+          encoded_user = KeyCodec.encode(user_id)
+          prefix = "#{encoded_user}#{KeyCodec::SEPARATOR}"
+          @store.keys
+            .select { |k| k.start_with?(prefix) }
+            .map { |k| KeyCodec.split_composite_key(k, parts: 2).last }
         end
 
         private
 
         def key_for(user_id, session_id)
-          "#{user_id}:#{session_id}"
+          KeyCodec.composite_key(user_id, session_id)
+        end
+
+        def deep_copy(value)
+          return nil if value.nil?
+
+          Marshal.load(Marshal.dump(value))
         end
       end
     end