import os import asyncio import random from dataclasses import dataclass from typing import final import configparser from ..utils import logger, validate_workspace from ..base import BaseGraphStorage from ..types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge from ..kg.shared_storage import get_data_init_lock import pipmaster as pm if not pm.is_installed("neo4j"): pm.install("neo4j") from neo4j import ( AsyncGraphDatabase, AsyncManagedTransaction, ) from neo4j.exceptions import TransientError, ResultFailedError from dotenv import load_dotenv # use the .env that is inside the current folder load_dotenv(dotenv_path=".env", override=False) MAX_GRAPH_NODES = int(os.getenv("MAX_GRAPH_NODES", 1000)) config = configparser.ConfigParser() config.read("config.ini", "utf-8") @final @dataclass class MemgraphStorage(BaseGraphStorage): def __init__(self, namespace, global_config, embedding_func, workspace=None): # Priority: 1) MEMGRAPH_WORKSPACE env 2) user arg 3) default 'base' memgraph_workspace = os.environ.get("MEMGRAPH_WORKSPACE") original_workspace = workspace # Save original value for logging if memgraph_workspace and memgraph_workspace.strip(): workspace = memgraph_workspace if not workspace or not str(workspace).strip(): workspace = "base" super().__init__( namespace=namespace, workspace=workspace, global_config=global_config, embedding_func=embedding_func, ) validate_workspace(self.workspace) # Log after super().__init__() to ensure self.workspace is initialized if memgraph_workspace and memgraph_workspace.strip(): logger.info( f"Using MEMGRAPH_WORKSPACE environment variable: '{memgraph_workspace}' (overriding '{original_workspace}/{namespace}')" ) self._driver = None def _get_workspace_label(self) -> str: """Return sanitized workspace label safe for use as a backtick-quoted identifier in Cypher queries. Escapes backticks by doubling them to prevent Cypher injection via the LIGHTRAG-WORKSPACE header, while preserving a 1-to-1 mapping for all other characters. The returned value is intended to be used inside backticks (for example, MATCH (n:`{label}`)) and is not validated as a standalone unquoted identifier. """ workspace = self.workspace.strip() if not workspace: return "base" return workspace.replace("`", "``") async def initialize(self): async with get_data_init_lock(): URI = os.environ.get( "MEMGRAPH_URI", config.get("memgraph", "uri", fallback="bolt://localhost:7687"), ) USERNAME = os.environ.get( "MEMGRAPH_USERNAME", config.get("memgraph", "username", fallback="") ) PASSWORD = os.environ.get( "MEMGRAPH_PASSWORD", config.get("memgraph", "password", fallback="") ) DATABASE = os.environ.get( "MEMGRAPH_DATABASE", config.get("memgraph", "database", fallback="memgraph"), ) self._driver = AsyncGraphDatabase.driver( URI, auth=(USERNAME, PASSWORD), ) self._DATABASE = DATABASE try: async with self._driver.session(database=DATABASE) as session: # Create index for base nodes on entity_id if it doesn't exist try: workspace_label = self._get_workspace_label() await session.run( f"""CREATE INDEX ON :{workspace_label}(entity_id)""" ) logger.info( f"[{self.workspace}] Created index on :{workspace_label}(entity_id) in Memgraph." ) except Exception as e: # Index may already exist, which is not an error logger.warning( f"[{self.workspace}] Index creation on :{workspace_label}(entity_id) may have failed or already exists: {e}" ) await session.run("RETURN 1") logger.info(f"[{self.workspace}] Connected to Memgraph at {URI}") except Exception as e: logger.error( f"[{self.workspace}] Failed to connect to Memgraph at {URI}: {e}" ) raise async def finalize(self): if self._driver is not None: await self._driver.close() self._driver = None async def __aexit__(self, exc_type, exc, tb): await self.finalize() async def index_done_callback(self): # Memgraph handles persistence automatically pass async def has_node(self, node_id: str) -> bool: """ Check if a node exists in the graph. Args: node_id: The ID of the node to check. Returns: bool: True if the node exists, False otherwise. Raises: Exception: If there is an error checking the node existence. """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: result = None try: workspace_label = self._get_workspace_label() query = f"MATCH (n:`{workspace_label}` {{entity_id: $entity_id}}) RETURN count(n) > 0 AS node_exists" result = await session.run(query, entity_id=node_id) single_result = await result.single() await result.consume() # Ensure result is fully consumed return ( single_result["node_exists"] if single_result is not None else False ) except Exception as e: logger.error( f"[{self.workspace}] Error checking node existence for {node_id}: {str(e)}" ) if result is not None: await ( result.consume() ) # Ensure the result is consumed even on error raise async def has_edge(self, source_node_id: str, target_node_id: str) -> bool: """ Check if an edge exists between two nodes in the graph. Args: source_node_id: The ID of the source node. target_node_id: The ID of the target node. Returns: bool: True if the edge exists, False otherwise. Raises: Exception: If there is an error checking the edge existence. """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: result = None try: workspace_label = self._get_workspace_label() query = ( f"MATCH (a:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(b:`{workspace_label}` {{entity_id: $target_entity_id}}) " "RETURN COUNT(r) > 0 AS edgeExists" ) result = await session.run( query, source_entity_id=source_node_id, target_entity_id=target_node_id, ) # type: ignore single_result = await result.single() await result.consume() # Ensure result is fully consumed return ( single_result["edgeExists"] if single_result is not None else False ) except Exception as e: logger.error( f"[{self.workspace}] Error checking edge existence between {source_node_id} and {target_node_id}: {str(e)}" ) if result is not None: await ( result.consume() ) # Ensure the result is consumed even on error raise async def get_node(self, node_id: str) -> dict[str, str] | None: """Get node by its label identifier, return only node properties Args: node_id: The node label to look up Returns: dict: Node properties if found None: If node not found Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: try: workspace_label = self._get_workspace_label() query = ( f"MATCH (n:`{workspace_label}` {{entity_id: $entity_id}}) RETURN n" ) result = await session.run(query, entity_id=node_id) try: records = await result.fetch( 2 ) # Get 2 records for duplication check if len(records) > 1: logger.warning( f"[{self.workspace}] Multiple nodes found with label '{node_id}'. Using first node." ) if records: node = records[0]["n"] node_dict = dict(node) # Remove workspace label from labels list if it exists if "labels" in node_dict: node_dict["labels"] = [ label for label in node_dict["labels"] if label != workspace_label ] return node_dict return None finally: await result.consume() # Ensure result is fully consumed except Exception as e: logger.error( f"[{self.workspace}] Error getting node for {node_id}: {str(e)}" ) raise async def node_degree(self, node_id: str) -> int: """Get the degree (number of relationships) of a node with the given label. If multiple nodes have the same label, returns the degree of the first node. If no node is found, returns 0. Args: node_id: The label of the node Returns: int: The number of relationships the node has, or 0 if no node found Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: try: workspace_label = self._get_workspace_label() query = f""" MATCH (n:`{workspace_label}` {{entity_id: $entity_id}}) OPTIONAL MATCH (n)-[r]-() RETURN COUNT(r) AS degree """ result = await session.run(query, entity_id=node_id) try: record = await result.single() if not record: logger.warning( f"[{self.workspace}] No node found with label '{node_id}'" ) return 0 degree = record["degree"] return degree finally: await result.consume() # Ensure result is fully consumed except Exception as e: logger.error( f"[{self.workspace}] Error getting node degree for {node_id}: {str(e)}" ) raise async def get_all_labels(self) -> list[str]: """ Get all existing node labels(entity names) in the database Returns: ["Person", "Company", ...] # Alphabetically sorted label list Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: result = None try: workspace_label = self._get_workspace_label() query = f""" MATCH (n:`{workspace_label}`) WHERE n.entity_id IS NOT NULL RETURN DISTINCT n.entity_id AS label ORDER BY label """ result = await session.run(query) labels = [] async for record in result: labels.append(record["label"]) await result.consume() return labels except Exception as e: logger.error(f"[{self.workspace}] Error getting all labels: {str(e)}") if result is not None: await ( result.consume() ) # Ensure the result is consumed even on error raise async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None: """Retrieves all edges (relationships) for a particular node identified by its label. Args: source_node_id: Label of the node to get edges for Returns: list[tuple[str, str]]: List of (source_label, target_label) tuples representing edges None: If no edges found Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) try: async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: results = None try: workspace_label = self._get_workspace_label() query = f"""MATCH (n:`{workspace_label}` {{entity_id: $entity_id}}) OPTIONAL MATCH (n)-[r]-(connected:`{workspace_label}`) WHERE connected.entity_id IS NOT NULL RETURN n.entity_id AS node_entity_id, connected.entity_id AS connected_entity_id, startNode(r).entity_id AS start_entity_id""" results = await session.run(query, entity_id=source_node_id) edges = [] async for record in results: node_entity_id = record["node_entity_id"] connected_entity_id = record["connected_entity_id"] start_entity_id = record["start_entity_id"] if not node_entity_id or not connected_entity_id: continue # Preserve the original edge direction via startNode(r) if start_entity_id == node_entity_id: edges.append((node_entity_id, connected_entity_id)) else: edges.append((connected_entity_id, node_entity_id)) await results.consume() # Ensure results are consumed return edges except Exception as e: logger.error( f"[{self.workspace}] Error getting edges for node {source_node_id}: {str(e)}" ) if results is not None: await ( results.consume() ) # Ensure results are consumed even on error raise except Exception as e: logger.error( f"[{self.workspace}] Error in get_node_edges for {source_node_id}: {str(e)}" ) raise async def get_edge( self, source_node_id: str, target_node_id: str ) -> dict[str, str] | None: """Get edge properties between two nodes. Args: source_node_id: Label of the source node target_node_id: Label of the target node Returns: dict: Edge properties if found, default properties if not found or on error Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: result = None try: workspace_label = self._get_workspace_label() query = f""" MATCH (start:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(end:`{workspace_label}` {{entity_id: $target_entity_id}}) RETURN properties(r) as edge_properties """ result = await session.run( query, source_entity_id=source_node_id, target_entity_id=target_node_id, ) records = await result.fetch(2) await result.consume() if records: edge_result = dict(records[0]["edge_properties"]) for key, default_value in { "weight": 1.0, "source_id": None, "description": None, "keywords": None, }.items(): if key not in edge_result: edge_result[key] = default_value logger.warning( f"[{self.workspace}] Edge between {source_node_id} and {target_node_id} is missing property: {key}. Using default value: {default_value}" ) return edge_result return None except Exception as e: logger.error( f"[{self.workspace}] Error getting edge between {source_node_id} and {target_node_id}: {str(e)}" ) if result is not None: await ( result.consume() ) # Ensure the result is consumed even on error raise async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None: """ Upsert a node in the Memgraph database with manual transaction-level retry logic for transient errors. Args: node_id: The unique identifier for the node (used as label) node_data: Dictionary of node properties """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) properties = node_data if "entity_id" not in properties: raise ValueError( "Memgraph: node properties must contain an 'entity_id' field" ) # Manual transaction-level retry following official Memgraph documentation max_retries = 100 initial_wait_time = 0.2 backoff_factor = 1.1 jitter_factor = 0.1 for attempt in range(max_retries): try: logger.debug( f"[{self.workspace}] Attempting node upsert, attempt {attempt + 1}/{max_retries}" ) async with self._driver.session(database=self._DATABASE) as session: workspace_label = self._get_workspace_label() async def execute_upsert(tx: AsyncManagedTransaction): query = f""" MERGE (n:`{workspace_label}` {{entity_id: $entity_id}}) SET n += $properties """ result = await tx.run( query, entity_id=node_id, properties=properties ) await result.consume() # Ensure result is fully consumed await session.execute_write(execute_upsert) break # Success - exit retry loop except (TransientError, ResultFailedError) as e: # Check if the root cause is a TransientError root_cause = e while hasattr(root_cause, "__cause__") and root_cause.__cause__: root_cause = root_cause.__cause__ # Check if this is a transient error that should be retried is_transient = ( isinstance(root_cause, TransientError) or isinstance(e, TransientError) or "TransientError" in str(e) or "Cannot resolve conflicting transactions" in str(e) ) if is_transient: if attempt < max_retries - 1: # Calculate wait time with exponential backoff and jitter jitter = random.uniform(0, jitter_factor) * initial_wait_time wait_time = ( initial_wait_time * (backoff_factor**attempt) + jitter ) logger.warning( f"[{self.workspace}] Node upsert failed. Attempt #{attempt + 1} retrying in {wait_time:.3f} seconds... Error: {str(e)}" ) await asyncio.sleep(wait_time) else: logger.error( f"[{self.workspace}] Memgraph transient error during node upsert after {max_retries} retries: {str(e)}" ) raise else: # Non-transient error, don't retry logger.error( f"[{self.workspace}] Non-transient error during node upsert: {str(e)}" ) raise except Exception as e: logger.error( f"[{self.workspace}] Unexpected error during node upsert: {str(e)}" ) raise async def upsert_edge( self, source_node_id: str, target_node_id: str, edge_data: dict[str, str] ) -> None: """ Upsert an edge and its properties between two nodes identified by their labels with manual transaction-level retry logic for transient errors. Ensures both source and target nodes exist and are unique before creating the edge. Uses entity_id property to uniquely identify nodes. Args: source_node_id (str): Label of the source node (used as identifier) target_node_id (str): Label of the target node (used as identifier) edge_data (dict): Dictionary of properties to set on the edge Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) edge_properties = edge_data # Manual transaction-level retry following official Memgraph documentation max_retries = 100 initial_wait_time = 0.2 backoff_factor = 1.1 jitter_factor = 0.1 for attempt in range(max_retries): try: logger.debug( f"[{self.workspace}] Attempting edge upsert, attempt {attempt + 1}/{max_retries}" ) async with self._driver.session(database=self._DATABASE) as session: async def execute_upsert(tx: AsyncManagedTransaction): workspace_label = self._get_workspace_label() query = f""" MATCH (source:`{workspace_label}` {{entity_id: $source_entity_id}}) WITH source MATCH (target:`{workspace_label}` {{entity_id: $target_entity_id}}) MERGE (source)-[r:DIRECTED]-(target) SET r += $properties RETURN r, source, target """ result = await tx.run( query, source_entity_id=source_node_id, target_entity_id=target_node_id, properties=edge_properties, ) try: await result.fetch(2) finally: await result.consume() # Ensure result is consumed await session.execute_write(execute_upsert) break # Success - exit retry loop except (TransientError, ResultFailedError) as e: # Check if the root cause is a TransientError root_cause = e while hasattr(root_cause, "__cause__") and root_cause.__cause__: root_cause = root_cause.__cause__ # Check if this is a transient error that should be retried is_transient = ( isinstance(root_cause, TransientError) or isinstance(e, TransientError) or "TransientError" in str(e) or "Cannot resolve conflicting transactions" in str(e) ) if is_transient: if attempt < max_retries - 1: # Calculate wait time with exponential backoff and jitter jitter = random.uniform(0, jitter_factor) * initial_wait_time wait_time = ( initial_wait_time * (backoff_factor**attempt) + jitter ) logger.warning( f"[{self.workspace}] Edge upsert failed. Attempt #{attempt + 1} retrying in {wait_time:.3f} seconds... Error: {str(e)}" ) await asyncio.sleep(wait_time) else: logger.error( f"[{self.workspace}] Memgraph transient error during edge upsert after {max_retries} retries: {str(e)}" ) raise else: # Non-transient error, don't retry logger.error( f"[{self.workspace}] Non-transient error during edge upsert: {str(e)}" ) raise except Exception as e: logger.error( f"[{self.workspace}] Unexpected error during edge upsert: {str(e)}" ) raise async def upsert_nodes_batch(self, nodes: list[tuple[str, dict[str, str]]]) -> None: """Batch insert/update multiple nodes using a single UNWIND Cypher query. Uses the same transient-error retry logic as upsert_node(). Args: nodes: List of (node_id, node_data) tuples. """ if not nodes: return if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) workspace_label = self._get_workspace_label() nodes_data = [] for node_id, node_data in nodes: if "entity_id" not in node_data: raise ValueError( "Memgraph: node properties must contain an 'entity_id' field" ) nodes_data.append({"entity_id": node_id, "props": node_data}) max_retries = 100 initial_wait_time = 0.2 backoff_factor = 1.1 jitter_factor = 0.1 for attempt in range(max_retries): try: async with self._driver.session(database=self._DATABASE) as session: async def execute_batch(tx: AsyncManagedTransaction): query = f""" UNWIND $nodes AS row MERGE (n:`{workspace_label}` {{entity_id: row.entity_id}}) SET n += row.props """ result = await tx.run(query, nodes=nodes_data) await result.consume() await session.execute_write(execute_batch) break except (TransientError, ResultFailedError) as e: root_cause = e while hasattr(root_cause, "__cause__") and root_cause.__cause__: root_cause = root_cause.__cause__ is_transient = ( isinstance(root_cause, TransientError) or isinstance(e, TransientError) or "TransientError" in str(e) or "Cannot resolve conflicting transactions" in str(e) ) if is_transient: if attempt < max_retries - 1: jitter = random.uniform(0, jitter_factor) * initial_wait_time wait_time = ( initial_wait_time * (backoff_factor**attempt) + jitter ) logger.warning( f"[{self.workspace}] Batch node upsert failed. Attempt #{attempt + 1} retrying in {wait_time:.3f}s... Error: {str(e)}" ) await asyncio.sleep(wait_time) else: logger.error( f"[{self.workspace}] Memgraph transient error during batch node upsert after {max_retries} retries: {str(e)}" ) raise else: logger.error( f"[{self.workspace}] Non-transient error during batch node upsert: {str(e)}" ) raise except Exception as e: logger.error( f"[{self.workspace}] Unexpected error during batch node upsert: {str(e)}" ) raise async def has_nodes_batch(self, node_ids: list[str]) -> set[str]: """Check existence of multiple nodes in a single UNWIND query. Args: node_ids: List of node IDs to check. Returns: Set of node_ids that exist in the graph. """ if not node_ids: return set() if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) workspace_label = self._get_workspace_label() try: async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: query = f""" UNWIND $ids AS id MATCH (n:`{workspace_label}` {{entity_id: id}}) RETURN n.entity_id AS entity_id """ result = await session.run(query, ids=node_ids) records = await result.data() await result.consume() return {r["entity_id"] for r in records} except Exception as e: logger.error( f"[{self.workspace}] Error during batch node existence check: {str(e)}" ) raise async def upsert_edges_batch( self, edges: list[tuple[str, str, dict[str, str]]] ) -> None: """Batch insert/update multiple edges using a single UNWIND Cypher query. Uses the same transient-error retry logic as upsert_edge(). Args: edges: List of (source_node_id, target_node_id, edge_data) tuples. """ if not edges: return if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) workspace_label = self._get_workspace_label() edges_data = [ {"src": src, "tgt": tgt, "props": edge_data} for src, tgt, edge_data in edges ] max_retries = 100 initial_wait_time = 0.2 backoff_factor = 1.1 jitter_factor = 0.1 for attempt in range(max_retries): try: async with self._driver.session(database=self._DATABASE) as session: async def execute_batch(tx: AsyncManagedTransaction): query = f""" UNWIND $edges AS row MATCH (source:`{workspace_label}` {{entity_id: row.src}}) WITH source, row MATCH (target:`{workspace_label}` {{entity_id: row.tgt}}) MERGE (source)-[r:DIRECTED]-(target) SET r += row.props RETURN r """ result = await tx.run(query, edges=edges_data) await result.consume() await session.execute_write(execute_batch) break except (TransientError, ResultFailedError) as e: root_cause = e while hasattr(root_cause, "__cause__") and root_cause.__cause__: root_cause = root_cause.__cause__ is_transient = ( isinstance(root_cause, TransientError) or isinstance(e, TransientError) or "TransientError" in str(e) or "Cannot resolve conflicting transactions" in str(e) ) if is_transient: if attempt < max_retries - 1: jitter = random.uniform(0, jitter_factor) * initial_wait_time wait_time = ( initial_wait_time * (backoff_factor**attempt) + jitter ) logger.warning( f"[{self.workspace}] Batch edge upsert failed. Attempt #{attempt + 1} retrying in {wait_time:.3f}s... Error: {str(e)}" ) await asyncio.sleep(wait_time) else: logger.error( f"[{self.workspace}] Memgraph transient error during batch edge upsert after {max_retries} retries: {str(e)}" ) raise else: logger.error( f"[{self.workspace}] Non-transient error during batch edge upsert: {str(e)}" ) raise except Exception as e: logger.error( f"[{self.workspace}] Unexpected error during batch edge upsert: {str(e)}" ) raise async def delete_node(self, node_id: str) -> None: """Delete a node with the specified label Args: node_id: The label of the node to delete Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) async def _do_delete(tx: AsyncManagedTransaction): workspace_label = self._get_workspace_label() query = f""" MATCH (n:`{workspace_label}` {{entity_id: $entity_id}}) DETACH DELETE n """ result = await tx.run(query, entity_id=node_id) logger.debug(f"[{self.workspace}] Deleted node with label {node_id}") await result.consume() try: async with self._driver.session(database=self._DATABASE) as session: await session.execute_write(_do_delete) except Exception as e: logger.error(f"[{self.workspace}] Error during node deletion: {str(e)}") raise async def remove_nodes(self, nodes: list[str]): """Delete multiple nodes Args: nodes: List of node labels to be deleted """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) for node in nodes: await self.delete_node(node) async def remove_edges(self, edges: list[tuple[str, str]]): """Delete multiple edges Args: edges: List of edges to be deleted, each edge is a (source, target) tuple Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) for source, target in edges: async def _do_delete_edge(tx: AsyncManagedTransaction): workspace_label = self._get_workspace_label() query = f""" MATCH (source:`{workspace_label}` {{entity_id: $source_entity_id}})-[r]-(target:`{workspace_label}` {{entity_id: $target_entity_id}}) DELETE r """ result = await tx.run( query, source_entity_id=source, target_entity_id=target ) logger.debug( f"[{self.workspace}] Deleted edge from '{source}' to '{target}'" ) await result.consume() # Ensure result is fully consumed try: async with self._driver.session(database=self._DATABASE) as session: await session.execute_write(_do_delete_edge) except Exception as e: logger.error(f"[{self.workspace}] Error during edge deletion: {str(e)}") raise async def drop(self) -> dict[str, str]: """Drop all data from the current workspace and clean up resources This method will delete all nodes and relationships in the Memgraph database. Returns: dict[str, str]: Operation status and message - On success: {"status": "success", "message": "data dropped"} - On failure: {"status": "error", "message": ""} Raises: Exception: If there is an error executing the query """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) try: async with self._driver.session(database=self._DATABASE) as session: workspace_label = self._get_workspace_label() query = f"MATCH (n:`{workspace_label}`) DETACH DELETE n" result = await session.run(query) await result.consume() logger.info( f"[{self.workspace}] Dropped workspace {workspace_label} from Memgraph database {self._DATABASE}" ) return {"status": "success", "message": "workspace data dropped"} except Exception as e: logger.error( f"[{self.workspace}] Error dropping workspace {workspace_label} from Memgraph database {self._DATABASE}: {e}" ) return {"status": "error", "message": str(e)} async def edge_degree(self, src_id: str, tgt_id: str) -> int: """Get the total degree (sum of relationships) of two nodes. Args: src_id: Label of the source node tgt_id: Label of the target node Returns: int: Sum of the degrees of both nodes """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) src_degree = await self.node_degree(src_id) trg_degree = await self.node_degree(tgt_id) # Convert None to 0 for addition src_degree = 0 if src_degree is None else src_degree trg_degree = 0 if trg_degree is None else trg_degree degrees = int(src_degree) + int(trg_degree) return degrees async def get_knowledge_graph( self, node_label: str, max_depth: int = 3, max_nodes: int = None, ) -> KnowledgeGraph: """ Retrieve a connected subgraph of nodes where the label includes the specified `node_label`. Args: node_label: Label of the starting node, * means all nodes max_depth: Maximum depth of the subgraph, Defaults to 3 max_nodes: Maximum nodes to return by BFS, Defaults to 1000 Returns: KnowledgeGraph object containing nodes and edges, with an is_truncated flag indicating whether the graph was truncated due to max_nodes limit """ # Get max_nodes from global_config if not provided if max_nodes is None: max_nodes = self.global_config.get("max_graph_nodes", 1000) else: # Limit max_nodes to not exceed global_config max_graph_nodes max_nodes = min(max_nodes, self.global_config.get("max_graph_nodes", 1000)) workspace_label = self._get_workspace_label() result = KnowledgeGraph() seen_nodes = set() seen_edges = set() async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: try: if node_label == "*": # First check total node count to determine if graph is truncated count_query = ( f"MATCH (n:`{workspace_label}`) RETURN count(n) as total" ) count_result = None try: count_result = await session.run(count_query) count_record = await count_result.single() if count_record and count_record["total"] > max_nodes: result.is_truncated = True logger.info( f"Graph truncated: {count_record['total']} nodes found, limited to {max_nodes}" ) finally: if count_result: await count_result.consume() # Run main query to get nodes with highest degree main_query = f""" MATCH (n:`{workspace_label}`) OPTIONAL MATCH (n)-[r]-() WITH n, COALESCE(count(r), 0) AS degree ORDER BY degree DESC LIMIT $max_nodes WITH collect({{node: n}}) AS filtered_nodes UNWIND filtered_nodes AS node_info WITH collect(node_info.node) AS kept_nodes, filtered_nodes OPTIONAL MATCH (a)-[r]-(b) WHERE a IN kept_nodes AND b IN kept_nodes RETURN filtered_nodes AS node_info, collect(DISTINCT r) AS relationships """ result_set = None try: result_set = await session.run( main_query, {"max_nodes": max_nodes}, ) record = await result_set.single() finally: if result_set: await result_set.consume() else: # Run subgraph query for specific node_label subgraph_query = f""" MATCH (start:`{workspace_label}`) WHERE start.entity_id = $entity_id OPTIONAL MATCH path = (start)-[*BFS 0..{max_depth}]-(end:`{workspace_label}`) WHERE path IS NULL OR ALL(n IN nodes(path) WHERE '{workspace_label}' IN labels(n)) WITH start, collect(DISTINCT end) AS discovered_nodes WITH start, [node IN discovered_nodes WHERE node IS NOT NULL AND node <> start] AS other_nodes WITH CASE WHEN 1 + size(other_nodes) <= $max_nodes THEN [start] + other_nodes ELSE [start] + other_nodes[0..$max_other_nodes] END AS limited_nodes, 1 + size(other_nodes) > $max_nodes AS is_truncated UNWIND limited_nodes AS n OPTIONAL MATCH (n)-[r]-(m) WHERE m IN limited_nodes WITH limited_nodes, collect(DISTINCT r) AS relationships, is_truncated RETURN [node IN limited_nodes | {{node: node}}] AS node_info, [rel IN relationships WHERE rel IS NOT NULL] AS relationships, is_truncated """ result_set = None try: result_set = await session.run( subgraph_query, { "entity_id": node_label, "max_nodes": max_nodes, "max_other_nodes": max(max_nodes - 1, 0), }, ) record = await result_set.single() # If no record found, return empty KnowledgeGraph if not record: logger.debug( f"[{self.workspace}] No nodes found for entity_id: {node_label}" ) return result # Check if the result was truncated if record.get("is_truncated"): result.is_truncated = True logger.info( f"[{self.workspace}] Graph truncated: breadth-first search limited to {max_nodes} nodes" ) finally: if result_set: await result_set.consume() if record: for node_info in record["node_info"]: node = node_info["node"] node_id = node.id if node_id not in seen_nodes: result.nodes.append( KnowledgeGraphNode( id=f"{node_id}", labels=[node.get("entity_id")], properties=dict(node), ) ) seen_nodes.add(node_id) for rel in record["relationships"]: edge_id = rel.id if edge_id not in seen_edges: start = rel.start_node end = rel.end_node result.edges.append( KnowledgeGraphEdge( id=f"{edge_id}", type=rel.type, source=f"{start.id}", target=f"{end.id}", properties=dict(rel), ) ) seen_edges.add(edge_id) logger.info( f"[{self.workspace}] Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}" ) except Exception as e: logger.warning( f"[{self.workspace}] Memgraph error during subgraph query: {str(e)}" ) return result async def get_all_nodes(self) -> list[dict]: """Get all nodes in the graph. Returns: A list of all nodes, where each node is a dictionary of its properties """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) workspace_label = self._get_workspace_label() async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: query = f""" MATCH (n:`{workspace_label}`) RETURN n """ result = await session.run(query) nodes = [] async for record in result: node = record["n"] node_dict = dict(node) # Add node id (entity_id) to the dictionary for easier access node_dict["id"] = node_dict.get("entity_id") nodes.append(node_dict) await result.consume() return nodes async def get_all_edges(self) -> list[dict]: """Get all edges in the graph. Returns: A list of all edges, where each edge is a dictionary of its properties """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) workspace_label = self._get_workspace_label() async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: query = f""" MATCH (a:`{workspace_label}`)-[r]-(b:`{workspace_label}`) RETURN DISTINCT a.entity_id AS source, b.entity_id AS target, properties(r) AS properties """ result = await session.run(query) edges = [] async for record in result: edge_properties = record["properties"] edge_properties["source"] = record["source"] edge_properties["target"] = record["target"] edges.append(edge_properties) await result.consume() return edges async def get_popular_labels(self, limit: int = 300) -> list[str]: """Get popular labels by node degree (most connected entities) Args: limit: Maximum number of labels(entity names) to return Returns: List of labels(entity names) sorted by degree (highest first) """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) result = None try: workspace_label = self._get_workspace_label() async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: query = f""" MATCH (n:`{workspace_label}`) WHERE n.entity_id IS NOT NULL OPTIONAL MATCH (n)-[r]-() WITH n.entity_id AS label, count(r) AS degree ORDER BY degree DESC, label ASC LIMIT {limit} RETURN label """ result = await session.run(query) labels = [] async for record in result: labels.append(record["label"]) await result.consume() logger.debug( f"[{self.workspace}] Retrieved {len(labels)} popular labels (limit: {limit})" ) return labels except Exception as e: logger.error(f"[{self.workspace}] Error getting popular labels: {str(e)}") if result is not None: await result.consume() return [] async def search_labels(self, query: str, limit: int = 50) -> list[str]: """Search labels(entity names) with fuzzy matching Args: query: Search query string limit: Maximum number of results to return Returns: List of matching labels(entity names) sorted by relevance """ if self._driver is None: raise RuntimeError( "Memgraph driver is not initialized. Call 'await initialize()' first." ) query_lower = query.lower().strip() if not query_lower: return [] result = None try: workspace_label = self._get_workspace_label() async with self._driver.session( database=self._DATABASE, default_access_mode="READ" ) as session: cypher_query = f""" MATCH (n:`{workspace_label}`) WHERE n.entity_id IS NOT NULL WITH n.entity_id AS label, toLower(n.entity_id) AS label_lower WHERE label_lower CONTAINS $query_lower WITH label, label_lower, CASE WHEN label_lower = $query_lower THEN 1000 WHEN label_lower STARTS WITH $query_lower THEN 500 ELSE 100 - size(label) END AS score ORDER BY score DESC, label ASC LIMIT {limit} RETURN label """ result = await session.run(cypher_query, query_lower=query_lower) labels = [] async for record in result: labels.append(record["label"]) await result.consume() logger.debug( f"[{self.workspace}] Search query '{query}' returned {len(labels)} results (limit: {limit})" ) return labels except Exception as e: logger.error(f"[{self.workspace}] Error searching labels: {str(e)}") if result is not None: await result.consume() return []