llmemory 0.1.1

data/lib/llmemory/long_term/graph_based/knowledge_graph.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+require_relative "node"
+require_relative "edge"
+require_relative "storages/base"
+require_relative "storages/memory_storage"
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      class KnowledgeGraph
+        def initialize(user_id:, storage: nil)
+          @user_id = user_id
+          @storage = storage || Storages::MemoryStorage.new
+        end
+
+        def add_node(entity_type:, name:, properties: {})
+          existing = @storage.find_node_by_name(@user_id, entity_type, name)
+          return existing.id if existing
+          node = Node.new(
+            id: nil,
+            user_id: @user_id,
+            entity_type: entity_type.to_s,
+            name: name.to_s,
+            properties: properties,
+            created_at: Time.now,
+            updated_at: Time.now
+          )
+          @storage.save_node(@user_id, node)
+        end
+
+        def find_node(name: nil, id: nil)
+          if id
+            @storage.find_node_by_id(@user_id, id)
+          elsif name
+            list_nodes.find { |n| n.name.to_s == name.to_s }
+          end
+        end
+
+        def find_node_by_id(id)
+          @storage.find_node_by_id(@user_id, id)
+        end
+
+        def add_edge(subject:, predicate:, object:, properties: {})
+          subject_id = subject.is_a?(Node) ? subject.id : subject.to_s
+          object_id = object.is_a?(Node) ? object.id : object.to_s
+          edge = Edge.new(
+            id: nil,
+            user_id: @user_id,
+            subject_id: subject_id,
+            predicate: predicate.to_s,
+            object_id: object_id,
+            properties: properties,
+            created_at: Time.now,
+            archived_at: nil
+          )
+          @storage.save_edge(@user_id, edge)
+        end
+
+        def find_edges(subject: nil, predicate: nil, object: nil, include_archived: false)
+          subject_id = subject.is_a?(Node) ? subject.id : subject
+          object_id = object.is_a?(Node) ? object.id : object
+          @storage.find_edges(
+            @user_id,
+            subject_id: subject_id,
+            predicate: predicate&.to_s,
+            object_id: object_id,
+            include_archived: include_archived
+          )
+        end
+
+        def traverse(start_node:, depth: 2)
+          start_id = start_node.is_a?(Node) ? start_node.id : start_node.to_s
+          visited = {}
+          queue = [[start_id, 0]]
+          result_nodes = []
+          result_edges = []
+
+          while queue.any?
+            node_id, d = queue.shift
+            next if d > depth
+            next if visited[node_id]
+            visited[node_id] = true
+            node = @storage.find_node_by_id(@user_id, node_id)
+            result_nodes << node if node
+
+            edges = @storage.find_edges(@user_id, subject_id: node_id, include_archived: false)
+            edges.each do |e|
+              result_edges << e
+              queue << [e.object_id, d + 1] unless visited[e.object_id]
+            end
+            edges_out = @storage.find_edges(@user_id, object_id: node_id, include_archived: false)
+            edges_out.each do |e|
+              result_edges << e
+              queue << [e.subject_id, d + 1] unless visited[e.subject_id]
+            end
+          end
+
+          { nodes: result_nodes.uniq, edges: result_edges.uniq }
+        end
+
+        def archive_edge(edge_id, reason: nil)
+          @storage.archive_edge(@user_id, edge_id, archived_at: Time.now)
+        end
+
+        def list_nodes
+          @storage.list_nodes(@user_id)
+        end
+
+        attr_reader :storage, :user_id
+      end
+    end
+  end
+end

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

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+require_relative "node"
+require_relative "edge"
+require_relative "knowledge_graph"
+require_relative "conflict_resolver"
+require_relative "storage"
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      class Memory
+        def initialize(user_id:, storage: nil, vector_store: nil, llm: nil, extractor: nil)
+          @user_id = user_id
+          @graph_storage = storage || Storages.build
+          @kg = KnowledgeGraph.new(user_id: user_id, storage: @graph_storage)
+          @conflict_resolver = ConflictResolver.new(@kg)
+          @vector_store = vector_store || build_vector_store
+          @llm = llm || Llmemory::LLM.client
+          @extractor = extractor || Llmemory::Extractors::EntityRelationExtractor.new(llm: @llm)
+        end
+
+        def memorize(conversation_text)
+          data = @extractor.extract(conversation_text)
+          name_to_id = {}
+
+          data[:entities].each do |e|
+            id = @kg.add_node(entity_type: e[:type], name: e[:name], properties: {})
+            name_to_id[e[:name]] ||= id
+          end
+
+          data[:relations].each do |r|
+            subject_id = name_to_id[r[:subject]] || @kg.add_node(entity_type: "concept", name: r[:subject], properties: {})
+            object_id = name_to_id[r[:object]] || @kg.add_node(entity_type: "concept", name: r[:object], properties: {})
+
+            edge = Edge.new(
+              id: nil,
+              user_id: @user_id,
+              subject_id: subject_id,
+              predicate: r[:predicate],
+              object_id: object_id,
+              properties: {},
+              created_at: Time.now,
+              archived_at: nil
+            )
+            @conflict_resolver.resolve(edge)
+            edge_id = @kg.add_edge(subject: subject_id, predicate: r[:predicate], object: object_id, properties: {})
+
+            text = "#{r[:subject]} #{r[:predicate]} #{r[:object]}"
+            embedding = @vector_store.respond_to?(:embed) ? @vector_store.embed(text) : nil
+            if embedding && @vector_store.respond_to?(:store)
+              @vector_store.store(id: "edge_#{edge_id}", embedding: embedding, metadata: { text: text, created_at: Time.now }, user_id: @user_id)
+            end
+          end
+
+          true
+        end
+
+        def retrieve(query, top_k: 10)
+          results = hybrid_search(query, top_k: top_k)
+          format_as_context(results)
+        end
+
+        def search_candidates(query, user_id: nil, top_k: 20)
+          uid = user_id || @user_id
+          return [] unless uid == @user_id
+          results = hybrid_search(query, top_k: top_k)
+          results.map do |r|
+            {
+              text: r[:text],
+              timestamp: r[:created_at] || r[:timestamp],
+              score: r[:score] || 1.0
+            }
+          end
+        end
+
+        attr_reader :user_id
+
+        def storage
+          @graph_storage
+        end
+
+        private
+
+        def build_vector_store
+          emb = Llmemory::VectorStore::OpenAIEmbeddings.new
+          Llmemory::VectorStore::MemoryStore.new(embedding_provider: emb)
+        end
+
+        def hybrid_search(query, top_k:)
+          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)
+          elsif @vector_store.respond_to?(:embed) && @vector_store.respond_to?(:search)
+            emb = @vector_store.embed(query.to_s)
+            vector_results = @vector_store.search(emb, top_k: top_k, user_id: @user_id)
+          end
+
+          out = vector_results.map do |v|
+            id = v[:id] || v["id"]
+            meta = v[:metadata] || v["metadata"] || {}
+            { text: meta["text"] || meta[:text] || id.to_s, score: v[:score] || v["score"] || 1.0, created_at: meta["created_at"] || meta[:created_at] }
+          end
+
+          node_ids = out.flat_map { |r| extract_node_ids_from_text(r[:text]) }.compact.uniq
+          node_ids.first(3).each do |node_id|
+            node = @kg.find_node_by_id(node_id)
+            next unless node
+            traversed = @kg.traverse(start_node: node, depth: 1)
+            traversed[:edges].each do |e|
+              subj = @kg.find_node_by_id(e.subject_id)
+              obj = @kg.find_node_by_id(e.object_id)
+              edge_text = "#{subj&.name} #{e.predicate} #{obj&.name}"
+              out << { text: edge_text, score: 0.85, created_at: e.created_at } unless out.any? { |o| o[:text] == edge_text }
+            end
+          end
+
+          out.sort_by { |r| -(r[:score] || 0) }.first(top_k)
+        end
+
+        def extract_node_ids_from_text(text)
+          return [] if text.to_s.empty?
+          ids = []
+          @kg.list_nodes.each do |n|
+            ids << n.id if text.to_s.include?(n.name.to_s)
+          end
+          ids
+        end
+
+        def format_as_context(results)
+          return "" if results.empty?
+          lines = ["=== RELEVANT MEMORIES (GRAPH) ===", ""]
+          results.each do |r|
+            lines << "- #{r[:text]}"
+          end
+          lines << ""
+          lines << "=== END MEMORIES ==="
+          lines.join("\n")
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/long_term/graph_based/node.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      Node = Struct.new(
+        :id,
+        :user_id,
+        :entity_type,
+        :name,
+        :properties,
+        :created_at,
+        :updated_at,
+        keyword_init: true
+      ) do
+        def self.from_h(hash)
+          new(
+            id: hash[:id] || hash["id"],
+            user_id: hash[:user_id] || hash["user_id"],
+            entity_type: (hash[:entity_type] || hash["entity_type"]).to_s,
+            name: (hash[:name] || hash["name"]).to_s,
+            properties: hash[:properties] || hash["properties"] || {},
+            created_at: hash[:created_at] || hash["created_at"],
+            updated_at: hash[:updated_at] || hash["updated_at"]
+          )
+        end
+
+        def to_h
+          {
+            id: id,
+            user_id: user_id,
+            entity_type: entity_type,
+            name: name,
+            properties: properties || {},
+            created_at: created_at,
+            updated_at: updated_at
+          }
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require_relative "storages/base"
+require_relative "storages/memory_storage"
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      module Storages
+        def self.build(store: nil)
+          case (store || :memory).to_s.to_sym
+          when :memory
+            MemoryStorage.new
+          when :active_record, :activerecord
+            require_relative "storages/active_record_storage"
+            ActiveRecordStorage.new
+          else
+            MemoryStorage.new
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/long_term/graph_based/storages/active_record_models.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      module Storages
+        class LlmemoryGraphNode < ::ActiveRecord::Base
+          self.table_name = "llmemory_graph_nodes"
+          self.primary_key = "id"
+          has_many :out_edges, class_name: "Llmemory::LongTerm::GraphBased::Storages::LlmemoryGraphEdge", foreign_key: :subject_id
+          has_many :in_edges, class_name: "Llmemory::LongTerm::GraphBased::Storages::LlmemoryGraphEdge", foreign_key: :object_id
+        end
+
+        class LlmemoryGraphEdge < ::ActiveRecord::Base
+          self.table_name = "llmemory_graph_edges"
+          self.primary_key = "id"
+          belongs_to :subject_node, class_name: "Llmemory::LongTerm::GraphBased::Storages::LlmemoryGraphNode", foreign_key: :subject_id, optional: true
+          belongs_to :object_node, class_name: "Llmemory::LongTerm::GraphBased::Storages::LlmemoryGraphNode", foreign_key: :object_id, optional: true
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require_relative "base"
+require_relative "active_record_models"
+require_relative "../node"
+require_relative "../edge"
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      module Storages
+        class ActiveRecordStorage < Base
+          def initialize
+            self.class.load_models!
+          end
+
+          def self.load_models!
+            return if @models_loaded
+            require "active_record"
+            require_relative "active_record_models"
+            @models_loaded = true
+          end
+
+          def save_node(user_id, node)
+            n = node.is_a?(Node) ? node : Node.from_h(node.to_h)
+            rec = if n.id
+              LlmemoryGraphNode.find_by(user_id: user_id, id: n.id)
+            else
+              LlmemoryGraphNode.find_by(user_id: user_id, entity_type: n.entity_type, name: n.name)
+            end
+            if rec
+              rec.update!(properties: n.properties || {}, updated_at: Time.current)
+              rec.id
+            else
+              id = n.id || "node_#{SecureRandom.hex(8)}"
+              LlmemoryGraphNode.create!(
+                id: id,
+                user_id: user_id,
+                entity_type: n.entity_type.to_s,
+                name: n.name.to_s,
+                properties: n.properties || {}
+              )
+              id
+            end
+          end
+
+          def find_node_by_id(user_id, id)
+            rec = LlmemoryGraphNode.find_by(user_id: user_id, id: id)
+            record_to_node(rec) if rec
+          end
+
+          def find_node_by_name(user_id, entity_type, name)
+            rec = LlmemoryGraphNode.find_by(user_id: user_id, entity_type: entity_type.to_s, name: name.to_s)
+            record_to_node(rec) if rec
+          end
+
+          def list_nodes(user_id)
+            LlmemoryGraphNode.where(user_id: user_id).map { |r| record_to_node(r) }
+          end
+
+          def save_edge(user_id, edge)
+            e = edge.is_a?(Edge) ? edge : Edge.from_h(edge.to_h)
+            id = e.id || "edge_#{SecureRandom.hex(8)}"
+            rec = LlmemoryGraphEdge.find_by(user_id: user_id, id: id)
+            if rec
+              rec.update!(
+                subject_id: e.subject_id,
+                predicate: e.predicate,
+                object_id: e.object_id,
+                properties: e.properties || {}
+              )
+            else
+              LlmemoryGraphEdge.create!(
+                id: id,
+                user_id: user_id,
+                subject_id: e.subject_id,
+                predicate: e.predicate,
+                object_id: e.object_id,
+                properties: e.properties || {}
+              )
+            end
+            id
+          end
+
+          def find_edges(user_id, subject_id: nil, predicate: nil, object_id: nil, include_archived: false)
+            scope = LlmemoryGraphEdge.where(user_id: user_id)
+            scope = scope.where(archived_at: nil) unless include_archived
+            scope = scope.where(subject_id: subject_id) if subject_id
+            scope = scope.where(predicate: predicate) if predicate
+            scope = scope.where(object_id: object_id) if object_id
+            scope.map { |r| record_to_edge(r) }
+          end
+
+          def archive_edge(user_id, edge_id, archived_at: nil)
+            rec = LlmemoryGraphEdge.find_by(user_id: user_id, id: edge_id)
+            return false unless rec
+            rec.update!(archived_at: archived_at || Time.current)
+            true
+          end
+
+          private
+
+          def record_to_node(r)
+            Node.new(
+              id: r.id,
+              user_id: r.user_id,
+              entity_type: r.entity_type,
+              name: r.name,
+              properties: r.properties || {},
+              created_at: r.created_at,
+              updated_at: r.updated_at
+            )
+          end
+
+          def record_to_edge(r)
+            Edge.new(
+              id: r.id,
+              user_id: r.user_id,
+              subject_id: r.subject_id,
+              predicate: r.predicate,
+              object_id: r.object_id,
+              properties: r.properties || {},
+              created_at: r.created_at,
+              archived_at: r.archived_at
+            )
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/long_term/graph_based/storages/base.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      module Storages
+        class Base
+          def save_node(user_id, node)
+            raise NotImplementedError, "#{self.class}#save_node must be implemented"
+          end
+
+          def find_node_by_id(user_id, id)
+            raise NotImplementedError, "#{self.class}#find_node_by_id must be implemented"
+          end
+
+          def find_node_by_name(user_id, entity_type, name)
+            raise NotImplementedError, "#{self.class}#find_node_by_name must be implemented"
+          end
+
+          def list_nodes(user_id)
+            raise NotImplementedError, "#{self.class}#list_nodes must be implemented"
+          end
+
+          def save_edge(user_id, edge)
+            raise NotImplementedError, "#{self.class}#save_edge must be implemented"
+          end
+
+          def find_edges(user_id, subject_id: nil, predicate: nil, object_id: nil, include_archived: false)
+            raise NotImplementedError, "#{self.class}#find_edges must be implemented"
+          end
+
+          def archive_edge(user_id, edge_id, archived_at: nil)
+            raise NotImplementedError, "#{self.class}#archive_edge must be implemented"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/long_term/graph_based/storages/memory_storage.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require_relative "base"
+require_relative "../node"
+require_relative "../edge"
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+      module Storages
+        class MemoryStorage < Base
+          def initialize
+            @nodes = Hash.new { |h, k| h[k] = {} }
+            @edges = Hash.new { |h, k| h[k] = [] }
+            @node_id_seq = 0
+            @edge_id_seq = 0
+          end
+
+          def save_node(user_id, node)
+            n = node.is_a?(Node) ? node : Node.from_h(node.to_h)
+            unless n.id
+              @node_id_seq += 1
+              n = Node.new(
+                id: "node_#{@node_id_seq}",
+                user_id: user_id,
+                entity_type: n.entity_type,
+                name: n.name,
+                properties: n.properties || {},
+                created_at: n.created_at || Time.now,
+                updated_at: Time.now
+              )
+            end
+            @nodes[user_id][n.id] = n
+            n.id
+          end
+
+          def find_node_by_id(user_id, id)
+            @nodes[user_id][id]
+          end
+
+          def find_node_by_name(user_id, entity_type, name)
+            @nodes[user_id].values.find { |n| n.entity_type == entity_type.to_s && n.name.to_s == name.to_s }
+          end
+
+          def list_nodes(user_id)
+            @nodes[user_id].values.to_a
+          end
+
+          def save_edge(user_id, edge)
+            e = edge.is_a?(Edge) ? edge : Edge.from_h(edge.to_h)
+            unless e.id
+              @edge_id_seq += 1
+              e = Edge.new(
+                id: "edge_#{@edge_id_seq}",
+                user_id: user_id,
+                subject_id: e.subject_id,
+                predicate: e.predicate,
+                object_id: e.object_id,
+                properties: e.properties || {},
+                created_at: e.created_at || Time.now,
+                archived_at: nil
+              )
+            end
+            idx = @edges[user_id].find_index { |x| x.id == e.id }
+            if idx
+              @edges[user_id][idx] = e
+            else
+              @edges[user_id] << e
+            end
+            e.id
+          end
+
+          def find_edges(user_id, subject_id: nil, predicate: nil, object_id: nil, include_archived: false)
+            list = @edges[user_id].select do |e|
+              next false unless include_archived || !e.archived?
+              next false if subject_id && e.subject_id != subject_id
+              next false if predicate && e.predicate != predicate.to_s
+              next false if object_id && e.object_id != object_id
+              true
+            end
+            list
+          end
+
+          def archive_edge(user_id, edge_id, archived_at: nil)
+            t = archived_at || Time.now
+            e = @edges[user_id].find { |x| x.id == edge_id }
+            return false unless e
+            @edges[user_id].delete(e)
+            @edges[user_id] << Edge.new(
+              id: e.id,
+              user_id: e.user_id,
+              subject_id: e.subject_id,
+              predicate: e.predicate,
+              object_id: e.object_id,
+              properties: e.properties,
+              created_at: e.created_at,
+              archived_at: t
+            )
+            true
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llmemory/long_term/graph_based.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require_relative "graph_based/node"
+require_relative "graph_based/edge"
+require_relative "graph_based/storage"
+require_relative "graph_based/knowledge_graph"
+require_relative "graph_based/conflict_resolver"
+require_relative "graph_based/memory"
+
+module Llmemory
+  module LongTerm
+    module GraphBased
+    end
+  end
+end

data/lib/llmemory/long_term.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require_relative "long_term/file_based"
+require_relative "long_term/graph_based"
+
+module Llmemory
+  module LongTerm
+  end
+end