kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

修复子图检索的bug

Browse Source
This commit is contained in:
Wenjie Zhang 2025-02-28 02:42:21 +08:00
parent 1bf0af44dd
commit 577b78821f
2 changed files with 22 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/core/graphbase.py
View File

@ -280,7 +280,7 @@ class GraphDatabase:
def query(tx, entity_name, hops): def query(tx, entity_name, hops):
result = tx.run(f""" result = tx.run(f"""
MATCH (n {{name: $entity_name}})-[r*1..{hops}]-(m) MATCH (n {{name: $entity_name}})-[r*1..{hops}]-(m)
RETURN n.name AS node_name, r, m.name AS neighbor_name RETURN n, r, m
""", entity_name=entity_name) """, entity_name=entity_name)
return result.values() return result.values()

31
src/core/retriever.py
View File

@ -175,7 +175,7 @@ class Retriever:
return entities return entities
def _extract_relationship_info(self, relationship, source_name, target_name): def _extract_relationship_info(self, relationship, source_name=None, target_name=None, node_dict=None):
""" """
提取关系信息并返回格式化的节点和边信息 提取关系信息并返回格式化的节点和边信息
""" """
@ -188,6 +188,9 @@ class Retriever:
source_id = source.element_id source_id = source.element_id
target_id = target.element_id target_id = target.element_id
source_name = node_dict[source_id]["name"] if source_name is None else source_name
target_name = node_dict[target_id]["name"] if target_name is None else target_name
relationship_type = relationship._properties.get("type", "unknown") relationship_type = relationship._properties.get("type", "unknown")
if relationship_type == "unknown": if relationship_type == "unknown":
relationship_type = relationship.type relationship_type = relationship.type
@ -234,20 +237,28 @@ class Retriever:
node_dict = {} node_dict = {}
for item in results: for item in results:
if not isinstance(item[1], list) or not item[1]: # 检查数据格式
if len(item) < 2 or not isinstance(item[1], list):
continue continue
relationship = item[1][0] node_dict[item[0].element_id] = dict(id=item[0].element_id, name=item[0]._properties.get("name", "Unknown"))
source_name = item[0] node_dict[item[2].element_id] = dict(id=item[2].element_id, name=item[2]._properties.get("name", "Unknown"))
target_name = item[2] if len(item) > 2 else "unknown"
node_info, edge_info = self._extract_relationship_info(relationship, source_name, target_name) # 处理关系列表中的每个关系
if node_info is None or edge_info is None: for i, relationship in enumerate(item[1]):
continue try:
# 提取关系信息
node_info, edge_info = self._extract_relationship_info(relationship, node_dict=node_dict)
if node_info is None or edge_info is None:
continue
node_dict.update({node["id"]: node for node in node_info}) # 添加边
formatted_results["edges"].append(edge_info) formatted_results["edges"].append(edge_info)
except Exception as e:
logger.error(f"处理关系时出错: {e}, 关系: {relationship}")
continue
# 将节点字典转换为列表
formatted_results["nodes"] = list(node_dict.values()) formatted_results["nodes"] = list(node_dict.values())
return formatted_results return formatted_results