Merge pull request #154 from xerrors/fix-graph-result-chongfu
修复图谱检索结果重复的问题，同时将检索部分的逻辑代码，迁移到graphbase 里面
This commit is contained in:
Wenjie Zhang
GitHub
commit
5a098e5ce2
@ -157,7 +157,7 @@ async def index_nodes(data: dict = Body(default={})):
|
||||
@data.get("/graph/node")
|
||||
async def get_graph_node(entity_name: str):
|
||||
result = graph_base.query_node(entity_name=entity_name)
|
||||
return {"result": retriever.format_query_results(result), "message": "success"}
|
||||
return {"result": graph_base.format_query_result_to_graph(result), "message": "success"}
|
||||
|
||||
@data.get("/graph/nodes")
|
||||
async def get_graph_nodes(kgdb_name: str, num: int):
|
||||
@ -166,7 +166,7 @@ async def get_graph_nodes(kgdb_name: str, num: int):
|
||||
|
||||
logger.debug(f"Get graph nodes in {kgdb_name} with {num} nodes")
|
||||
result = graph_base.get_sample_nodes(kgdb_name, num)
|
||||
return {"result": retriever.format_general_results(result), "message": "success"}
|
||||
return {"result": graph_base.format_general_results(result), "message": "success"}
|
||||
|
||||
@data.post("/graph/add-by-jsonl")
|
||||
async def add_graph_entity(file_path: str = Body(...), kgdb_name: Optional[str] = Body(None)):
|
||||
|
||||
@ -498,6 +498,125 @@ class GraphDatabase:
|
||||
|
||||
return count
|
||||
|
||||
|
||||
def _extract_relationship_info(self, relationship, source_name=None, target_name=None, node_dict=None):
|
||||
"""
|
||||
提取关系信息并返回格式化的节点和边信息
|
||||
"""
|
||||
rel_id = relationship.element_id
|
||||
nodes = relationship.nodes
|
||||
if len(nodes) != 2:
|
||||
return None, None
|
||||
|
||||
source, target = nodes
|
||||
source_id = source.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")
|
||||
if relationship_type == "unknown":
|
||||
relationship_type = relationship.type
|
||||
|
||||
edge_info = {
|
||||
"id": rel_id,
|
||||
"type": relationship_type,
|
||||
"source_id": source_id,
|
||||
"target_id": target_id,
|
||||
"source_name": source_name,
|
||||
"target_name": target_name,
|
||||
}
|
||||
|
||||
node_info = [
|
||||
{"id": source_id, "name": source_name},
|
||||
{"id": target_id, "name": target_name},
|
||||
]
|
||||
|
||||
return node_info, edge_info
|
||||
|
||||
def format_general_results(self, results):
|
||||
formatted_results = {"nodes": [], "edges": []}
|
||||
|
||||
for item in results:
|
||||
relationship = item[1]
|
||||
source_name = item[0]._properties.get("name", "unknown")
|
||||
target_name = item[2]._properties.get("name", "unknown") 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:
|
||||
continue
|
||||
|
||||
for node in node_info:
|
||||
if node["id"] not in [n["id"] for n in formatted_results["nodes"]]:
|
||||
formatted_results["nodes"].append(node)
|
||||
|
||||
formatted_results["edges"].append(edge_info)
|
||||
|
||||
return formatted_results
|
||||
|
||||
def format_query_result_to_graph(self, query_results):
|
||||
"""将检索到的结果转换为 {"nodes": [], "edges": []} 的格式
|
||||
|
||||
例如：
|
||||
{
|
||||
"nodes": [
|
||||
{
|
||||
"id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:103",
|
||||
"name": "张若锦"
|
||||
},
|
||||
{
|
||||
"id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:20",
|
||||
"name": "贾宝玉"
|
||||
},
|
||||
....
|
||||
],
|
||||
"edges": [
|
||||
{
|
||||
"id": "5:5efbff88-72ef-44f9-b867-6c0e164a4a13:71",
|
||||
"type": "奴仆",
|
||||
"source_id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:88",
|
||||
"target_id": "4:5efbff88-72ef-44f9-b867-6c0e164a4a13:20",
|
||||
"source_name": "宋嬷嬷",
|
||||
"target_name": "贾宝玉"
|
||||
},
|
||||
....
|
||||
]
|
||||
}
|
||||
"""
|
||||
formatted_results = {"nodes": [], "edges": []}
|
||||
node_dict = {}
|
||||
edge_dict = {}
|
||||
|
||||
for item in query_results:
|
||||
# 检查数据格式
|
||||
if len(item) < 2 or not isinstance(item[1], list):
|
||||
continue
|
||||
|
||||
node_dict[item[0].element_id] = dict(id=item[0].element_id, name=item[0]._properties.get("name", "Unknown"))
|
||||
node_dict[item[2].element_id] = dict(id=item[2].element_id, name=item[2]._properties.get("name", "Unknown"))
|
||||
|
||||
# 处理关系列表中的每个关系
|
||||
for i, relationship in enumerate(item[1]):
|
||||
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
|
||||
|
||||
# 添加边
|
||||
edge_dict[edge_info["id"]] = edge_info
|
||||
except Exception as e:
|
||||
logger.error(f"处理关系时出错: {e}, 关系: {relationship}, {traceback.format_exc()}")
|
||||
continue
|
||||
|
||||
# 将节点字典转换为列表
|
||||
formatted_results["nodes"] = list(node_dict.values())
|
||||
formatted_results["edges"] = list(edge_dict.values())
|
||||
|
||||
|
||||
return formatted_results
|
||||
|
||||
def clean_triples_embedding(triples):
|
||||
for item in triples:
|
||||
if hasattr(item[0], '_properties'):
|
||||
|
||||
@ -84,7 +84,7 @@ class Retriever:
|
||||
result = graph_base.query_node(entity)
|
||||
if result != []:
|
||||
results.extend(result)
|
||||
return {"results": self.format_query_results(results)}
|
||||
return {"results": graph_base.format_query_result_to_graph(results)}
|
||||
|
||||
|
||||
def query_knowledgebase(self, query, history, refs):
|
||||
@ -177,94 +177,6 @@ class Retriever:
|
||||
|
||||
return entities
|
||||
|
||||
def _extract_relationship_info(self, relationship, source_name=None, target_name=None, node_dict=None):
|
||||
"""
|
||||
提取关系信息并返回格式化的节点和边信息
|
||||
"""
|
||||
rel_id = relationship.element_id
|
||||
nodes = relationship.nodes
|
||||
if len(nodes) != 2:
|
||||
return None, None
|
||||
|
||||
source, target = nodes
|
||||
source_id = source.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")
|
||||
if relationship_type == "unknown":
|
||||
relationship_type = relationship.type
|
||||
|
||||
edge_info = {
|
||||
"id": rel_id,
|
||||
"type": relationship_type,
|
||||
"source_id": source_id,
|
||||
"target_id": target_id,
|
||||
"source_name": source_name,
|
||||
"target_name": target_name,
|
||||
}
|
||||
|
||||
node_info = [
|
||||
{"id": source_id, "name": source_name},
|
||||
{"id": target_id, "name": target_name},
|
||||
]
|
||||
|
||||
return node_info, edge_info
|
||||
|
||||
def format_general_results(self, results):
|
||||
formatted_results = {"nodes": [], "edges": []}
|
||||
|
||||
for item in results:
|
||||
relationship = item[1]
|
||||
source_name = item[0]._properties.get("name", "unknown")
|
||||
target_name = item[2]._properties.get("name", "unknown") 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:
|
||||
continue
|
||||
|
||||
for node in node_info:
|
||||
if node["id"] not in [n["id"] for n in formatted_results["nodes"]]:
|
||||
formatted_results["nodes"].append(node)
|
||||
|
||||
formatted_results["edges"].append(edge_info)
|
||||
|
||||
return formatted_results
|
||||
|
||||
def format_query_results(self, results):
|
||||
# logger.debug(f"Graph Query Results: {results}")
|
||||
formatted_results = {"nodes": [], "edges": []}
|
||||
node_dict = {}
|
||||
|
||||
for item in results:
|
||||
# 检查数据格式
|
||||
if len(item) < 2 or not isinstance(item[1], list):
|
||||
continue
|
||||
|
||||
node_dict[item[0].element_id] = dict(id=item[0].element_id, name=item[0]._properties.get("name", "Unknown"))
|
||||
node_dict[item[2].element_id] = dict(id=item[2].element_id, name=item[2]._properties.get("name", "Unknown"))
|
||||
|
||||
# 处理关系列表中的每个关系
|
||||
for i, relationship in enumerate(item[1]):
|
||||
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
|
||||
|
||||
# 添加边
|
||||
formatted_results["edges"].append(edge_info)
|
||||
except Exception as e:
|
||||
logger.error(f"处理关系时出错: {e}, 关系: {relationship}, {traceback.format_exc()}")
|
||||
continue
|
||||
|
||||
# 将节点字典转换为列表
|
||||
formatted_results["nodes"] = list(node_dict.values())
|
||||
|
||||
return formatted_results
|
||||
|
||||
def __call__(self, query, history, meta):
|
||||
refs = self.retrieval(query, history, meta)
|
||||
query = self.construct_query(query, refs, meta)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user