wip fix Memgraph get_knowledge_graph issues
This commit is contained in:
DavIvek
parent
63d554bad8
commit
8d295bd294
1 changed files with 109 additions and 80 deletions
|
|
@ -734,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`.
|
||||
|
|
@ -747,137 +747,166 @@ class MemgraphStorage(BaseGraphStorage):
|
|||
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,
|
||||
collect(DISTINCT r) AS relationships
|
||||
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"""
|
||||
# return await self._robust_fallback(node_label, max_depth, max_nodes)
|
||||
# First try without limit to check if we need to truncate
|
||||
full_query = f"""
|
||||
MATCH (start:`{workspace_label}`)
|
||||
WHERE start.entity_id = $entity_id
|
||||
WITH start
|
||||
CALL {{
|
||||
WITH start
|
||||
MATCH path = (start)-[*BFS 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_found
|
||||
WITH
|
||||
CASE
|
||||
WHEN total_nodes_found <= {max_nodes} THEN nodes
|
||||
ELSE nodes[0..{max_nodes}]
|
||||
END AS limited_nodes,
|
||||
relationships,
|
||||
total_nodes_found,
|
||||
total_nodes_found > {max_nodes} AS is_truncated
|
||||
|
||||
UNWIND relationships AS rel
|
||||
WITH limited_nodes, rel, total_nodes_found, is_truncated
|
||||
WHERE startNode(rel) IN limited_nodes AND endNode(rel) IN limited_nodes
|
||||
WITH limited_nodes, collect(DISTINCT rel) AS limited_relationships, total_nodes_found, is_truncated
|
||||
MATCH path = (start)-[*BFS ..{max_depth}]-(node:`{workspace_label}`)
|
||||
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
|
||||
[node IN limited_nodes | {{node: node}}] AS node_info,
|
||||
limited_relationships AS relationships,
|
||||
total_nodes_found,
|
||||
is_truncated
|
||||
[node IN all_nodes | {{node: node}}] AS node_info,
|
||||
all_rels AS relationships,
|
||||
size(all_nodes) AS total_nodes
|
||||
"""
|
||||
|
||||
result_set = None
|
||||
# Try to get full result
|
||||
full_result = None
|
||||
try:
|
||||
result_set = await session.run(
|
||||
bfs_query,
|
||||
full_result = await session.run(
|
||||
full_query,
|
||||
{
|
||||
"entity_id": node_label,
|
||||
},
|
||||
)
|
||||
record = await result_set.single()
|
||||
if not record:
|
||||
full_record = await full_result.single()
|
||||
|
||||
# If no record found, return empty KnowledgeGraph
|
||||
if not full_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"]:
|
||||
# If record found, check node count
|
||||
total_nodes = full_record["total_nodes"]
|
||||
|
||||
if total_nodes <= max_nodes:
|
||||
# If node count is within limit, use full result directly
|
||||
logger.debug(
|
||||
f"Using full result with {total_nodes} nodes (no truncation needed)"
|
||||
)
|
||||
record = full_record
|
||||
else:
|
||||
# If node count exceeds limit, set truncated flag and run limited query
|
||||
result.is_truncated = True
|
||||
logger.info(
|
||||
f"Graph truncated: breadth-first search limited to {max_nodes} nodes"
|
||||
f"Graph truncated: {total_nodes} nodes found, breadth-first search limited to {max_nodes}"
|
||||
)
|
||||
|
||||
finally:
|
||||
if result_set:
|
||||
await result_set.consume()
|
||||
# Run limited query
|
||||
limited_query = f"""
|
||||
MATCH (start:`{workspace_label}`)
|
||||
WHERE start.entity_id = $entity_id
|
||||
WITH start
|
||||
MATCH path = (start)-[*BFS ..{max_depth}]-(node:`{workspace_label}`)
|
||||
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
|
||||
WITH
|
||||
CASE
|
||||
WHEN size(all_nodes) <= $max_nodes THEN all_nodes
|
||||
ELSE all_nodes[0..$max_nodes]
|
||||
END AS limited_nodes,
|
||||
all_rels
|
||||
UNWIND all_rels AS rel
|
||||
WITH limited_nodes, rel
|
||||
WHERE startNode(rel) IN limited_nodes AND endNode(rel) IN limited_nodes
|
||||
WITH limited_nodes, collect(DISTINCT rel) AS limited_relationships
|
||||
RETURN
|
||||
[node IN limited_nodes | {{node: node}}] AS node_info,
|
||||
limited_relationships AS relationships
|
||||
"""
|
||||
|
||||
# Process the record if it exists
|
||||
if record and record["node_info"]:
|
||||
result_set = None
|
||||
try:
|
||||
result_set = await session.run(
|
||||
limited_query,
|
||||
{
|
||||
"entity_id": node_label,
|
||||
"max_nodes": max_nodes,
|
||||
},
|
||||
)
|
||||
record = await result_set.single()
|
||||
finally:
|
||||
if result_set:
|
||||
await result_set.consume()
|
||||
finally:
|
||||
if full_result:
|
||||
await full_result.consume()
|
||||
|
||||
if record:
|
||||
# Handle nodes (compatible with multi-label cases)
|
||||
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}",
|
||||
|
|
@ -885,11 +914,12 @@ class MemgraphStorage(BaseGraphStorage):
|
|||
properties=dict(node),
|
||||
)
|
||||
)
|
||||
seen_nodes.add(node_id)
|
||||
|
||||
# Handle relationships (including direction information)
|
||||
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(
|
||||
|
|
@ -901,14 +931,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.error(f"Error during subgraph query for {node_label}: {str(e)}")
|
||||
|
||||
return result
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue