Class: Llmemory::LongTerm::GraphBased::Storages::MemoryStorage

Inherits:
Base
  • Object
show all
Defined in:
lib/llmemory/long_term/graph_based/storages/memory_storage.rb

Instance Method Summary collapse

Constructor Details

#initialize ⇒ MemoryStorage

Returns a new instance of MemoryStorage.



13
14
15
16
17
18
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 13

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

Instance Method Details

#archive_edge(user_id, edge_id, archived_at: nil) ⇒ Object 



87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 87

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,
    target_id: e.target_id,
    properties: e.properties,
    created_at: e.created_at,
    archived_at: t
  )
  true
end

#count_edges(user_id) ⇒ Object 



118
119
120
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 118

def count_edges(user_id)
  @edges[user_id].count { |e| !e.archived? }
end

#count_nodes(user_id) ⇒ Object 



114
115
116
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 114

def count_nodes(user_id)
  @nodes[user_id].size
end

#find_edges(user_id, subject_id: nil, predicate: nil, object_id: nil, include_archived: false) ⇒ Object 



76
77
78
79
80
81
82
83
84
85
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 76

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.target_id != object_id
    true
  end
  list
end

#find_node_by_id(user_id, id) ⇒ Object 



38
39
40
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 38

def find_node_by_id(user_id, id)
  @nodes[user_id][id]
end

#find_node_by_name(user_id, entity_type, name) ⇒ Object 



42
43
44
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 42

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

#get_edges_around(user_id, reference, before: 5, after: 5) ⇒ Object 



122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 122

def get_edges_around(user_id, reference, before: 5, after: 5)
  edges = @edges[user_id].reject(&:archived?).sort_by(&:created_at)
  return { before: [], target: nil, after: [] } if edges.empty?

  idx = if reference.is_a?(String) && reference.match?(/^\d{4}-/)
    target_time = Time.parse(reference)
    edges.index { |e| e.created_at >= target_time } || edges.size
  else
    edges.index { |e| e.id == reference }
  end

  return { before: [], target: nil, after: [] } unless idx

  start_idx = [idx - before, 0].max
  end_idx = [idx + after, edges.size - 1].min

  {
    before: edges[start_idx...idx] || [],
    target: edges[idx],
    after: edges[(idx + 1)..end_idx] || []
  }
end

#list_edges(user_id, subject_id: nil, predicate: nil, limit: nil) ⇒ Object 



109
110
111
112
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 109

def list_edges(user_id, subject_id: nil, predicate: nil, limit: nil)
  list = find_edges(user_id, subject_id: subject_id, predicate: predicate, object_id: nil, include_archived: false)
  limit ? list.take(limit) : list
end

#list_nodes(user_id, entity_type: nil, limit: nil) ⇒ Object 



46
47
48
49
50
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 46

def list_nodes(user_id, entity_type: nil, limit: nil)
  list = @nodes[user_id].values
  list = list.select { |n| n.entity_type.to_s == entity_type.to_s } if entity_type
  limit ? list.take(limit) : list
end

#list_users ⇒ Object 



105
106
107
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 105

def list_users
  (@nodes.keys + @edges.keys).uniq
end

#save_edge(user_id, edge) ⇒ Object 



52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 52

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,
      target_id: e.target_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

#save_node(user_id, node) ⇒ Object 



20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
 
# File 'lib/llmemory/long_term/graph_based/storages/memory_storage.rb', line 20

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