wehub-resource-sync
1348 lines
55 KiB
Python
1348 lines
55 KiB
Python
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": "<error details>"}
|
|
|
|
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 []
|