llmemory 0.2.2 → 0.2.4

data/lib/llmemory/long_term/procedural/storages/database_storage.rb

@@ -4,6 +4,7 @@ require "json"
 require "securerandom"
 require "time"
 require_relative "base"
+require_relative "../../../crypto/field_helpers"
 
 module Llmemory
   module LongTerm
@@ -13,9 +14,12 @@ module Llmemory
         # (plus id/user_id/created_at and a denormalized search_text), mirroring
         # the file-based DatabaseStorage pattern.
         class DatabaseStorage < Base
-          def initialize(database_url: nil)
+          include Llmemory::Crypto::FieldHelpers
+
+          def initialize(database_url: nil, cipher: nil)
             @database_url = database_url || Llmemory.configuration.database_url
             @connection = nil
+            @cipher = cipher || Llmemory.build_cipher
           end
 
           def save_skill(user_id, skill)
@@ -27,7 +31,7 @@ module Llmemory
               "INSERT INTO llmemory_skills (id, user_id, data, search_text, created_at) " \
               "VALUES ($1, $2, $3::jsonb, $4, $5) " \
               "ON CONFLICT (id) DO UPDATE SET data = $3::jsonb, search_text = $4",
-              [id, user_id, JSON.generate(data), searchable_text(data), created_at_value(data)]
+              [id, user_id, store_data(data), enc(searchable_text(data)), created_at_value(data)]
             )
             id
           end
@@ -72,7 +76,7 @@ module Llmemory
             data[:updated_at] = Time.now.utc.iso8601
             conn.exec_params(
               "UPDATE llmemory_skills SET data = $3::jsonb, search_text = $4 WHERE user_id = $1 AND id = $2",
-              [user_id, skill_id, JSON.generate(data), searchable_text(data)]
+              [user_id, skill_id, store_data(data), enc(searchable_text(data))]
             )
             data
           end
@@ -144,11 +148,28 @@ module Llmemory
           end
 
           def parse_data(value)
-            JSON.parse(value.to_s, symbolize_names: true)
+            if value.is_a?(Hash)
+              return value.transform_keys(&:to_sym)
+            end
+
+            str = value.to_s
+            if cipher.encrypted?(str)
+              cipher.decrypt_json(str)
+            else
+              JSON.parse(str, symbolize_names: true)
+            end
           rescue JSON::ParserError
             {}
           end
 
+          def store_data(data)
+            if cipher.enabled?
+              JSON.generate(enc_json(data))
+            else
+              JSON.generate(data)
+            end
+          end
+
           def symbolize(hash)
             hash.each_with_object({}) { |(k, v), acc| acc[k.to_sym] = v }
           end

data/lib/llmemory/long_term/procedural/storages/file_storage.rb

@@ -4,22 +4,26 @@ require "fileutils"
 require "json"
 require "time"
 require_relative "base"
+require_relative "../../../crypto/field_helpers"
 
 module Llmemory
   module LongTerm
     module Procedural
       module Storages
         class FileStorage < Base
-          def initialize(base_path: nil)
+          include Llmemory::Crypto::FieldHelpers
+
+          def initialize(base_path: nil, cipher: nil)
             @base_path = base_path || Llmemory.configuration.long_term_storage_path || "./llmemory_data"
             @base_path = File.expand_path(@base_path)
+            @cipher = cipher || Llmemory.build_cipher
           end
 
           def save_skill(user_id, skill)
             id = skill[:id] || skill["id"] || "skill_#{next_seq(user_id)}"
             data = stringify_for_json(skill).merge("id" => id, "user_id" => user_id)
             data["created_at"] ||= Time.now.iso8601(6)
-            File.write(skill_path(user_id, id), JSON.generate(data))
+            write_skill_file(skill_path(user_id, id), data)
             id
           end
 
@@ -50,7 +54,7 @@ module Llmemory
             key = success ? :success_count : :failure_count
             skill[key] = (skill[key] || 0).to_i + 1
             skill[:updated_at] = Time.now.iso8601(6)
-            File.write(skill_path(user_id, skill_id), JSON.generate(stringify_for_json(skill)))
+            write_skill_file(skill_path(user_id, skill_id), stringify_for_json(skill))
             skill
           end
 
@@ -71,10 +75,10 @@ module Llmemory
             Array(ids).map(&:to_s).count do |id|
               path = skill_path(user_id, id)
               next false unless File.file?(path)
-              data = JSON.parse(File.read(path))
+              data = load_skill_raw(path)
               next false if data["archived_at"]
               data["archived_at"] = Time.now.iso8601
-              File.write(path, JSON.generate(data))
+              write_skill_file(path, data)
               true
             end
           end
@@ -103,7 +107,9 @@ module Llmemory
           end
 
           def load_skill(path)
-            data = JSON.parse(File.read(path), symbolize_names: true)
+            data = load_skill_raw(path)
+            return nil unless data
+
             data[:created_at] = parse_time(data[:created_at])
             data[:updated_at] = parse_time(data[:updated_at]) if data[:updated_at]
             data
@@ -111,6 +117,17 @@ module Llmemory
             nil
           end
 
+          def load_skill_raw(path)
+            raw = File.read(path)
+            json = cipher.enabled? && cipher.encrypted?(raw) ? cipher.decrypt(raw) : raw
+            JSON.parse(json, symbolize_names: true)
+          end
+
+          def write_skill_file(path, data)
+            payload = JSON.generate(data)
+            File.write(path, cipher.enabled? ? cipher.encrypt(payload) : payload)
+          end
+
           def skill_text(skill)
             [skill[:name], skill[:description], skill[:body]].compact.join("\n")
           end

data/lib/llmemory/mcp/tools/memory_stats.rb

@@ -52,6 +52,8 @@ module Llmemory
               stats[:long_term] = { error: e.message }
             end
 
+            stats[:llm_usage] = Llmemory::LLM::UsageLedger.new(store: store).totals(user_id)
+
             ::MCP::Tool::Response.new([{
               type: "text",
               text: format_stats(stats)
@@ -102,8 +104,19 @@ module Llmemory
               output << "  Resources: #{stats[:long_term][:resources]}"
             end
 
+            output << ""
+            output << Llmemory::LLM::UsageLedger.format_text(stats[:llm_usage] || default_llm_usage)
+
             output.join("\n")
           end
+
+          def default_llm_usage
+            {
+              invoke: { input_tokens: 0, output_tokens: 0, total_tokens: 0, calls: 0 },
+              embed: { total_tokens: 0, calls: 0 },
+              updated_at: nil
+            }
+          end
         end
       end
     end

data/lib/llmemory/memory.rb

@@ -10,47 +10,76 @@ module Llmemory
     DEFAULT_SESSION_ID = "default"
     STATE_KEY_MESSAGES = :messages
 
-    def initialize(user_id:, session_id: DEFAULT_SESSION_ID, checkpoint: nil, long_term: nil, long_term_type: nil, retrieval_engine: nil, working_memory: nil, episodic: nil, procedural: nil, api_key: nil)
+    def initialize(user_id:, session_id: DEFAULT_SESSION_ID, checkpoint: nil, long_term: nil, long_term_type: nil, retrieval_engine: nil, working_memory: nil, episodic: nil, procedural: nil, api_key: nil, encryption_key: :inherit)
       @user_id = user_id
       @session_id = session_id
-      @checkpoint = checkpoint || ShortTerm::Checkpoint.new(user_id: user_id, session_id: session_id)
+      resolved_key = encryption_key == :inherit ? nil : encryption_key
+      @cipher = Llmemory.build_cipher(resolved_key)
+      if checkpoint
+        @checkpoint = checkpoint
+        @short_term_store = checkpoint.store
+      else
+        @short_term_store = build_short_term_store(@cipher)
+        @checkpoint = ShortTerm::Checkpoint.new(
+          user_id: user_id,
+          session_id: session_id,
+          store: @short_term_store,
+          cipher: @cipher
+        )
+      end
       @working_memory = working_memory
       @episodic = episodic
       @procedural = procedural
-      @llm = api_key.to_s.empty? ? nil : Llmemory::LLM.client(api_key: api_key)
+      @api_key = api_key unless api_key.to_s.empty?
       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)
+      @retrieval_engine = retrieval_engine || Retrieval::Engine.new(
+        @long_term,
+        llm: tracked_llm_client,
+        feedback: Retrieval::FeedbackStore.new(store: @short_term_store)
+      )
     end
 
     # Structured working memory for this session (CoALA working memory),
     # parallel to the message checkpoint. Lazily built.
     def working_memory
-      @working_memory ||= WorkingMemory.new(user_id: @user_id, session_id: @session_id)
+      @working_memory ||= WorkingMemory.new(
+        user_id: @user_id,
+        session_id: @session_id,
+        store: build_short_term_store(@cipher)
+      )
     end
 
     # Episodic long-term memory (CoALA): records and retrieves agent trajectories.
     # Additive — coexists with the semantic store (file/graph). Lazily built.
     def episodic
-      @episodic ||= LongTerm::Episodic::Memory.new(user_id: @user_id)
+      @episodic ||= LongTerm::Episodic::Memory.new(
+        user_id: @user_id,
+        storage: LongTerm::Episodic::Storages.build(cipher: @cipher),
+        cipher: @cipher
+      )
     end
 
     # Procedural long-term memory (Voyager-style skill library). Lazily built.
     def procedural
-      @procedural ||= LongTerm::Procedural::Memory.new(user_id: @user_id)
+      @procedural ||= LongTerm::Procedural::Memory.new(
+        user_id: @user_id,
+        storage: LongTerm::Procedural::Storages.build(cipher: @cipher),
+        cipher: @cipher
+      )
     end
 
     # Reflects over recent episodes and writes distilled insights to the
     # semantic store (file/graph) with provenance back to source episodes.
     def reflect!(window: 10, category: "insights")
-      Reflection::Reflector.new(episodic: episodic, semantic: @long_term, llm: @llm)
+      Reflection::Reflector.new(episodic: episodic, semantic: @long_term, llm: tracked_llm_client)
         .reflect(window: window, category: category)
     end
 
     # Reasoning action: render a prompt from working memory, call the LLM, write
     # the result back. Composable; does not touch long-term memory.
     def reason(template:, into: Actions::Reason::DEFAULT_SLOT, parse: nil)
-      Actions::Reason.call(working_memory: working_memory, template: template, into: into, parse: parse, llm: @llm)
+      Actions::Reason.call(working_memory: working_memory, template: template, into: into, parse: parse, llm: tracked_llm_client)
     end
 
     # Mines recent episodes for reusable skills (Voyager-style). Human-in-the-loop
@@ -58,7 +87,7 @@ module Llmemory
     # `auto_register: true`, registers them in procedural memory (with provenance
     # back to the source episodes) and returns the new skill ids.
     def mine_skills!(window: SkillMining::Miner::DEFAULT_WINDOW, outcomes: nil, auto_register: false)
-      SkillMining::Miner.new(episodic: episodic, procedural: procedural, llm: @llm)
+      SkillMining::Miner.new(episodic: episodic, procedural: procedural, llm: tracked_llm_client)
         .mine(window: window, outcomes: outcomes, auto_register: auto_register)
     end
 
@@ -68,7 +97,7 @@ module Llmemory
     def maintain!(**opts)
       Maintenance::CognitivePass.run!(
         @user_id,
-        memory: self, episodic: episodic, procedural: procedural, semantic: @long_term, llm: @llm,
+        memory: self, episodic: episodic, procedural: procedural, semantic: @long_term, llm: tracked_llm_client,
         **opts
       )
     end
@@ -222,6 +251,10 @@ module Llmemory
       @user_id
     end
 
+    def llm_usage
+      Llmemory::LLM::UsageLedger.new(store: @short_term_store).totals(@user_id)
+    end
+
     private
 
     def summarize_messages(msgs)
@@ -240,7 +273,16 @@ module Llmemory
     end
 
     def llm_client
-      @llm ||= Llmemory::LLM.client
+      tracked_llm_client
+    end
+
+    def tracked_llm_client
+      @tracked_llm_client ||= Llmemory::LLM::TrackingClient.new(
+        nil,
+        user_id: @user_id,
+        store: @short_term_store,
+        api_key: @api_key
+      )
     end
 
     def flush_memory_before_compaction!(msgs)
@@ -316,19 +358,28 @@ module Llmemory
     end
 
     def build_long_term(long_term_type)
-      llm_opts = @llm ? { llm: @llm } : {}
+      llm_opts = { llm: tracked_llm_client }
       case long_term_type.to_s.to_sym
       when :graph_based
         LongTerm::GraphBased::Memory.new(
           user_id: @user_id,
-          storage: LongTerm::GraphBased::Storages.build,
+          storage: LongTerm::GraphBased::Storages.build(cipher: @cipher),
+          cipher: @cipher,
           **llm_opts
         )
       else
-        LongTerm::FileBased::Memory.new(user_id: @user_id, storage: LongTerm::FileBased::Storages.build, **llm_opts)
+        LongTerm::FileBased::Memory.new(
+          user_id: @user_id,
+          storage: LongTerm::FileBased::Storages.build(cipher: @cipher),
+          **llm_opts
+        )
       end
     end
 
+    def build_short_term_store(cipher)
+      ShortTerm::Stores.build(cipher: cipher)
+    end
+
     def save_state(messages:, last_flush_at: nil, last_compact_at: nil)
       state = { STATE_KEY_MESSAGES => messages, last_activity_at: Time.now }
       state[:last_flush_at] = last_flush_at if last_flush_at

data/lib/llmemory/short_term/checkpoint.rb

@@ -7,12 +7,15 @@ module Llmemory
     class Checkpoint
       DEFAULT_SESSION_ID = "default"
 
-      def initialize(user_id:, session_id: DEFAULT_SESSION_ID, store: nil)
+      def initialize(user_id:, session_id: DEFAULT_SESSION_ID, store: nil, cipher: nil)
         @user_id = user_id
         @session_id = session_id
+        @cipher = cipher
         @store = store || build_store
       end
 
+      attr_reader :store
+
       def save_state(state)
         @store.save(@user_id, @session_id, state)
       end
@@ -28,7 +31,7 @@ module Llmemory
       private
 
       def build_store
-        Stores.build
+        Stores.build(cipher: @cipher)
       end
     end
   end

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

@@ -1,12 +1,16 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "../../crypto/field_helpers"
 
 module Llmemory
   module ShortTerm
     module Stores
       class ActiveRecordStore < Base
-        def initialize
+        include Llmemory::Crypto::FieldHelpers
+
+        def initialize(cipher: nil)
+          @cipher = cipher || Llmemory.build_cipher
           self.class.load_model!
         end
 
@@ -22,7 +26,7 @@ module Llmemory
             user_id: user_id,
             session_id: session_id
           )
-          record.state = state
+          record.state = cipher.enabled? ? serialize_state(state) : state
           record.updated_at = Time.current
           record.save!
           true
@@ -34,8 +38,13 @@ module Llmemory
             session_id: session_id
           )
           return nil unless record
+
           raw = record.state
-          raw.is_a?(Hash) ? raw.transform_keys(&:to_sym) : deserialize(raw)
+          if raw.is_a?(Hash)
+            raw.transform_keys(&:to_sym)
+          else
+            deserialize_state(raw)
+          end
         end
 
         def delete(user_id, session_id)
@@ -53,13 +62,6 @@ module Llmemory
         def list_sessions(user_id:)
           Llmemory::ShortTerm::Stores::ActiveRecordCheckpoint.where(user_id: user_id).pluck(:session_id)
         end
-
-        private
-
-        def deserialize(data)
-          return data if data.is_a?(Hash)
-          JSON.parse(data.to_s, symbolize_names: true)
-        end
       end
     end
   end

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

@@ -6,7 +6,7 @@ module Llmemory
   module ShortTerm
     module Stores
       class MemoryStore < Base
-        def initialize
+        def initialize(cipher: nil)
           @store = {}
         end
 

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

@@ -1,14 +1,18 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "../../crypto/field_helpers"
 
 module Llmemory
   module ShortTerm
     module Stores
       class PostgresStore < Base
-        def initialize(database_url: nil)
+        include Llmemory::Crypto::FieldHelpers
+
+        def initialize(database_url: nil, cipher: nil)
           @database_url = database_url || Llmemory.configuration.database_url
           @connection = nil
+          @cipher = cipher || Llmemory.build_cipher
         end
 
         def save(user_id, session_id, state)
@@ -81,13 +85,15 @@ module Llmemory
         end
 
         def serialize(state)
-          require "json"
-          JSON.generate(state)
+          payload = serialize_state(state)
+          cipher.enabled? ? JSON.generate(payload) : payload
         end
 
         def deserialize(data)
-          require "json"
-          JSON.parse(data, symbolize_names: true)
+          if data.is_a?(String) && !cipher.encrypted?(data)
+            data = JSON.parse(data)
+          end
+          deserialize_state(data)
         end
       end
     end

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

@@ -1,14 +1,18 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "../../crypto/field_helpers"
 
 module Llmemory
   module ShortTerm
     module Stores
       class RedisStore < Base
-        def initialize(redis_url: nil)
+        include Llmemory::Crypto::FieldHelpers
+
+        def initialize(redis_url: nil, cipher: nil)
           @redis_url = redis_url || Llmemory.configuration.redis_url
           @redis = nil
+          @cipher = cipher || Llmemory.build_cipher
         end
 
         def save(user_id, session_id, state)
@@ -50,13 +54,11 @@ module Llmemory
         end
 
         def serialize(state)
-          require "json"
-          JSON.generate(state)
+          serialize_state(state)
         end
 
         def deserialize(data)
-          require "json"
-          JSON.parse(data, symbolize_names: true)
+          deserialize_state(data)
         end
       end
     end

data/lib/llmemory/short_term/stores.rb

@@ -10,16 +10,17 @@ module Llmemory
     module Stores
       # Single source of truth for selecting a short-term store backend.
       # Shared by Checkpoint, SessionLifecycle and WorkingMemory.
-      def self.build(store_type = nil)
+      def self.build(store_type = nil, cipher: nil)
+        resolved_cipher = cipher || Llmemory.build_cipher
         case (store_type || Llmemory.configuration.short_term_store).to_sym
-        when :memory then MemoryStore.new
-        when :redis then RedisStore.new
-        when :postgres then PostgresStore.new
+        when :memory then MemoryStore.new(cipher: resolved_cipher)
+        when :redis then RedisStore.new(cipher: resolved_cipher)
+        when :postgres then PostgresStore.new(cipher: resolved_cipher)
         when :active_record, :activerecord
           require_relative "stores/active_record_store"
-          ActiveRecordStore.new
+          ActiveRecordStore.new(cipher: resolved_cipher)
         else
-          MemoryStore.new
+          MemoryStore.new(cipher: resolved_cipher)
         end
       end
     end

data/lib/llmemory/vector_store/active_record_store.rb

@@ -7,10 +7,11 @@ module Llmemory
     # Persists embeddings in llmemory_embeddings (pgvector).
     # Use when long_term_store is :active_record so hybrid search finds persisted embeddings.
     class ActiveRecordStore < Base
-      def initialize(embedding_provider: nil, source_type: "edge")
+      def initialize(embedding_provider: nil, source_type: "edge", cipher: nil)
         self.class.load_model!
         @embedding_provider = embedding_provider
         @source_type = source_type.to_s
+        @cipher = cipher || Llmemory.build_cipher
       end
 
       def self.load_model!
@@ -25,6 +26,12 @@ module Llmemory
         @embedding_provider.embed(text)
       end
 
+      def last_usage
+        return @embedding_provider.last_usage if @embedding_provider&.respond_to?(:last_usage)
+
+        Llmemory::LLM::Usage.zero
+      end
+
       def store(id:, embedding:, metadata: {}, user_id: nil)
         return id if user_id.nil? || user_id.to_s.empty?
         text_content = (metadata || {}).dig("text") || (metadata || {}).dig(:text)
@@ -34,7 +41,7 @@ module Llmemory
           source_id: id.to_s
         )
         rec.embedding = embedding.to_a.map(&:to_f)
-        rec.text_content = text_content
+        rec.text_content = encrypt_text_content(text_content)
         rec.save!
         id
       end
@@ -58,7 +65,11 @@ module Llmemory
           {
             id: r.source_id,
             score: score,
-            metadata: { "text" => r.text_content, "created_at" => r.created_at, "user_id" => r.user_id }
+            metadata: {
+              "text" => decrypt_text_content(r.text_content),
+              "created_at" => r.created_at,
+              "user_id" => r.user_id
+            }
           }
         end
       end
@@ -69,6 +80,22 @@ module Llmemory
         query_embedding = @embedding_provider.embed(query_text)
         search(query_embedding, top_k: top_k, user_id: user_id)
       end
+
+      private
+
+      def encrypt_text_content(text)
+        return text if text.nil? || text.to_s.empty?
+        return text unless @cipher.enabled?
+
+        @cipher.encrypt(text.to_s)
+      end
+
+      def decrypt_text_content(text)
+        return text if text.nil?
+        return text unless text.is_a?(String) && @cipher.enabled? && @cipher.encrypted?(text)
+
+        @cipher.decrypt(text)
+      end
     end
   end
 end

data/lib/llmemory/vector_store/memory_store.rb

@@ -5,9 +5,10 @@ require_relative "base"
 module Llmemory
   module VectorStore
     class MemoryStore < Base
-      def initialize(embedding_provider: nil)
+      def initialize(embedding_provider: nil, cipher: nil)
         @entries = {}
         @embedding_provider = embedding_provider
+        @cipher = cipher || Llmemory.build_cipher
       end
 
       def embed(text)
@@ -15,9 +16,21 @@ module Llmemory
         @embedding_provider.embed(text)
       end
 
+      def last_usage
+        return @embedding_provider.last_usage if @embedding_provider&.respond_to?(:last_usage)
+
+        Llmemory::LLM::Usage.zero
+      end
+
       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: (metadata || {}).merge("user_id" => user_id) }
+        meta = (metadata || {}).dup
+        if meta["text"] && @cipher.enabled?
+          meta["text"] = @cipher.encrypt(meta["text"].to_s)
+        elsif meta[:text] && @cipher.enabled?
+          meta[:text] = @cipher.encrypt(meta[:text].to_s)
+        end
+        @entries[key] = { embedding: embedding.to_a.map(&:to_f), metadata: meta.merge("user_id" => user_id) }
         id
       end
 
@@ -27,7 +40,7 @@ module Llmemory
         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] }
+          { id: id, score: sim, metadata: decrypt_metadata(data[:metadata]) }
         end
         scores.sort_by { |s| -s[:score] }.first(top_k)
       end
@@ -40,6 +53,19 @@ module Llmemory
 
       private
 
+      def decrypt_metadata(meta)
+        return meta unless meta.is_a?(Hash) && @cipher.enabled?
+
+        out = meta.dup
+        text = out["text"] || out[:text]
+        if text.is_a?(String) && @cipher.encrypted?(text)
+          decrypted = @cipher.decrypt(text)
+          out["text"] = decrypted
+          out[:text] = decrypted
+        end
+        out
+      end
+
       def cosine_similarity(a, b)
         return 0.0 if a.size != b.size || a.empty?
         dot = a.zip(b).sum { |x, y| x * y }