llmemory 0.2.3 → 0.2.5

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

@@ -31,7 +31,11 @@ module Llmemory
             end
 
             storage = build_storage
-            result = storage.get_items_around(user_id, reference, before: before_count, after: after_count)
+            result = if graph_based?
+              storage.get_edges_around(user_id, reference, before: before_count, after: after_count)
+            else
+              storage.get_items_around(user_id, reference, before: before_count, after: after_count)
+            end
 
             ::MCP::Tool::Response.new([{
               type: "text",
@@ -47,13 +51,17 @@ module Llmemory
           private
 
           def build_storage
-            if Llmemory.configuration.long_term_type.to_s == "graph_based"
+            if graph_based?
               LongTerm::GraphBased::Storages.build
             else
               LongTerm::FileBased::Storages.build
             end
           end
 
+          def graph_based?
+            Llmemory.configuration.long_term_type.to_s == "graph_based"
+          end
+
           def format_context(result, reference)
             output = []
             output << "Timeline Context around '#{reference}':\n"
@@ -122,10 +130,12 @@ module Llmemory
           end
 
           def extract_category(item)
-            if item.respond_to?(:category)
-              item.category
-            else
+            if item.is_a?(Hash)
               item[:category] || item["category"]
+            elsif item.respond_to?(:members) && item.members.include?(:category)
+              item.category
+            elsif item.respond_to?(:category)
+              item.category
             end
           end
 

data/lib/llmemory/memory.rb

@@ -15,22 +15,28 @@ module Llmemory
       @session_id = session_id
       resolved_key = encryption_key == :inherit ? nil : encryption_key
       @cipher = Llmemory.build_cipher(resolved_key)
-      @checkpoint = checkpoint || ShortTerm::Checkpoint.new(
-        user_id: user_id,
-        session_id: session_id,
-        cipher: @cipher
-      )
+      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)
-      short_term_store = build_short_term_store(@cipher)
       @retrieval_engine = retrieval_engine || Retrieval::Engine.new(
         @long_term,
-        llm: @llm,
-        feedback: Retrieval::FeedbackStore.new(store: short_term_store)
+        llm: tracked_llm_client,
+        feedback: Retrieval::FeedbackStore.new(store: @short_term_store)
       )
     end
 
@@ -66,14 +72,14 @@ module Llmemory
     # 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
@@ -81,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
 
@@ -91,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
@@ -245,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)
@@ -263,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)
@@ -339,7 +358,7 @@ 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(

data/lib/llmemory/short_term/checkpoint.rb

@@ -14,6 +14,8 @@ module Llmemory
         @store = store || build_store
       end
 
+      attr_reader :store
+
       def save_state(state)
         @store.save(@user_id, @session_id, state)
       end

data/lib/llmemory/vector_store/active_record_store.rb

@@ -26,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)

data/lib/llmemory/vector_store/memory_store.rb

@@ -16,6 +16,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)
         key = user_id ? "#{user_id}:#{id}" : id.to_s
         meta = (metadata || {}).dup
@@ -24,7 +30,11 @@ module Llmemory
         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) }
+        @entries[key] = {
+          source_id: id.to_s,
+          embedding: embedding.to_a.map(&:to_f),
+          metadata: meta.merge("user_id" => user_id)
+        }
         id
       end
 
@@ -32,9 +42,9 @@ module Llmemory
         query = query_embedding.to_a.map(&:to_f)
         return [] if query.empty?
         entries = user_id ? @entries.select { |k, _| k.to_s.start_with?("#{user_id}:") } : @entries
-        scores = entries.map do |id, data|
+        scores = entries.map do |_key, data|
           sim = cosine_similarity(query, data[:embedding])
-          { id: id, score: sim, metadata: decrypt_metadata(data[:metadata]) }
+          { id: data[:source_id], score: sim, metadata: decrypt_metadata(data[:metadata]) }
         end
         scores.sort_by { |s| -s[:score] }.first(top_k)
       end

data/lib/llmemory/vector_store/openai_embeddings.rb

@@ -4,6 +4,7 @@ require "faraday"
 require "json"
 require "digest"
 require_relative "base"
+require_relative "../llm/usage"
 
 module Llmemory
   module VectorStore
@@ -11,11 +12,14 @@ module Llmemory
       DEFAULT_MODEL = "text-embedding-3-small"
       DEFAULT_DIMS = 1536
 
+      attr_reader :last_usage
+
       def initialize(api_key: nil, model: nil)
         @api_key = api_key || Llmemory.configuration.llm_api_key
         @model = model || DEFAULT_MODEL
         @cache = {}
         @cache_order = []
+        @last_usage = Llmemory::LLM::Usage.zero
       end
 
       def embed(text)
@@ -23,7 +27,10 @@ module Llmemory
 
         if Llmemory.configuration.embedding_cache_enabled
           key = cache_key(text)
-          return @cache[key].dup if @cache.key?(key)
+          if @cache.key?(key)
+            @last_usage = Llmemory::LLM::Usage.zero
+            return @cache[key].dup
+          end
         end
 
         result = fetch_embedding(text)
@@ -55,7 +62,8 @@ module Llmemory
 
       def fetch_embedding(text)
         result = nil
-        Llmemory::Instrumentation.instrument(:llm_embed, provider: :openai, model: @model, text_chars: text.to_s.length) do
+        payload = { provider: :openai, model: @model, text_chars: text.to_s.length }
+        Llmemory::Instrumentation.instrument(:llm_embed, payload) do
           response = connection.post("embeddings") do |req|
             req.headers["Authorization"] = "Bearer #{@api_key}"
             req.headers["Content-Type"] = "application/json"
@@ -63,11 +71,24 @@ module Llmemory
           end
           raise Llmemory::LLMError, "OpenAI Embeddings API error: #{response.body}" unless response.success?
           body = response.body.is_a?(Hash) ? response.body : JSON.parse(response.body.to_s)
+          @last_usage = parse_embed_usage(body["usage"])
+          payload.merge!(
+            input_tokens: @last_usage.input_tokens,
+            output_tokens: @last_usage.output_tokens,
+            total_tokens: @last_usage.total_tokens
+          )
           result = body.dig("data", 0, "embedding")&.map(&:to_f) || Array.new(DEFAULT_DIMS, 0.0)
         end
         result
       end
 
+      def parse_embed_usage(raw)
+        return Llmemory::LLM::Usage.zero unless raw.is_a?(Hash)
+
+        total = raw["total_tokens"] || raw[:total_tokens] || 0
+        Llmemory::LLM::Usage.new(input_tokens: 0, output_tokens: 0, total_tokens: total)
+      end
+
       def connection
         @connection ||= Faraday.new(url: "https://api.openai.com/v1") do |f|
           f.request :json

data/lib/llmemory/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Llmemory
-  VERSION = "0.2.3"
+  VERSION = "0.2.5"
 end

data/lib/tasks/llmemory.rake

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+namespace :llmemory do
+  desc "Backfill search_tokens blind index (and name_det on skills) for encrypted keyword search. " \
+       "Env: FORCE=1 (re-backfill all rows), DRY_RUN=1 (count only), STORE=active_record|postgres"
+  task :backfill_search_tokens, [:user_id] do |_t, args|
+    require "llmemory"
+    require_relative "../llmemory/maintenance/search_tokens_backfill"
+
+    dry_run = ENV["DRY_RUN"].to_s == "1"
+    force = ENV["FORCE"].to_s == "1"
+    store = ENV["STORE"]&.to_sym
+
+    if store.nil? && !Llmemory.configuration.encryption_enabled
+      warn "Note: encryption_enabled is false; backfill still aligns schema but keyword search uses plaintext LIKE."
+    end
+
+    backfill = Llmemory::Maintenance::SearchTokensBackfill.new(
+      dry_run: dry_run,
+      force: force,
+      store: store
+    )
+
+    user_id = args[:user_id].presence
+    puts "Backfilling search_tokens#{user_id ? " for user_id=#{user_id}" : " (all users)"} " \
+         "[store=#{store || Llmemory.configuration.long_term_store}]" \
+         "#{dry_run ? " [DRY RUN]" : ""}#{force ? " [FORCE]" : ""}"
+
+    result = backfill.run(user_id: user_id)
+    summary = result.to_h
+
+    puts "Updated: items=#{summary[:items]} resources=#{summary[:resources]} " \
+         "episodes=#{summary[:episodes]} skills=#{summary[:skills]} " \
+         "(total=#{summary[:total]})"
+    puts "Dry run — no rows written." if dry_run
+  end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: llmemory
 version: !ruby/object:Gem::Version
-  version: 0.2.3
+  version: 0.2.5
 platform: ruby
 authors:
 - llmemory
@@ -175,6 +175,11 @@ files:
 - lib/llmemory/llm/anthropic.rb
 - lib/llmemory/llm/base.rb
 - lib/llmemory/llm/openai.rb
+- lib/llmemory/llm/response.rb
+- lib/llmemory/llm/tracking_client.rb
+- lib/llmemory/llm/usage.rb
+- lib/llmemory/llm/usage_ledger.rb
+- lib/llmemory/llm/usage_recorder.rb
 - lib/llmemory/long_term.rb
 - lib/llmemory/long_term/episodic.rb
 - lib/llmemory/long_term/episodic/episode.rb
@@ -225,6 +230,7 @@ files:
 - lib/llmemory/maintenance/consolidator.rb
 - lib/llmemory/maintenance/reindexer.rb
 - lib/llmemory/maintenance/runner.rb
+- lib/llmemory/maintenance/search_tokens_backfill.rb
 - lib/llmemory/maintenance/summarizer.rb
 - lib/llmemory/maintenance/ttl_expiry.rb
 - lib/llmemory/mcp.rb
@@ -283,6 +289,7 @@ files:
 - lib/llmemory/vector_store/openai_embeddings.rb
 - lib/llmemory/version.rb
 - lib/llmemory/working_memory.rb
+- lib/tasks/llmemory.rake
 - lib/tasks/release.rake
 homepage: https://github.com/entaina/llmemory
 licenses:
@@ -304,7 +311,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 4.0.10
+rubygems_version: 4.0.16
 specification_version: 4
 summary: Persistent memory system for LLM agents
 test_files: []