llmemory 0.2.2 → 0.2.4

data/lib/llmemory/llm/usage.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module LLM
+    class Usage
+      attr_reader :input_tokens, :output_tokens, :total_tokens
+
+      def initialize(input_tokens:, output_tokens:, total_tokens: nil)
+        @input_tokens = input_tokens.to_i
+        @output_tokens = output_tokens.to_i
+        @total_tokens = total_tokens.nil? ? (@input_tokens + @output_tokens) : total_tokens.to_i
+      end
+
+      def self.zero
+        new(input_tokens: 0, output_tokens: 0, total_tokens: 0)
+      end
+
+      def +(other)
+        self.class.new(
+          input_tokens: @input_tokens + other.input_tokens,
+          output_tokens: @output_tokens + other.output_tokens,
+          total_tokens: @total_tokens + other.total_tokens
+        )
+      end
+
+      def to_h
+        { input_tokens: @input_tokens, output_tokens: @output_tokens, total_tokens: @total_tokens }
+      end
+    end
+  end
+end

data/lib/llmemory/llm/usage_ledger.rb

@@ -0,0 +1,118 @@
+# frozen_string_literal: true
+
+require "time"
+require_relative "../short_term/stores"
+
+module Llmemory
+  module LLM
+    # Cumulative LLM token usage per user, persisted in the short-term store
+    # under a pseudo-session key (same pattern as ForgetLog).
+    class UsageLedger
+      SESSION_KEY = "__llm_usage__"
+
+      def initialize(store: nil)
+        @store = store || ShortTerm::Stores.build
+      end
+
+      def record(user_id, usage, operation:)
+        state = load_raw(user_id)
+        case operation.to_sym
+        when :invoke
+          bucket = symbolize_bucket(state[:invoke] || state["invoke"])
+          state = state.merge(
+            invoke: {
+              input_tokens: bucket[:input_tokens] + usage.input_tokens,
+              output_tokens: bucket[:output_tokens] + usage.output_tokens,
+              total_tokens: bucket[:total_tokens] + usage.total_tokens,
+              calls: bucket[:calls] + 1
+            }
+          )
+        when :embed
+          bucket = symbolize_bucket(state[:embed] || state["embed"], embed: true)
+          state = state.merge(
+            embed: {
+              total_tokens: bucket[:total_tokens] + usage.total_tokens,
+              calls: bucket[:calls] + 1
+            }
+          )
+        else
+          return totals(user_id)
+        end
+        state[:updated_at] = Time.now.iso8601
+        @store.save(user_id, SESSION_KEY, stringify(state))
+        totals(user_id)
+      end
+
+      def totals(user_id)
+        normalize(load_raw(user_id))
+      end
+
+      def reset!(user_id)
+        empty = default_state
+        @store.save(user_id, SESSION_KEY, stringify(empty))
+        empty
+      end
+
+      def self.format_text(totals)
+        inv = totals[:invoke]
+        emb = totals[:embed]
+        lines = [
+          "LLM TOKEN USAGE:",
+          "  Chat/completions: #{inv[:total_tokens]} total (#{inv[:input_tokens]} in, #{inv[:output_tokens]} out, #{inv[:calls]} calls)",
+          "  Embeddings: #{emb[:total_tokens]} total (#{emb[:calls]} calls)"
+        ]
+        lines << "  Last updated: #{totals[:updated_at]}" if totals[:updated_at]
+        lines.join("\n")
+      end
+
+      private
+
+      def load_raw(user_id)
+        state = @store.load(user_id, SESSION_KEY)
+        return default_state unless state.is_a?(Hash)
+        normalize(state)
+      end
+
+      def default_state
+        {
+          invoke: { input_tokens: 0, output_tokens: 0, total_tokens: 0, calls: 0 },
+          embed: { total_tokens: 0, calls: 0 },
+          updated_at: nil
+        }
+      end
+
+      def normalize(state)
+        invoke = symbolize_bucket(state[:invoke] || state["invoke"])
+        embed = symbolize_bucket(state[:embed] || state["embed"], embed: true)
+        {
+          invoke: invoke,
+          embed: embed,
+          updated_at: state[:updated_at] || state["updated_at"]
+        }
+      end
+
+      def symbolize_bucket(bucket, embed: false)
+        bucket = {} unless bucket.is_a?(Hash)
+        if embed
+          {
+            total_tokens: (bucket[:total_tokens] || bucket["total_tokens"] || 0).to_i,
+            calls: (bucket[:calls] || bucket["calls"] || 0).to_i
+          }
+        else
+          {
+            input_tokens: (bucket[:input_tokens] || bucket["input_tokens"] || 0).to_i,
+            output_tokens: (bucket[:output_tokens] || bucket["output_tokens"] || 0).to_i,
+            total_tokens: (bucket[:total_tokens] || bucket["total_tokens"] || 0).to_i,
+            calls: (bucket[:calls] || bucket["calls"] || 0).to_i
+          }
+        end
+      end
+
+      def stringify(state)
+        state.transform_keys(&:to_s).transform_values do |v|
+          v.is_a?(Hash) ? v.transform_keys(&:to_s) : v
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/llm/usage_recorder.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require_relative "usage_ledger"
+
+module Llmemory
+  module LLM
+    module UsageRecorder
+      module_function
+
+      def record(user_id:, usage:, operation:, store: nil)
+        return if user_id.nil? || user_id.to_s.empty?
+        return if usage.nil?
+
+        UsageLedger.new(store: store).record(user_id, usage, operation: operation)
+      end
+
+      def record_embed_from_store(user_id:, vector_store:, store: nil)
+        usage = embed_usage_from(vector_store)
+        return unless usage
+
+        record(user_id: user_id, usage: usage, operation: :embed, store: store)
+      end
+
+      def embed_usage_from(vector_store)
+        return nil unless vector_store
+
+        if vector_store.respond_to?(:last_usage)
+          usage = vector_store.last_usage
+          return usage unless usage.nil?
+        end
+
+        provider = vector_store.instance_variable_get(:@embedding_provider) if vector_store.instance_variable_defined?(:@embedding_provider)
+        provider&.last_usage if provider&.respond_to?(:last_usage)
+      end
+    end
+  end
+end

data/lib/llmemory/llm.rb

@@ -1,6 +1,11 @@
 # frozen_string_literal: true
 
 require_relative "llm/base"
+require_relative "llm/usage"
+require_relative "llm/response"
+require_relative "llm/usage_ledger"
+require_relative "llm/usage_recorder"
+require_relative "llm/tracking_client"
 require_relative "llm/openai"
 require_relative "llm/anthropic"
 

data/lib/llmemory/long_term/episodic/memory.rb

@@ -21,9 +21,10 @@ module Llmemory
 
         attr_reader :user_id, :storage
 
-        def initialize(user_id:, storage: nil, vector_store: nil)
+        def initialize(user_id:, storage: nil, vector_store: nil, cipher: nil)
           @user_id = user_id
-          @storage = storage || Storages.build
+          @cipher = cipher || Llmemory.build_cipher
+          @storage = storage || Storages.build(cipher: @cipher)
           @vector_store = vector_store
           @vector_explicit = !vector_store.nil?
         end
@@ -125,7 +126,7 @@ module Llmemory
           if @vector_explicit
             @vector_store
           elsif Llmemory.configuration.episodic_vector_enabled
-            @vector_store ||= Llmemory::VectorStore.build(source_type: "episode")
+            @vector_store ||= Llmemory::VectorStore.build(source_type: "episode", cipher: @cipher)
           end
         end
 
@@ -134,6 +135,7 @@ module Llmemory
           vs = vector_store
           return if vs.nil? || text.to_s.strip.empty?
           embedding = vs.embed(text)
+          record_embed_usage(vs)
           return unless embedding
           vs.store(id: id, embedding: embedding, metadata: { text: text, created_at: Time.now }, user_id: @user_id)
         rescue StandardError
@@ -141,7 +143,9 @@ module Llmemory
         end
 
         def vector_candidates(query, top_k, vs)
-          vs.search_by_text(query.to_s, top_k: top_k, user_id: @user_id).filter_map do |r|
+          results = vs.search_by_text(query.to_s, top_k: top_k, user_id: @user_id)
+          record_embed_usage(vs)
+          results.filter_map do |r|
             raw = @storage.get_episode(@user_id, r[:id] || r["id"])
             raw && candidate_for(raw, (r[:score] || r["score"] || 1.0).to_f)
           end
@@ -182,6 +186,14 @@ module Llmemory
           return nil if actions.empty?
           "Episode with #{normalized.size} step(s): #{actions.join(' -> ')}"
         end
+
+        def record_embed_usage(vector_store)
+          Llmemory::LLM::UsageRecorder.record_embed_from_store(
+            user_id: @user_id,
+            vector_store: vector_store,
+            store: Llmemory::ShortTerm::Stores.build(cipher: @cipher)
+          )
+        end
       end
     end
   end

data/lib/llmemory/long_term/episodic/storage.rb

@@ -12,17 +12,24 @@ module Llmemory
       Storage = Storages::MemoryStorage
 
       module Storages
-        def self.build(store: nil, base_path: nil, database_url: nil)
+        def self.build(store: nil, base_path: nil, database_url: nil, cipher: nil)
+          resolved_cipher = cipher || Llmemory.build_cipher
           case (store || Llmemory.configuration.long_term_store).to_s.to_sym
           when :memory
             MemoryStorage.new
           when :file
-            FileStorage.new(base_path: base_path || Llmemory.configuration.long_term_storage_path)
+            FileStorage.new(
+              base_path: base_path || Llmemory.configuration.long_term_storage_path,
+              cipher: resolved_cipher
+            )
           when :postgres, :database
-            DatabaseStorage.new(database_url: database_url || Llmemory.configuration.database_url)
+            DatabaseStorage.new(
+              database_url: database_url || Llmemory.configuration.database_url,
+              cipher: resolved_cipher
+            )
           when :active_record, :activerecord
             require_relative "storages/active_record_storage"
-            ActiveRecordStorage.new
+            ActiveRecordStorage.new(cipher: resolved_cipher)
           else
             MemoryStorage.new
           end

data/lib/llmemory/long_term/episodic/storages/active_record_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,7 +14,10 @@ module Llmemory
         # AR auto-deserializes jsonb to a Hash (string keys), which Episode.from_h
         # handles. Mirrors the file-based ActiveRecordStorage pattern.
         class ActiveRecordStorage < Base
-          def initialize
+          include Llmemory::Crypto::FieldHelpers
+
+          def initialize(cipher: nil)
+            @cipher = cipher || Llmemory.build_cipher
             self.class.load_models!
           end
 
@@ -30,8 +34,8 @@ module Llmemory
             data["created_at"] ||= Time.now.utc.iso8601
             rec = LlmemoryEpisode.find_or_initialize_by(id: id)
             rec.user_id = user_id
-            rec.data = data
-            rec.search_text = searchable_text(data)
+            rec.data = cipher.enabled? ? enc_json(data) : data
+            rec.search_text = enc(searchable_text(data))
             rec.created_at ||= Time.current
             rec.save!
             id
@@ -39,19 +43,21 @@ module Llmemory
 
           def get_episode(user_id, id)
             rec = LlmemoryEpisode.find_by(user_id: user_id, id: id)
-            rec&.data
+            return nil unless rec
+
+            decode_data(rec.data)
           end
 
           def list_episodes(user_id, limit: nil, offset: nil)
             scope = LlmemoryEpisode.where(user_id: user_id, archived_at: nil).order(created_at: :desc)
             scope = scope.limit(limit) if limit && limit.to_i.positive?
             scope = scope.offset(offset) if offset && offset.to_i.positive?
-            scope.map(&:data)
+            scope.map { |r| decode_data(r.data) }
           end
 
           def search_episodes(user_id, query)
             token_scope(LlmemoryEpisode.where(user_id: user_id, archived_at: nil), "search_text", query)
-              .order(created_at: :desc).map(&:data)
+              .order(created_at: :desc).map { |r| decode_data(r.data) }
           end
 
           def count_episodes(user_id)
@@ -96,6 +102,13 @@ module Llmemory
             end
             parts.compact.join("\n")
           end
+
+          def decode_data(raw)
+            return raw.transform_keys(&:to_sym) if raw.is_a?(Hash)
+            return dec_json(raw) if raw.is_a?(String) && cipher.encrypted?(raw)
+
+            raw
+          end
         end
       end
     end

data/lib/llmemory/long_term/episodic/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 for keyword
         # search), 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_episode(user_id, episode)
@@ -23,11 +27,12 @@ module Llmemory
             id = episode[:id] || episode["id"] || "ep_#{SecureRandom.hex(8)}"
             data = symbolize(episode).merge(id: id, user_id: user_id)
             data[:created_at] ||= Time.now.utc.iso8601
+            search = searchable_text(data)
             conn.exec_params(
               "INSERT INTO llmemory_episodes (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(search), created_at_value(data)]
             )
             id
           end
@@ -124,11 +129,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/episodic/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 Episodic
       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_episode(user_id, episode)
             id = episode[:id] || episode["id"] || "ep_#{next_seq(user_id)}"
             data = stringify_for_json(episode).merge("id" => id, "user_id" => user_id)
             data["created_at"] ||= Time.now.iso8601
-            File.write(episode_path(user_id, id), JSON.generate(data))
+            write_episode_file(episode_path(user_id, id), data)
             id
           end
 
@@ -57,10 +61,10 @@ module Llmemory
             Array(ids).map(&:to_s).count do |id|
               path = episode_path(user_id, id)
               next false unless File.file?(path)
-              data = JSON.parse(File.read(path))
+              data = load_episode_raw(path)
               next false if data["archived_at"]
               data["archived_at"] = Time.now.iso8601
-              File.write(path, JSON.generate(data))
+              write_episode_file(path, data)
               true
             end
           end
@@ -89,13 +93,26 @@ module Llmemory
           end
 
           def load_episode(path)
-            data = JSON.parse(File.read(path), symbolize_names: true)
+            data = load_episode_raw(path)
+            return nil unless data
+
             data[:created_at] = parse_time(data[:created_at])
             data
           rescue JSON::ParserError
             nil
           end
 
+          def load_episode_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_episode_file(path, data)
+            payload = JSON.generate(data)
+            File.write(path, cipher.enabled? ? cipher.encrypt(payload) : payload)
+          end
+
           def episode_text(episode)
             parts = [episode[:summary], episode[:outcome]]
             Array(episode[:steps]).each do |s|

data/lib/llmemory/long_term/file_based/storage.rb

@@ -14,17 +14,24 @@ module Llmemory
       Storage = Storages::MemoryStorage
 
       module Storages
-        def self.build(store: nil, base_path: nil, database_url: nil)
+        def self.build(store: nil, base_path: nil, database_url: nil, cipher: nil)
+          resolved_cipher = cipher || Llmemory.build_cipher
           case (store || Llmemory.configuration.long_term_store).to_s.to_sym
           when :memory
             MemoryStorage.new
           when :file
-            FileStorage.new(base_path: base_path || Llmemory.configuration.long_term_storage_path)
+            FileStorage.new(
+              base_path: base_path || Llmemory.configuration.long_term_storage_path,
+              cipher: resolved_cipher
+            )
           when :postgres, :database
-            DatabaseStorage.new(database_url: database_url || Llmemory.configuration.database_url)
+            DatabaseStorage.new(
+              database_url: database_url || Llmemory.configuration.database_url,
+              cipher: resolved_cipher
+            )
           when :active_record, :activerecord
             require_relative "storages/active_record_storage"
-            ActiveRecordStorage.new
+            ActiveRecordStorage.new(cipher: resolved_cipher)
           else
             MemoryStorage.new
           end

data/lib/llmemory/long_term/file_based/storages/active_record_storage.rb

@@ -2,13 +2,17 @@
 
 require "securerandom"
 require_relative "base"
+require_relative "../../../crypto/field_helpers"
 
 module Llmemory
   module LongTerm
     module FileBased
       module Storages
         class ActiveRecordStorage < Base
-          def initialize
+          include Llmemory::Crypto::FieldHelpers
+
+          def initialize(cipher: nil)
+            @cipher = cipher || Llmemory.build_cipher
             self.class.load_models!
           end
 
@@ -24,7 +28,7 @@ module Llmemory
             LlmemoryResource.create!(
               id: id,
               user_id: user_id,
-              text: text,
+              text: enc(text),
               created_at: Time.current
             )
             id
@@ -36,24 +40,26 @@ module Llmemory
               id: id,
               user_id: user_id,
               category: category,
-              content: content,
+              content: enc(content),
               source_resource_id: source_resource_id,
               created_at: Time.current
             }
             attrs[:importance] = importance if LlmemoryItem.column_names.include?("importance")
-            attrs[:provenance] = provenance if provenance && LlmemoryItem.column_names.include?("provenance")
+            if provenance && LlmemoryItem.column_names.include?("provenance")
+              attrs[:provenance] = cipher.enabled? ? enc_json(provenance) : provenance
+            end
             LlmemoryItem.create!(attrs)
             id
           end
 
           def load_category(user_id, category_name)
             rec = LlmemoryCategory.find_by(user_id: user_id, category_name: category_name)
-            rec ? rec.content.to_s : ""
+            rec ? dec(rec.content.to_s) : ""
           end
 
           def save_category(user_id, category_name, content)
             rec = LlmemoryCategory.find_or_initialize_by(user_id: user_id, category_name: category_name)
-            rec.content = content
+            rec.content = enc(content)
             rec.updated_at = Time.current
             rec.save!
             true
@@ -101,7 +107,7 @@ module Llmemory
               id: "item_#{SecureRandom.hex(8)}",
               user_id: user_id,
               category: merged_item[:category],
-              content: merged_item[:content],
+              content: enc(merged_item[:content]),
               source_resource_id: merged_item[:source_resource_id],
               created_at: created_at
             }
@@ -194,19 +200,19 @@ module Llmemory
             h = {
               id: r.id,
               category: r.category,
-              content: r.content,
+              content: dec(r.content),
               source_resource_id: r.source_resource_id,
               created_at: r.created_at
             }
             h[:importance] = r.respond_to?(:importance) ? (r.importance || 0.7).to_f : 0.7
-            h[:provenance] = r.provenance if r.respond_to?(:provenance)
+            h[:provenance] = parse_provenance(r.provenance) if r.respond_to?(:provenance)
             h
           end
 
           def row_to_resource(r)
             {
               id: r.id,
-              text: r.text,
+              text: dec(r.text),
               created_at: r.created_at
             }
           end