feat: support nebular database

Author: tianxing02

src/memos/graph_dbs/nebular.py

@@ -281,29 +281,25 @@ def get_memory_count(self, memory_type: str) -> int:
         query += "\nRETURN COUNT(n) AS count"
 
         try:
-            print(f"\n ======> query: {query}\n")
             result = self.client.execute(query)
-            print(result.one_or_none()["count"].value)
             return result.one_or_none()["count"].value
         except Exception as e:
             logger.error(f"[get_memory_count] Failed: {e}")
             return -1
 
-    # TODO
     def count_nodes(self, scope: str) -> int:
         query = f"""
                 MATCH (n@Memory)
                 WHERE n.memory_type = {scope}
                 """
         if not self.config.use_multi_db and self.config.user_name:
             user_name = self.config.user_name
-            query += f"\nAND n.user_name = {user_name}"
+            query += f"\nAND n.user_name = '{user_name}'"
         query += "\nRETURN count(n) AS count"
 
         result = self.client.execute(query)
-        return result.one_or_none().values()["count"]
+        return result.one_or_none()["count"].value
 
-    # TODO
     def edge_exists(
         self, source_id: str, target_id: str, type: str = "ANY", direction: str = "OUTGOING"
     ) -> bool:
@@ -319,26 +315,27 @@ def edge_exists(
             True if the edge exists, otherwise False.
         """
         # Prepare the relationship pattern
-        rel = "r" if type == "ANY" else f"r:{type}"
+        rel = "r" if type == "ANY" else f"r@{type}"
 
         # Prepare the match pattern with direction
         if direction == "OUTGOING":
-            pattern = f"(a@Memory {{id: {source_id}}})-[{rel}]->(b@Memory {{id: {target_id}}})"
+            pattern = f"(a@Memory {{id: '{source_id}'}})-[{rel}]->(b@Memory {{id: '{target_id}'}})"
         elif direction == "INCOMING":
-            pattern = f"(a@Memory {{id: {source_id}}})<-[{rel}]-(b@Memory {{id: {target_id}}})"
+            pattern = f"(a@Memory {{id: '{source_id}'}})<-[{rel}]-(b@Memory {{id: '{target_id}'}})"
         elif direction == "ANY":
-            pattern = f"(a@Memory {{id: {source_id}}})-[{rel}]-(b@Memory {{id: {target_id}}})"
+            pattern = f"(a@Memory {{id: '{source_id}'}})-[{rel}]-(b@Memory {{id: '{target_id}'}})"
         else:
             raise ValueError(
                 f"Invalid direction: {direction}. Must be 'OUTGOING', 'INCOMING', or 'ANY'."
             )
         query = f"MATCH {pattern}"
         if not self.config.use_multi_db and self.config.user_name:
             user_name = self.config.user_name
-            query += f"\nWHERE a.user_name = {user_name} AND b.user_name = {user_name}"
+            query += f"\nWHERE a.user_name = '{user_name}' AND b.user_name = '{user_name}'"
         query += "\nRETURN r"
 
         # Run the Cypher query
+        print("\n ======> query: ", query)
         result = self.client.execute(query)
         return result.one_or_none().values() is not None
 
@@ -374,7 +371,6 @@ def get_node(self, id: str) -> dict[str, Any] | None:
             logger.error(f"[get_node] Failed to retrieve node '{id}': {e}")
             return None
 
-    # TODO
     def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]:
         """
         Retrieve the metadata and memory of a list of nodes.
@@ -392,14 +388,13 @@ def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]:
 
         where_user = ""
         if not self.config.use_multi_db and self.config.user_name:
-            where_user = f" AND n.user_name = {self.config.user_name}"
+            where_user = f" AND n.user_name = '{self.config.user_name}'"
 
         query = f"MATCH (n@Memory) WHERE n.id IN {ids} {where_user} RETURN n"
 
         results = self.client.execute(query)
-        return [self._parse_node(dict(record.values()["n"])) for record in results]
+        return [self._parse_node(record["n"]) for record in results]
 
-    # TODO
     def get_edges(self, id: str, type: str = "ANY", direction: str = "ANY") -> list[dict[str, str]]:
         """
         Get edges connected to a node, with optional type and direction filter.
@@ -417,28 +412,28 @@ def get_edges(self, id: str, type: str = "ANY", direction: str = "ANY") -> list[
             ]
         """
         # Build relationship type filter
-        rel_type = "" if type == "ANY" else f":{type}"
+        rel_type = "" if type == "ANY" else f"@{type}"
 
         # Build Cypher pattern based on direction
         if direction == "OUTGOING":
             pattern = f"(a@Memory)-[r{rel_type}]->(b@Memory)"
-            where_clause = f"a.id = {id}"
+            where_clause = f"a.id = '{id}'"
         elif direction == "INCOMING":
             pattern = f"(a@Memory)<-[r{rel_type}]-(b@Memory)"
-            where_clause = f"a.id = {id}"
+            where_clause = f"a.id = '{id}'"
         elif direction == "ANY":
             pattern = f"(a@Memory)-[r{rel_type}]-(b@Memory)"
             where_clause = f"a.id = {id} OR b.id = {id}"
         else:
             raise ValueError("Invalid direction. Must be 'OUTGOING', 'INCOMING', or 'ANY'.")
 
         if not self.config.use_multi_db and self.config.user_name:
-            where_clause += f" AND a.user_name = {self.config.user_name} AND b.user_name = {self.config.user_name}"
+            where_clause += f" AND a.user_name = '{self.config.user_name}' AND b.user_name = '{self.config.user_name}'"
 
         query = f"""
                     MATCH {pattern}
                     WHERE {where_clause}
-                    RETURN a.id AS from_id, b.id AS to_id, type(r) AS type
+                    RETURN a.id AS from_id, b.id AS to_id, type(r) AS edge_type
                 """
 
         result = self.client.execute(query)
@@ -448,11 +443,12 @@ def get_edges(self, id: str, type: str = "ANY", direction: str = "ANY") -> list[
                 {
                     "from": record["from_id"].value,
                     "to": record["to_id"].value,
-                    "type": record["type"].value,
+                    "type": record["edge_type"].value,
                 }
             )
         return edges
 
+    # TODO
     def get_neighbors_by_tag(
         self,
         tags: list[str],
@@ -544,7 +540,7 @@ def get_subgraph(
                    collect(EDGES(p)) AS edge_chains
             """
 
-        result = self.client.execute(gql).one_or_none()  # 执行查询
+        result = self.client.execute(gql).one_or_none()
         if not result or result.size == 0:
             return {"core_node": None, "neighbors": [], "edges": []}
 
@@ -734,7 +730,6 @@ def get_grouped_counts(
 
         return output
 
-    # TODO
     def clear(self) -> None:
         """
         Clear the entire graph if the target database exists.
@@ -745,13 +740,11 @@ def clear(self) -> None:
             else:
                 query = "MATCH (n) DETACH DELETE n"
 
-            # Step 2: Clear the graph in that database
             self.client.execute(query)
             logger.info("Cleared all nodes from database.")
 
         except Exception as e:
             logger.error(f"[ERROR] Failed to clear database: {e}")
-            raise
 
     def export_graph(self) -> dict[str, Any]:
         """

tests/graph_dbs/test_nebular.py

@@ -25,7 +25,6 @@
         "base_url": os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
     },
 }
-
 nebular_config = {
     "hosts": json.loads(os.getenv("NEBULAR_HOSTS", "localhost")),
     "user_name": os.getenv("NEBULAR_USER", "root"),
@@ -35,6 +34,8 @@
     "embedding_dimension": 3072,
     "use_multi_db": False,
 }
+
+
 embedder_config = EmbedderConfigFactory.model_validate(gpt_config)
 embedder = EmbedderFactory.from_config(embedder_config)
 
@@ -46,60 +47,95 @@ def embed_memory_item(memory: str) -> list[float]:
     return embedding_list
 
 
+now = datetime.now(timezone.utc).isoformat()
+test_node1 = TextualMemoryItem(
+    memory="This is a test node",
+    metadata=TreeNodeTextualMemoryMetadata(
+        memory_type="LongTermMemory",
+        key="Research Topic",
+        hierarchy_level="topic",
+        type="fact",
+        memory_time="2024-01-01",
+        status="activated",
+        visibility="public",
+        updated_at=now,
+        embedding=embed_memory_item("This is a test node"),
+    ),
+)
+
+test_node2 = TextualMemoryItem(
+    memory="This is another test node",
+    metadata=TreeNodeTextualMemoryMetadata(
+        memory_type="LongTermMemory",
+        key="Research Topic",
+        hierarchy_level="topic",
+        type="fact",
+        memory_time="2024-01-01",
+        status="activated",
+        visibility="public",
+        updated_at=now,
+        embedding=embed_memory_item("This is another test node"),
+    ),
+)
+
+
 def test_get_memory_count():
     config = GraphDBConfigFactory(backend="nebular", config=nebular_config)
-
     graph = GraphStoreFactory.from_config(config)
     graph.clear()
-    now = datetime.now(timezone.utc).isoformat()
 
-    # Insert memory
-    mem = TextualMemoryItem(
-        memory="Example A",
-        metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="LongTermMemory",
-            key="Research Topic",
-            hierarchy_level="topic",
-            type="fact",
-            memory_time="2024-01-01",
-            status="activated",
-            visibility="public",
-            updated_at=now,
-            embedding=embed_memory_item("Example A"),
-        ),
-    )
+    mem = test_node1
     graph.add_node(mem.id, mem.memory, mem.metadata.model_dump(exclude_none=True))
 
     count = graph.get_memory_count('"LongTermMemory"')  # quoting string literal for Cypher
     print("Memory Count:", count)
     assert count == 1
 
 
+def test_count_nodes():
+    graph = GraphStoreFactory.from_config(
+        GraphDBConfigFactory(
+            backend="nebular",
+            config=nebular_config,
+        )
+    )
+    graph.clear()
+
+    # Insert two nodes
+    for i in range(2):
+        mem = TextualMemoryItem(
+            memory=f"Memory {i}",
+            metadata=TreeNodeTextualMemoryMetadata(
+                memory_type="LongTermMemory",
+                key="Research Topic",
+                hierarchy_level="topic",
+                type="fact",
+                memory_time="2024-01-01",
+                status="activated",
+                visibility="public",
+                updated_at=now,
+                embedding=embed_memory_item(f"Memory {i}"),
+            ),
+        )
+        graph.add_node(mem.id, mem.memory, mem.metadata.model_dump(exclude_none=True))
+
+    count = graph.count_nodes('"LongTermMemory"')
+    print("Node Count:", count)
+    assert count == 2
+
+
 def test_get_nodes():
     graph = GraphStoreFactory.from_config(
         GraphDBConfigFactory(backend="nebular", config=nebular_config)
     )
     graph.clear()
 
-    now = datetime.now(timezone.utc).isoformat()
-    mem = TextualMemoryItem(
-        memory="Test node",
-        metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="LongTermMemory",
-            key="Research Topic",
-            hierarchy_level="topic",
-            type="fact",
-            memory_time="2024-01-01",
-            status="activated",
-            visibility="public",
-            updated_at=now,
-            embedding=embed_memory_item("Test node"),
-        ),
-    )
+    mem = test_node1
     graph.add_node(mem.id, mem.memory, mem.metadata.model_dump(exclude_none=True))
+
     nodes = graph.get_nodes([mem.id])
     assert len(nodes) == 1
-    assert nodes[0]["id"] == mem.id
+    assert nodes[0]["properties"]["id"] == mem.id
 
 
 def test_edge_exists():
@@ -112,8 +148,13 @@ def test_edge_exists():
         memory="Edge topic",
         metadata=TreeNodeTextualMemoryMetadata(
             memory_type="LongTermMemory",
-            key="topic",
-            updated_at=datetime.now().isoformat(),
+            key="Research Topic",
+            hierarchy_level="topic",
+            type="fact",
+            memory_time="2024-01-01",
+            status="activated",
+            visibility="public",
+            updated_at=now,
             embedding=embed_memory_item("Edge topic"),
         ),
     )
@@ -122,8 +163,13 @@ def test_edge_exists():
         memory="Edge concept",
         metadata=TreeNodeTextualMemoryMetadata(
             memory_type="LongTermMemory",
-            key="concept",
-            updated_at=datetime.now().isoformat(),
+            key="Research Topic",
+            hierarchy_level="topic",
+            type="fact",
+            memory_time="2024-01-01",
+            status="activated",
+            visibility="public",
+            updated_at=now,
             embedding=embed_memory_item("Edge concept"),
         ),
     )
@@ -144,18 +190,28 @@ def test_get_edges():
     source = TextualMemoryItem(
         memory="Source",
         metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="WorkingMemory",
-            key="src",
-            updated_at=datetime.now().isoformat(),
+            memory_type="LongTermMemory",
+            key="Research Topic",
+            hierarchy_level="topic",
+            type="fact",
+            memory_time="2024-01-01",
+            status="activated",
+            visibility="public",
+            updated_at=now,
             embedding=embed_memory_item("Source"),
         ),
     )
     target = TextualMemoryItem(
         memory="Target",
         metadata=TreeNodeTextualMemoryMetadata(
-            memory_type="WorkingMemory",
-            key="tgt",
-            updated_at=datetime.now().isoformat(),
+            memory_type="LongTermMemory",
+            key="Research Topic",
+            hierarchy_level="topic",
+            type="fact",
+            memory_time="2024-01-01",
+            status="activated",
+            visibility="public",
+            updated_at=now,
             embedding=embed_memory_item("Target"),
         ),
     )