From a0941a9298c319746a9f46700f31f0b3ade26db6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Thu, 17 Jul 2025 21:26:19 +0800 Subject: [PATCH 01/12] feat: init neo4j-community --- src/memos/graph_dbs/neo4j_community.py | 1096 ++++++++++++++++++++++++ 1 file changed, 1096 insertions(+) create mode 100644 src/memos/graph_dbs/neo4j_community.py diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py new file mode 100644 index 000000000..0dc0b48fd --- /dev/null +++ b/src/memos/graph_dbs/neo4j_community.py @@ -0,0 +1,1096 @@ +import time + +from datetime import datetime +from typing import Any, Literal + +from neo4j import GraphDatabase +from neo4j.exceptions import ClientError + +from memos.configs.graph_db import Neo4jGraphDBConfig +from memos.graph_dbs.base import BaseGraphDB +from memos.log import get_logger + + +logger = get_logger(__name__) + + +def _parse_node(node_data: dict[str, Any]) -> dict[str, Any]: + node = node_data.copy() + + # Convert Neo4j datetime to string + for time_field in ("created_at", "updated_at"): + if time_field in node and hasattr(node[time_field], "isoformat"): + node[time_field] = node[time_field].isoformat() + node.pop("user_name", None) + + return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node} + + +def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]: + node_id = item["id"] + memory = item["memory"] + metadata = item.get("metadata", {}) + return node_id, memory, metadata + + +def _prepare_node_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + """ + Ensure metadata has proper datetime fields and normalized types. + + - Fill `created_at` and `updated_at` if missing (in ISO 8601 format). + - Convert embedding to list of float if present. + """ + now = datetime.utcnow().isoformat() + + # Fill timestamps if missing + metadata.setdefault("created_at", now) + metadata.setdefault("updated_at", now) + + # Normalize embedding type + embedding = metadata.get("embedding") + if embedding and isinstance(embedding, list): + metadata["embedding"] = [float(x) for x in embedding] + + return metadata + + +class Neo4jGraphDB(BaseGraphDB): + """Neo4j-based implementation of a graph memory store.""" + + def __init__(self, config: Neo4jGraphDBConfig): + """Neo4j-based implementation of a graph memory store. + + Tenant Modes: + - use_multi_db = True: + Dedicated Database Mode (Multi-Database Multi-Tenant). + Each tenant or logical scope uses a separate Neo4j database. + `db_name` is the specific tenant database. + `user_name` can be None (optional). + + - use_multi_db = False: + Shared Database Multi-Tenant Mode. + All tenants share a single Neo4j database. + `db_name` is the shared database. + `user_name` is required to isolate each tenant's data at the node level. + All node queries will enforce `user_name` in WHERE conditions and store it in metadata, + but it will be removed automatically before returning to external consumers. + """ + + self.config = config + self.driver = GraphDatabase.driver(config.uri, auth=(config.user, config.password)) + self.db_name = config.db_name + self.user_name = config.user_name + + self.system_db_name = "system" if config.use_multi_db else config.db_name + if config.auto_create: + self._ensure_database_exists() + + # Create only if not exists + self.create_index(dimensions=config.embedding_dimension) + + def create_index( + self, + label: str = "Memory", + vector_property: str = "embedding", + dimensions: int = 1536, + index_name: str = "memory_vector_index", + ) -> None: + """ + Create the vector index for embedding and datetime indexes for created_at and updated_at fields. + """ + # Create vector index if it doesn't exist + if not self._vector_index_exists(index_name): + self._create_vector_index(label, vector_property, dimensions, index_name) + # Create indexes + self._create_basic_property_indexes() + + def get_memory_count(self, memory_type: str) -> int: + query = """ + MATCH (n:Memory) + WHERE n.memory_type = $memory_type + """ + if not self.config.use_multi_db and self.config.user_name: + query += "\nAND n.user_name = $user_name" + query += "\nRETURN COUNT(n) AS count" + with self.driver.session(database=self.db_name) as session: + result = session.run( + query, + { + "memory_type": memory_type, + "user_name": self.config.user_name if self.config.user_name else None, + }, + ) + return result.single()["count"] + + def count_nodes(self, scope: str) -> int: + query = """ + MATCH (n:Memory) + WHERE n.memory_type = $scope + """ + if not self.config.use_multi_db and self.config.user_name: + query += "\nAND n.user_name = $user_name" + query += "\nRETURN count(n) AS count" + + with self.driver.session(database=self.db_name) as session: + result = session.run( + query, + { + "scope": scope, + "user_name": self.config.user_name if self.config.user_name else None, + }, + ) + return result.single()["count"] + + def remove_oldest_memory(self, memory_type: str, keep_latest: int) -> None: + """ + Remove all WorkingMemory nodes except the latest `keep_latest` entries. + + Args: + memory_type (str): Memory type (e.g., 'WorkingMemory', 'LongTermMemory'). + keep_latest (int): Number of latest WorkingMemory entries to keep. + """ + query = f""" + MATCH (n:Memory) + WHERE n.memory_type = '{memory_type}' + """ + if not self.config.use_multi_db and self.config.user_name: + query += f"\nAND n.user_name = '{self.config.user_name}'" + + query += f""" + WITH n ORDER BY n.updated_at DESC + SKIP {keep_latest} + DETACH DELETE n + """ + with self.driver.session(database=self.db_name) as session: + session.run(query) + + def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None: + if not self.config.use_multi_db and self.config.user_name: + metadata["user_name"] = self.config.user_name + + # Safely process metadata + metadata = _prepare_node_metadata(metadata) + + # Merge node and set metadata + created_at = metadata.pop("created_at") + updated_at = metadata.pop("updated_at") + + query = """ + MERGE (n:Memory {id: $id}) + SET n.memory = $memory, + n.created_at = datetime($created_at), + n.updated_at = datetime($updated_at), + n += $metadata + """ + with self.driver.session(database=self.db_name) as session: + session.run( + query, + id=id, + memory=memory, + created_at=created_at, + updated_at=updated_at, + metadata=metadata, + ) + + def update_node(self, id: str, fields: dict[str, Any]) -> None: + """ + Update node fields in Neo4j, auto-converting `created_at` and `updated_at` to datetime type if present. + """ + fields = fields.copy() # Avoid mutating external dict + set_clauses = [] + params = {"id": id, "fields": fields} + + for time_field in ("created_at", "updated_at"): + if time_field in fields: + # Set clause like: n.created_at = datetime($created_at) + set_clauses.append(f"n.{time_field} = datetime(${time_field})") + params[time_field] = fields.pop(time_field) + + set_clauses.append("n += $fields") # Merge remaining fields + set_clause_str = ",\n ".join(set_clauses) + + query = """ + MATCH (n:Memory {id: $id}) + """ + if not self.config.use_multi_db and self.config.user_name: + query += "\nWHERE n.user_name = $user_name" + params["user_name"] = self.config.user_name + + query += f"\nSET {set_clause_str}" + + with self.driver.session(database=self.db_name) as session: + session.run(query, **params) + + def delete_node(self, id: str) -> None: + """ + Delete a node from the graph. + Args: + id: Node identifier to delete. + """ + query = "MATCH (n:Memory {id: $id})" + + params = {"id": id} + if not self.config.use_multi_db and self.config.user_name: + query += " WHERE n.user_name = $user_name" + params["user_name"] = self.config.user_name + + query += " DETACH DELETE n" + + with self.driver.session(database=self.db_name) as session: + session.run(query, **params) + + # Edge (Relationship) Management + def add_edge(self, source_id: str, target_id: str, type: str) -> None: + """ + Create an edge from source node to target node. + Args: + source_id: ID of the source node. + target_id: ID of the target node. + type: Relationship type (e.g., 'RELATE_TO', 'PARENT'). + """ + query = """ + MATCH (a:Memory {id: $source_id}) + MATCH (b:Memory {id: $target_id}) + """ + params = {"source_id": source_id, "target_id": target_id} + if not self.config.use_multi_db and self.config.user_name: + query += """ + WHERE a.user_name = $user_name AND b.user_name = $user_name + """ + params["user_name"] = self.config.user_name + + query += f"\nMERGE (a)-[:{type}]->(b)" + + with self.driver.session(database=self.db_name) as session: + session.run(query, params) + + def delete_edge(self, source_id: str, target_id: str, type: str) -> None: + """ + Delete a specific edge between two nodes. + Args: + source_id: ID of the source node. + target_id: ID of the target node. + type: Relationship type to remove. + """ + query = f""" + MATCH (a:Memory {{id: $source}}) + -[r:{type}]-> + (b:Memory {{id: $target}}) + """ + params = {"source": source_id, "target": target_id} + + if not self.config.use_multi_db and self.config.user_name: + query += "\nWHERE a.user_name = $user_name AND b.user_name = $user_name" + params["user_name"] = self.config.user_name + + query += "\nDELETE r" + + with self.driver.session(database=self.db_name) as session: + session.run(query, params) + + def edge_exists( + self, source_id: str, target_id: str, type: str = "ANY", direction: str = "OUTGOING" + ) -> bool: + """ + Check if an edge exists between two nodes. + Args: + source_id: ID of the source node. + target_id: ID of the target node. + type: Relationship type. Use "ANY" to match any relationship type. + direction: Direction of the edge. + Use "OUTGOING" (default), "INCOMING", or "ANY". + Returns: + True if the edge exists, otherwise False. + """ + # Prepare the relationship pattern + rel = "r" if type == "ANY" else f"r:{type}" + + # Prepare the match pattern with direction + if direction == "OUTGOING": + pattern = f"(a:Memory {{id: $source}})-[{rel}]->(b:Memory {{id: $target}})" + elif direction == "INCOMING": + pattern = f"(a:Memory {{id: $source}})<-[{rel}]-(b:Memory {{id: $target}})" + elif direction == "ANY": + pattern = f"(a:Memory {{id: $source}})-[{rel}]-(b:Memory {{id: $target}})" + else: + raise ValueError( + f"Invalid direction: {direction}. Must be 'OUTGOING', 'INCOMING', or 'ANY'." + ) + query = f"MATCH {pattern}" + params = {"source": source_id, "target": target_id} + + if not self.config.use_multi_db and self.config.user_name: + query += "\nWHERE a.user_name = $user_name AND b.user_name = $user_name" + params["user_name"] = self.config.user_name + + query += "\nRETURN r" + + # Run the Cypher query + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + return result.single() is not None + + # Graph Query & Reasoning + def get_node(self, id: str) -> dict[str, Any] | None: + """ + Retrieve the metadata and memory of a node. + Args: + id: Node identifier. + Returns: + Dictionary of node fields, or None if not found. + """ + where_user = "" + params = {"id": id} + if not self.config.use_multi_db and self.config.user_name: + where_user = " AND n.user_name = $user_name" + params["user_name"] = self.config.user_name + + query = f"MATCH (n:Memory) WHERE n.id = $id {where_user} RETURN n" + + with self.driver.session(database=self.db_name) as session: + record = session.run(query, params).single() + return _parse_node(dict(record["n"])) if record else None + + def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]: + """ + Retrieve the metadata and memory of a list of nodes. + Args: + ids: List of Node identifier. + Returns: + list[dict]: Parsed node records containing 'id', 'memory', and 'metadata'. + + Notes: + - Assumes all provided IDs are valid and exist. + - Returns empty list if input is empty. + """ + if not ids: + return [] + + where_user = "" + params = {"ids": ids} + + if not self.config.use_multi_db and self.config.user_name: + where_user = " AND n.user_name = $user_name" + params["user_name"] = self.config.user_name + + query = f"MATCH (n:Memory) WHERE n.id IN $ids{where_user} RETURN n" + + with self.driver.session(database=self.db_name) as session: + results = session.run(query, params) + return [_parse_node(dict(record["n"])) for record in results] + + 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. + + Args: + id: Node ID to retrieve edges for. + type: Relationship type to match, or 'ANY' to match all. + direction: 'OUTGOING', 'INCOMING', or 'ANY'. + + Returns: + List of edges: + [ + {"from": "source_id", "to": "target_id", "type": "RELATE"}, + ... + ] + """ + # Build relationship type filter + 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 = "a.id = $id" + elif direction == "INCOMING": + pattern = f"(a:Memory)<-[r{rel_type}]-(b:Memory)" + where_clause = "a.id = $id" + elif direction == "ANY": + pattern = f"(a:Memory)-[r{rel_type}]-(b:Memory)" + where_clause = "a.id = $id OR b.id = $id" + else: + raise ValueError("Invalid direction. Must be 'OUTGOING', 'INCOMING', or 'ANY'.") + + params = {"id": id} + + if not self.config.use_multi_db and self.config.user_name: + where_clause += " AND a.user_name = $user_name AND b.user_name = $user_name" + params["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 + """ + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + edges = [] + for record in result: + edges.append( + {"from": record["from_id"], "to": record["to_id"], "type": record["type"]} + ) + return edges + + def get_neighbors( + self, id: str, type: str, direction: Literal["in", "out", "both"] = "out" + ) -> list[str]: + """ + Get connected node IDs in a specific direction and relationship type. + Args: + id: Source node ID. + type: Relationship type. + direction: Edge direction to follow ('out', 'in', or 'both'). + Returns: + List of neighboring node IDs. + """ + raise NotImplementedError + + def get_neighbors_by_tag( + self, + tags: list[str], + exclude_ids: list[str], + top_k: int = 5, + min_overlap: int = 1, + ) -> list[dict[str, Any]]: + """ + Find top-K neighbor nodes with maximum tag overlap. + + Args: + tags: The list of tags to match. + exclude_ids: Node IDs to exclude (e.g., local cluster). + top_k: Max number of neighbors to return. + min_overlap: Minimum number of overlapping tags required. + + Returns: + List of dicts with node details and overlap count. + """ + where_user = "" + params = { + "tags": tags, + "exclude_ids": exclude_ids, + "min_overlap": min_overlap, + "top_k": top_k, + } + + if not self.config.use_multi_db and self.config.user_name: + where_user = "AND n.user_name = $user_name" + params["user_name"] = self.config.user_name + + query = f""" + MATCH (n:Memory) + WHERE NOT n.id IN $exclude_ids + AND n.status = 'activated' + AND n.type <> 'reasoning' + AND n.memory_type <> 'WorkingMemory' + {where_user} + WITH n, [tag IN n.tags WHERE tag IN $tags] AS overlap_tags + WHERE size(overlap_tags) >= $min_overlap + RETURN n, size(overlap_tags) AS overlap_count + ORDER BY overlap_count DESC + LIMIT $top_k + """ + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + return [_parse_node(dict(record["n"])) for record in result] + + def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]: + where_user = "" + params = {"id": id} + + if not self.config.use_multi_db and self.config.user_name: + where_user = "AND p.user_name = $user_name AND c.user_name = $user_name" + params["user_name"] = self.config.user_name + + query = f""" + MATCH (p:Memory)-[:PARENT]->(c:Memory) + WHERE p.id = $id {where_user} + RETURN c.id AS id, c.embedding AS embedding, c.memory AS memory + """ + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + return [ + {"id": r["id"], "embedding": r["embedding"], "memory": r["memory"]} for r in result + ] + + def get_path(self, source_id: str, target_id: str, max_depth: int = 3) -> list[str]: + """ + Get the path of nodes from source to target within a limited depth. + Args: + source_id: Starting node ID. + target_id: Target node ID. + max_depth: Maximum path length to traverse. + Returns: + Ordered list of node IDs along the path. + """ + raise NotImplementedError + + def get_subgraph( + self, center_id: str, depth: int = 2, center_status: str = "activated" + ) -> dict[str, Any]: + """ + Retrieve a local subgraph centered at a given node. + Args: + center_id: The ID of the center node. + depth: The hop distance for neighbors. + center_status: Required status for center node. + Returns: + { + "core_node": {...}, + "neighbors": [...], + "edges": [...] + } + """ + with self.driver.session(database=self.db_name) as session: + params = {"center_id": center_id} + center_user_clause = "" + neighbor_user_clause = "" + + if not self.config.use_multi_db and self.config.user_name: + center_user_clause = " AND center.user_name = $user_name" + neighbor_user_clause = " WHERE neighbor.user_name = $user_name" + params["user_name"] = self.config.user_name + status_clause = f" AND center.status = '{center_status}'" if center_status else "" + + query = f""" + MATCH (center:Memory) + WHERE center.id = $center_id{status_clause}{center_user_clause} + + OPTIONAL MATCH (center)-[r*1..{depth}]-(neighbor:Memory) + {neighbor_user_clause} + + WITH collect(DISTINCT center) AS centers, + collect(DISTINCT neighbor) AS neighbors, + collect(DISTINCT r) AS rels + RETURN centers, neighbors, rels + """ + record = session.run(query, params).single() + + if not record: + return {"core_node": None, "neighbors": [], "edges": []} + + centers = record["centers"] + if not centers or centers[0] is None: + return {"core_node": None, "neighbors": [], "edges": []} + + core_node = _parse_node(dict(centers[0])) + neighbors = [_parse_node(dict(n)) for n in record["neighbors"] if n] + edges = [] + for rel_chain in record["rels"]: + for rel in rel_chain: + edges.append( + { + "type": rel.type, + "source": rel.start_node["id"], + "target": rel.end_node["id"], + } + ) + + return {"core_node": core_node, "neighbors": neighbors, "edges": edges} + + def get_context_chain(self, id: str, type: str = "FOLLOWS") -> list[str]: + """ + Get the ordered context chain starting from a node, following a relationship type. + Args: + id: Starting node ID. + type: Relationship type to follow (e.g., 'FOLLOWS'). + Returns: + List of ordered node IDs in the chain. + """ + raise NotImplementedError + + # Search / recall operations + def search_by_embedding( + self, + vector: list[float], + top_k: int = 5, + scope: str | None = None, + status: str | None = None, + threshold: float | None = None, + ) -> list[dict]: + """ + Retrieve node IDs based on vector similarity. + + Args: + vector (list[float]): The embedding vector representing query semantics. + top_k (int): Number of top similar nodes to retrieve. + scope (str, optional): Memory type filter (e.g., 'WorkingMemory', 'LongTermMemory'). + status (str, optional): Node status filter (e.g., 'active', 'archived'). + If provided, restricts results to nodes with matching status. + threshold (float, optional): Minimum similarity score threshold (0 ~ 1). + + Returns: + list[dict]: A list of dicts with 'id' and 'score', ordered by similarity. + + Notes: + - This method uses Neo4j native vector indexing to search for similar nodes. + - If scope is provided, it restricts results to nodes with matching memory_type. + - If 'status' is provided, only nodes with the matching status will be returned. + - If threshold is provided, only results with score >= threshold will be returned. + - Typical use case: restrict to 'status = activated' to avoid + matching archived or merged nodes. + """ + # Build WHERE clause dynamically + where_clauses = [] + if scope: + where_clauses.append("node.memory_type = $scope") + if status: + where_clauses.append("node.status = $status") + if not self.config.use_multi_db and self.config.user_name: + where_clauses.append("node.user_name = $user_name") + + where_clause = "" + if where_clauses: + where_clause = "WHERE " + " AND ".join(where_clauses) + + query = f""" + CALL db.index.vector.queryNodes('memory_vector_index', $k, $embedding) + YIELD node, score + {where_clause} + RETURN node.id AS id, score + """ + + parameters = {"embedding": vector, "k": top_k, "scope": scope} + if scope: + parameters["scope"] = scope + if status: + parameters["status"] = status + if not self.config.use_multi_db and self.config.user_name: + parameters["user_name"] = self.config.user_name + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, parameters) + records = [{"id": record["id"], "score": record["score"]} for record in result] + + # Threshold filtering after retrieval + if threshold is not None: + records = [r for r in records if r["score"] >= threshold] + + return records + + def get_by_metadata(self, filters: list[dict[str, Any]]) -> list[str]: + """ + TODO: + 1. ADD logic: "AND" vs "OR"(support logic combination); + 2. Support nested conditional expressions; + + Retrieve node IDs that match given metadata filters. + Supports exact match. + + Args: + filters: List of filter dicts like: + [ + {"field": "key", "op": "in", "value": ["A", "B"]}, + {"field": "confidence", "op": ">=", "value": 80}, + {"field": "tags", "op": "contains", "value": "AI"}, + ... + ] + + Returns: + list[str]: Node IDs whose metadata match the filter conditions. (AND logic). + + Notes: + - Supports structured querying such as tag/category/importance/time filtering. + - Can be used for faceted recall or prefiltering before embedding rerank. + """ + where_clauses = [] + params = {} + + for i, f in enumerate(filters): + field = f["field"] + op = f.get("op", "=") + value = f["value"] + param_key = f"val{i}" + + # Build WHERE clause + if op == "=": + where_clauses.append(f"n.{field} = ${param_key}") + params[param_key] = value + elif op == "in": + where_clauses.append(f"n.{field} IN ${param_key}") + params[param_key] = value + elif op == "contains": + where_clauses.append(f"ANY(x IN ${param_key} WHERE x IN n.{field})") + params[param_key] = value + elif op == "starts_with": + where_clauses.append(f"n.{field} STARTS WITH ${param_key}") + params[param_key] = value + elif op == "ends_with": + where_clauses.append(f"n.{field} ENDS WITH ${param_key}") + params[param_key] = value + elif op in [">", ">=", "<", "<="]: + where_clauses.append(f"n.{field} {op} ${param_key}") + params[param_key] = value + else: + raise ValueError(f"Unsupported operator: {op}") + + if not self.config.use_multi_db and self.config.user_name: + where_clauses.append("n.user_name = $user_name") + params["user_name"] = self.config.user_name + + where_str = " AND ".join(where_clauses) + query = f"MATCH (n:Memory) WHERE {where_str} RETURN n.id AS id" + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + return [record["id"] for record in result] + + def get_grouped_counts( + self, + group_fields: list[str], + where_clause: str = "", + params: dict[str, Any] | None = None, + ) -> list[dict[str, Any]]: + """ + Count nodes grouped by any fields. + + Args: + group_fields (list[str]): Fields to group by, e.g., ["memory_type", "status"] + where_clause (str, optional): Extra WHERE condition. E.g., + "WHERE n.status = 'activated'" + params (dict, optional): Parameters for WHERE clause. + + Returns: + list[dict]: e.g., [{ 'memory_type': 'WorkingMemory', 'status': 'active', 'count': 10 }, ...] + """ + if not group_fields: + raise ValueError("group_fields cannot be empty") + + final_params = params.copy() if params else {} + + if not self.config.use_multi_db and self.config.user_name: + user_clause = "n.user_name = $user_name" + final_params["user_name"] = self.config.user_name + if where_clause: + where_clause = where_clause.strip() + if where_clause.upper().startswith("WHERE"): + where_clause += f" AND {user_clause}" + else: + where_clause = f"WHERE {where_clause} AND {user_clause}" + else: + where_clause = f"WHERE {user_clause}" + + # Force RETURN field AS field to guarantee key match + group_fields_cypher = ", ".join([f"n.{field} AS {field}" for field in group_fields]) + + query = f""" + MATCH (n:Memory) + {where_clause} + RETURN {group_fields_cypher}, COUNT(n) AS count + """ + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, final_params) + return [ + {**{field: record[field] for field in group_fields}, "count": record["count"]} + for record in result + ] + + # Structure Maintenance + def deduplicate_nodes(self) -> None: + """ + Deduplicate redundant or semantically similar nodes. + This typically involves identifying nodes with identical or near-identical memory. + """ + raise NotImplementedError + + def detect_conflicts(self) -> list[tuple[str, str]]: + """ + Detect conflicting nodes based on logical or semantic inconsistency. + Returns: + A list of (node_id1, node_id2) tuples that conflict. + """ + raise NotImplementedError + + def merge_nodes(self, id1: str, id2: str) -> str: + """ + Merge two similar or duplicate nodes into one. + Args: + id1: First node ID. + id2: Second node ID. + Returns: + ID of the resulting merged node. + """ + raise NotImplementedError + + # Utilities + def clear(self) -> None: + """ + Clear the entire graph if the target database exists. + """ + try: + if not self.config.use_multi_db and self.config.user_name: + query = "MATCH (n:Memory) WHERE n.user_name = $user_name DETACH DELETE n" + params = {"user_name": self.config.user_name} + else: + query = "MATCH (n) DETACH DELETE n" + params = {} + + # Step 2: Clear the graph in that database + with self.driver.session(database=self.db_name) as session: + session.run(query, params) + logger.info(f"Cleared all nodes from database '{self.db_name}'.") + + except Exception as e: + logger.error(f"[ERROR] Failed to clear database '{self.db_name}': {e}") + raise + + def export_graph(self) -> dict[str, Any]: + """ + Export all graph nodes and edges in a structured form. + + Returns: + { + "nodes": [ { "id": ..., "memory": ..., "metadata": {...} }, ... ], + "edges": [ { "source": ..., "target": ..., "type": ... }, ... ] + } + """ + with self.driver.session(database=self.db_name) as session: + # Export nodes + node_query = "MATCH (n:Memory)" + edge_query = "MATCH (a:Memory)-[r]->(b:Memory)" + params = {} + + if not self.config.use_multi_db and self.config.user_name: + node_query += " WHERE n.user_name = $user_name" + edge_query += " WHERE a.user_name = $user_name AND b.user_name = $user_name" + params["user_name"] = self.config.user_name + + node_result = session.run(f"{node_query} RETURN n", params) + nodes = [_parse_node(dict(record["n"])) for record in node_result] + + # Export edges + edge_result = session.run( + f"{edge_query} RETURN a.id AS source, b.id AS target, type(r) AS type", params + ) + edges = [ + {"source": record["source"], "target": record["target"], "type": record["type"]} + for record in edge_result + ] + + return {"nodes": nodes, "edges": edges} + + def import_graph(self, data: dict[str, Any]) -> None: + """ + Import the entire graph from a serialized dictionary. + + Args: + data: A dictionary containing all nodes and edges to be loaded. + """ + with self.driver.session(database=self.db_name) as session: + for node in data.get("nodes", []): + id, memory, metadata = _compose_node(node) + + if not self.config.use_multi_db and self.config.user_name: + metadata["user_name"] = self.config.user_name + + metadata = _prepare_node_metadata(metadata) + + # Merge node and set metadata + created_at = metadata.pop("created_at") + updated_at = metadata.pop("updated_at") + + session.run( + """ + MERGE (n:Memory {id: $id}) + SET n.memory = $memory, + n.created_at = datetime($created_at), + n.updated_at = datetime($updated_at), + n += $metadata + """, + id=id, + memory=memory, + created_at=created_at, + updated_at=updated_at, + metadata=metadata, + ) + + for edge in data.get("edges", []): + session.run( + f""" + MATCH (a:Memory {{id: $source_id}}) + MATCH (b:Memory {{id: $target_id}}) + MERGE (a)-[:{edge["type"]}]->(b) + """, + source_id=edge["source"], + target_id=edge["target"], + ) + + def get_all_memory_items(self, scope: str) -> list[dict]: + """ + Retrieve all memory items of a specific memory_type. + + Args: + scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'. + + Returns: + list[dict]: Full list of memory items under this scope. + """ + if scope not in {"WorkingMemory", "LongTermMemory", "UserMemory"}: + raise ValueError(f"Unsupported memory type scope: {scope}") + + where_clause = "WHERE n.memory_type = $scope" + params = {"scope": scope} + + if not self.config.use_multi_db and self.config.user_name: + where_clause += " AND n.user_name = $user_name" + params["user_name"] = self.config.user_name + + query = f""" + MATCH (n:Memory) + {where_clause} + RETURN n + """ + + with self.driver.session(database=self.db_name) as session: + results = session.run(query, params) + return [_parse_node(dict(record["n"])) for record in results] + + def get_structure_optimization_candidates(self, scope: str) -> list[dict]: + """ + Find nodes that are likely candidates for structure optimization: + - Isolated nodes, nodes with empty background, or nodes with exactly one child. + - Plus: the child of any parent node that has exactly one child. + """ + where_clause = """ + WHERE n.memory_type = $scope + AND n.status = 'activated' + AND NOT ( (n)-[:PARENT]->() OR ()-[:PARENT]->(n) ) + """ + params = {"scope": scope} + + if not self.config.use_multi_db and self.config.user_name: + where_clause += " AND n.user_name = $user_name" + params["user_name"] = self.config.user_name + + query = f""" + MATCH (n:Memory) + {where_clause} + RETURN n.id AS id, n AS node + """ + + with self.driver.session(database=self.db_name) as session: + results = session.run(query, params) + return [_parse_node({"id": record["id"], **dict(record["node"])}) for record in results] + + def drop_database(self) -> None: + """ + Permanently delete the entire database this instance is using. + WARNING: This operation is destructive and cannot be undone. + """ + if self.config.use_multi_db: + if self.db_name in ("system", "neo4j"): + raise ValueError(f"Refusing to drop protected database: {self.db_name}") + + with self.driver.session(database=self.system_db_name) as session: + session.run(f"DROP DATABASE {self.db_name} IF EXISTS") + print(f"Database '{self.db_name}' has been dropped.") + else: + raise ValueError( + f"Refusing to drop protected database: {self.db_name} in " + f"Shared Database Multi-Tenant mode" + ) + + def _ensure_database_exists(self): + try: + with self.driver.session(database="system") as session: + session.run(f"CREATE DATABASE `{self.db_name}` IF NOT EXISTS") + except ClientError as e: + if "ExistingDatabaseFound" in str(e): + pass # Ignore, database already exists + else: + raise + + # Wait until the database is available + for _ in range(10): + with self.driver.session(database=self.system_db_name) as session: + result = session.run( + "SHOW DATABASES YIELD name, currentStatus RETURN name, currentStatus" + ) + status_map = {r["name"]: r["currentStatus"] for r in result} + if self.db_name in status_map and status_map[self.db_name] == "online": + return + time.sleep(1) + + raise RuntimeError(f"Database {self.db_name} not ready after waiting.") + + def _vector_index_exists(self, index_name: str = "memory_vector_index") -> bool: + query = "SHOW INDEXES YIELD name WHERE name = $name RETURN name" + with self.driver.session(database=self.db_name) as session: + result = session.run(query, name=index_name) + return result.single() is not None + + def _create_vector_index( + self, label: str, vector_property: str, dimensions: int, index_name: str + ) -> None: + """ + Create a vector index for the specified property in the label. + """ + try: + query = f""" + CREATE VECTOR INDEX {index_name} IF NOT EXISTS + FOR (n:{label}) ON (n.{vector_property}) + OPTIONS {{ + indexConfig: {{ + `vector.dimensions`: {dimensions}, + `vector.similarity_function`: 'cosine' + }} + }} + """ + with self.driver.session(database=self.db_name) as session: + session.run(query) + logger.debug(f"Vector index '{index_name}' ensured.") + except Exception as e: + logger.warning(f"Failed to create vector index '{index_name}': {e}") + + def _create_basic_property_indexes(self) -> None: + """ + Create standard B-tree indexes on memory_type, created_at, + and updated_at fields. + Create standard B-tree indexes on user_name when use Shared Database + Multi-Tenant Mode + """ + try: + with self.driver.session(database=self.db_name) as session: + session.run(""" + CREATE INDEX memory_type_index IF NOT EXISTS + FOR (n:Memory) ON (n.memory_type) + """) + logger.debug("Index 'memory_type_index' ensured.") + + session.run(""" + CREATE INDEX memory_created_at_index IF NOT EXISTS + FOR (n:Memory) ON (n.created_at) + """) + logger.debug("Index 'memory_created_at_index' ensured.") + + session.run(""" + CREATE INDEX memory_updated_at_index IF NOT EXISTS + FOR (n:Memory) ON (n.updated_at) + """) + logger.debug("Index 'memory_updated_at_index' ensured.") + + if not self.config.use_multi_db and self.config.user_name: + session.run( + """ + CREATE INDEX memory_user_name_index IF NOT EXISTS + FOR (n:Memory) ON (n.user_name) + """ + ) + logger.debug("Index 'memory_user_name_index' ensured.") + except Exception as e: + logger.warning(f"Failed to create basic property indexes: {e}") + + def _index_exists(self, index_name: str) -> bool: + """ + Check if an index with the given name exists. + """ + query = "SHOW INDEXES" + with self.driver.session(database=self.db_name) as session: + result = session.run(query) + for record in result: + if record["name"] == index_name: + return True + return False From ba299d2b8c9dadb1480b0f41652cbd23c054d1fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Thu, 17 Jul 2025 21:55:26 +0800 Subject: [PATCH 02/12] feat: Inherit existing class for Community Edition database support --- src/memos/graph_dbs/neo4j_community.py | 743 +------------------------ 1 file changed, 12 insertions(+), 731 deletions(-) diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index 0dc0b48fd..3144a8734 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -1,60 +1,19 @@ import time -from datetime import datetime -from typing import Any, Literal +from typing import Any from neo4j import GraphDatabase from neo4j.exceptions import ClientError from memos.configs.graph_db import Neo4jGraphDBConfig -from memos.graph_dbs.base import BaseGraphDB +from memos.graph_dbs.neo4j import Neo4jGraphDB, _parse_node from memos.log import get_logger logger = get_logger(__name__) -def _parse_node(node_data: dict[str, Any]) -> dict[str, Any]: - node = node_data.copy() - - # Convert Neo4j datetime to string - for time_field in ("created_at", "updated_at"): - if time_field in node and hasattr(node[time_field], "isoformat"): - node[time_field] = node[time_field].isoformat() - node.pop("user_name", None) - - return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node} - - -def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]: - node_id = item["id"] - memory = item["memory"] - metadata = item.get("metadata", {}) - return node_id, memory, metadata - - -def _prepare_node_metadata(metadata: dict[str, Any]) -> dict[str, Any]: - """ - Ensure metadata has proper datetime fields and normalized types. - - - Fill `created_at` and `updated_at` if missing (in ISO 8601 format). - - Convert embedding to list of float if present. - """ - now = datetime.utcnow().isoformat() - - # Fill timestamps if missing - metadata.setdefault("created_at", now) - metadata.setdefault("updated_at", now) - - # Normalize embedding type - embedding = metadata.get("embedding") - if embedding and isinstance(embedding, list): - metadata["embedding"] = [float(x) for x in embedding] - - return metadata - - -class Neo4jGraphDB(BaseGraphDB): +class Neo4jCommunityGraphDB(Neo4jGraphDB): """Neo4j-based implementation of a graph memory store.""" def __init__(self, config: Neo4jGraphDBConfig): @@ -80,11 +39,7 @@ def __init__(self, config: Neo4jGraphDBConfig): self.driver = GraphDatabase.driver(config.uri, auth=(config.user, config.password)) self.db_name = config.db_name self.user_name = config.user_name - - self.system_db_name = "system" if config.use_multi_db else config.db_name - if config.auto_create: - self._ensure_database_exists() - + self.system_db_name = config.db_name # Create only if not exists self.create_index(dimensions=config.embedding_dimension) @@ -98,9 +53,6 @@ def create_index( """ Create the vector index for embedding and datetime indexes for created_at and updated_at fields. """ - # Create vector index if it doesn't exist - if not self._vector_index_exists(index_name): - self._create_vector_index(label, vector_property, dimensions, index_name) # Create indexes self._create_basic_property_indexes() @@ -164,337 +116,6 @@ def remove_oldest_memory(self, memory_type: str, keep_latest: int) -> None: with self.driver.session(database=self.db_name) as session: session.run(query) - def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None: - if not self.config.use_multi_db and self.config.user_name: - metadata["user_name"] = self.config.user_name - - # Safely process metadata - metadata = _prepare_node_metadata(metadata) - - # Merge node and set metadata - created_at = metadata.pop("created_at") - updated_at = metadata.pop("updated_at") - - query = """ - MERGE (n:Memory {id: $id}) - SET n.memory = $memory, - n.created_at = datetime($created_at), - n.updated_at = datetime($updated_at), - n += $metadata - """ - with self.driver.session(database=self.db_name) as session: - session.run( - query, - id=id, - memory=memory, - created_at=created_at, - updated_at=updated_at, - metadata=metadata, - ) - - def update_node(self, id: str, fields: dict[str, Any]) -> None: - """ - Update node fields in Neo4j, auto-converting `created_at` and `updated_at` to datetime type if present. - """ - fields = fields.copy() # Avoid mutating external dict - set_clauses = [] - params = {"id": id, "fields": fields} - - for time_field in ("created_at", "updated_at"): - if time_field in fields: - # Set clause like: n.created_at = datetime($created_at) - set_clauses.append(f"n.{time_field} = datetime(${time_field})") - params[time_field] = fields.pop(time_field) - - set_clauses.append("n += $fields") # Merge remaining fields - set_clause_str = ",\n ".join(set_clauses) - - query = """ - MATCH (n:Memory {id: $id}) - """ - if not self.config.use_multi_db and self.config.user_name: - query += "\nWHERE n.user_name = $user_name" - params["user_name"] = self.config.user_name - - query += f"\nSET {set_clause_str}" - - with self.driver.session(database=self.db_name) as session: - session.run(query, **params) - - def delete_node(self, id: str) -> None: - """ - Delete a node from the graph. - Args: - id: Node identifier to delete. - """ - query = "MATCH (n:Memory {id: $id})" - - params = {"id": id} - if not self.config.use_multi_db and self.config.user_name: - query += " WHERE n.user_name = $user_name" - params["user_name"] = self.config.user_name - - query += " DETACH DELETE n" - - with self.driver.session(database=self.db_name) as session: - session.run(query, **params) - - # Edge (Relationship) Management - def add_edge(self, source_id: str, target_id: str, type: str) -> None: - """ - Create an edge from source node to target node. - Args: - source_id: ID of the source node. - target_id: ID of the target node. - type: Relationship type (e.g., 'RELATE_TO', 'PARENT'). - """ - query = """ - MATCH (a:Memory {id: $source_id}) - MATCH (b:Memory {id: $target_id}) - """ - params = {"source_id": source_id, "target_id": target_id} - if not self.config.use_multi_db and self.config.user_name: - query += """ - WHERE a.user_name = $user_name AND b.user_name = $user_name - """ - params["user_name"] = self.config.user_name - - query += f"\nMERGE (a)-[:{type}]->(b)" - - with self.driver.session(database=self.db_name) as session: - session.run(query, params) - - def delete_edge(self, source_id: str, target_id: str, type: str) -> None: - """ - Delete a specific edge between two nodes. - Args: - source_id: ID of the source node. - target_id: ID of the target node. - type: Relationship type to remove. - """ - query = f""" - MATCH (a:Memory {{id: $source}}) - -[r:{type}]-> - (b:Memory {{id: $target}}) - """ - params = {"source": source_id, "target": target_id} - - if not self.config.use_multi_db and self.config.user_name: - query += "\nWHERE a.user_name = $user_name AND b.user_name = $user_name" - params["user_name"] = self.config.user_name - - query += "\nDELETE r" - - with self.driver.session(database=self.db_name) as session: - session.run(query, params) - - def edge_exists( - self, source_id: str, target_id: str, type: str = "ANY", direction: str = "OUTGOING" - ) -> bool: - """ - Check if an edge exists between two nodes. - Args: - source_id: ID of the source node. - target_id: ID of the target node. - type: Relationship type. Use "ANY" to match any relationship type. - direction: Direction of the edge. - Use "OUTGOING" (default), "INCOMING", or "ANY". - Returns: - True if the edge exists, otherwise False. - """ - # Prepare the relationship pattern - rel = "r" if type == "ANY" else f"r:{type}" - - # Prepare the match pattern with direction - if direction == "OUTGOING": - pattern = f"(a:Memory {{id: $source}})-[{rel}]->(b:Memory {{id: $target}})" - elif direction == "INCOMING": - pattern = f"(a:Memory {{id: $source}})<-[{rel}]-(b:Memory {{id: $target}})" - elif direction == "ANY": - pattern = f"(a:Memory {{id: $source}})-[{rel}]-(b:Memory {{id: $target}})" - else: - raise ValueError( - f"Invalid direction: {direction}. Must be 'OUTGOING', 'INCOMING', or 'ANY'." - ) - query = f"MATCH {pattern}" - params = {"source": source_id, "target": target_id} - - if not self.config.use_multi_db and self.config.user_name: - query += "\nWHERE a.user_name = $user_name AND b.user_name = $user_name" - params["user_name"] = self.config.user_name - - query += "\nRETURN r" - - # Run the Cypher query - with self.driver.session(database=self.db_name) as session: - result = session.run(query, params) - return result.single() is not None - - # Graph Query & Reasoning - def get_node(self, id: str) -> dict[str, Any] | None: - """ - Retrieve the metadata and memory of a node. - Args: - id: Node identifier. - Returns: - Dictionary of node fields, or None if not found. - """ - where_user = "" - params = {"id": id} - if not self.config.use_multi_db and self.config.user_name: - where_user = " AND n.user_name = $user_name" - params["user_name"] = self.config.user_name - - query = f"MATCH (n:Memory) WHERE n.id = $id {where_user} RETURN n" - - with self.driver.session(database=self.db_name) as session: - record = session.run(query, params).single() - return _parse_node(dict(record["n"])) if record else None - - def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]: - """ - Retrieve the metadata and memory of a list of nodes. - Args: - ids: List of Node identifier. - Returns: - list[dict]: Parsed node records containing 'id', 'memory', and 'metadata'. - - Notes: - - Assumes all provided IDs are valid and exist. - - Returns empty list if input is empty. - """ - if not ids: - return [] - - where_user = "" - params = {"ids": ids} - - if not self.config.use_multi_db and self.config.user_name: - where_user = " AND n.user_name = $user_name" - params["user_name"] = self.config.user_name - - query = f"MATCH (n:Memory) WHERE n.id IN $ids{where_user} RETURN n" - - with self.driver.session(database=self.db_name) as session: - results = session.run(query, params) - return [_parse_node(dict(record["n"])) for record in results] - - 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. - - Args: - id: Node ID to retrieve edges for. - type: Relationship type to match, or 'ANY' to match all. - direction: 'OUTGOING', 'INCOMING', or 'ANY'. - - Returns: - List of edges: - [ - {"from": "source_id", "to": "target_id", "type": "RELATE"}, - ... - ] - """ - # Build relationship type filter - 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 = "a.id = $id" - elif direction == "INCOMING": - pattern = f"(a:Memory)<-[r{rel_type}]-(b:Memory)" - where_clause = "a.id = $id" - elif direction == "ANY": - pattern = f"(a:Memory)-[r{rel_type}]-(b:Memory)" - where_clause = "a.id = $id OR b.id = $id" - else: - raise ValueError("Invalid direction. Must be 'OUTGOING', 'INCOMING', or 'ANY'.") - - params = {"id": id} - - if not self.config.use_multi_db and self.config.user_name: - where_clause += " AND a.user_name = $user_name AND b.user_name = $user_name" - params["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 - """ - - with self.driver.session(database=self.db_name) as session: - result = session.run(query, params) - edges = [] - for record in result: - edges.append( - {"from": record["from_id"], "to": record["to_id"], "type": record["type"]} - ) - return edges - - def get_neighbors( - self, id: str, type: str, direction: Literal["in", "out", "both"] = "out" - ) -> list[str]: - """ - Get connected node IDs in a specific direction and relationship type. - Args: - id: Source node ID. - type: Relationship type. - direction: Edge direction to follow ('out', 'in', or 'both'). - Returns: - List of neighboring node IDs. - """ - raise NotImplementedError - - def get_neighbors_by_tag( - self, - tags: list[str], - exclude_ids: list[str], - top_k: int = 5, - min_overlap: int = 1, - ) -> list[dict[str, Any]]: - """ - Find top-K neighbor nodes with maximum tag overlap. - - Args: - tags: The list of tags to match. - exclude_ids: Node IDs to exclude (e.g., local cluster). - top_k: Max number of neighbors to return. - min_overlap: Minimum number of overlapping tags required. - - Returns: - List of dicts with node details and overlap count. - """ - where_user = "" - params = { - "tags": tags, - "exclude_ids": exclude_ids, - "min_overlap": min_overlap, - "top_k": top_k, - } - - if not self.config.use_multi_db and self.config.user_name: - where_user = "AND n.user_name = $user_name" - params["user_name"] = self.config.user_name - - query = f""" - MATCH (n:Memory) - WHERE NOT n.id IN $exclude_ids - AND n.status = 'activated' - AND n.type <> 'reasoning' - AND n.memory_type <> 'WorkingMemory' - {where_user} - WITH n, [tag IN n.tags WHERE tag IN $tags] AS overlap_tags - WHERE size(overlap_tags) >= $min_overlap - RETURN n, size(overlap_tags) AS overlap_count - ORDER BY overlap_count DESC - LIMIT $top_k - """ - - with self.driver.session(database=self.db_name) as session: - result = session.run(query, params) - return [_parse_node(dict(record["n"])) for record in result] - def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]: where_user = "" params = {"id": id} @@ -515,18 +136,6 @@ def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]: {"id": r["id"], "embedding": r["embedding"], "memory": r["memory"]} for r in result ] - def get_path(self, source_id: str, target_id: str, max_depth: int = 3) -> list[str]: - """ - Get the path of nodes from source to target within a limited depth. - Args: - source_id: Starting node ID. - target_id: Target node ID. - max_depth: Maximum path length to traverse. - Returns: - Ordered list of node IDs along the path. - """ - raise NotImplementedError - def get_subgraph( self, center_id: str, depth: int = 2, center_status: str = "activated" ) -> dict[str, Any]: @@ -590,17 +199,6 @@ def get_subgraph( return {"core_node": core_node, "neighbors": neighbors, "edges": edges} - def get_context_chain(self, id: str, type: str = "FOLLOWS") -> list[str]: - """ - Get the ordered context chain starting from a node, following a relationship type. - Args: - id: Starting node ID. - type: Relationship type to follow (e.g., 'FOLLOWS'). - Returns: - List of ordered node IDs in the chain. - """ - raise NotImplementedError - # Search / recall operations def search_by_embedding( self, @@ -632,291 +230,11 @@ def search_by_embedding( - Typical use case: restrict to 'status = activated' to avoid matching archived or merged nodes. """ - # Build WHERE clause dynamically - where_clauses = [] - if scope: - where_clauses.append("node.memory_type = $scope") - if status: - where_clauses.append("node.status = $status") - if not self.config.use_multi_db and self.config.user_name: - where_clauses.append("node.user_name = $user_name") - - where_clause = "" - if where_clauses: - where_clause = "WHERE " + " AND ".join(where_clauses) - - query = f""" - CALL db.index.vector.queryNodes('memory_vector_index', $k, $embedding) - YIELD node, score - {where_clause} - RETURN node.id AS id, score - """ - - parameters = {"embedding": vector, "k": top_k, "scope": scope} - if scope: - parameters["scope"] = scope - if status: - parameters["status"] = status - if not self.config.use_multi_db and self.config.user_name: - parameters["user_name"] = self.config.user_name - - with self.driver.session(database=self.db_name) as session: - result = session.run(query, parameters) - records = [{"id": record["id"], "score": record["score"]} for record in result] + # TODO + from your_vector_index import vector_index - # Threshold filtering after retrieval - if threshold is not None: - records = [r for r in records if r["score"] >= threshold] - - return records - - def get_by_metadata(self, filters: list[dict[str, Any]]) -> list[str]: - """ - TODO: - 1. ADD logic: "AND" vs "OR"(support logic combination); - 2. Support nested conditional expressions; - - Retrieve node IDs that match given metadata filters. - Supports exact match. - - Args: - filters: List of filter dicts like: - [ - {"field": "key", "op": "in", "value": ["A", "B"]}, - {"field": "confidence", "op": ">=", "value": 80}, - {"field": "tags", "op": "contains", "value": "AI"}, - ... - ] - - Returns: - list[str]: Node IDs whose metadata match the filter conditions. (AND logic). - - Notes: - - Supports structured querying such as tag/category/importance/time filtering. - - Can be used for faceted recall or prefiltering before embedding rerank. - """ - where_clauses = [] - params = {} - - for i, f in enumerate(filters): - field = f["field"] - op = f.get("op", "=") - value = f["value"] - param_key = f"val{i}" - - # Build WHERE clause - if op == "=": - where_clauses.append(f"n.{field} = ${param_key}") - params[param_key] = value - elif op == "in": - where_clauses.append(f"n.{field} IN ${param_key}") - params[param_key] = value - elif op == "contains": - where_clauses.append(f"ANY(x IN ${param_key} WHERE x IN n.{field})") - params[param_key] = value - elif op == "starts_with": - where_clauses.append(f"n.{field} STARTS WITH ${param_key}") - params[param_key] = value - elif op == "ends_with": - where_clauses.append(f"n.{field} ENDS WITH ${param_key}") - params[param_key] = value - elif op in [">", ">=", "<", "<="]: - where_clauses.append(f"n.{field} {op} ${param_key}") - params[param_key] = value - else: - raise ValueError(f"Unsupported operator: {op}") - - if not self.config.use_multi_db and self.config.user_name: - where_clauses.append("n.user_name = $user_name") - params["user_name"] = self.config.user_name - - where_str = " AND ".join(where_clauses) - query = f"MATCH (n:Memory) WHERE {where_str} RETURN n.id AS id" - - with self.driver.session(database=self.db_name) as session: - result = session.run(query, params) - return [record["id"] for record in result] - - def get_grouped_counts( - self, - group_fields: list[str], - where_clause: str = "", - params: dict[str, Any] | None = None, - ) -> list[dict[str, Any]]: - """ - Count nodes grouped by any fields. - - Args: - group_fields (list[str]): Fields to group by, e.g., ["memory_type", "status"] - where_clause (str, optional): Extra WHERE condition. E.g., - "WHERE n.status = 'activated'" - params (dict, optional): Parameters for WHERE clause. - - Returns: - list[dict]: e.g., [{ 'memory_type': 'WorkingMemory', 'status': 'active', 'count': 10 }, ...] - """ - if not group_fields: - raise ValueError("group_fields cannot be empty") - - final_params = params.copy() if params else {} - - if not self.config.use_multi_db and self.config.user_name: - user_clause = "n.user_name = $user_name" - final_params["user_name"] = self.config.user_name - if where_clause: - where_clause = where_clause.strip() - if where_clause.upper().startswith("WHERE"): - where_clause += f" AND {user_clause}" - else: - where_clause = f"WHERE {where_clause} AND {user_clause}" - else: - where_clause = f"WHERE {user_clause}" - - # Force RETURN field AS field to guarantee key match - group_fields_cypher = ", ".join([f"n.{field} AS {field}" for field in group_fields]) - - query = f""" - MATCH (n:Memory) - {where_clause} - RETURN {group_fields_cypher}, COUNT(n) AS count - """ - - with self.driver.session(database=self.db_name) as session: - result = session.run(query, final_params) - return [ - {**{field: record[field] for field in group_fields}, "count": record["count"]} - for record in result - ] - - # Structure Maintenance - def deduplicate_nodes(self) -> None: - """ - Deduplicate redundant or semantically similar nodes. - This typically involves identifying nodes with identical or near-identical memory. - """ - raise NotImplementedError - - def detect_conflicts(self) -> list[tuple[str, str]]: - """ - Detect conflicting nodes based on logical or semantic inconsistency. - Returns: - A list of (node_id1, node_id2) tuples that conflict. - """ - raise NotImplementedError - - def merge_nodes(self, id1: str, id2: str) -> str: - """ - Merge two similar or duplicate nodes into one. - Args: - id1: First node ID. - id2: Second node ID. - Returns: - ID of the resulting merged node. - """ - raise NotImplementedError - - # Utilities - def clear(self) -> None: - """ - Clear the entire graph if the target database exists. - """ - try: - if not self.config.use_multi_db and self.config.user_name: - query = "MATCH (n:Memory) WHERE n.user_name = $user_name DETACH DELETE n" - params = {"user_name": self.config.user_name} - else: - query = "MATCH (n) DETACH DELETE n" - params = {} - - # Step 2: Clear the graph in that database - with self.driver.session(database=self.db_name) as session: - session.run(query, params) - logger.info(f"Cleared all nodes from database '{self.db_name}'.") - - except Exception as e: - logger.error(f"[ERROR] Failed to clear database '{self.db_name}': {e}") - raise - - def export_graph(self) -> dict[str, Any]: - """ - Export all graph nodes and edges in a structured form. - - Returns: - { - "nodes": [ { "id": ..., "memory": ..., "metadata": {...} }, ... ], - "edges": [ { "source": ..., "target": ..., "type": ... }, ... ] - } - """ - with self.driver.session(database=self.db_name) as session: - # Export nodes - node_query = "MATCH (n:Memory)" - edge_query = "MATCH (a:Memory)-[r]->(b:Memory)" - params = {} - - if not self.config.use_multi_db and self.config.user_name: - node_query += " WHERE n.user_name = $user_name" - edge_query += " WHERE a.user_name = $user_name AND b.user_name = $user_name" - params["user_name"] = self.config.user_name - - node_result = session.run(f"{node_query} RETURN n", params) - nodes = [_parse_node(dict(record["n"])) for record in node_result] - - # Export edges - edge_result = session.run( - f"{edge_query} RETURN a.id AS source, b.id AS target, type(r) AS type", params - ) - edges = [ - {"source": record["source"], "target": record["target"], "type": record["type"]} - for record in edge_result - ] - - return {"nodes": nodes, "edges": edges} - - def import_graph(self, data: dict[str, Any]) -> None: - """ - Import the entire graph from a serialized dictionary. - - Args: - data: A dictionary containing all nodes and edges to be loaded. - """ - with self.driver.session(database=self.db_name) as session: - for node in data.get("nodes", []): - id, memory, metadata = _compose_node(node) - - if not self.config.use_multi_db and self.config.user_name: - metadata["user_name"] = self.config.user_name - - metadata = _prepare_node_metadata(metadata) - - # Merge node and set metadata - created_at = metadata.pop("created_at") - updated_at = metadata.pop("updated_at") - - session.run( - """ - MERGE (n:Memory {id: $id}) - SET n.memory = $memory, - n.created_at = datetime($created_at), - n.updated_at = datetime($updated_at), - n += $metadata - """, - id=id, - memory=memory, - created_at=created_at, - updated_at=updated_at, - metadata=metadata, - ) - - for edge in data.get("edges", []): - session.run( - f""" - MATCH (a:Memory {{id: $source_id}}) - MATCH (b:Memory {{id: $target_id}}) - MERGE (a)-[:{edge["type"]}]->(b) - """, - source_id=edge["source"], - target_id=edge["target"], - ) + results = vector_index.query(vector, top_k=top_k) + return [{"id": item.id, "score": item.score} for item in results] def get_all_memory_items(self, scope: str) -> list[dict]: """ @@ -980,18 +298,10 @@ def drop_database(self) -> None: Permanently delete the entire database this instance is using. WARNING: This operation is destructive and cannot be undone. """ - if self.config.use_multi_db: - if self.db_name in ("system", "neo4j"): - raise ValueError(f"Refusing to drop protected database: {self.db_name}") - - with self.driver.session(database=self.system_db_name) as session: - session.run(f"DROP DATABASE {self.db_name} IF EXISTS") - print(f"Database '{self.db_name}' has been dropped.") - else: - raise ValueError( - f"Refusing to drop protected database: {self.db_name} in " - f"Shared Database Multi-Tenant mode" - ) + raise ValueError( + f"Refusing to drop protected database: {self.db_name} in " + f"Shared Database Multi-Tenant mode" + ) def _ensure_database_exists(self): try: @@ -1016,35 +326,6 @@ def _ensure_database_exists(self): raise RuntimeError(f"Database {self.db_name} not ready after waiting.") - def _vector_index_exists(self, index_name: str = "memory_vector_index") -> bool: - query = "SHOW INDEXES YIELD name WHERE name = $name RETURN name" - with self.driver.session(database=self.db_name) as session: - result = session.run(query, name=index_name) - return result.single() is not None - - def _create_vector_index( - self, label: str, vector_property: str, dimensions: int, index_name: str - ) -> None: - """ - Create a vector index for the specified property in the label. - """ - try: - query = f""" - CREATE VECTOR INDEX {index_name} IF NOT EXISTS - FOR (n:{label}) ON (n.{vector_property}) - OPTIONS {{ - indexConfig: {{ - `vector.dimensions`: {dimensions}, - `vector.similarity_function`: 'cosine' - }} - }} - """ - with self.driver.session(database=self.db_name) as session: - session.run(query) - logger.debug(f"Vector index '{index_name}' ensured.") - except Exception as e: - logger.warning(f"Failed to create vector index '{index_name}': {e}") - def _create_basic_property_indexes(self) -> None: """ Create standard B-tree indexes on memory_type, created_at, From 7094fb2526a77afdf30891b066cf4f2adf50f3f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Fri, 18 Jul 2025 12:17:59 +0800 Subject: [PATCH 03/12] feat: add outer vector db --- src/memos/graph_dbs/neo4j.py | 41 ++-- src/memos/graph_dbs/neo4j_community.py | 186 +++++++++++------- .../tree_text_memory/retrieve/recall.py | 1 - 3 files changed, 138 insertions(+), 90 deletions(-) diff --git a/src/memos/graph_dbs/neo4j.py b/src/memos/graph_dbs/neo4j.py index 0dc0b48fd..ad26b182e 100644 --- a/src/memos/graph_dbs/neo4j.py +++ b/src/memos/graph_dbs/neo4j.py @@ -14,18 +14,6 @@ logger = get_logger(__name__) -def _parse_node(node_data: dict[str, Any]) -> dict[str, Any]: - node = node_data.copy() - - # Convert Neo4j datetime to string - for time_field in ("created_at", "updated_at"): - if time_field in node and hasattr(node[time_field], "isoformat"): - node[time_field] = node[time_field].isoformat() - node.pop("user_name", None) - - return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node} - - def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]: node_id = item["id"] memory = item["memory"] @@ -349,7 +337,7 @@ def get_node(self, id: str) -> dict[str, Any] | None: with self.driver.session(database=self.db_name) as session: record = session.run(query, params).single() - return _parse_node(dict(record["n"])) if record else None + return self._parse_node(dict(record["n"])) if record else None def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]: """ @@ -377,7 +365,7 @@ def get_nodes(self, ids: list[str]) -> list[dict[str, Any]]: with self.driver.session(database=self.db_name) as session: results = session.run(query, params) - return [_parse_node(dict(record["n"])) for record in results] + return [self._parse_node(dict(record["n"])) for record in results] def get_edges(self, id: str, type: str = "ANY", direction: str = "ANY") -> list[dict[str, str]]: """ @@ -493,7 +481,7 @@ def get_neighbors_by_tag( with self.driver.session(database=self.db_name) as session: result = session.run(query, params) - return [_parse_node(dict(record["n"])) for record in result] + return [self._parse_node(dict(record["n"])) for record in result] def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]: where_user = "" @@ -575,8 +563,8 @@ def get_subgraph( if not centers or centers[0] is None: return {"core_node": None, "neighbors": [], "edges": []} - core_node = _parse_node(dict(centers[0])) - neighbors = [_parse_node(dict(n)) for n in record["neighbors"] if n] + core_node = self._parse_node(dict(centers[0])) + neighbors = [self._parse_node(dict(n)) for n in record["neighbors"] if n] edges = [] for rel_chain in record["rels"]: for rel in rel_chain: @@ -859,7 +847,7 @@ def export_graph(self) -> dict[str, Any]: params["user_name"] = self.config.user_name node_result = session.run(f"{node_query} RETURN n", params) - nodes = [_parse_node(dict(record["n"])) for record in node_result] + nodes = [self._parse_node(dict(record["n"])) for record in node_result] # Export edges edge_result = session.run( @@ -946,7 +934,7 @@ def get_all_memory_items(self, scope: str) -> list[dict]: with self.driver.session(database=self.db_name) as session: results = session.run(query, params) - return [_parse_node(dict(record["n"])) for record in results] + return [self._parse_node(dict(record["n"])) for record in results] def get_structure_optimization_candidates(self, scope: str) -> list[dict]: """ @@ -973,7 +961,9 @@ def get_structure_optimization_candidates(self, scope: str) -> list[dict]: with self.driver.session(database=self.db_name) as session: results = session.run(query, params) - return [_parse_node({"id": record["id"], **dict(record["node"])}) for record in results] + return [ + self._parse_node({"id": record["id"], **dict(record["node"])}) for record in results + ] def drop_database(self) -> None: """ @@ -1094,3 +1084,14 @@ def _index_exists(self, index_name: str) -> bool: if record["name"] == index_name: return True return False + + def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]: + node = node_data.copy() + + # Convert Neo4j datetime to string + for time_field in ("created_at", "updated_at"): + if time_field in node and hasattr(node[time_field], "isoformat"): + node[time_field] = node[time_field].isoformat() + node.pop("user_name", None) + + return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node} diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index 3144a8734..b770d498c 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -1,47 +1,34 @@ -import time - from typing import Any -from neo4j import GraphDatabase -from neo4j.exceptions import ClientError - from memos.configs.graph_db import Neo4jGraphDBConfig -from memos.graph_dbs.neo4j import Neo4jGraphDB, _parse_node +from memos.graph_dbs.neo4j import Neo4jGraphDB, _prepare_node_metadata from memos.log import get_logger +from memos.vec_dbs.factory import VecDBFactory +from memos.vec_dbs.item import VecDBItem logger = get_logger(__name__) class Neo4jCommunityGraphDB(Neo4jGraphDB): - """Neo4j-based implementation of a graph memory store.""" + """ + Neo4j Community Edition graph memory store. + + Note: + This class avoids Enterprise-only features: + - No multi-database support + - No vector index + - No CREATE DATABASE + """ def __init__(self, config: Neo4jGraphDBConfig): - """Neo4j-based implementation of a graph memory store. - - Tenant Modes: - - use_multi_db = True: - Dedicated Database Mode (Multi-Database Multi-Tenant). - Each tenant or logical scope uses a separate Neo4j database. - `db_name` is the specific tenant database. - `user_name` can be None (optional). - - - use_multi_db = False: - Shared Database Multi-Tenant Mode. - All tenants share a single Neo4j database. - `db_name` is the shared database. - `user_name` is required to isolate each tenant's data at the node level. - All node queries will enforce `user_name` in WHERE conditions and store it in metadata, - but it will be removed automatically before returning to external consumers. - """ + assert config.auto_create is False + assert config.use_multi_db is False + # Call parent init + super().__init__(config) - self.config = config - self.driver = GraphDatabase.driver(config.uri, auth=(config.user, config.password)) - self.db_name = config.db_name - self.user_name = config.user_name - self.system_db_name = config.db_name - # Create only if not exists - self.create_index(dimensions=config.embedding_dimension) + # Init vector database + self.vec_db = VecDBFactory.from_config(config.vec_config) def create_index( self, @@ -116,6 +103,54 @@ def remove_oldest_memory(self, memory_type: str, keep_latest: int) -> None: with self.driver.session(database=self.db_name) as session: session.run(query) + def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None: + # Safely process metadata + metadata = _prepare_node_metadata(metadata) + + # Extract required fields + embedding = metadata.pop("embedding", None) + if embedding is None: + raise ValueError(f"Missing 'embedding' in metadata for node {id}") + + # Merge node and set metadata + created_at = metadata.pop("created_at") + updated_at = metadata.pop("updated_at") + vector_sync_status = "success" + + try: + # Write to Vector DB + item = VecDBItem( + id=id, + vector=embedding, + payload={ + "memory": memory, + "metadata": metadata, + "vector_sync": vector_sync_status, + }, + ) + self.vec_db.add([item]) + except Exception as e: + logger.warning(f"[VecDB] Vector insert failed for node {id}: {e}") + vector_sync_status = "failed" + + metadata["vector_sync"] = vector_sync_status + query = """ + MERGE (n:Memory {id: $id}) + SET n.memory = $memory, + n.created_at = datetime($created_at), + n.updated_at = datetime($updated_at), + n += $metadata + """ + with self.driver.session(database=self.db_name) as session: + session.run( + query, + id=id, + memory=memory, + created_at=created_at, + updated_at=updated_at, + metadata=metadata, + ) + def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]: where_user = "" params = {"id": id} @@ -184,8 +219,8 @@ def get_subgraph( if not centers or centers[0] is None: return {"core_node": None, "neighbors": [], "edges": []} - core_node = _parse_node(dict(centers[0])) - neighbors = [_parse_node(dict(n)) for n in record["neighbors"] if n] + core_node = self._parse_node(dict(centers[0])) + neighbors = [self._parse_node(dict(n)) for n in record["neighbors"] if n] edges = [] for rel_chain in record["rels"]: for rel in rel_chain: @@ -209,32 +244,42 @@ def search_by_embedding( threshold: float | None = None, ) -> list[dict]: """ - Retrieve node IDs based on vector similarity. + Retrieve node IDs based on vector similarity using external vector DB. Args: vector (list[float]): The embedding vector representing query semantics. top_k (int): Number of top similar nodes to retrieve. scope (str, optional): Memory type filter (e.g., 'WorkingMemory', 'LongTermMemory'). - status (str, optional): Node status filter (e.g., 'active', 'archived'). - If provided, restricts results to nodes with matching status. + status (str, optional): Node status filter (e.g., 'activated', 'archived'). threshold (float, optional): Minimum similarity score threshold (0 ~ 1). Returns: list[dict]: A list of dicts with 'id' and 'score', ordered by similarity. Notes: - - This method uses Neo4j native vector indexing to search for similar nodes. - - If scope is provided, it restricts results to nodes with matching memory_type. - - If 'status' is provided, only nodes with the matching status will be returned. - - If threshold is provided, only results with score >= threshold will be returned. - - Typical use case: restrict to 'status = activated' to avoid - matching archived or merged nodes. + - This method uses an external vector database (not Neo4j) to perform the search. + - If 'scope' is provided, it restricts results to nodes with matching memory_type. + - If 'status' is provided, it further filters nodes by status. + - If 'threshold' is provided, only results with score >= threshold will be returned. + - The returned IDs can be used to fetch full node data from Neo4j if needed. """ - # TODO - from your_vector_index import vector_index + # Build VecDB filter + vec_filter = {} + if scope: + vec_filter["metadata.memory_type"] = scope + if status: + vec_filter["metadata.status"] = status + vec_filter["metadata.vector_sync"] = "success" - results = vector_index.query(vector, top_k=top_k) - return [{"id": item.id, "score": item.score} for item in results] + # Perform vector search + results = self.vec_db.search(query_vector=vector, top_k=top_k, filter=vec_filter) + + # Filter by threshold + if threshold is not None: + results = [r for r in results if r.score is None or r.score >= threshold] + + # Return consistent format + return [{"id": r.id, "score": r.score} for r in results] def get_all_memory_items(self, scope: str) -> list[dict]: """ @@ -264,7 +309,7 @@ def get_all_memory_items(self, scope: str) -> list[dict]: with self.driver.session(database=self.db_name) as session: results = session.run(query, params) - return [_parse_node(dict(record["n"])) for record in results] + return [self._parse_node(dict(record["n"])) for record in results] def get_structure_optimization_candidates(self, scope: str) -> list[dict]: """ @@ -291,7 +336,9 @@ def get_structure_optimization_candidates(self, scope: str) -> list[dict]: with self.driver.session(database=self.db_name) as session: results = session.run(query, params) - return [_parse_node({"id": record["id"], **dict(record["node"])}) for record in results] + return [ + self._parse_node({"id": record["id"], **dict(record["node"])}) for record in results + ] def drop_database(self) -> None: """ @@ -303,28 +350,9 @@ def drop_database(self) -> None: f"Shared Database Multi-Tenant mode" ) + # Avoid enterprise feature def _ensure_database_exists(self): - try: - with self.driver.session(database="system") as session: - session.run(f"CREATE DATABASE `{self.db_name}` IF NOT EXISTS") - except ClientError as e: - if "ExistingDatabaseFound" in str(e): - pass # Ignore, database already exists - else: - raise - - # Wait until the database is available - for _ in range(10): - with self.driver.session(database=self.system_db_name) as session: - result = session.run( - "SHOW DATABASES YIELD name, currentStatus RETURN name, currentStatus" - ) - status_map = {r["name"]: r["currentStatus"] for r in result} - if self.db_name in status_map and status_map[self.db_name] == "online": - return - time.sleep(1) - - raise RuntimeError(f"Database {self.db_name} not ready after waiting.") + pass def _create_basic_property_indexes(self) -> None: """ @@ -375,3 +403,23 @@ def _index_exists(self, index_name: str) -> bool: if record["name"] == index_name: return True return False + + def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]: + """Parse Neo4j node and optionally fetch embedding from vector DB.""" + node = node_data.copy() + + # Convert Neo4j datetime to string + for time_field in ("created_at", "updated_at"): + if time_field in node and hasattr(node[time_field], "isoformat"): + node[time_field] = node[time_field].isoformat() + node.pop("user_name", None) + + new_node = {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node} + try: + vec_item = self.vec_db.get_by_id(new_node["id"]) + if vec_item and vec_item.vector: + new_node["embedding"] = vec_item.vector + except Exception as e: + logger.warning(f"Failed to fetch vector for node {new_node['id']}: {e}") + new_node["embedding"] = None + return new_node diff --git a/src/memos/memories/textual/tree_text_memory/retrieve/recall.py b/src/memos/memories/textual/tree_text_memory/retrieve/recall.py index 8a3fc4b5c..36a0b5fee 100644 --- a/src/memos/memories/textual/tree_text_memory/retrieve/recall.py +++ b/src/memos/memories/textual/tree_text_memory/retrieve/recall.py @@ -56,7 +56,6 @@ def retrieve( # Step 3: Merge and deduplicate results combined = {item.id: item for item in graph_results + vector_results} - # Debug: 打印在 graph_results 中但不在 combined 中的 id graph_ids = {item.id for item in graph_results} combined_ids = set(combined.keys()) lost_ids = graph_ids - combined_ids From e520770044e3a177b342b6636d666ea1bdfc6b03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Sat, 19 Jul 2025 18:49:05 +0800 Subject: [PATCH 04/12] feat: finish neo4j community --- examples/basic_modules/neo4j_example.py | 198 +++++++++++++++++++++++- src/memos/configs/graph_db.py | 25 +++ src/memos/graph_dbs/factory.py | 2 + src/memos/graph_dbs/neo4j_community.py | 196 ++++------------------- src/memos/vec_dbs/base.py | 4 + 5 files changed, 250 insertions(+), 175 deletions(-) diff --git a/examples/basic_modules/neo4j_example.py b/examples/basic_modules/neo4j_example.py index bf7bdf9c5..60d28d139 100644 --- a/examples/basic_modules/neo4j_example.py +++ b/examples/basic_modules/neo4j_example.py @@ -42,7 +42,6 @@ def example_multi_db(db_name: str = "paper"): metadata=TreeNodeTextualMemoryMetadata( memory_type="LongTermMemory", key="Multi-UAV Long-Term Coverage", - value="Research topic on distributed multi-agent UAV navigation and coverage", hierarchy_level="topic", type="fact", memory_time="2024-01-01", @@ -74,7 +73,6 @@ def example_multi_db(db_name: str = "paper"): metadata=TreeNodeTextualMemoryMetadata( memory_type="LongTermMemory", key="Reward Function Design", - value="Combines coverage, energy efficiency, and overlap penalty", hierarchy_level="concept", type="fact", memory_time="2024-01-01", @@ -99,7 +97,6 @@ def example_multi_db(db_name: str = "paper"): metadata=TreeNodeTextualMemoryMetadata( memory_type="LongTermMemory", key="Energy Model", - value="Includes communication and motion energy consumption", hierarchy_level="concept", type="fact", memory_time="2024-01-01", @@ -122,7 +119,6 @@ def example_multi_db(db_name: str = "paper"): metadata=TreeNodeTextualMemoryMetadata( memory_type="LongTermMemory", key="Coverage Metrics", - value="CT and FT used for long-term area and fairness evaluation", hierarchy_level="concept", type="fact", memory_time="2024-01-01", @@ -161,7 +157,6 @@ def example_multi_db(db_name: str = "paper"): 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", @@ -186,7 +181,6 @@ def example_multi_db(db_name: str = "paper"): metadata=TreeNodeTextualMemoryMetadata( memory_type="LongTermMemory", key="Energy Cost Components", - value="Includes movement and communication energy", hierarchy_level="fact", type="fact", memory_time="2024-01-01", @@ -211,7 +205,6 @@ def example_multi_db(db_name: str = "paper"): metadata=TreeNodeTextualMemoryMetadata( memory_type="LongTermMemory", key="CT and FT Definition", - value="CT: total coverage duration; FT: fairness index", hierarchy_level="fact", type="fact", memory_time="2024-01-01", @@ -347,9 +340,198 @@ def example_shared_db(db_name: str = "shared-traval-group"): print(graph_alice.get_node(node["id"])) +def run_user_session( + user_name: str, + db_name: str, + topic_text: str, + concept_texts: list[str], + fact_texts: list[str], + community: bool = False, +): + print(f"\n=== {user_name} starts building their memory graph ===") + + # Manually initialize correct GraphDB class + if community: + config = GraphDBConfigFactory( + backend="neo4j-community", + config={ + "uri": "bolt://localhost:7687", + "user": "neo4j", + "password": "12345678", + "db_name": "neo4j", + "user_name": user_name, + "use_multi_db": False, + "auto_create": False, # Neo4j Community does not allow auto DB creation + "embedding_dimension": 768, + "vec_config": { # Pass nested config to initialize external vector DB + "backend": "qdrant", + "config": { + "collection_name": f"{user_name}_memory", + "vector_dimension": 768, + "distance_metric": "cosine", + "path": "./qdrant_data", # or any valid path + }, + }, + }, + ) + else: + 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, + }, + ) + graph = GraphStoreFactory.from_config(config) + + # Start with a clean slate for this user + graph.clear() + + now = datetime.utcnow().isoformat() + + # === Step 1: Create a root topic node (e.g., user's research focus) === + topic = TextualMemoryItem( + memory=topic_text, + 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(topic_text), + ), + ) + graph.add_node(topic.id, topic.memory, topic.metadata.model_dump(exclude_none=True)) + + # === Step 2: Create two concept nodes linked to the topic === + concept_items = [] + for i, text in enumerate(concept_texts): + concept = TextualMemoryItem( + memory=text, + metadata=TreeNodeTextualMemoryMetadata( + memory_type="LongTermMemory", + key=f"Concept {i + 1}", + hierarchy_level="concept", + type="fact", + memory_time="2024-01-01", + status="activated", + visibility="public", + updated_at=now, + embedding=embed_memory_item(text), + tags=["concept"], + confidence=90 + i, + ), + ) + graph.add_node(concept.id, concept.memory, concept.metadata.model_dump(exclude_none=True)) + graph.add_edge(topic.id, concept.id, type="PARENT") + concept_items.append(concept) + + # === Step 3: Create supporting facts under each concept === + for i, text in enumerate(fact_texts): + fact = TextualMemoryItem( + memory=text, + metadata=TreeNodeTextualMemoryMetadata( + memory_type="WorkingMemory", + key=f"Fact {i + 1}", + hierarchy_level="fact", + type="fact", + memory_time="2024-01-01", + status="activated", + visibility="public", + updated_at=now, + embedding=embed_memory_item(text), + confidence=85.0, + tags=["fact"], + ), + ) + graph.add_node(fact.id, fact.memory, fact.metadata.model_dump(exclude_none=True)) + graph.add_edge(concept_items[i % len(concept_items)].id, fact.id, type="PARENT") + + # === Step 4: Retrieve memory using semantic search === + vector = embed_memory_item("How is memory retrieved?") + search_result = graph.search_by_embedding(vector, top_k=2) + for r in search_result: + print("🔍 Search result:", graph.get_node(r["id"])["memory"]) + + # === Step 5: Tag-based neighborhood discovery === + neighbors = graph.get_neighbors_by_tag(["concept"], exclude_ids=[], top_k=2) + print("📎 Tag-related nodes:", [neighbor["memory"] for neighbor in neighbors]) + + # === Step 6: Retrieve children (facts) of first concept === + children = graph.get_children_with_embeddings(concept_items[0].id) + print("📍 Children of concept:", [child["memory"] for child in children]) + + # === Step 7: Export a local subgraph and grouped statistics === + subgraph = graph.get_subgraph(topic.id, depth=2) + print("📌 Subgraph node count:", len(subgraph["neighbors"])) + + stats = graph.get_grouped_counts(["memory_type", "status"]) + print("📊 Grouped counts:", stats) + + # === Step 8: Demonstrate updates and cleanup === + graph.update_node(concept_items[0].id, {"confidence": 99.0}) + graph.remove_oldest_memory("WorkingMemory", keep_latest=1) + graph.delete_edge(topic.id, concept_items[0].id, type="PARENT") + graph.delete_node(concept_items[1].id) + + # === Step 9: Export and re-import the entire graph structure === + exported = graph.export_graph() + graph.import_graph(exported) + print("📦 Graph exported and re-imported, total nodes:", len(exported["nodes"])) + + +def example_complex_shared_db(db_name: str = "shared-traval-group-complex", community=False): + # User 1: Alice explores structured memory for LLMs + run_user_session( + user_name="alice", + db_name=db_name, + topic_text="Alice studies structured memory and long-term memory optimization in LLMs.", + concept_texts=[ + "Short-term memory can be simulated using WorkingMemory blocks.", + "A structured memory graph improves retrieval precision for agents.", + ], + fact_texts=[ + "Embedding search is used to find semantically similar memory items.", + "User memories are stored as node-edge structures that support hierarchical reasoning.", + ], + community=community, + ) + + # User 2: Bob focuses on GNN-based reasoning + run_user_session( + user_name="bob", + db_name=db_name, + topic_text="Bob investigates how graph neural networks can support knowledge reasoning.", + concept_texts=[ + "GNNs can learn high-order relations among entities.", + "Attention mechanisms in graphs improve inference precision.", + ], + fact_texts=[ + "GAT outperforms GCN in graph classification tasks.", + "Multi-hop reasoning helps answer complex queries.", + ], + community=community, + ) + + if __name__ == "__main__": print("\n=== Example: Multi-DB ===") example_multi_db(db_name="paper") print("\n=== Example: Single-DB ===") - example_shared_db(db_name="shared-traval-group11") + example_shared_db(db_name="shared-traval-group") + + print("\n=== Example: Single-DB-Complex ===") + example_complex_shared_db(db_name="shared-traval-group-complex-new") + + print("\n=== Example: Single-Community-DB-Complex ===") + example_complex_shared_db(db_name="neo4j", community=True) diff --git a/src/memos/configs/graph_db.py b/src/memos/configs/graph_db.py index 21baf4cd9..cca93fede 100644 --- a/src/memos/configs/graph_db.py +++ b/src/memos/configs/graph_db.py @@ -3,6 +3,7 @@ from pydantic import BaseModel, Field, field_validator, model_validator from memos.configs.base import BaseConfig +from memos.configs.vec_db import VectorDBConfigFactory class BaseGraphDBConfig(BaseConfig): @@ -79,12 +80,36 @@ def validate_config(self): return self +class Neo4jCommunityGraphDBConfig(Neo4jGraphDBConfig): + """ + Community edition config for Neo4j. + + Notes: + - Must set `use_multi_db = False` + - Must provide `user_name` for logical isolation + - Embedding vector DB config is required + """ + + vec_config: VectorDBConfigFactory = Field( + ..., description="Vector DB config for embedding search" + ) + + @model_validator(mode="after") + def validate_community(self): + if self.use_multi_db: + raise ValueError("Neo4j Community Edition does not support use_multi_db=True.") + if not self.user_name: + raise ValueError("Neo4j Community config requires user_name for logical isolation.") + return self + + class GraphDBConfigFactory(BaseModel): backend: str = Field(..., description="Backend for graph database") config: dict[str, Any] = Field(..., description="Configuration for the graph database backend") backend_to_class: ClassVar[dict[str, Any]] = { "neo4j": Neo4jGraphDBConfig, + "neo4j-community": Neo4jCommunityGraphDBConfig, } @field_validator("backend") diff --git a/src/memos/graph_dbs/factory.py b/src/memos/graph_dbs/factory.py index c100270da..c4365d16a 100644 --- a/src/memos/graph_dbs/factory.py +++ b/src/memos/graph_dbs/factory.py @@ -3,6 +3,7 @@ from memos.configs.graph_db import GraphDBConfigFactory from memos.graph_dbs.base import BaseGraphDB from memos.graph_dbs.neo4j import Neo4jGraphDB +from memos.graph_dbs.neo4j_community import Neo4jCommunityGraphDB class GraphStoreFactory(BaseGraphDB): @@ -10,6 +11,7 @@ class GraphStoreFactory(BaseGraphDB): backend_to_class: ClassVar[dict[str, Any]] = { "neo4j": Neo4jGraphDB, + "neo4j-community": Neo4jCommunityGraphDB, } @classmethod diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index b770d498c..533cb13cb 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -43,67 +43,10 @@ def create_index( # Create indexes self._create_basic_property_indexes() - def get_memory_count(self, memory_type: str) -> int: - query = """ - MATCH (n:Memory) - WHERE n.memory_type = $memory_type - """ - if not self.config.use_multi_db and self.config.user_name: - query += "\nAND n.user_name = $user_name" - query += "\nRETURN COUNT(n) AS count" - with self.driver.session(database=self.db_name) as session: - result = session.run( - query, - { - "memory_type": memory_type, - "user_name": self.config.user_name if self.config.user_name else None, - }, - ) - return result.single()["count"] - - def count_nodes(self, scope: str) -> int: - query = """ - MATCH (n:Memory) - WHERE n.memory_type = $scope - """ - if not self.config.use_multi_db and self.config.user_name: - query += "\nAND n.user_name = $user_name" - query += "\nRETURN count(n) AS count" - - with self.driver.session(database=self.db_name) as session: - result = session.run( - query, - { - "scope": scope, - "user_name": self.config.user_name if self.config.user_name else None, - }, - ) - return result.single()["count"] - - def remove_oldest_memory(self, memory_type: str, keep_latest: int) -> None: - """ - Remove all WorkingMemory nodes except the latest `keep_latest` entries. - - Args: - memory_type (str): Memory type (e.g., 'WorkingMemory', 'LongTermMemory'). - keep_latest (int): Number of latest WorkingMemory entries to keep. - """ - query = f""" - MATCH (n:Memory) - WHERE n.memory_type = '{memory_type}' - """ + def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None: if not self.config.use_multi_db and self.config.user_name: - query += f"\nAND n.user_name = '{self.config.user_name}'" + metadata["user_name"] = self.config.user_name - query += f""" - WITH n ORDER BY n.updated_at DESC - SKIP {keep_latest} - DETACH DELETE n - """ - with self.driver.session(database=self.db_name) as session: - session.run(query) - - def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None: # Safely process metadata metadata = _prepare_node_metadata(metadata) @@ -124,8 +67,8 @@ def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None: vector=embedding, payload={ "memory": memory, - "metadata": metadata, "vector_sync": vector_sync_status, + **metadata, # unpack all metadata keys to top-level }, ) self.vec_db.add([item]) @@ -162,77 +105,22 @@ def get_children_with_embeddings(self, id: str) -> list[dict[str, Any]]: query = f""" MATCH (p:Memory)-[:PARENT]->(c:Memory) WHERE p.id = $id {where_user} - RETURN c.id AS id, c.embedding AS embedding, c.memory AS memory + RETURN c.id AS id, c.memory AS memory """ with self.driver.session(database=self.db_name) as session: result = session.run(query, params) - return [ - {"id": r["id"], "embedding": r["embedding"], "memory": r["memory"]} for r in result - ] + child_nodes = [{"id": r["id"], "memory": r["memory"]} for r in result] - def get_subgraph( - self, center_id: str, depth: int = 2, center_status: str = "activated" - ) -> dict[str, Any]: - """ - Retrieve a local subgraph centered at a given node. - Args: - center_id: The ID of the center node. - depth: The hop distance for neighbors. - center_status: Required status for center node. - Returns: - { - "core_node": {...}, - "neighbors": [...], - "edges": [...] - } - """ - with self.driver.session(database=self.db_name) as session: - params = {"center_id": center_id} - center_user_clause = "" - neighbor_user_clause = "" - - if not self.config.use_multi_db and self.config.user_name: - center_user_clause = " AND center.user_name = $user_name" - neighbor_user_clause = " WHERE neighbor.user_name = $user_name" - params["user_name"] = self.config.user_name - status_clause = f" AND center.status = '{center_status}'" if center_status else "" - - query = f""" - MATCH (center:Memory) - WHERE center.id = $center_id{status_clause}{center_user_clause} - - OPTIONAL MATCH (center)-[r*1..{depth}]-(neighbor:Memory) - {neighbor_user_clause} - - WITH collect(DISTINCT center) AS centers, - collect(DISTINCT neighbor) AS neighbors, - collect(DISTINCT r) AS rels - RETURN centers, neighbors, rels - """ - record = session.run(query, params).single() - - if not record: - return {"core_node": None, "neighbors": [], "edges": []} - - centers = record["centers"] - if not centers or centers[0] is None: - return {"core_node": None, "neighbors": [], "edges": []} - - core_node = self._parse_node(dict(centers[0])) - neighbors = [self._parse_node(dict(n)) for n in record["neighbors"] if n] - edges = [] - for rel_chain in record["rels"]: - for rel in rel_chain: - edges.append( - { - "type": rel.type, - "source": rel.start_node["id"], - "target": rel.end_node["id"], - } - ) + # Get embeddings from vector DB + ids = [n["id"] for n in child_nodes] + vec_items = {v.id: v.vector for v in self.vec_db.get_by_ids(ids)} - return {"core_node": core_node, "neighbors": neighbors, "edges": edges} + # Merge results + for node in child_nodes: + node["embedding"] = vec_items.get(node["id"]) + + return child_nodes # Search / recall operations def search_by_embedding( @@ -266,10 +154,10 @@ def search_by_embedding( # Build VecDB filter vec_filter = {} if scope: - vec_filter["metadata.memory_type"] = scope + vec_filter["memory_type"] = scope if status: - vec_filter["metadata.status"] = status - vec_filter["metadata.vector_sync"] = "success" + vec_filter["status"] = status + vec_filter["vector_sync"] = "success" # Perform vector search results = self.vec_db.search(query_vector=vector, top_k=top_k, filter=vec_filter) @@ -311,34 +199,20 @@ def get_all_memory_items(self, scope: str) -> list[dict]: results = session.run(query, params) return [self._parse_node(dict(record["n"])) for record in results] - def get_structure_optimization_candidates(self, scope: str) -> list[dict]: + def clear(self) -> None: """ - Find nodes that are likely candidates for structure optimization: - - Isolated nodes, nodes with empty background, or nodes with exactly one child. - - Plus: the child of any parent node that has exactly one child. + Clear the entire graph if the target database exists. """ - where_clause = """ - WHERE n.memory_type = $scope - AND n.status = 'activated' - AND NOT ( (n)-[:PARENT]->() OR ()-[:PARENT]->(n) ) - """ - params = {"scope": scope} - - if not self.config.use_multi_db and self.config.user_name: - where_clause += " AND n.user_name = $user_name" - params["user_name"] = self.config.user_name + # Step 1: clear Neo4j part via parent logic + super().clear() - query = f""" - MATCH (n:Memory) - {where_clause} - RETURN n.id AS id, n AS node - """ - - with self.driver.session(database=self.db_name) as session: - results = session.run(query, params) - return [ - self._parse_node({"id": record["id"], **dict(record["node"])}) for record in results - ] + # Step2: Clear the vector db + try: + items = self.vec_db.get_by_filter({"user_name": self.config.user_name}) + self.vec_db.delete([item.id for item in items]) + logger.info(f"Cleared {len(items)} vectors for user '{self.config.user_name}'.") + except Exception as e: + logger.warning(f"Failed to clear vector DB for user '{self.config.user_name}': {e}") def drop_database(self) -> None: """ @@ -392,18 +266,6 @@ def _create_basic_property_indexes(self) -> None: except Exception as e: logger.warning(f"Failed to create basic property indexes: {e}") - def _index_exists(self, index_name: str) -> bool: - """ - Check if an index with the given name exists. - """ - query = "SHOW INDEXES" - with self.driver.session(database=self.db_name) as session: - result = session.run(query) - for record in result: - if record["name"] == index_name: - return True - return False - def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]: """Parse Neo4j node and optionally fetch embedding from vector DB.""" node = node_data.copy() @@ -418,8 +280,8 @@ def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]: try: vec_item = self.vec_db.get_by_id(new_node["id"]) if vec_item and vec_item.vector: - new_node["embedding"] = vec_item.vector + new_node["metadata"]["embedding"] = vec_item.vector except Exception as e: logger.warning(f"Failed to fetch vector for node {new_node['id']}: {e}") - new_node["embedding"] = None + new_node["metadata"]["embedding"] = None return new_node diff --git a/src/memos/vec_dbs/base.py b/src/memos/vec_dbs/base.py index 2ffa3b539..08919823d 100644 --- a/src/memos/vec_dbs/base.py +++ b/src/memos/vec_dbs/base.py @@ -55,6 +55,10 @@ def search( def get_by_id(self, id: str) -> VecDBItem | None: """Get an item from the vector database.""" + @abstractmethod + def get_by_ids(self, ids: list[str]) -> list[VecDBItem]: + """Get multiple items by their IDs.""" + @abstractmethod def get_by_filter(self, filter: dict[str, Any]) -> list[VecDBItem]: """ From ce4f3b5042b2b470f62b1ebda84448e5aa00076c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Sat, 19 Jul 2025 19:50:54 +0800 Subject: [PATCH 05/12] feat: add tree community config example --- .../data/config/tree_config_community.json | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 examples/data/config/tree_config_community.json diff --git a/examples/data/config/tree_config_community.json b/examples/data/config/tree_config_community.json new file mode 100644 index 000000000..bd6325a6b --- /dev/null +++ b/examples/data/config/tree_config_community.json @@ -0,0 +1,49 @@ +{ + "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-community", + "config": { + "uri": "bolt://localhost:7687", + "user": "neo4j", + "password": "12345678", + "db_name": "neo4j", + "user_name": "alice", + "use_multi_db": false, + "auto_create": false, + "embedding_dimension": 768, + "vec_config": { + "backend": "qdrant", + "config": { + "collection_name": "shared-db-community-20250719", + "vector_dimension": 768, + "distance_metric": "cosine", + "path": "./qdrant_data" + } + } + } + }, + "reorganize": false +} From c22ed164a44e3e88237c21b4ec31fd8a3558d608 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Sat, 19 Jul 2025 23:15:02 +0800 Subject: [PATCH 06/12] fix: ensure outer db for neo4j-community is server mode --- examples/basic_modules/neo4j_example.py | 16 ++++++++++------ src/memos/graph_dbs/neo4j_community.py | 1 + 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/examples/basic_modules/neo4j_example.py b/examples/basic_modules/neo4j_example.py index 60d28d139..082ad8c3e 100644 --- a/examples/basic_modules/neo4j_example.py +++ b/examples/basic_modules/neo4j_example.py @@ -358,18 +358,21 @@ def run_user_session( "uri": "bolt://localhost:7687", "user": "neo4j", "password": "12345678", - "db_name": "neo4j", + "db_name": db_name, "user_name": user_name, "use_multi_db": False, "auto_create": False, # Neo4j Community does not allow auto DB creation "embedding_dimension": 768, - "vec_config": { # Pass nested config to initialize external vector DB + "vec_config": { + # Pass nested config to initialize external vector DB + # If you use qdrant, please use Server instead of local mode. "backend": "qdrant", "config": { - "collection_name": f"{user_name}_memory", + "collection_name": "neo4j_vec_db", "vector_dimension": 768, "distance_metric": "cosine", - "path": "./qdrant_data", # or any valid path + "host": "localhost", + "port": 6333, }, }, }, @@ -460,7 +463,8 @@ def run_user_session( vector = embed_memory_item("How is memory retrieved?") search_result = graph.search_by_embedding(vector, top_k=2) for r in search_result: - print("🔍 Search result:", graph.get_node(r["id"])["memory"]) + node = graph.get_node(r["id"]) + print("🔍 Search result:", node["memory"]) # === Step 5: Tag-based neighborhood discovery === neighbors = graph.get_neighbors_by_tag(["concept"], exclude_ids=[], top_k=2) @@ -534,4 +538,4 @@ def example_complex_shared_db(db_name: str = "shared-traval-group-complex", comm example_complex_shared_db(db_name="shared-traval-group-complex-new") print("\n=== Example: Single-Community-DB-Complex ===") - example_complex_shared_db(db_name="neo4j", community=True) + example_complex_shared_db(db_name="paper", community=True) diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index 533cb13cb..3718f152d 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -158,6 +158,7 @@ def search_by_embedding( if status: vec_filter["status"] = status vec_filter["vector_sync"] = "success" + vec_filter["user_name"] = self.config.user_name # Perform vector search results = self.vec_db.search(query_vector=vector, top_k=top_k, filter=vec_filter) From 435b9764734dc24482da73c14ee4c0731f7fd0da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Sat, 19 Jul 2025 23:18:15 +0800 Subject: [PATCH 07/12] feat: update tree-config-community --- examples/data/config/tree_config_community.json | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/data/config/tree_config_community.json b/examples/data/config/tree_config_community.json index bd6325a6b..cbd2c8a07 100644 --- a/examples/data/config/tree_config_community.json +++ b/examples/data/config/tree_config_community.json @@ -37,13 +37,14 @@ "vec_config": { "backend": "qdrant", "config": { - "collection_name": "shared-db-community-20250719", + "collection_name": "neo4j_vec_db", "vector_dimension": 768, "distance_metric": "cosine", - "path": "./qdrant_data" + "host": "localhost", + "port": 6333 } } } }, - "reorganize": false + "reorganize": true } From 835051fd398039c1a17ed7b09de97886c98633de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Sat, 19 Jul 2025 23:26:29 +0800 Subject: [PATCH 08/12] fix: bug for qdrant delete --- src/memos/graph_dbs/neo4j_community.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index 3718f152d..605853a75 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -210,8 +210,11 @@ def clear(self) -> None: # Step2: Clear the vector db try: items = self.vec_db.get_by_filter({"user_name": self.config.user_name}) - self.vec_db.delete([item.id for item in items]) - logger.info(f"Cleared {len(items)} vectors for user '{self.config.user_name}'.") + if items: + self.vec_db.delete([item.id for item in items]) + logger.info(f"Cleared {len(items)} vectors for user '{self.config.user_name}'.") + else: + logger.info(f"No vectors to clear for user '{self.config.user_name}'.") except Exception as e: logger.warning(f"Failed to clear vector DB for user '{self.config.user_name}': {e}") From 521b27f50acd7c15f97b589a0b4b98a65d196a1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Sun, 20 Jul 2025 00:17:05 +0800 Subject: [PATCH 09/12] fix: create index for outer_db --- src/memos/graph_dbs/neo4j_community.py | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index 605853a75..a33abb5d9 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -24,11 +24,10 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB): def __init__(self, config: Neo4jGraphDBConfig): assert config.auto_create is False assert config.use_multi_db is False - # Call parent init - super().__init__(config) - # Init vector database self.vec_db = VecDBFactory.from_config(config.vec_config) + # Call parent init + super().__init__(config) def create_index( self, @@ -239,6 +238,7 @@ def _create_basic_property_indexes(self) -> None: Create standard B-tree indexes on user_name when use Shared Database Multi-Tenant Mode """ + # Step 1: Neo4j indexes try: with self.driver.session(database=self.db_name) as session: session.run(""" @@ -270,6 +270,25 @@ def _create_basic_property_indexes(self) -> None: except Exception as e: logger.warning(f"Failed to create basic property indexes: {e}") + # Step 2: Qdrant payload indexes + try: + # Qdrant supports `create_payload_index`, which is idempotent + self.vec_db.client.create_payload_index( + collection_name=self.vec_db.config.collection_name, + field_name="user_name", + field_schema="keyword", + ) + logger.debug("Qdrant payload index on 'user_name' ensured.") + + self.vec_db.client.create_payload_index( + collection_name=self.vec_db.config.collection_name, + field_name="memory_type", + field_schema="keyword", + ) + logger.debug("Qdrant payload index on 'memory_type' ensured.") + except Exception as e: + logger.warning(f"Failed to create Qdrant payload indexes: {e}") + def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]: """Parse Neo4j node and optionally fetch embedding from vector DB.""" node = node_data.copy() From 2b3af237f00a71ab55bad25e8f86a58438131829 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Sun, 20 Jul 2025 00:22:06 +0800 Subject: [PATCH 10/12] feat: add ensure payload index exists in qdrant --- src/memos/graph_dbs/neo4j_community.py | 22 ++++++---------------- src/memos/vec_dbs/base.py | 8 ++++++++ src/memos/vec_dbs/qdrant.py | 19 +++++++++++++++++++ 3 files changed, 33 insertions(+), 16 deletions(-) diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index a33abb5d9..2b6e672c5 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -270,24 +270,14 @@ def _create_basic_property_indexes(self) -> None: except Exception as e: logger.warning(f"Failed to create basic property indexes: {e}") - # Step 2: Qdrant payload indexes + # Step 2: VectorDB indexes try: - # Qdrant supports `create_payload_index`, which is idempotent - self.vec_db.client.create_payload_index( - collection_name=self.vec_db.config.collection_name, - field_name="user_name", - field_schema="keyword", - ) - logger.debug("Qdrant payload index on 'user_name' ensured.") - - self.vec_db.client.create_payload_index( - collection_name=self.vec_db.config.collection_name, - field_name="memory_type", - field_schema="keyword", - ) - logger.debug("Qdrant payload index on 'memory_type' ensured.") + if hasattr(self.vec_db, "ensure_payload_indexes"): + self.vec_db.ensure_payload_indexes(["user_name", "memory_type"]) + else: + logger.debug("VecDB does not support payload index creation; skipping.") except Exception as e: - logger.warning(f"Failed to create Qdrant payload indexes: {e}") + logger.warning(f"Failed to create VecDB payload indexes: {e}") def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]: """Parse Neo4j node and optionally fetch embedding from vector DB.""" diff --git a/src/memos/vec_dbs/base.py b/src/memos/vec_dbs/base.py index 08919823d..ee1bfb3ca 100644 --- a/src/memos/vec_dbs/base.py +++ b/src/memos/vec_dbs/base.py @@ -107,3 +107,11 @@ def upsert(self, data: list[VecDBItem | dict[str, Any]]) -> None: @abstractmethod def delete(self, ids: list[str]) -> None: """Delete items from the vector database.""" + + @abstractmethod + def ensure_payload_indexes(self, fields: list[str]) -> None: + """ + Create payload indexes for specified fields in the collection. + Args: + fields (list[str]): List of field names to index (as keyword). + """ diff --git a/src/memos/vec_dbs/qdrant.py b/src/memos/vec_dbs/qdrant.py index c37f113bd..a0ebf1d80 100644 --- a/src/memos/vec_dbs/qdrant.py +++ b/src/memos/vec_dbs/qdrant.py @@ -278,6 +278,25 @@ def update(self, id: str, data: VecDBItem | dict[str, Any]) -> None: collection_name=self.config.collection_name, payload=data.payload, points=[id] ) + def ensure_payload_indexes(self, fields: list[str]) -> None: + """ + Create payload indexes for specified fields in the collection. + This is idempotent: it will skip if index already exists. + + Args: + fields (list[str]): List of field names to index (as keyword). + """ + for field in fields: + try: + self.client.create_payload_index( + collection_name=self.config.collection_name, + field_name=field, + field_schema="keyword", # Could be extended in future + ) + logger.debug(f"Qdrant payload index on '{field}' ensured.") + except Exception as e: + logger.warning(f"Failed to create payload index on '{field}': {e}") + def upsert(self, data: list[VecDBItem | dict[str, Any]]) -> None: """ Add or update data in the vector database. From 3c8f991c995661de04fa1c3893908bb034f66ab9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Mon, 21 Jul 2025 11:51:57 +0800 Subject: [PATCH 11/12] feat: create index for qdrant --- src/memos/graph_dbs/neo4j_community.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index 2b6e672c5..98d9723bb 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -273,7 +273,7 @@ def _create_basic_property_indexes(self) -> None: # Step 2: VectorDB indexes try: if hasattr(self.vec_db, "ensure_payload_indexes"): - self.vec_db.ensure_payload_indexes(["user_name", "memory_type"]) + self.vec_db.ensure_payload_indexes(["user_name", "memory_type", "status"]) else: logger.debug("VecDB does not support payload index creation; skipping.") except Exception as e: From b18b9c79f4f835ef1db84b121d0c9ec3b7466c4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B8=AD=E9=98=B3=E9=98=B3?= Date: Mon, 21 Jul 2025 12:08:42 +0800 Subject: [PATCH 12/12] feat: add simple_openapi_memos_neo4j_community --- .../simple_openapi_memos_neo4j_community.py | 315 ++++++++++++++++++ 1 file changed, 315 insertions(+) create mode 100644 examples/mem_os/simple_openapi_memos_neo4j_community.py diff --git a/examples/mem_os/simple_openapi_memos_neo4j_community.py b/examples/mem_os/simple_openapi_memos_neo4j_community.py new file mode 100644 index 000000000..aad1b8c77 --- /dev/null +++ b/examples/mem_os/simple_openapi_memos_neo4j_community.py @@ -0,0 +1,315 @@ +import os +import time +import uuid + +from datetime import datetime + +from dotenv import load_dotenv + +from memos.configs.mem_cube import GeneralMemCubeConfig +from memos.configs.mem_os import MOSConfig +from memos.mem_cube.general import GeneralMemCube +from memos.mem_os.main import MOS + + +load_dotenv() + +# 1. Create MOS Config and set openai config +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to create MOS configuration...") +start_time = time.time() + +user_name = str(uuid.uuid4()) +print(user_name) + +# 1.1 Set openai config +openapi_config = { + "model_name_or_path": "gpt-4o-mini", + "temperature": 0.8, + "max_tokens": 1024, + "top_p": 0.9, + "top_k": 50, + "remove_think_prefix": True, + "api_key": os.getenv("OPENAI_API_KEY", "sk-xxxxx"), + "api_base": os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"), +} +embedder_config = { + "backend": "universal_api", + "config": { + "provider": "openai", + "api_key": os.getenv("OPENAI_API_KEY", "sk-xxxxx"), + "model_name_or_path": "text-embedding-3-large", + "base_url": os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"), + }, +} +EMBEDDING_DIMENSION = 3072 + +# 1.2 Set neo4j config +neo4j_uri = os.getenv("NEO4J_URI", "bolt://localhost:7687") + +# 1.3 Create MOS Config +config = { + "user_id": user_name, + "chat_model": { + "backend": "openai", + "config": openapi_config, + }, + "mem_reader": { + "backend": "simple_struct", + "config": { + "llm": { + "backend": "openai", + "config": openapi_config, + }, + "embedder": embedder_config, + "chunker": { + "backend": "sentence", + "config": { + "tokenizer_or_token_counter": "gpt2", + "chunk_size": 512, + "chunk_overlap": 128, + "min_sentences_per_chunk": 1, + }, + }, + }, + }, + "max_turns_window": 20, + "top_k": 5, + "enable_textual_memory": True, + "enable_activation_memory": False, + "enable_parametric_memory": False, +} + +mos_config = MOSConfig(**config) +# you can set PRO_MODE to True to enable CoT enhancement mos_config.PRO_MODE = True +mos = MOS(mos_config) + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] MOS configuration created successfully, time elapsed: {time.time() - start_time:.2f}s\n" +) + +# 2. Initialize memory cube +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to initialize MemCube configuration...") +start_time = time.time() + +config = GeneralMemCubeConfig.model_validate( + { + "user_id": user_name, + "cube_id": f"{user_name}", + "text_mem": { + "backend": "tree_text", + "config": { + "extractor_llm": { + "backend": "openai", + "config": openapi_config, + }, + "dispatcher_llm": { + "backend": "openai", + "config": openapi_config, + }, + "embedder": embedder_config, + "graph_db": { + "backend": "neo4j-community", + "config": { + "uri": neo4j_uri, + "user": "neo4j", + "password": "12345678", + "db_name": "neo4j", + "user_name": "alice", + "use_multi_db": False, + "auto_create": False, + "embedding_dimension": EMBEDDING_DIMENSION, + "vec_config": { + "backend": "qdrant", + "config": { + "collection_name": "neo4j_vec_db", + "vector_dimension": EMBEDDING_DIMENSION, + "distance_metric": "cosine", + "host": "localhost", + "port": 6333, + }, + }, + }, + }, + "reorganize": True, + }, + }, + "act_mem": {}, + "para_mem": {}, + }, +) + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] MemCube configuration initialization completed, time elapsed: {time.time() - start_time:.2f}s\n" +) + +# 3. Initialize the MemCube with the configuration +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to create MemCube instance...") +start_time = time.time() + +mem_cube = GeneralMemCube(config) +try: + mem_cube.dump(f"/tmp/{user_name}/") + print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] MemCube created and saved successfully, time elapsed: {time.time() - start_time:.2f}s\n" + ) +except Exception as e: + print( + f"❌ [{datetime.now().strftime('%H:%M:%S')}] MemCube save failed: {e}, time elapsed: {time.time() - start_time:.2f}s\n" + ) + +# 4. Register the MemCube +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to register MemCube...") +start_time = time.time() + +mos.register_mem_cube(f"/tmp/{user_name}", mem_cube_id=user_name) + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] MemCube registration completed, time elapsed: {time.time() - start_time:.2f}s\n" +) + +# 5. Add, get, search memory +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to add single memory...") +start_time = time.time() + +mos.add(memory_content="I like playing football.") + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] Single memory added successfully, time elapsed: {time.time() - start_time:.2f}s" +) + +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to get all memories...") +start_time = time.time() + +get_all_results = mos.get_all() + + +# Filter out embedding fields, keeping only necessary fields +def filter_memory_data(memories_data): + filtered_data = {} + for key, value in memories_data.items(): + if key == "text_mem": + filtered_data[key] = [] + for mem_group in value: + # Check if it's the new data structure (list of TextualMemoryItem objects) + if "memories" in mem_group and isinstance(mem_group["memories"], list): + # New data structure: directly a list of TextualMemoryItem objects + filtered_memories = [] + for memory_item in mem_group["memories"]: + # Create filtered dictionary + filtered_item = { + "id": memory_item.id, + "memory": memory_item.memory, + "metadata": {}, + } + # Filter metadata, excluding embedding + if hasattr(memory_item, "metadata") and memory_item.metadata: + for attr_name in dir(memory_item.metadata): + if not attr_name.startswith("_") and attr_name != "embedding": + attr_value = getattr(memory_item.metadata, attr_name) + if not callable(attr_value): + filtered_item["metadata"][attr_name] = attr_value + filtered_memories.append(filtered_item) + + filtered_group = { + "cube_id": mem_group.get("cube_id", ""), + "memories": filtered_memories, + } + filtered_data[key].append(filtered_group) + else: + # Old data structure: dictionary with nodes and edges + filtered_group = { + "memories": {"nodes": [], "edges": mem_group["memories"].get("edges", [])} + } + for node in mem_group["memories"].get("nodes", []): + filtered_node = { + "id": node.get("id"), + "memory": node.get("memory"), + "metadata": { + k: v + for k, v in node.get("metadata", {}).items() + if k != "embedding" + }, + } + filtered_group["memories"]["nodes"].append(filtered_node) + filtered_data[key].append(filtered_group) + else: + filtered_data[key] = value + return filtered_data + + +filtered_results = filter_memory_data(get_all_results) +print(f"Get all results after add memory: {filtered_results['text_mem'][0]['memories']}") + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] Get all memories completed, time elapsed: {time.time() - start_time:.2f}s\n" +) + +# 6. Add messages +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to add conversation messages...") +start_time = time.time() + +messages = [ + {"role": "user", "content": "I like playing football."}, + {"role": "assistant", "content": "yes football is my favorite game."}, +] +mos.add(messages) + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] Conversation messages added successfully, time elapsed: {time.time() - start_time:.2f}s" +) + +print( + f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to get all memories (after adding messages)..." +) +start_time = time.time() + +get_all_results = mos.get_all() +filtered_results = filter_memory_data(get_all_results) +print(f"Get all results after add messages: {filtered_results}") + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] Get all memories completed, time elapsed: {time.time() - start_time:.2f}s\n" +) + +# 7. Add document +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to add document...") +start_time = time.time() +## 7.1 add pdf for ./tmp/data if use doc mem mos.add(doc_path="./tmp/data/") +start_time = time.time() + +get_all_results = mos.get_all() +filtered_results = filter_memory_data(get_all_results) +print(f"Get all results after add doc: {filtered_results}") + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] Get all memories completed, time elapsed: {time.time() - start_time:.2f}s\n" +) + +# 8. Search +print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to search memories...") +start_time = time.time() + +search_results = mos.search(query="my favorite football game", user_id=user_name) +filtered_search_results = filter_memory_data(search_results) +print(f"Search results: {filtered_search_results}") + +print( + f"✅ [{datetime.now().strftime('%H:%M:%S')}] Memory search completed, time elapsed: {time.time() - start_time:.2f}s\n" +) + +# 9. Chat +print(f"🎯 [{datetime.now().strftime('%H:%M:%S')}] Starting chat mode...") +while True: + user_input = input("👤 [You] ").strip() + if user_input.lower() in ["quit", "exit"]: + break + + print() + chat_start_time = time.time() + response = mos.chat(user_input) + chat_duration = time.time() - chat_start_time + + print(f"🤖 [Assistant] {response}") + print(f"⏱️ [Response time: {chat_duration:.2f}s]\n") + +print("📢 [System] MemChat has stopped.")