feat: add neo4j share-db example (#59)

* feat: add neo4j share-db example

* feat: add Shared Database Multi-Tenant Mode

* fix: neo4j auto create bug

* fix: tree text memory reorganizer bug

* feat: change 'RELATE_TO' to 'RELATE'

* fix: database exist bug

* feat: delete useless relation edge

* feat: finish Shared Database Multi-Tenant Mode

* feat: add share database multi tenant tree config

* feat: add shared-db mode config

* fix: file name bug in example tree text memory]

* style: add user_name from config

Author: CaralHsi

examples/basic_modules/neo4j_example.py

@@ -8,12 +8,7 @@
 
 
 embedder_config = EmbedderConfigFactory.model_validate(
-    {
-        "backend": "sentence_transformer",
-        "config": {
-            "model_name_or_path": "nomic-ai/nomic-embed-text-v1.5",
-        },
-    }
+    {"backend": "ollama", "config": {"model_name_or_path": "nomic-embed-text:latest"}}
 )
 embedder = EmbedderFactory.from_config(embedder_config)
 
@@ -22,7 +17,7 @@ def embed_memory_item(memory: str) -> list[float]:
     return embedder.embed([memory])[0]
 
 
-def example_1_paper(db_name: str = "paper"):
+def example_multi_db(db_name: str = "paper"):
     # Step 1: Build factory config
     config = GraphDBConfigFactory(
         backend="neo4j",
@@ -33,6 +28,7 @@ def example_1_paper(db_name: str = "paper"):
             "db_name": db_name,
             "auto_create": True,
             "embedding_dimension": 768,
+            "use_multi_db": True,
         },
     )
 
@@ -68,7 +64,7 @@ def example_1_paper(db_name: str = "paper"):
     )
 
     graph.add_node(
-        id=topic.id, content=topic.memory, metadata=topic.metadata.model_dump(exclude_none=True)
+        id=topic.id, memory=topic.memory, metadata=topic.metadata.model_dump(exclude_none=True)
     )
 
     # Step 4: Define and write concept nodes
@@ -150,7 +146,7 @@ def example_1_paper(db_name: str = "paper"):
     for concept in concepts:
         graph.add_node(
             id=concept.id,
-            content=concept.memory,
+            memory=concept.memory,
             metadata=concept.metadata.model_dump(exclude_none=True),
         )
         graph.add_edge(source_id=concept.id, target_id=topic.id, type="RELATED")
@@ -259,113 +255,101 @@ def example_1_paper(db_name: str = "paper"):
         print(graph.get_node(node_i["id"]))
 
 
-def example_2_travel(db_name: str = "travel"):
-    # Step 1: Build factory config
-    config = GraphDBConfigFactory(
-        backend="neo4j",
-        config={
-            "uri": "bolt://localhost:7687",
-            "user": "neo4j",
-            "password": "12345678",
-            "db_name": db_name,
-            "auto_create": True,
-            "embedding_dimension": 768,
-        },
-    )
-
-    # Step 2: Instantiate the graph store
-    graph = GraphStoreFactory.from_config(config)
-    graph.clear()
-
-    # Step 3: Create topic node
-    topic = TextualMemoryItem(
-        memory="Travel",
-        metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="LongTermMemory",
-            hierarchy_level="topic",
-            status="activated",
-            visibility="public",
-            embedding=embed_memory_item("Travel"),
-        ),
-    )
-
-    graph.add_node(
-        id=topic.id, content=topic.memory, metadata=topic.metadata.model_dump(exclude_none=True)
-    )
-
-    concept1 = TextualMemoryItem(
-        memory="Travel in Italy",
-        metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="LongTermMemory",
-            hierarchy_level="concept",
-            status="activated",
-            visibility="public",
-            embedding=embed_memory_item("Travel in Italy"),
-        ),
-    )
+def example_shared_db(db_name: str = "shared-traval-group"):
+    """
+    Example: Single(Shared)-DB multi-tenant (logical isolation)
+    Multiple users' data in the same Neo4j DB with user_name as a tag.
+    """
+    # users
+    user_list = ["travel_member_alice", "travel_member_bob"]
+
+    for user_name in user_list:
+        # Step 1: Build factory config
+        config = GraphDBConfigFactory(
+            backend="neo4j",
+            config={
+                "uri": "bolt://localhost:7687",
+                "user": "neo4j",
+                "password": "12345678",
+                "db_name": db_name,
+                "user_name": user_name,
+                "use_multi_db": False,
+                "auto_create": True,
+                "embedding_dimension": 768,
+            },
+        )
+        # Step 2: Instantiate graph store
+        graph = GraphStoreFactory.from_config(config)
+        print(f"\n[INFO] Working in shared DB: {db_name}, for user: {user_name}")
+        graph.clear()
+
+        # Step 3: Create topic node
+        topic = TextualMemoryItem(
+            memory=f"Travel notes for {user_name}",
+            metadata=TreeNodeTextualMemoryMetadata(
+                memory_type="LongTermMemory",
+                hierarchy_level="topic",
+                status="activated",
+                visibility="public",
+                embedding=embed_memory_item(f"Travel notes for {user_name}"),
+            ),
+        )
 
-    graph.add_node(
-        id=concept1.id,
-        content=concept1.memory,
-        metadata=concept1.metadata.model_dump(exclude_none=True),
-    )
-    graph.add_edge(source_id=topic.id, target_id=concept1.id, type="INCLUDE")
+        graph.add_node(
+            id=topic.id, memory=topic.memory, metadata=topic.metadata.model_dump(exclude_none=True)
+        )
 
-    concept2 = TextualMemoryItem(
-        memory="Traval plan",
-        metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="LongTermMemory",
-            hierarchy_level="concept",
-            status="activated",
-            visibility="public",
-            embedding=embed_memory_item("Traval plan"),
-        ),
-    )
+        # Step 4: Add a concept for each user
+        concept = TextualMemoryItem(
+            memory=f"Itinerary plan for {user_name}",
+            metadata=TreeNodeTextualMemoryMetadata(
+                memory_type="LongTermMemory",
+                hierarchy_level="concept",
+                status="activated",
+                visibility="public",
+                embedding=embed_memory_item(f"Itinerary plan for {user_name}"),
+            ),
+        )
 
-    graph.add_node(
-        id=concept2.id,
-        content=concept2.memory,
-        metadata=concept2.metadata.model_dump(exclude_none=True),
-    )
-    graph.add_edge(source_id=concept1.id, target_id=concept2.id, type="INCLUDE")
+        graph.add_node(
+            id=concept.id,
+            memory=concept.memory,
+            metadata=concept.metadata.model_dump(exclude_none=True),
+        )
 
-    fact1 = TextualMemoryItem(
-        memory="10-Day Itinerary for Traveling in Italy",
-        metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="WorkingMemory",
-            key="Reward Components",
-            value="Coverage gain, energy usage penalty, overlap penalty",
-            hierarchy_level="fact",
-            type="fact",
-            memory_time="2024-01-01",
-            source="file",
-            sources=["paper://multi-uav-coverage/reward-details"],
-            status="activated",
-            confidence=90.0,
-            tags=["reward", "overlap", "multi-agent"],
-            entities=["coverage", "energy", "overlap"],
-            visibility="public",
-            embedding=embed_memory_item("10-Day Itinerary for Traveling in Italy"),
-            updated_at=datetime.now().isoformat(),
-        ),
-    )
+        # Link concept to topic
+        graph.add_edge(source_id=concept.id, target_id=topic.id, type="INCLUDE")
 
-    graph.add_node(
-        id=fact1.id, content=fact1.memory, metadata=fact1.metadata.model_dump(exclude_none=True)
-    )
-    graph.add_edge(source_id=concept2.id, target_id=fact1.id, type="INCLUDE")
+        print(f"[INFO] Added nodes for {user_name}")
 
+    # Step 5: Query and print ALL for verification
+    print("\n=== Export entire DB (for verification, includes ALL users) ===")
+    graph = GraphStoreFactory.from_config(config)
     all_graph_data = graph.export_graph()
     print(all_graph_data)
 
-    nodes = graph.search_by_embedding(vector=embed_memory_item("what does FT reflect?"), top_k=1)
-
-    for node_i in nodes:
-        print(graph.get_node(node_i["id"]))
+    # Step 6: Search for alice's data only
+    print("\n=== Search for travel_member_alice ===")
+    config_alice = GraphDBConfigFactory(
+        backend="neo4j",
+        config={
+            "uri": "bolt://localhost:7687",
+            "user": "neo4j",
+            "password": "12345678",
+            "db_name": db_name,
+            "user_name": user_list[0],
+            "embedding_dimension": 768,
+        },
+    )
+    graph_alice = GraphStoreFactory.from_config(config_alice)
+    nodes = graph_alice.search_by_embedding(vector=embed_memory_item("travel itinerary"), top_k=1)
+    for node in nodes:
+        print(graph_alice.get_node(node["id"]))
 
 
 if __name__ == "__main__":
-    example_1_paper(db_name="paper")
+    print("\n=== Example: Multi-DB ===")
+    example_multi_db(db_name="paper")
 
-if __name__ == "__main__":
-    example_2_travel(db_name="traval")
+    print("\n=== Example: Single-DB ===")
+    example_shared_db(db_name="shared-traval-group11")

examples/core_memories/tree_textual_memory.py

@@ -1,3 +1,5 @@
+import time
+
 from memos import log
 from memos.configs.embedder import EmbedderConfigFactory
 from memos.configs.mem_reader import SimpleStructMemReaderConfig
@@ -25,7 +27,9 @@ def embed_memory_item(memory: str) -> list[float]:
     return embedder.embed([memory])[0]
 
 
-tree_config = TreeTextMemoryConfig.from_json_file("examples/data/config/tree_config.json")
+tree_config = TreeTextMemoryConfig.from_json_file(
+    "examples/data/config/tree_config_shared_database.json"
+)
 my_tree_textual_memory = TreeTextMemory(tree_config)
 my_tree_textual_memory.delete_all()
 
@@ -185,6 +189,8 @@ def embed_memory_item(memory: str) -> list[float]:
     my_tree_textual_memory.add(m_list)
     my_tree_textual_memory.memory_manager.wait_reorganizer()
 
+time.sleep(60)
+
 results = my_tree_textual_memory.search(
     "Talk about the user's childhood story?",
     top_k=10,

examples/data/config/tree_config_shared_database.json

@@ -0,0 +1,40 @@
+{
+  "extractor_llm": {
+    "backend": "ollama",
+    "config": {
+      "model_name_or_path": "qwen3:0.6b",
+      "temperature": 0.0,
+      "remove_think_prefix": true,
+      "max_tokens": 8192
+    }
+  },
+  "dispatcher_llm": {
+    "backend": "ollama",
+    "config": {
+      "model_name_or_path": "qwen3:0.6b",
+      "temperature": 0.0,
+      "remove_think_prefix": true,
+      "max_tokens": 8192
+    }
+  },
+  "embedder": {
+    "backend": "ollama",
+    "config": {
+      "model_name_or_path": "nomic-embed-text:latest"
+    }
+  },
+  "graph_db": {
+    "backend": "neo4j",
+    "config": {
+      "uri": "bolt://localhost:7687",
+      "user": "neo4j",
+      "password": "12345678",
+      "db_name": "shared-tree-textual-memory",
+      "user_name": "alice",
+      "auto_create": true,
+      "use_multi_db": false,
+      "embedding_dimension": 768
+    }
+  },
+  "reorganize": true
+}

src/memos/configs/graph_db.py

@@ -14,14 +14,70 @@ class BaseGraphDBConfig(BaseConfig):
 
 
 class Neo4jGraphDBConfig(BaseGraphDBConfig):
-    """Neo4j-specific configuration."""
+    """
+    Neo4j-specific configuration.
+
+    This config supports:
+    1) Physical isolation (multi-db) — each user gets a dedicated Neo4j database.
+    2) Logical isolation (single-db) — all users share one or more databases, but each node is tagged with `user_name`.
+
+    How to use:
+    - If `use_multi_db=True`, then `db_name` should usually be the same as `user_name`.
+      Each user gets a separate database for physical isolation.
+      Example: db_name = "alice", user_name = None or "alice".
+
+    - If `use_multi_db=False`, then `db_name` is your shared database (e.g., "neo4j" or "shared_db").
+      You must provide `user_name` to logically isolate each user's data.
+      All nodes and queries must respect this tag.
+
+    Example configs:
+    ---
+    # Physical isolation:
+    db_name = "alice"
+    use_multi_db = True
+    user_name = None
+
+    # Logical isolation:
+    db_name = "shared_db_student_group"
+    use_multi_db = False
+    user_name = "alice"
+    """
 
     db_name: str = Field(..., description="The name of the target Neo4j database")
     auto_create: bool = Field(
-        default=False, description="Whether to create the DB if it doesn't exist"
+        default=False,
+        description="If True, automatically create the target db_name in multi-db mode if it does not exist.",
+    )
+
+    use_multi_db: bool = Field(
+        default=True,
+        description=(
+            "If True: use Neo4j's multi-database feature for physical isolation; "
+            "each user typically gets a separate database. "
+            "If False: use a single shared database with logical isolation by user_name."
+        ),
     )
+
+    user_name: str | None = Field(
+        default=None,
+        description=(
+            "Logical user or tenant ID for data isolation. "
+            "Required if use_multi_db is False. "
+            "All nodes must be tagged with this and all queries must filter by this."
+        ),
+    )
+
     embedding_dimension: int = Field(default=768, description="Dimension of vector embedding")
 
+    @model_validator(mode="after")
+    def validate_config(self):
+        """Validate logical constraints to avoid misconfiguration."""
+        if not self.use_multi_db and not self.user_name:
+            raise ValueError(
+                "In single-database mode (use_multi_db=False), `user_name` must be provided for logical isolation."
+            )
+        return self
+
 
 class GraphDBConfigFactory(BaseModel):
     backend: str = Field(..., description="Backend for graph database")