llmemory 0.2.3 → 0.2.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8e44ccb1c23fc659d9607e1eb3181e598a57edb95f47b94df4060ad46bfe7c31
-  data.tar.gz: 5d80e1fefb1dd77cbdc8c0b4d24e588abaf6f18c570db10e3fc0b2c808b09461
+  metadata.gz: 296b9d61d6c474145ecaa607653b37438b2491c846aac602f65d5fd850dae9ef
+  data.tar.gz: 521fd05b577c6c17a7dbc5d3771ff9fb3f7cddeaeef31938efabdcfd40db74a5
 SHA512:
-  metadata.gz: edb2f849efc9a6d1dbabc0be0a1210bf21e77f33174b965e3b22e9df3ca90aa536e47152c75891e548d6b43a08d82cfda29afa59d3b2e777aa144f4a3900570f
-  data.tar.gz: a8c84411d9c262dee5b5a0c07835c05816283a17749c1b6b282c75694dc56a497d8fb9f2b18eeccd242a82d5dce36234342ede75a13e2739ad11f99fb23ae005
+  metadata.gz: edf6ee6b41cb366f2ddef5ac2deec9c3c03090d4920d2386f9578e31999d135b34ad7f995ee9a9e803cf43627f969a3fab9ee191c68e2362a4d62fa8bade0729
+  data.tar.gz: 83266918faeb2bb4f7b57e89007b642283d7d64600bf83513948d5378a625fa911aefa6823c54115961bafe34dfdd147f4fb297a8226caad49aaceb6a01372e0

data/README.md

@@ -51,6 +51,39 @@ memory.clear_session!
 - **`prune!(mode: nil)`** — Prunes oversized tool results (soft-trim or hard-clear). Only when `prune_tool_results_enabled` is true.
 - **`check_context_window!`** — Triggers consolidate and compact when context exceeds configured thresholds.
 - **`clear_session!`** — Clears short-term only.
+- **`llm_usage`** — Returns cumulative LLM token usage for this `user_id` (chat/completions + embeddings), persisted in the short-term store.
+
+## LLM token usage
+
+llmemory captures **real token counts** from OpenAI and Anthropic API responses (chat and embeddings), accumulates them per `user_id`, and exposes them for cost monitoring.
+
+```ruby
+memory = Llmemory::Memory.new(user_id: "user_123")
+memory.consolidate!
+memory.maintain!
+
+usage = memory.llm_usage
+# => {
+#      invoke: { input_tokens: 1200, output_tokens: 400, total_tokens: 1600, calls: 3 },
+#      embed:  { total_tokens: 48, calls: 2 },
+#      updated_at: "2026-07-02T12:00:00Z"
+#    }
+```
+
+| What | Details |
+|------|---------|
+| **Counted** | `consolidate!`, reflection, skill mining, compaction summaries, iterative retrieval, graph/file extraction, OpenAI embeddings (index + search) |
+| **Scope** | Cumulative per `user_id` (not per session); stored under pseudo-session `__llm_usage__` |
+| **Not counted** | `context_tokens` (local byte estimate), retrieval context budget, MCP auth tokens |
+| **Cache** | Embedding cache hits record zero tokens |
+
+**Other surfaces:**
+
+- **CLI:** `llmemory stats USER_ID` prints an `LLM TOKEN USAGE` section.
+- **MCP:** `memory_stats` includes the same totals.
+- **Rails metrics:** subscribe to `llm_invoke.llmemory` and `llm_embed.llmemory` (payload includes `input_tokens`, `output_tokens`, `total_tokens`, `response_chars`).
+
+Dollar cost is not computed — multiply tokens by your model pricing externally. For lower-level access, `Llmemory::LLM::OpenAI#invoke` returns a `Response` with `#content` (via `#to_s`) and `#usage`.
 
 ## Configuration
 
@@ -685,7 +718,7 @@ MCP_TOKEN=your-secret-token llmemory mcp serve --http --port 443 \
 | `memory_timeline_context` | Get N items before/after a specific memory |
 | `memory_add_message` | Add message to short-term conversation (roles: user, assistant, system, tool, tool_result) |
 | `memory_consolidate` | Extract facts from conversation to long-term |
-| `memory_stats` | Get memory statistics for a user |
+| `memory_stats` | Get memory statistics for a user (includes LLM token usage) |
 | `memory_info` | Documentation on how to use the tools |
 | `memory_episode_record` / `memory_episodes` | Record / list episodic trajectories |
 | `memory_skill_register` / `memory_skill_report` / `memory_skills` | Register / outcome-track / list procedural skills |

data/lib/llmemory/cli/commands/stats.rb

@@ -41,6 +41,11 @@ module Llmemory
             puts "Long-term (file) categories: #{storage.list_categories(user_id).size}"
             puts "Long-term (file) resources: #{storage.list_resources(user_id: user_id).size}"
           end
+
+          puts "---"
+          puts Llmemory::LLM::UsageLedger.format_text(
+            Llmemory::LLM::UsageLedger.new(store: short_store).totals(user_id)
+          )
         end
 
         def print_global_stats(short_store, long_type)

data/lib/llmemory/instrumentation.rb

@@ -10,8 +10,10 @@ module Llmemory
   # Events (payload keys are best-effort; subscribers should treat them as
   # optional):
   #
-  #   llm_invoke.llmemory       provider:, model:, prompt_chars:, response_chars:
-  #   llm_embed.llmemory        provider:, model:, text_chars:, dimensions:
+  #   llm_invoke.llmemory       provider:, model:, prompt_chars:, response_chars:,
+  #                             input_tokens:, output_tokens:, total_tokens:
+  #   llm_embed.llmemory        provider:, model:, text_chars:, input_tokens:,
+  #                             output_tokens:, total_tokens:
   #   memory_write.llmemory     memory_type:, user_id:
   #   memory_forget.llmemory    memory_type:, user_id:, count:
   #   retrieve.llmemory         query_chars:, candidates:, results:

data/lib/llmemory/llm/anthropic.rb

@@ -11,6 +11,7 @@ module Llmemory
       DEFAULT_MODEL = "claude-sonnet-4-6"
 
       def initialize(api_key: nil, model: nil, base_url: nil)
+        super()
         @api_key = api_key || config.llm_api_key || ENV["ANTHROPIC_API_KEY"]
         @model = model || config.llm_model || DEFAULT_MODEL
         @base_url = base_url || config.llm_base_url || DEFAULT_BASE_URL
@@ -18,7 +19,8 @@ module Llmemory
 
       def invoke(prompt)
         result = nil
-        Llmemory::Instrumentation.instrument(:llm_invoke, provider: :anthropic, model: @model, prompt_chars: prompt.to_s.length) do
+        payload = { provider: :anthropic, model: @model, prompt_chars: prompt.to_s.length }
+        Llmemory::Instrumentation.instrument(:llm_invoke, payload) do
           response = connection.post("v1/messages") do |req|
             req.body = {
               model: @model,
@@ -33,8 +35,11 @@ module Llmemory
           raise Llmemory::LLMError, "Anthropic API error: #{response.body}" unless response.success?
 
           body = response.body.is_a?(Hash) ? response.body : JSON.parse(response.body.to_s)
-          content = body.dig("content", 0, "text")
-          result = content&.strip || ""
+          content = body.dig("content", 0, "text")&.strip || ""
+          usage = parse_anthropic_usage(body["usage"])
+          record_usage(usage)
+          payload.merge!(instrumentation_payload(usage, content))
+          result = Response.new(content, usage: usage)
         end
         result
       end

data/lib/llmemory/llm/base.rb

@@ -1,8 +1,17 @@
 # frozen_string_literal: true
 
+require_relative "usage"
+require_relative "response"
+
 module Llmemory
   module LLM
     class Base
+      attr_reader :last_usage
+
+      def initialize(*)
+        @last_usage = Usage.zero
+      end
+
       def invoke(prompt)
         raise NotImplementedError, "#{self.class}#invoke must be implemented"
       end
@@ -18,6 +27,39 @@ module Llmemory
       def config
         Llmemory.configuration
       end
+
+      def parse_openai_chat_usage(raw)
+        return Usage.zero unless raw.is_a?(Hash)
+
+        Usage.new(
+          input_tokens: raw["prompt_tokens"] || raw[:prompt_tokens] || 0,
+          output_tokens: raw["completion_tokens"] || raw[:completion_tokens] || 0,
+          total_tokens: raw["total_tokens"] || raw[:total_tokens]
+        )
+      end
+
+      def parse_anthropic_usage(raw)
+        return Usage.zero unless raw.is_a?(Hash)
+
+        input = raw["input_tokens"] || raw[:input_tokens] || 0
+        output = raw["output_tokens"] || raw[:output_tokens] || 0
+        Usage.new(input_tokens: input, output_tokens: output)
+      end
+
+      def parse_openai_embed_usage(raw)
+        return Usage.zero unless raw.is_a?(Hash)
+
+        total = raw["total_tokens"] || raw[:total_tokens] || 0
+        Usage.new(input_tokens: 0, output_tokens: 0, total_tokens: total)
+      end
+
+      def record_usage(usage)
+        @last_usage = usage
+      end
+
+      def instrumentation_payload(usage, content, extra = {})
+        usage.to_h.merge(response_chars: content.to_s.length).merge(extra)
+      end
     end
   end
 end

data/lib/llmemory/llm/openai.rb

@@ -11,6 +11,7 @@ module Llmemory
       DEFAULT_MODEL = "gpt-4"
 
       def initialize(api_key: nil, model: nil, base_url: nil)
+        super()
         @api_key = api_key || config.llm_api_key
         @model = model || config.llm_model || DEFAULT_MODEL
         @base_url = base_url || config.llm_base_url || DEFAULT_BASE_URL
@@ -18,7 +19,8 @@ module Llmemory
 
       def invoke(prompt)
         result = nil
-        Llmemory::Instrumentation.instrument(:llm_invoke, provider: :openai, model: @model, prompt_chars: prompt.to_s.length) do
+        payload = { provider: :openai, model: @model, prompt_chars: prompt.to_s.length }
+        Llmemory::Instrumentation.instrument(:llm_invoke, payload) do
           response = connection.post("chat/completions") do |req|
             req.body = {
               model: @model,
@@ -32,7 +34,11 @@ module Llmemory
           raise Llmemory::LLMError, "OpenAI API error: #{response.body}" unless response.success?
 
           body = response.body.is_a?(Hash) ? response.body : JSON.parse(response.body.to_s)
-          result = body.dig("choices", 0, "message", "content")&.strip || ""
+          content = body.dig("choices", 0, "message", "content")&.strip || ""
+          usage = parse_openai_chat_usage(body["usage"])
+          record_usage(usage)
+          payload.merge!(instrumentation_payload(usage, content))
+          result = Response.new(content, usage: usage)
         end
         result
       end
@@ -54,18 +60,27 @@ module Llmemory
             }
           }
         }
-        response = connection.post("chat/completions") do |req|
-          req.body = payload.to_json
-          req.headers["Content-Type"] = "application/json"
-          req.headers["Authorization"] = "Bearer #{@api_key}"
-        end
+        parsed = nil
+        instrument_payload = { provider: :openai, model: @model, prompt_chars: prompt.to_s.length }
+        Llmemory::Instrumentation.instrument(:llm_invoke, instrument_payload) do
+          response = connection.post("chat/completions") do |req|
+            req.body = payload.to_json
+            req.headers["Content-Type"] = "application/json"
+            req.headers["Authorization"] = "Bearer #{@api_key}"
+          end
 
-        raise Llmemory::LLMError, "OpenAI API error: #{response.body}" unless response.success?
+          raise Llmemory::LLMError, "OpenAI API error: #{response.body}" unless response.success?
+
+          body = response.body.is_a?(Hash) ? response.body : JSON.parse(response.body.to_s)
+          content = body.dig("choices", 0, "message", "content")&.strip
+          usage = parse_openai_chat_usage(body["usage"])
+          record_usage(usage)
+          instrument_payload.merge!(instrumentation_payload(usage, content.to_s))
+          return {} if content.nil? || content.empty?
 
-        body = response.body.is_a?(Hash) ? response.body : JSON.parse(response.body.to_s)
-        content = body.dig("choices", 0, "message", "content")&.strip
-        return {} if content.nil? || content.empty?
-        JSON.parse(content)
+          parsed = JSON.parse(content)
+        end
+        parsed
       rescue JSON::ParserError => e
         raise Llmemory::LLMError, "Failed to parse JSON response: #{e.message}"
       end

data/lib/llmemory/llm/response.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module LLM
+    class Response
+      attr_reader :content, :usage
+
+      def initialize(content, usage: Usage.zero)
+        @content = content.to_s
+        @usage = usage
+      end
+
+      def to_s
+        @content
+      end
+    end
+  end
+end

data/lib/llmemory/llm/tracking_client.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require_relative "usage_recorder"
+
+module Llmemory
+  module LLM
+    # Transparent wrapper that records token usage to the per-user ledger.
+    class TrackingClient
+      def initialize(inner, user_id:, store: nil, api_key: nil)
+        @inner = inner
+        @user_id = user_id
+        @store = store
+        @api_key = api_key
+      end
+
+      def invoke(prompt)
+        response = inner_client.invoke(prompt)
+        usage = if response.respond_to?(:usage)
+                  response.usage
+                elsif inner_client.respond_to?(:last_usage)
+                  inner_client.last_usage
+                else
+                  Usage.zero
+                end
+        UsageRecorder.record(user_id: @user_id, usage: usage, operation: :invoke, store: @store)
+        response
+      end
+
+      def invoke_with_json_schema(prompt, json_schema)
+        result = inner_client.invoke_with_json_schema(prompt, json_schema)
+        usage = inner_client.respond_to?(:last_usage) ? inner_client.last_usage : Usage.zero
+        UsageRecorder.record(user_id: @user_id, usage: usage, operation: :invoke, store: @store)
+        result
+      end
+
+      def last_usage
+        return inner_client.last_usage if inner_client.respond_to?(:last_usage)
+
+        Usage.zero
+      end
+
+      def respond_to?(method, include_private = false)
+        inner_client.respond_to?(method, include_private) || super
+      end
+
+      def method_missing(method, *args, &block)
+        if inner_client.respond_to?(method)
+          inner_client.public_send(method, *args, &block)
+        else
+          super
+        end
+      end
+
+      private
+
+      def inner_client
+        @inner_client ||= @inner || Llmemory::LLM.client(api_key: @api_key)
+      end
+    end
+  end
+end

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

@@ -135,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
@@ -142,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
@@ -183,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/graph_based/memory.rb

@@ -162,6 +162,7 @@ module Llmemory
 
             edge_text = "#{subject} #{predicate} #{object}"
             embedding = @vector_store.respond_to?(:embed) ? @vector_store.embed(edge_text) : nil
+            record_embed_usage(@vector_store) if embedding
             if embedding && @vector_store.respond_to?(:store)
               @vector_store.store(id: edge_id, embedding: embedding, metadata: { text: edge_text, created_at: Time.now }, user_id: @user_id)
             end
@@ -172,8 +173,10 @@ module Llmemory
           vector_results = []
           if @vector_store.respond_to?(:search_by_text)
             vector_results = @vector_store.search_by_text(query.to_s, top_k: top_k, user_id: @user_id)
+            record_embed_usage(@vector_store)
           elsif @vector_store.respond_to?(:embed) && @vector_store.respond_to?(:search)
             emb = @vector_store.embed(query.to_s)
+            record_embed_usage(@vector_store)
             vector_results = @vector_store.search(emb, top_k: top_k, user_id: @user_id)
           end
 
@@ -231,6 +234,16 @@ module Llmemory
           lines << "=== END MEMORIES ==="
           lines.join("\n")
         end
+
+        def record_embed_usage(vector_store)
+          return unless 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/procedural/memory.rb

@@ -141,6 +141,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
@@ -148,7 +149,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_skill(@user_id, r[:id] || r["id"])
             raw && candidate_for(raw, (r[:score] || r["score"] || 1.0).to_f)
           end
@@ -178,6 +181,14 @@ module Llmemory
           end
           by_id.values.sort_by { |c| -c[:score].to_f }.first(top_k)
         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/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

@@ -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

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.4"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: llmemory
 version: !ruby/object:Gem::Version
-  version: 0.2.3
+  version: 0.2.4
 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