llmemory 0.2.4 → 0.2.5

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

@@ -35,7 +35,11 @@ module Llmemory
             rec = LlmemorySkill.find_or_initialize_by(id: id)
             rec.user_id = user_id
             rec.data = cipher.enabled? ? enc_json(data) : data
-            rec.search_text = enc(searchable_text(data))
+            text = searchable_text(data)
+            rec.search_text = enc(text)
+            rec.search_tokens = search_tokens_for(text) if LlmemorySkill.column_names.include?("search_tokens")
+            name = (data["name"] || data[:name]).to_s
+            rec.name_det = enc_det(name) if LlmemorySkill.column_names.include?("name_det")
             rec.created_at ||= Time.current
             rec.save!
             id
@@ -56,12 +60,21 @@ module Llmemory
           end
 
           def search_skills(user_id, query)
-            token_scope(LlmemorySkill.where(user_id: user_id, archived_at: nil), "search_text", query)
-              .order(created_at: :desc).map { |r| decode_data(r.data) }
+            scope = LlmemorySkill.where(user_id: user_id, archived_at: nil)
+            token_scope(scope, "search_text", query, model: LlmemorySkill)
+              .sort_by { |r| -r.created_at.to_i }
+              .map { |r| decode_data(r.data) }
           end
 
           def find_skills_by_name(user_id, name)
-            if cipher.enabled?
+            if cipher.enabled? && LlmemorySkill.column_names.include?("name_det")
+              indexed = LlmemorySkill.where(user_id: user_id, archived_at: nil, name_det: enc_det(name.to_s))
+                .map { |r| decode_data(r.data) }
+              legacy = LlmemorySkill.where(user_id: user_id, archived_at: nil, name_det: nil)
+                .map { |r| decode_data(r.data) }
+                .select { |s| (s[:name] || s["name"]).to_s == name.to_s }
+              (indexed + legacy).uniq { |s| s[:id] || s["id"] }
+            elsif cipher.enabled?
               LlmemorySkill.where(user_id: user_id, archived_at: nil).map { |r| decode_data(r.data) }
                 .select { |s| (s[:name] || s["name"]).to_s == name.to_s }
             else
@@ -77,7 +90,11 @@ module Llmemory
             data[key] = (data[key] || 0).to_i + 1
             data[:updated_at] = Time.now.utc.iso8601
             rec.data = cipher.enabled? ? enc_json(data) : data
-            rec.search_text = enc(searchable_text(data))
+            text = searchable_text(data)
+            rec.search_text = enc(text)
+            rec.search_tokens = search_tokens_for(text) if LlmemorySkill.column_names.include?("search_tokens")
+            name = (data[:name] || data["name"]).to_s
+            rec.name_det = enc_det(name) if LlmemorySkill.column_names.include?("name_det")
             rec.save!
             data
           end
@@ -105,11 +122,25 @@ module Llmemory
 
           private
 
-          def token_scope(scope, column, query)
+          def token_scope(scope, column, query, model:)
             tokens = Llmemory::Tokenizer.tokenize(query)
-            return scope if tokens.empty?
-            clause = tokens.map { "LOWER(#{column}) LIKE LOWER(?)" }.join(" OR ")
-            scope.where(clause, *tokens.map { |t| "%#{t}%" })
+            return scope.to_a if tokens.empty?
+
+            if cipher.enabled? && model.column_names.include?("search_tokens")
+              digests = tokens.map { |t| cipher.blind_index(t) }
+              clause = digests.map { "search_tokens LIKE ?" }.join(" OR ")
+              indexed = scope.where(clause, *digests.map { |d| "% #{d} %" })
+              legacy_scope = scope.where(search_tokens: nil)
+              indexed_records = indexed.to_a
+              legacy_records = legacy_scope.to_a.select do |record|
+                plaintext = dec(record.public_send(column))
+                Llmemory::Tokenizer.matches?(plaintext, query)
+              end
+              (indexed_records + legacy_records).uniq { |r| r.id }
+            else
+              clause = tokens.map { "LOWER(#{column}) LIKE LOWER(?)" }.join(" OR ")
+              scope.where(clause, *tokens.map { |t| "%#{t}%" }).to_a
+            end
           end
 
           def stringify(hash)

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

@@ -27,11 +27,13 @@ module Llmemory
             id = skill[:id] || skill["id"] || "skill_#{SecureRandom.hex(8)}"
             data = symbolize(skill).merge(id: id, user_id: user_id)
             data[:created_at] ||= Time.now.utc.iso8601
+            search = searchable_text(data)
+            name = (data[:name] || data["name"]).to_s
             conn.exec_params(
-              "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, store_data(data), enc(searchable_text(data)), created_at_value(data)]
+              "INSERT INTO llmemory_skills (id, user_id, data, search_text, search_tokens, name_det, created_at) " \
+              "VALUES ($1, $2, $3::jsonb, $4, $5, $6, $7) " \
+              "ON CONFLICT (id) DO UPDATE SET data = $3::jsonb, search_text = $4, search_tokens = $5, name_det = $6",
+              [id, user_id, store_data(data), enc(search), search_tokens_for(search), enc_det(name), created_at_value(data)]
             )
             id
           end
@@ -52,19 +54,56 @@ module Llmemory
 
           def search_skills(user_id, query)
             ensure_tables!
-            suffix, params = token_filter("search_text", query, 2)
-            conn.exec_params(
+            tokens = Llmemory::Tokenizer.tokenize(query)
+            return conn.exec_params(
+              "SELECT data FROM llmemory_skills WHERE user_id = $1 AND archived_at IS NULL ORDER BY created_at DESC",
+              [user_id]
+            ).map { |r| parse_data(r["data"]) } if tokens.empty?
+
+            suffix, params = blind_token_filter("search_text", query, 2, search_tokens_column: cipher.enabled? ? "search_tokens" : nil)
+            rows = conn.exec_params(
               "SELECT data FROM llmemory_skills WHERE user_id = $1 AND archived_at IS NULL#{suffix} ORDER BY created_at DESC",
               [user_id, *params]
-            ).map { |r| parse_data(r["data"]) }
+            )
+            skills = rows.map { |r| parse_data(r["data"]) }
+            return skills unless cipher.enabled?
+
+            legacy_rows = conn.exec_params(
+              "SELECT data FROM llmemory_skills WHERE user_id = $1 AND archived_at IS NULL AND search_tokens IS NULL",
+              [user_id]
+            )
+            legacy = legacy_rows.map { |r| parse_data(r["data"]) }.select do |skill|
+              Llmemory::Tokenizer.matches?(searchable_text(skill), query)
+            end
+            by_id = {}
+            (skills + legacy).each { |skill| by_id[skill[:id] || skill["id"]] = skill }
+            by_id.values
           end
 
           def find_skills_by_name(user_id, name)
             ensure_tables!
-            conn.exec_params(
-              "SELECT data FROM llmemory_skills WHERE user_id = $1 AND archived_at IS NULL AND data->>'name' = $2",
-              [user_id, name.to_s]
-            ).map { |r| parse_data(r["data"]) }
+            if cipher.enabled?
+              rows = conn.exec_params(
+                "SELECT data FROM llmemory_skills WHERE user_id = $1 AND archived_at IS NULL AND name_det = $2",
+                [user_id, enc_det(name.to_s)]
+              )
+              indexed = rows.map { |r| parse_data(r["data"]) }
+              legacy_rows = conn.exec_params(
+                "SELECT data FROM llmemory_skills WHERE user_id = $1 AND archived_at IS NULL AND name_det IS NULL",
+                [user_id]
+              )
+              legacy = legacy_rows.map { |r| parse_data(r["data"]) }.select do |skill|
+                (skill[:name] || skill["name"]).to_s == name.to_s
+              end
+              by_id = {}
+              (indexed + legacy).each { |skill| by_id[skill[:id] || skill["id"]] = skill }
+              by_id.values
+            else
+              conn.exec_params(
+                "SELECT data FROM llmemory_skills WHERE user_id = $1 AND archived_at IS NULL AND data->>'name' = $2",
+                [user_id, name.to_s]
+              ).map { |r| parse_data(r["data"]) }
+            end
           end
 
           def record_outcome(user_id, skill_id, success:)
@@ -74,9 +113,11 @@ module Llmemory
             key = success ? :success_count : :failure_count
             data[key] = (data[key] || 0).to_i + 1
             data[:updated_at] = Time.now.utc.iso8601
+            search = searchable_text(data)
+            name = (data[:name] || data["name"]).to_s
             conn.exec_params(
-              "UPDATE llmemory_skills SET data = $3::jsonb, search_text = $4 WHERE user_id = $1 AND id = $2",
-              [user_id, skill_id, store_data(data), enc(searchable_text(data))]
+              "UPDATE llmemory_skills SET data = $3::jsonb, search_text = $4, search_tokens = $5, name_det = $6 WHERE user_id = $1 AND id = $2",
+              [user_id, skill_id, store_data(data), enc(search), search_tokens_for(search), enc_det(name)]
             )
             data
           end
@@ -137,14 +178,13 @@ module Llmemory
               );
               CREATE INDEX IF NOT EXISTS idx_llmemory_skills_user_id ON llmemory_skills(user_id);
               ALTER TABLE llmemory_skills ADD COLUMN IF NOT EXISTS archived_at TIMESTAMPTZ;
+              ALTER TABLE llmemory_skills ADD COLUMN IF NOT EXISTS search_tokens TEXT;
+              ALTER TABLE llmemory_skills ADD COLUMN IF NOT EXISTS name_det TEXT;
             SQL
           end
 
           def token_filter(column, query, start_index)
-            tokens = Llmemory::Tokenizer.tokenize(query)
-            return ["", []] if tokens.empty?
-            likes = tokens.each_index.map { |i| "LOWER(#{column}) LIKE $#{start_index + i}" }
-            [" AND (#{likes.join(' OR ')})", tokens.map { |t| "%#{t}%" }]
+            blind_token_filter(column, query, start_index, search_tokens_column: cipher.enabled? ? "search_tokens" : nil)
           end
 
           def parse_data(value)

data/lib/llmemory/maintenance/search_tokens_backfill.rb

@@ -0,0 +1,339 @@
+# frozen_string_literal: true
+
+require "json"
+require_relative "../crypto/field_helpers"
+
+module Llmemory
+  module Maintenance
+    # Populates search_tokens (blind index) and procedural name_det for rows
+    # written before encrypted keyword search support. In-place UPDATE only.
+    class SearchTokensBackfill
+      include Crypto::FieldHelpers
+
+      Result = Struct.new(
+        :items, :resources, :episodes, :skills, :skipped, :dry_run,
+        keyword_init: true
+      ) do
+        def total
+          items + resources + episodes + skills
+        end
+
+        def to_h
+          { items: items, resources: resources, episodes: episodes, skills: skills,
+            skipped: skipped, dry_run: dry_run, total: total }
+        end
+      end
+
+      def initialize(cipher: nil, store: nil, dry_run: false, force: false)
+        @cipher = cipher || Llmemory.build_cipher
+        @store = (store || Llmemory.configuration.long_term_store).to_s.to_sym
+        @dry_run = dry_run
+        @force = force
+      end
+
+      def run(user_id: nil)
+        case @store
+        when :active_record, :activerecord
+          run_active_record(user_id: user_id)
+        when :postgres
+          run_postgres(user_id: user_id)
+        else
+          raise ConfigurationError,
+                "backfill_search_tokens requires long_term_store :active_record or :postgres (got #{@store.inspect})"
+        end
+      end
+
+      private
+
+      def run_active_record(user_id: nil)
+        require "active_record"
+        load_active_record_models!
+
+        result = empty_result
+        result.items = backfill_ar_items(user_id)
+        result.resources = backfill_ar_resources(user_id)
+        result.episodes = backfill_ar_episodes(user_id)
+        result.skills = backfill_ar_skills(user_id)
+        result
+      end
+
+      def run_postgres(user_id: nil)
+        result = empty_result
+        result.items = backfill_pg_items(user_id)
+        result.resources = backfill_pg_resources(user_id)
+        result.episodes = backfill_pg_episodes(user_id)
+        result.skills = backfill_pg_skills(user_id)
+        result
+      end
+
+      def empty_result
+        Result.new(items: 0, resources: 0, episodes: 0, skills: 0, skipped: 0, dry_run: @dry_run)
+      end
+
+      def load_active_record_models!
+        LongTerm::FileBased::Storages::ActiveRecordStorage.load_models!
+        LongTerm::Episodic::Storages::ActiveRecordStorage.load_models!
+        LongTerm::Procedural::Storages::ActiveRecordStorage.load_models!
+      end
+
+      # --- ActiveRecord ---
+
+      def backfill_ar_items(user_id)
+        model = LongTerm::FileBased::Storages::LlmemoryItem
+        return 0 unless model.column_names.include?("search_tokens")
+
+        count = 0
+        ar_scope(model, user_id).find_each do |rec|
+          tokens = search_tokens_for(dec(rec.content))
+          next if tokens.nil?
+
+          rec.update_column(:search_tokens, tokens) unless @dry_run
+          count += 1
+        end
+        count
+      end
+
+      def backfill_ar_resources(user_id)
+        model = LongTerm::FileBased::Storages::LlmemoryResource
+        return 0 unless model.column_names.include?("search_tokens")
+
+        count = 0
+        ar_scope(model, user_id).find_each do |rec|
+          tokens = search_tokens_for(dec(rec.text))
+          next if tokens.nil?
+
+          rec.update_column(:search_tokens, tokens) unless @dry_run
+          count += 1
+        end
+        count
+      end
+
+      def backfill_ar_episodes(user_id)
+        model = LongTerm::Episodic::Storages::LlmemoryEpisode
+        return 0 unless model.column_names.include?("search_tokens")
+
+        count = 0
+        ar_scope(model, user_id).find_each do |rec|
+          plain = dec(rec.search_text.to_s)
+          tokens = search_tokens_for(plain)
+          next if tokens.nil?
+
+          rec.update_column(:search_tokens, tokens) unless @dry_run
+          count += 1
+        end
+        count
+      end
+
+      def backfill_ar_skills(user_id)
+        model = LongTerm::Procedural::Storages::LlmemorySkill
+        has_tokens = model.column_names.include?("search_tokens")
+        has_name_det = model.column_names.include?("name_det")
+        return 0 unless has_tokens || has_name_det
+
+        count = 0
+        ar_scope_skills(model, user_id, has_tokens: has_tokens, has_name_det: has_name_det).find_each do |rec|
+          updates = {}
+          if has_tokens
+            plain = dec(rec.search_text.to_s)
+            tokens = search_tokens_for(plain)
+            updates[:search_tokens] = tokens if tokens
+          end
+          if has_name_det
+            data = decode_json_document(rec.data)
+            name = (data[:name] || data["name"]).to_s
+            updates[:name_det] = enc_det(name) unless name.empty?
+          end
+          next if updates.empty?
+
+          if @dry_run
+            count += 1
+          else
+            rec.update_columns(updates)
+            count += 1
+          end
+        end
+        count
+      end
+
+      def ar_scope(model, user_id)
+        scope = model.all
+        scope = scope.where(user_id: user_id) if user_id
+        scope = scope.where(search_tokens: nil) unless @force
+        scope
+      end
+
+      def ar_scope_skills(model, user_id, has_tokens:, has_name_det:)
+        scope = model.all
+        scope = scope.where(user_id: user_id) if user_id
+        return scope if @force
+
+        if has_tokens && has_name_det
+          scope.where("search_tokens IS NULL OR name_det IS NULL")
+        elsif has_tokens
+          scope.where(search_tokens: nil)
+        elsif has_name_det
+          scope.where(name_det: nil)
+        else
+          scope
+        end
+      end
+
+      # --- Postgres ---
+
+      def backfill_pg_items(user_id)
+        return 0 unless pg_column_exists?("llmemory_items", "search_tokens")
+
+        count = 0
+        pg_each_row("llmemory_items", "id, user_id, content", user_id) do |row|
+          tokens = search_tokens_for(dec(row["content"]))
+          next if tokens.nil?
+
+          pg_update("llmemory_items", row["id"], search_tokens: tokens) unless @dry_run
+          count += 1
+        end
+        count
+      end
+
+      def backfill_pg_resources(user_id)
+        return 0 unless pg_column_exists?("llmemory_resources", "search_tokens")
+
+        count = 0
+        pg_each_row("llmemory_resources", "id, user_id, text", user_id) do |row|
+          tokens = search_tokens_for(dec(row["text"]))
+          next if tokens.nil?
+
+          pg_update("llmemory_resources", row["id"], search_tokens: tokens) unless @dry_run
+          count += 1
+        end
+        count
+      end
+
+      def backfill_pg_episodes(user_id)
+        return 0 unless pg_column_exists?("llmemory_episodes", "search_tokens")
+
+        count = 0
+        pg_each_row("llmemory_episodes", "id, user_id, search_text", user_id) do |row|
+          tokens = search_tokens_for(dec(row["search_text"].to_s))
+          next if tokens.nil?
+
+          pg_update("llmemory_episodes", row["id"], search_tokens: tokens) unless @dry_run
+          count += 1
+        end
+        count
+      end
+
+      def backfill_pg_skills(user_id)
+        has_tokens = pg_column_exists?("llmemory_skills", "search_tokens")
+        has_name_det = pg_column_exists?("llmemory_skills", "name_det")
+        return 0 unless has_tokens || has_name_det
+
+        count = 0
+        sql = "SELECT id, user_id, data, search_text FROM llmemory_skills"
+        params = []
+        clauses = []
+        if user_id
+          params << user_id
+          clauses << "user_id = $#{params.size}"
+        end
+        unless @force
+          parts = []
+          parts << "search_tokens IS NULL" if has_tokens
+          parts << "name_det IS NULL" if has_name_det
+          clauses << "(#{parts.join(' OR ')})" if parts.any?
+        end
+        sql += " WHERE #{clauses.join(' AND ')}" unless clauses.empty?
+
+        pg_conn.exec_params(sql, params).each do |row|
+          sets = {}
+          if has_tokens
+            tokens = search_tokens_for(dec(row["search_text"].to_s))
+            sets["search_tokens"] = tokens if tokens
+          end
+          if has_name_det
+            data = decode_json_document(row["data"])
+            name = (data[:name] || data["name"]).to_s
+            sets["name_det"] = enc_det(name) unless name.empty?
+          end
+          next if sets.empty?
+
+          pg_update_columns("llmemory_skills", row["id"], sets) unless @dry_run
+          count += 1
+        end
+        count
+      end
+
+      def pg_each_row(table, columns, user_id)
+        sql = "SELECT #{columns} FROM #{table}"
+        params = []
+        clauses = []
+        if user_id
+          params << user_id
+          clauses << "user_id = $#{params.size}"
+        end
+        unless @force
+          clauses << "search_tokens IS NULL" if pg_column_exists?(table, "search_tokens")
+        end
+        sql += " WHERE #{clauses.join(' AND ')}" unless clauses.empty?
+
+        pg_conn.exec_params(sql, params).each { |row| yield row }
+      end
+
+      def pg_update(table, id, search_tokens:)
+        pg_conn.exec_params(
+          "UPDATE #{table} SET search_tokens = $1 WHERE id = $2",
+          [search_tokens, id]
+        )
+      end
+
+      def pg_update_columns(table, id, sets)
+        return if sets.empty?
+
+        idx = 1
+        params = []
+        assignments = sets.each_key.map do |col|
+          params << sets[col]
+          clause = "#{col} = $#{idx}"
+          idx += 1
+          clause
+        end
+        params << id
+        pg_conn.exec_params(
+          "UPDATE #{table} SET #{assignments.join(', ')} WHERE id = $#{idx}",
+          params
+        )
+      end
+
+      def pg_column_exists?(table, column)
+        @pg_columns ||= {}
+        cache_key = "#{table}.#{column}"
+        return @pg_columns[cache_key] if @pg_columns.key?(cache_key)
+
+        rows = pg_conn.exec_params(
+          "SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = $2 LIMIT 1",
+          [table, column]
+        )
+        @pg_columns[cache_key] = rows.any?
+      end
+
+      def pg_conn
+        @pg_conn ||= begin
+          require "pg"
+          PG.connect(Llmemory.configuration.database_url)
+        end
+      end
+
+      def decode_json_document(raw)
+        return raw.transform_keys(&:to_sym) if raw.is_a?(Hash)
+
+        str = raw.to_s
+        if cipher.encrypted?(str)
+          dec_json(str)
+        else
+          JSON.parse(str, symbolize_names: true)
+        end
+      rescue JSON::ParserError
+        {}
+      end
+    end
+  end
+end

data/lib/llmemory/maintenance.rb

@@ -3,6 +3,7 @@
 require_relative "maintenance/runner"
 require_relative "maintenance/ttl_expiry"
 require_relative "maintenance/cognitive_pass"
+require_relative "maintenance/search_tokens_backfill"
 
 module Llmemory
   module Maintenance

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

@@ -55,6 +55,8 @@ module Llmemory
           private
 
           def fetch_timeline_context(user_id, query, window)
+            return "" if Llmemory.configuration.long_term_type.to_s == "graph_based"
+
             storage = build_storage
             items = storage.search_items(user_id, query)
             return "" if items.empty?

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

@@ -23,19 +23,10 @@ module Llmemory
             timeline = []
 
             # Get recent items from long-term memory
-            storage = build_long_term_storage
-            begin
-              items = storage.get_items_since(user_id, hours: hours)
-              items.each do |item|
-                timeline << {
-                  type: "fact",
-                  timestamp: item[:created_at] || item["created_at"],
-                  category: item[:category] || item["category"],
-                  content: item[:content] || item["content"]
-                }
-              end
-            rescue NotImplementedError
-              # Some storages may not implement get_items_since
+            if graph_based?
+              append_graph_timeline(timeline, user_id, hours)
+            else
+              append_file_timeline(timeline, user_id, hours)
             end
 
             # Get recent messages from short-term
@@ -89,7 +80,48 @@ module Llmemory
           end
 
           def build_long_term_storage
-            LongTerm::FileBased::Storages.build
+            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 append_file_timeline(timeline, user_id, hours)
+            storage = build_long_term_storage
+            items = storage.get_items_since(user_id, hours: hours)
+            items.each do |item|
+              timeline << {
+                type: "fact",
+                timestamp: item[:created_at] || item["created_at"],
+                category: item[:category] || item["category"],
+                content: item[:content] || item["content"]
+              }
+            end
+          rescue NotImplementedError
+            # Some storages may not implement get_items_since
+          end
+
+          def append_graph_timeline(timeline, user_id, hours)
+            storage = build_long_term_storage
+            cutoff = Time.now - (hours * 3600)
+            storage.list_edges(user_id).each do |edge|
+              created_at = edge.respond_to?(:created_at) ? edge.created_at : edge[:created_at]
+              next unless created_at && created_at >= cutoff
+
+              subject = edge.respond_to?(:subject_id) ? edge.subject_id : edge[:subject_id]
+              predicate = edge.respond_to?(:predicate) ? edge.predicate : edge[:predicate]
+              target = edge.respond_to?(:target_id) ? edge.target_id : edge[:target_id]
+              timeline << {
+                type: "relation",
+                timestamp: created_at,
+                content: "#{subject} -> #{predicate} -> #{target}"
+              }
+            end
           end
 
           def format_timeline(timeline, hours)
@@ -101,6 +133,8 @@ module Llmemory
               when "fact"
                 cat_info = entry[:category] ? "[#{entry[:category]}]" : ""
                 output << "- [FACT] #{cat_info} #{truncate(entry[:content], 150)}"
+              when "relation"
+                output << "- [REL] #{truncate(entry[:content], 150)}"
               when "message"
                 output << "- [MSG #{entry[:session_id]}] [#{entry[:role]}] #{truncate(entry[:content], 150)}"
               end

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