MemoryOS 0.1.13__py3-none-any.whl → 0.2.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of MemoryOS might be problematic. Click here for more details.

Files changed (84) hide show
  1. {memoryos-0.1.13.dist-info → memoryos-0.2.1.dist-info}/METADATA +78 -49
  2. memoryos-0.2.1.dist-info/RECORD +152 -0
  3. memoryos-0.2.1.dist-info/entry_points.txt +3 -0
  4. memos/__init__.py +1 -1
  5. memos/api/config.py +471 -0
  6. memos/api/exceptions.py +28 -0
  7. memos/api/mcp_serve.py +502 -0
  8. memos/api/product_api.py +35 -0
  9. memos/api/product_models.py +159 -0
  10. memos/api/routers/__init__.py +1 -0
  11. memos/api/routers/product_router.py +358 -0
  12. memos/chunkers/sentence_chunker.py +8 -2
  13. memos/cli.py +113 -0
  14. memos/configs/embedder.py +27 -0
  15. memos/configs/graph_db.py +83 -2
  16. memos/configs/llm.py +48 -0
  17. memos/configs/mem_cube.py +1 -1
  18. memos/configs/mem_reader.py +4 -0
  19. memos/configs/mem_scheduler.py +91 -5
  20. memos/configs/memory.py +10 -4
  21. memos/dependency.py +52 -0
  22. memos/embedders/ark.py +92 -0
  23. memos/embedders/factory.py +4 -0
  24. memos/embedders/sentence_transformer.py +8 -2
  25. memos/embedders/universal_api.py +32 -0
  26. memos/graph_dbs/base.py +2 -2
  27. memos/graph_dbs/factory.py +2 -0
  28. memos/graph_dbs/item.py +46 -0
  29. memos/graph_dbs/neo4j.py +377 -101
  30. memos/graph_dbs/neo4j_community.py +300 -0
  31. memos/llms/base.py +9 -0
  32. memos/llms/deepseek.py +54 -0
  33. memos/llms/factory.py +10 -1
  34. memos/llms/hf.py +170 -13
  35. memos/llms/hf_singleton.py +114 -0
  36. memos/llms/ollama.py +4 -0
  37. memos/llms/openai.py +68 -1
  38. memos/llms/qwen.py +63 -0
  39. memos/llms/vllm.py +153 -0
  40. memos/mem_cube/general.py +77 -16
  41. memos/mem_cube/utils.py +102 -0
  42. memos/mem_os/core.py +131 -41
  43. memos/mem_os/main.py +93 -11
  44. memos/mem_os/product.py +1098 -35
  45. memos/mem_os/utils/default_config.py +352 -0
  46. memos/mem_os/utils/format_utils.py +1154 -0
  47. memos/mem_reader/simple_struct.py +13 -8
  48. memos/mem_scheduler/base_scheduler.py +467 -36
  49. memos/mem_scheduler/general_scheduler.py +125 -244
  50. memos/mem_scheduler/modules/base.py +9 -0
  51. memos/mem_scheduler/modules/dispatcher.py +68 -2
  52. memos/mem_scheduler/modules/misc.py +39 -0
  53. memos/mem_scheduler/modules/monitor.py +228 -49
  54. memos/mem_scheduler/modules/rabbitmq_service.py +317 -0
  55. memos/mem_scheduler/modules/redis_service.py +32 -22
  56. memos/mem_scheduler/modules/retriever.py +250 -23
  57. memos/mem_scheduler/modules/schemas.py +189 -7
  58. memos/mem_scheduler/mos_for_test_scheduler.py +143 -0
  59. memos/mem_scheduler/utils.py +51 -2
  60. memos/mem_user/persistent_user_manager.py +260 -0
  61. memos/memories/activation/item.py +25 -0
  62. memos/memories/activation/kv.py +10 -3
  63. memos/memories/activation/vllmkv.py +219 -0
  64. memos/memories/factory.py +2 -0
  65. memos/memories/textual/general.py +7 -5
  66. memos/memories/textual/item.py +3 -1
  67. memos/memories/textual/tree.py +14 -6
  68. memos/memories/textual/tree_text_memory/organize/conflict.py +198 -0
  69. memos/memories/textual/tree_text_memory/organize/manager.py +72 -23
  70. memos/memories/textual/tree_text_memory/organize/redundancy.py +193 -0
  71. memos/memories/textual/tree_text_memory/organize/relation_reason_detector.py +233 -0
  72. memos/memories/textual/tree_text_memory/organize/reorganizer.py +606 -0
  73. memos/memories/textual/tree_text_memory/retrieve/recall.py +0 -1
  74. memos/memories/textual/tree_text_memory/retrieve/reranker.py +2 -2
  75. memos/memories/textual/tree_text_memory/retrieve/searcher.py +6 -5
  76. memos/parsers/markitdown.py +8 -2
  77. memos/templates/mem_reader_prompts.py +105 -36
  78. memos/templates/mem_scheduler_prompts.py +96 -47
  79. memos/templates/tree_reorganize_prompts.py +223 -0
  80. memos/vec_dbs/base.py +12 -0
  81. memos/vec_dbs/qdrant.py +46 -20
  82. memoryos-0.1.13.dist-info/RECORD +0 -122
  83. {memoryos-0.1.13.dist-info → memoryos-0.2.1.dist-info}/LICENSE +0 -0
  84. {memoryos-0.1.13.dist-info → memoryos-0.2.1.dist-info}/WHEEL +0 -0
memos/graph_dbs/neo4j.py CHANGED
@@ -3,9 +3,8 @@ import time
3
3
  from datetime import datetime
4
4
  from typing import Any, Literal
5
5
 
6
- from neo4j import GraphDatabase
7
-
8
6
  from memos.configs.graph_db import Neo4jGraphDBConfig
7
+ from memos.dependency import require_python_package
9
8
  from memos.graph_dbs.base import BaseGraphDB
10
9
  from memos.log import get_logger
11
10
 
@@ -13,17 +12,6 @@ from memos.log import get_logger
13
12
  logger = get_logger(__name__)
14
13
 
15
14
 
16
- def _parse_node(node_data: dict[str, Any]) -> dict[str, Any]:
17
- node = node_data.copy()
18
-
19
- # Convert Neo4j datetime to string
20
- for time_field in ("created_at", "updated_at"):
21
- if time_field in node and hasattr(node[time_field], "isoformat"):
22
- node[time_field] = node[time_field].isoformat()
23
-
24
- return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node}
25
-
26
-
27
15
  def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]:
28
16
  node_id = item["id"]
29
17
  memory = item["memory"]
@@ -55,11 +43,37 @@ def _prepare_node_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
55
43
  class Neo4jGraphDB(BaseGraphDB):
56
44
  """Neo4j-based implementation of a graph memory store."""
57
45
 
46
+ @require_python_package(
47
+ import_name="neo4j",
48
+ install_command="pip install neo4j",
49
+ install_link="https://neo4j.com/docs/python-manual/current/install/",
50
+ )
58
51
  def __init__(self, config: Neo4jGraphDBConfig):
52
+ """Neo4j-based implementation of a graph memory store.
53
+
54
+ Tenant Modes:
55
+ - use_multi_db = True:
56
+ Dedicated Database Mode (Multi-Database Multi-Tenant).
57
+ Each tenant or logical scope uses a separate Neo4j database.
58
+ `db_name` is the specific tenant database.
59
+ `user_name` can be None (optional).
60
+
61
+ - use_multi_db = False:
62
+ Shared Database Multi-Tenant Mode.
63
+ All tenants share a single Neo4j database.
64
+ `db_name` is the shared database.
65
+ `user_name` is required to isolate each tenant's data at the node level.
66
+ All node queries will enforce `user_name` in WHERE conditions and store it in metadata,
67
+ but it will be removed automatically before returning to external consumers.
68
+ """
69
+ from neo4j import GraphDatabase
70
+
59
71
  self.config = config
60
72
  self.driver = GraphDatabase.driver(config.uri, auth=(config.user, config.password))
61
73
  self.db_name = config.db_name
74
+ self.user_name = config.user_name
62
75
 
76
+ self.system_db_name = "system" if config.use_multi_db else config.db_name
63
77
  if config.auto_create:
64
78
  self._ensure_database_exists()
65
79
 
@@ -86,10 +100,37 @@ class Neo4jGraphDB(BaseGraphDB):
86
100
  query = """
87
101
  MATCH (n:Memory)
88
102
  WHERE n.memory_type = $memory_type
89
- RETURN COUNT(n) AS count
90
103
  """
104
+ if not self.config.use_multi_db and self.config.user_name:
105
+ query += "\nAND n.user_name = $user_name"
106
+ query += "\nRETURN COUNT(n) AS count"
107
+ with self.driver.session(database=self.db_name) as session:
108
+ result = session.run(
109
+ query,
110
+ {
111
+ "memory_type": memory_type,
112
+ "user_name": self.config.user_name if self.config.user_name else None,
113
+ },
114
+ )
115
+ return result.single()["count"]
116
+
117
+ def count_nodes(self, scope: str) -> int:
118
+ query = """
119
+ MATCH (n:Memory)
120
+ WHERE n.memory_type = $scope
121
+ """
122
+ if not self.config.use_multi_db and self.config.user_name:
123
+ query += "\nAND n.user_name = $user_name"
124
+ query += "\nRETURN count(n) AS count"
125
+
91
126
  with self.driver.session(database=self.db_name) as session:
92
- result = session.run(query, memory_type=memory_type)
127
+ result = session.run(
128
+ query,
129
+ {
130
+ "scope": scope,
131
+ "user_name": self.config.user_name if self.config.user_name else None,
132
+ },
133
+ )
93
134
  return result.single()["count"]
94
135
 
95
136
  def remove_oldest_memory(self, memory_type: str, keep_latest: int) -> None:
@@ -103,14 +144,22 @@ class Neo4jGraphDB(BaseGraphDB):
103
144
  query = f"""
104
145
  MATCH (n:Memory)
105
146
  WHERE n.memory_type = '{memory_type}'
106
- WITH n ORDER BY n.updated_at DESC
107
- SKIP {keep_latest}
108
- DETACH DELETE n
147
+ """
148
+ if not self.config.use_multi_db and self.config.user_name:
149
+ query += f"\nAND n.user_name = '{self.config.user_name}'"
150
+
151
+ query += f"""
152
+ WITH n ORDER BY n.updated_at DESC
153
+ SKIP {keep_latest}
154
+ DETACH DELETE n
109
155
  """
110
156
  with self.driver.session(database=self.db_name) as session:
111
157
  session.run(query)
112
158
 
113
159
  def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None:
160
+ if not self.config.use_multi_db and self.config.user_name:
161
+ metadata["user_name"] = self.config.user_name
162
+
114
163
  # Safely process metadata
115
164
  metadata = _prepare_node_metadata(metadata)
116
165
 
@@ -152,10 +201,14 @@ class Neo4jGraphDB(BaseGraphDB):
152
201
  set_clauses.append("n += $fields") # Merge remaining fields
153
202
  set_clause_str = ",\n ".join(set_clauses)
154
203
 
155
- query = f"""
156
- MATCH (n:Memory {{id: $id}})
157
- SET {set_clause_str}
204
+ query = """
205
+ MATCH (n:Memory {id: $id})
158
206
  """
207
+ if not self.config.use_multi_db and self.config.user_name:
208
+ query += "\nWHERE n.user_name = $user_name"
209
+ params["user_name"] = self.config.user_name
210
+
211
+ query += f"\nSET {set_clause_str}"
159
212
 
160
213
  with self.driver.session(database=self.db_name) as session:
161
214
  session.run(query, **params)
@@ -166,8 +219,17 @@ class Neo4jGraphDB(BaseGraphDB):
166
219
  Args:
167
220
  id: Node identifier to delete.
168
221
  """
222
+ query = "MATCH (n:Memory {id: $id})"
223
+
224
+ params = {"id": id}
225
+ if not self.config.use_multi_db and self.config.user_name:
226
+ query += " WHERE n.user_name = $user_name"
227
+ params["user_name"] = self.config.user_name
228
+
229
+ query += " DETACH DELETE n"
230
+
169
231
  with self.driver.session(database=self.db_name) as session:
170
- session.run("MATCH (n:Memory {id: $id}) DETACH DELETE n", id=id)
232
+ session.run(query, **params)
171
233
 
172
234
  # Edge (Relationship) Management
173
235
  def add_edge(self, source_id: str, target_id: str, type: str) -> None:
@@ -178,15 +240,21 @@ class Neo4jGraphDB(BaseGraphDB):
178
240
  target_id: ID of the target node.
179
241
  type: Relationship type (e.g., 'RELATE_TO', 'PARENT').
180
242
  """
243
+ query = """
244
+ MATCH (a:Memory {id: $source_id})
245
+ MATCH (b:Memory {id: $target_id})
246
+ """
247
+ params = {"source_id": source_id, "target_id": target_id}
248
+ if not self.config.use_multi_db and self.config.user_name:
249
+ query += """
250
+ WHERE a.user_name = $user_name AND b.user_name = $user_name
251
+ """
252
+ params["user_name"] = self.config.user_name
253
+
254
+ query += f"\nMERGE (a)-[:{type}]->(b)"
255
+
181
256
  with self.driver.session(database=self.db_name) as session:
182
- session.run(
183
- f"""
184
- MATCH (a:Memory {{id: $source_id}})
185
- MATCH (b:Memory {{id: $target_id}})
186
- MERGE (a)-[:{type}]->(b)
187
- """,
188
- {"source_id": source_id, "target_id": target_id},
189
- )
257
+ session.run(query, params)
190
258
 
191
259
  def delete_edge(self, source_id: str, target_id: str, type: str) -> None:
192
260
  """
@@ -196,12 +264,21 @@ class Neo4jGraphDB(BaseGraphDB):
196
264
  target_id: ID of the target node.
197
265
  type: Relationship type to remove.
198
266
  """
267
+ query = f"""
268
+ MATCH (a:Memory {{id: $source}})
269
+ -[r:{type}]->
270
+ (b:Memory {{id: $target}})
271
+ """
272
+ params = {"source": source_id, "target": target_id}
273
+
274
+ if not self.config.use_multi_db and self.config.user_name:
275
+ query += "\nWHERE a.user_name = $user_name AND b.user_name = $user_name"
276
+ params["user_name"] = self.config.user_name
277
+
278
+ query += "\nDELETE r"
279
+
199
280
  with self.driver.session(database=self.db_name) as session:
200
- session.run(
201
- f"MATCH (a:Memory {{id: $source}})-[r:{type}]->(b:Memory {{id: $target}})\nDELETE r",
202
- source=source_id,
203
- target=target_id,
204
- )
281
+ session.run(query, params)
205
282
 
206
283
  def edge_exists(
207
284
  self, source_id: str, target_id: str, type: str = "ANY", direction: str = "OUTGOING"
@@ -231,14 +308,18 @@ class Neo4jGraphDB(BaseGraphDB):
231
308
  raise ValueError(
232
309
  f"Invalid direction: {direction}. Must be 'OUTGOING', 'INCOMING', or 'ANY'."
233
310
  )
311
+ query = f"MATCH {pattern}"
312
+ params = {"source": source_id, "target": target_id}
313
+
314
+ if not self.config.use_multi_db and self.config.user_name:
315
+ query += "\nWHERE a.user_name = $user_name AND b.user_name = $user_name"
316
+ params["user_name"] = self.config.user_name
317
+
318
+ query += "\nRETURN r"
234
319
 
235
320
  # Run the Cypher query
236
321
  with self.driver.session(database=self.db_name) as session:
237
- result = session.run(
238
- f"MATCH {pattern} RETURN r",
239
- source=source_id,
240
- target=target_id,
241
- )
322
+ result = session.run(query, params)
242
323
  return result.single() is not None
243
324
 
244
325
  # Graph Query & Reasoning
@@ -250,10 +331,17 @@ class Neo4jGraphDB(BaseGraphDB):
250
331
  Returns:
251
332
  Dictionary of node fields, or None if not found.
252
333
  """
334
+ where_user = ""
335
+ params = {"id": id}
336
+ if not self.config.use_multi_db and self.config.user_name:
337
+ where_user = " AND n.user_name = $user_name"
338
+ params["user_name"] = self.config.user_name
339
+
340
+ query = f"MATCH (n:Memory) WHERE n.id = $id {where_user} RETURN n"
341
+
253
342
  with self.driver.session(database=self.db_name) as session:
254
- result = session.run("MATCH (n:Memory {id: $id}) RETURN n", id=id)
255
- record = result.single()
256
- return _parse_node(dict(record["n"])) if record else None
343
+ record = session.run(query, params).single()
344
+ return self._parse_node(dict(record["n"])) if record else None
257
345
 
258
346
  def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]:
259
347
  """
@@ -270,10 +358,18 @@ class Neo4jGraphDB(BaseGraphDB):
270
358
  if not ids:
271
359
  return []
272
360
 
273
- query = "MATCH (n:Memory) WHERE n.id IN $ids RETURN n"
361
+ where_user = ""
362
+ params = {"ids": ids}
363
+
364
+ if not self.config.use_multi_db and self.config.user_name:
365
+ where_user = " AND n.user_name = $user_name"
366
+ params["user_name"] = self.config.user_name
367
+
368
+ query = f"MATCH (n:Memory) WHERE n.id IN $ids{where_user} RETURN n"
369
+
274
370
  with self.driver.session(database=self.db_name) as session:
275
- results = session.run(query, {"ids": ids})
276
- return [_parse_node(dict(record["n"])) for record in results]
371
+ results = session.run(query, params)
372
+ return [self._parse_node(dict(record["n"])) for record in results]
277
373
 
278
374
  def get_edges(self, id: str, type: str = "ANY", direction: str = "ANY") -> list[dict[str, str]]:
279
375
  """
@@ -307,14 +403,20 @@ class Neo4jGraphDB(BaseGraphDB):
307
403
  else:
308
404
  raise ValueError("Invalid direction. Must be 'OUTGOING', 'INCOMING', or 'ANY'.")
309
405
 
406
+ params = {"id": id}
407
+
408
+ if not self.config.use_multi_db and self.config.user_name:
409
+ where_clause += " AND a.user_name = $user_name AND b.user_name = $user_name"
410
+ params["user_name"] = self.config.user_name
411
+
310
412
  query = f"""
311
- MATCH {pattern}
312
- WHERE {where_clause}
313
- RETURN a.id AS from_id, b.id AS to_id, type(r) AS type
314
- """
413
+ MATCH {pattern}
414
+ WHERE {where_clause}
415
+ RETURN a.id AS from_id, b.id AS to_id, type(r) AS type
416
+ """
315
417
 
316
418
  with self.driver.session(database=self.db_name) as session:
317
- result = session.run(query, id=id)
419
+ result = session.run(query, params)
318
420
  edges = []
319
421
  for record in result:
320
422
  edges.append(
@@ -336,14 +438,74 @@ class Neo4jGraphDB(BaseGraphDB):
336
438
  """
337
439
  raise NotImplementedError
338
440
 
339
- def get_children_with_embeddings(self, id: str) -> list[str]:
340
- query = """
341
- MATCH (p:Memory)-[:PARENT]->(c:Memory)
342
- WHERE p.id = $id
343
- RETURN c.id AS id, c.embedding AS embedding, c.memory AS memory
441
+ def get_neighbors_by_tag(
442
+ self,
443
+ tags: list[str],
444
+ exclude_ids: list[str],
445
+ top_k: int = 5,
446
+ min_overlap: int = 1,
447
+ ) -> list[dict[str, Any]]:
448
+ """
449
+ Find top-K neighbor nodes with maximum tag overlap.
450
+
451
+ Args:
452
+ tags: The list of tags to match.
453
+ exclude_ids: Node IDs to exclude (e.g., local cluster).
454
+ top_k: Max number of neighbors to return.
455
+ min_overlap: Minimum number of overlapping tags required.
456
+
457
+ Returns:
458
+ List of dicts with node details and overlap count.
344
459
  """
460
+ where_user = ""
461
+ params = {
462
+ "tags": tags,
463
+ "exclude_ids": exclude_ids,
464
+ "min_overlap": min_overlap,
465
+ "top_k": top_k,
466
+ }
467
+
468
+ if not self.config.use_multi_db and self.config.user_name:
469
+ where_user = "AND n.user_name = $user_name"
470
+ params["user_name"] = self.config.user_name
471
+
472
+ query = f"""
473
+ MATCH (n:Memory)
474
+ WHERE NOT n.id IN $exclude_ids
475
+ AND n.status = 'activated'
476
+ AND n.type <> 'reasoning'
477
+ AND n.memory_type <> 'WorkingMemory'
478
+ {where_user}
479
+ WITH n, [tag IN n.tags WHERE tag IN $tags] AS overlap_tags
480
+ WHERE size(overlap_tags) >= $min_overlap
481
+ RETURN n, size(overlap_tags) AS overlap_count
482
+ ORDER BY overlap_count DESC
483
+ LIMIT $top_k
484
+ """
485
+
345
486
  with self.driver.session(database=self.db_name) as session:
346
- return list(session.run(query, id=id))
487
+ result = session.run(query, params)
488
+ return [self._parse_node(dict(record["n"])) for record in result]
489
+
490
+ def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]:
491
+ where_user = ""
492
+ params = {"id": id}
493
+
494
+ if not self.config.use_multi_db and self.config.user_name:
495
+ where_user = "AND p.user_name = $user_name AND c.user_name = $user_name"
496
+ params["user_name"] = self.config.user_name
497
+
498
+ query = f"""
499
+ MATCH (p:Memory)-[:PARENT]->(c:Memory)
500
+ WHERE p.id = $id {where_user}
501
+ RETURN c.id AS id, c.embedding AS embedding, c.memory AS memory
502
+ """
503
+
504
+ with self.driver.session(database=self.db_name) as session:
505
+ result = session.run(query, params)
506
+ return [
507
+ {"id": r["id"], "embedding": r["embedding"], "memory": r["memory"]} for r in result
508
+ ]
347
509
 
348
510
  def get_path(self, source_id: str, target_id: str, max_depth: int = 3) -> list[str]:
349
511
  """
@@ -374,30 +536,39 @@ class Neo4jGraphDB(BaseGraphDB):
374
536
  }
375
537
  """
376
538
  with self.driver.session(database=self.db_name) as session:
377
- status_clause = f", status: '{center_status}'" if center_status else ""
539
+ params = {"center_id": center_id}
540
+ center_user_clause = ""
541
+ neighbor_user_clause = ""
542
+
543
+ if not self.config.use_multi_db and self.config.user_name:
544
+ center_user_clause = " AND center.user_name = $user_name"
545
+ neighbor_user_clause = " WHERE neighbor.user_name = $user_name"
546
+ params["user_name"] = self.config.user_name
547
+ status_clause = f" AND center.status = '{center_status}'" if center_status else ""
548
+
378
549
  query = f"""
379
- MATCH (center:Memory {{id: $center_id{status_clause}}})
380
- OPTIONAL MATCH (center)-[r*1..{depth}]-(neighbor:Memory)
381
- WITH collect(DISTINCT center) AS centers,
382
- collect(DISTINCT neighbor) AS neighbors,
383
- collect(DISTINCT r) AS rels
384
- RETURN centers, neighbors, rels
550
+ MATCH (center:Memory)
551
+ WHERE center.id = $center_id{status_clause}{center_user_clause}
552
+
553
+ OPTIONAL MATCH (center)-[r*1..{depth}]-(neighbor:Memory)
554
+ {neighbor_user_clause}
555
+
556
+ WITH collect(DISTINCT center) AS centers,
557
+ collect(DISTINCT neighbor) AS neighbors,
558
+ collect(DISTINCT r) AS rels
559
+ RETURN centers, neighbors, rels
385
560
  """
386
- record = session.run(query, {"center_id": center_id}).single()
561
+ record = session.run(query, params).single()
387
562
 
388
563
  if not record:
389
- logger.warning(
390
- f"No active node found for center_id={center_id} with status={center_status}"
391
- )
392
564
  return {"core_node": None, "neighbors": [], "edges": []}
393
565
 
394
566
  centers = record["centers"]
395
567
  if not centers or centers[0] is None:
396
- logger.warning(f"Center node not found or inactive for id={center_id}")
397
568
  return {"core_node": None, "neighbors": [], "edges": []}
398
569
 
399
- core_node = _parse_node(dict(centers[0]))
400
- neighbors = [_parse_node(dict(n)) for n in record["neighbors"] if n]
570
+ core_node = self._parse_node(dict(centers[0]))
571
+ neighbors = [self._parse_node(dict(n)) for n in record["neighbors"] if n]
401
572
  edges = []
402
573
  for rel_chain in record["rels"]:
403
574
  for rel in rel_chain:
@@ -459,6 +630,8 @@ class Neo4jGraphDB(BaseGraphDB):
459
630
  where_clauses.append("node.memory_type = $scope")
460
631
  if status:
461
632
  where_clauses.append("node.status = $status")
633
+ if not self.config.use_multi_db and self.config.user_name:
634
+ where_clauses.append("node.user_name = $user_name")
462
635
 
463
636
  where_clause = ""
464
637
  if where_clauses:
@@ -476,6 +649,8 @@ class Neo4jGraphDB(BaseGraphDB):
476
649
  parameters["scope"] = scope
477
650
  if status:
478
651
  parameters["status"] = status
652
+ if not self.config.use_multi_db and self.config.user_name:
653
+ parameters["user_name"] = self.config.user_name
479
654
 
480
655
  with self.driver.session(database=self.db_name) as session:
481
656
  result = session.run(query, parameters)
@@ -543,6 +718,10 @@ class Neo4jGraphDB(BaseGraphDB):
543
718
  else:
544
719
  raise ValueError(f"Unsupported operator: {op}")
545
720
 
721
+ if not self.config.use_multi_db and self.config.user_name:
722
+ where_clauses.append("n.user_name = $user_name")
723
+ params["user_name"] = self.config.user_name
724
+
546
725
  where_str = " AND ".join(where_clauses)
547
726
  query = f"MATCH (n:Memory) WHERE {where_str} RETURN n.id AS id"
548
727
 
@@ -571,6 +750,20 @@ class Neo4jGraphDB(BaseGraphDB):
571
750
  if not group_fields:
572
751
  raise ValueError("group_fields cannot be empty")
573
752
 
753
+ final_params = params.copy() if params else {}
754
+
755
+ if not self.config.use_multi_db and self.config.user_name:
756
+ user_clause = "n.user_name = $user_name"
757
+ final_params["user_name"] = self.config.user_name
758
+ if where_clause:
759
+ where_clause = where_clause.strip()
760
+ if where_clause.upper().startswith("WHERE"):
761
+ where_clause += f" AND {user_clause}"
762
+ else:
763
+ where_clause = f"WHERE {where_clause} AND {user_clause}"
764
+ else:
765
+ where_clause = f"WHERE {user_clause}"
766
+
574
767
  # Force RETURN field AS field to guarantee key match
575
768
  group_fields_cypher = ", ".join([f"n.{field} AS {field}" for field in group_fields])
576
769
 
@@ -581,7 +774,7 @@ class Neo4jGraphDB(BaseGraphDB):
581
774
  """
582
775
 
583
776
  with self.driver.session(database=self.db_name) as session:
584
- result = session.run(query, params or {})
777
+ result = session.run(query, final_params)
585
778
  return [
586
779
  {**{field: record[field] for field in group_fields}, "count": record["count"]}
587
780
  for record in result
@@ -620,17 +813,16 @@ class Neo4jGraphDB(BaseGraphDB):
620
813
  Clear the entire graph if the target database exists.
621
814
  """
622
815
  try:
623
- # Step 1: Check if the database exists
624
- with self.driver.session(database="system") as session:
625
- result = session.run("SHOW DATABASES YIELD name RETURN name")
626
- db_names = [record["name"] for record in result]
627
- if self.db_name not in db_names:
628
- logger.info(f"[Skip] Database '{self.db_name}' does not exist.")
629
- return
816
+ if not self.config.use_multi_db and self.config.user_name:
817
+ query = "MATCH (n:Memory) WHERE n.user_name = $user_name DETACH DELETE n"
818
+ params = {"user_name": self.config.user_name}
819
+ else:
820
+ query = "MATCH (n) DETACH DELETE n"
821
+ params = {}
630
822
 
631
823
  # Step 2: Clear the graph in that database
632
824
  with self.driver.session(database=self.db_name) as session:
633
- session.run("MATCH (n) DETACH DELETE n")
825
+ session.run(query, params)
634
826
  logger.info(f"Cleared all nodes from database '{self.db_name}'.")
635
827
 
636
828
  except Exception as e:
@@ -649,14 +841,22 @@ class Neo4jGraphDB(BaseGraphDB):
649
841
  """
650
842
  with self.driver.session(database=self.db_name) as session:
651
843
  # Export nodes
652
- node_result = session.run("MATCH (n:Memory) RETURN n")
653
- nodes = [_parse_node(dict(record["n"])) for record in node_result]
844
+ node_query = "MATCH (n:Memory)"
845
+ edge_query = "MATCH (a:Memory)-[r]->(b:Memory)"
846
+ params = {}
847
+
848
+ if not self.config.use_multi_db and self.config.user_name:
849
+ node_query += " WHERE n.user_name = $user_name"
850
+ edge_query += " WHERE a.user_name = $user_name AND b.user_name = $user_name"
851
+ params["user_name"] = self.config.user_name
852
+
853
+ node_result = session.run(f"{node_query} RETURN n", params)
854
+ nodes = [self._parse_node(dict(record["n"])) for record in node_result]
654
855
 
655
856
  # Export edges
656
- edge_result = session.run("""
657
- MATCH (a:Memory)-[r]->(b:Memory)
658
- RETURN a.id AS source, b.id AS target, type(r) AS type
659
- """)
857
+ edge_result = session.run(
858
+ f"{edge_query} RETURN a.id AS source, b.id AS target, type(r) AS type", params
859
+ )
660
860
  edges = [
661
861
  {"source": record["source"], "target": record["target"], "type": record["type"]}
662
862
  for record in edge_result
@@ -675,6 +875,9 @@ class Neo4jGraphDB(BaseGraphDB):
675
875
  for node in data.get("nodes", []):
676
876
  id, memory, metadata = _compose_node(node)
677
877
 
878
+ if not self.config.use_multi_db and self.config.user_name:
879
+ metadata["user_name"] = self.config.user_name
880
+
678
881
  metadata = _prepare_node_metadata(metadata)
679
882
 
680
883
  # Merge node and set metadata
@@ -720,35 +923,85 @@ class Neo4jGraphDB(BaseGraphDB):
720
923
  if scope not in {"WorkingMemory", "LongTermMemory", "UserMemory"}:
721
924
  raise ValueError(f"Unsupported memory type scope: {scope}")
722
925
 
723
- query = """
724
- MATCH (n:Memory)
725
- WHERE n.memory_type = $scope
726
- RETURN n
926
+ where_clause = "WHERE n.memory_type = $scope"
927
+ params = {"scope": scope}
928
+
929
+ if not self.config.use_multi_db and self.config.user_name:
930
+ where_clause += " AND n.user_name = $user_name"
931
+ params["user_name"] = self.config.user_name
932
+
933
+ query = f"""
934
+ MATCH (n:Memory)
935
+ {where_clause}
936
+ RETURN n
937
+ """
938
+
939
+ with self.driver.session(database=self.db_name) as session:
940
+ results = session.run(query, params)
941
+ return [self._parse_node(dict(record["n"])) for record in results]
942
+
943
+ def get_structure_optimization_candidates(self, scope: str) -> list[dict]:
727
944
  """
945
+ Find nodes that are likely candidates for structure optimization:
946
+ - Isolated nodes, nodes with empty background, or nodes with exactly one child.
947
+ - Plus: the child of any parent node that has exactly one child.
948
+ """
949
+ where_clause = """
950
+ WHERE n.memory_type = $scope
951
+ AND n.status = 'activated'
952
+ AND NOT ( (n)-[:PARENT]->() OR ()-[:PARENT]->(n) )
953
+ """
954
+ params = {"scope": scope}
955
+
956
+ if not self.config.use_multi_db and self.config.user_name:
957
+ where_clause += " AND n.user_name = $user_name"
958
+ params["user_name"] = self.config.user_name
959
+
960
+ query = f"""
961
+ MATCH (n:Memory)
962
+ {where_clause}
963
+ RETURN n.id AS id, n AS node
964
+ """
728
965
 
729
966
  with self.driver.session(database=self.db_name) as session:
730
- results = session.run(query, {"scope": scope})
731
- return [_parse_node(dict(record["n"])) for record in results]
967
+ results = session.run(query, params)
968
+ return [
969
+ self._parse_node({"id": record["id"], **dict(record["node"])}) for record in results
970
+ ]
732
971
 
733
972
  def drop_database(self) -> None:
734
973
  """
735
974
  Permanently delete the entire database this instance is using.
736
975
  WARNING: This operation is destructive and cannot be undone.
737
976
  """
738
- if self.db_name in ("system", "neo4j"):
739
- raise ValueError(f"Refusing to drop protected database: {self.db_name}")
977
+ if self.config.use_multi_db:
978
+ if self.db_name in ("system", "neo4j"):
979
+ raise ValueError(f"Refusing to drop protected database: {self.db_name}")
740
980
 
741
- with self.driver.session(database="system") as session:
742
- session.run(f"DROP DATABASE {self.db_name} IF EXISTS")
743
- print(f"Database '{self.db_name}' has been dropped.")
981
+ with self.driver.session(database=self.system_db_name) as session:
982
+ session.run(f"DROP DATABASE {self.db_name} IF EXISTS")
983
+ print(f"Database '{self.db_name}' has been dropped.")
984
+ else:
985
+ raise ValueError(
986
+ f"Refusing to drop protected database: {self.db_name} in "
987
+ f"Shared Database Multi-Tenant mode"
988
+ )
744
989
 
745
990
  def _ensure_database_exists(self):
746
- with self.driver.session(database="system") as session:
747
- session.run(f"CREATE DATABASE {self.db_name} IF NOT EXISTS")
991
+ from neo4j.exceptions import ClientError
992
+
993
+ try:
994
+ with self.driver.session(database="system") as session:
995
+ session.run(f"CREATE DATABASE `{self.db_name}` IF NOT EXISTS")
996
+ except ClientError as e:
997
+ if "ExistingDatabaseFound" in str(e):
998
+ pass # Ignore, database already exists
999
+ else:
1000
+ raise
748
1001
 
749
1002
  # Wait until the database is available
750
1003
  for _ in range(10):
751
- with self.driver.session(database="system") as session:
1004
+ with self.driver.session(database=self.system_db_name) as session:
752
1005
  result = session.run(
753
1006
  "SHOW DATABASES YIELD name, currentStatus RETURN name, currentStatus"
754
1007
  )
@@ -790,7 +1043,10 @@ class Neo4jGraphDB(BaseGraphDB):
790
1043
 
791
1044
  def _create_basic_property_indexes(self) -> None:
792
1045
  """
793
- Create standard B-tree indexes on memory_type, created_at, and updated_at fields.
1046
+ Create standard B-tree indexes on memory_type, created_at,
1047
+ and updated_at fields.
1048
+ Create standard B-tree indexes on user_name when use Shared Database
1049
+ Multi-Tenant Mode
794
1050
  """
795
1051
  try:
796
1052
  with self.driver.session(database=self.db_name) as session:
@@ -811,6 +1067,15 @@ class Neo4jGraphDB(BaseGraphDB):
811
1067
  FOR (n:Memory) ON (n.updated_at)
812
1068
  """)
813
1069
  logger.debug("Index 'memory_updated_at_index' ensured.")
1070
+
1071
+ if not self.config.use_multi_db and self.config.user_name:
1072
+ session.run(
1073
+ """
1074
+ CREATE INDEX memory_user_name_index IF NOT EXISTS
1075
+ FOR (n:Memory) ON (n.user_name)
1076
+ """
1077
+ )
1078
+ logger.debug("Index 'memory_user_name_index' ensured.")
814
1079
  except Exception as e:
815
1080
  logger.warning(f"Failed to create basic property indexes: {e}")
816
1081
 
@@ -825,3 +1090,14 @@ class Neo4jGraphDB(BaseGraphDB):
825
1090
  if record["name"] == index_name:
826
1091
  return True
827
1092
  return False
1093
+
1094
+ def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]:
1095
+ node = node_data.copy()
1096
+
1097
+ # Convert Neo4j datetime to string
1098
+ for time_field in ("created_at", "updated_at"):
1099
+ if time_field in node and hasattr(node[time_field], "isoformat"):
1100
+ node[time_field] = node[time_field].isoformat()
1101
+ node.pop("user_name", None)
1102
+
1103
+ return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node}