gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #1766 from HKUDS/fix-memgraph-max-nodes-issue

Browse source 
Fix Memgraph get_knowledge_graph issues
This commit is contained in:
Daniel.y 2025-07-15 16:07:04 +08:00 • committed by GitHub
commit 6d1260aafa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 64 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
lightrag/api/README.md
View file

@ -181,9 +181,9 @@ The command-line `workspace` argument and the `WORKSPACE` environment variable i
- **For local file-based databases, data isolation is achieved through workspace subdirectories:** `JsonKVStorage`, `JsonDocStatusStorage`, `NetworkXStorage`, `NanoVectorDBStorage`, `FaissVectorDBStorage`.
- **For databases that store data in collections, it's done by adding a workspace prefix to the collection name:** `RedisKVStorage`, `RedisDocStatusStorage`, `MilvusVectorDBStorage`, `QdrantVectorDBStorage`, `MongoKVStorage`, `MongoDocStatusStorage`, `MongoVectorDBStorage`, `MongoGraphStorage`, `PGGraphStorage`.
- **For relational databases, data isolation is achieved by adding a `workspace` field to the tables for logical data separation:** `PGKVStorage`, `PGVectorStorage`, `PGDocStatusStorage`.
- **For the Neo4j graph database, logical data isolation is achieved through labels:** `Neo4JStorage`
- **For graph databases, logical data isolation is achieved through labels:** `Neo4JStorage`, `MemgraphStorage`
To maintain compatibility with legacy data, the default workspace for PostgreSQL is `default` and for Neo4j is `base` when no workspace is configured. For all external storages, the system provides dedicated workspace environment variables to override the common `WORKSPACE` environment variable configuration. These storage-specific workspace environment variables are: `REDIS_WORKSPACE`, `MILVUS_WORKSPACE`, `QDRANT_WORKSPACE`, `MONGODB_WORKSPACE`, `POSTGRES_WORKSPACE`, `NEO4J_WORKSPACE`.
To maintain compatibility with legacy data, the default workspace for PostgreSQL is `default` and for Neo4j is `base` when no workspace is configured. For all external storages, the system provides dedicated workspace environment variables to override the common `WORKSPACE` environment variable configuration. These storage-specific workspace environment variables are: `REDIS_WORKSPACE`, `MILVUS_WORKSPACE`, `QDRANT_WORKSPACE`, `MONGODB_WORKSPACE`, `POSTGRES_WORKSPACE`, `NEO4J_WORKSPACE`, `MEMGRAPH_WORKSPACE`.
### Multiple workers for Gunicorn + Uvicorn
@ -396,6 +396,7 @@ MongoKVStorage MongoDB
NetworkXStorage NetworkX (default)
Neo4JStorage Neo4J
PGGraphStorage PostgreSQL with AGE plugin
MemgraphStorage. Memgraph
```
> Testing has shown that Neo4J delivers superior performance in production environments compared to PostgreSQL with AGE plugin.

125
lightrag/kg/memgraph_impl.py
View file

@ -435,7 +435,7 @@ class MemgraphStorage(BaseGraphStorage):
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
"""
Upsert a node in the Neo4j database.
Upsert a node in the Memgraph database.
Args:
node_id: The unique identifier for the node (used as label)
@ -448,7 +448,9 @@ class MemgraphStorage(BaseGraphStorage):
properties = node_data
entity_type = properties["entity_type"]
if "entity_id" not in properties:
raise ValueError("Neo4j: node properties must contain an 'entity_id' field")
raise ValueError(
"Memgraph: node properties must contain an 'entity_id' field"
)
try:
async with self._driver.session(database=self._DATABASE) as session:
@ -732,7 +734,7 @@ class MemgraphStorage(BaseGraphStorage):
self,
node_label: str,
max_depth: int = 3,
max_nodes: int = MAX_GRAPH_NODES,
max_nodes: int = None,
) -> KnowledgeGraph:
"""
Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
@ -740,120 +742,118 @@ class MemgraphStorage(BaseGraphStorage):
Args:
node_label: Label of the starting node, * means all nodes
max_depth: Maximum depth of the subgraph, Defaults to 3
max_nodes: Maxiumu nodes to return by BFS, Defaults to 1000
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
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."
)
# 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()
workspace_label = self._get_workspace_label()
async with self._driver.session(
database=self._DATABASE, default_access_mode="READ"
) as session:
try:
if node_label == "*":
# First check if database has any nodes
count_query = "MATCH (n) RETURN count(n) as total"
# 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
total_count = 0
try:
count_result = await session.run(count_query)
count_record = await count_result.single()
if count_record:
total_count = count_record["total"]
if total_count == 0:
logger.debug("No nodes found in database")
return result
if total_count > max_nodes:
result.is_truncated = True
logger.info(
f"Graph truncated: {total_count} nodes found, limited to {max_nodes}"
)
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 the main query to get nodes with highest degree
# 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(n) AS kept_nodes
MATCH (a)-[r]-(b)
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 [node IN kept_nodes | {{node: node}}] AS node_info,
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}
main_query,
{"max_nodes": max_nodes},
)
record = await result_set.single()
if not record:
logger.debug("No record returned from main query")
return result
finally:
if result_set:
await result_set.consume()
else:
bfs_query = f"""
# Run subgraph query for specific node_label
subgraph_query = f"""
MATCH (start:`{workspace_label}`)
WHERE start.entity_id = $entity_id
WITH start
CALL {{
WITH start
MATCH path = (start)-[*0..{max_depth}]-(node)
WITH nodes(path) AS path_nodes, relationships(path) AS path_rels
UNWIND path_nodes AS n
WITH collect(DISTINCT n) AS all_nodes, collect(DISTINCT path_rels) AS all_rel_lists
WITH all_nodes, reduce(r = [], x IN all_rel_lists | r + x) AS all_rels
RETURN all_nodes, all_rels
}}
WITH all_nodes AS nodes, all_rels AS relationships, size(all_nodes) AS total_nodes
MATCH path = (start)-[*BFS 0..{max_depth}]-(end:`{workspace_label}`)
WHERE ALL(n IN nodes(path) WHERE '{workspace_label}' IN labels(n))
WITH collect(DISTINCT end) + start AS all_nodes_unlimited
WITH
CASE
WHEN total_nodes <= {max_nodes} THEN nodes
ELSE nodes[0..{max_nodes}]
WHEN size(all_nodes_unlimited) <= $max_nodes THEN all_nodes_unlimited
ELSE all_nodes_unlimited[0..$max_nodes]
END AS limited_nodes,
relationships,
total_nodes,
total_nodes > {max_nodes} AS is_truncated
size(all_nodes_unlimited) > $max_nodes AS is_truncated
UNWIND limited_nodes AS n
MATCH (n)-[r]-(m)
WHERE m IN limited_nodes
WITH collect(DISTINCT n) AS limited_nodes, collect(DISTINCT r) AS relationships, is_truncated
RETURN
[node IN limited_nodes | {{node: node}}] AS node_info,
relationships,
total_nodes,
is_truncated
"""
result_set = None
try:
result_set = await session.run(
bfs_query,
subgraph_query,
{
"entity_id": node_label,
"max_nodes": max_nodes,
},
)
record = await result_set.single()
# If no record found, return empty KnowledgeGraph
if not record:
logger.debug(f"No nodes found for entity_id: {node_label}")
return result
# Check if the query indicates truncation
if "is_truncated" in record and record["is_truncated"]:
# Check if the result was truncated
if record.get("is_truncated"):
result.is_truncated = True
logger.info(
f"Graph truncated: breadth-first search limited to {max_nodes} nodes"
@ -863,13 +863,11 @@ class MemgraphStorage(BaseGraphStorage):
if result_set:
await result_set.consume()
# Process the record if it exists
if record and record["node_info"]:
if record:
for node_info in record["node_info"]:
node = node_info["node"]
node_id = node.id
if node_id not in seen_nodes:
seen_nodes.add(node_id)
result.nodes.append(
KnowledgeGraphNode(
id=f"{node_id}",
@ -877,11 +875,11 @@ class MemgraphStorage(BaseGraphStorage):
properties=dict(node),
)
)
seen_nodes.add(node_id)
for rel in record["relationships"]:
edge_id = rel.id
if edge_id not in seen_edges:
seen_edges.add(edge_id)
start = rel.start_node
end = rel.end_node
result.edges.append(
@ -893,14 +891,13 @@ class MemgraphStorage(BaseGraphStorage):
properties=dict(rel),
)
)
seen_edges.add(edge_id)
logger.info(
f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
)
logger.info(
f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
)
except Exception as e:
logger.error(f"Error getting knowledge graph: {str(e)}")
# Return empty but properly initialized KnowledgeGraph on error
return KnowledgeGraph()
logger.warning(f"Memgraph error during subgraph query: {str(e)}")
return result