kgrag/kg_build_v2/utils.py
2026-07-29 18:10:19 +08:00

76 lines
2.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import json
from collections import defaultdict
def remove_duplicates(records):
# 使用字典分组key 为 (type, name)
grouped = defaultdict(list)
for record in records:
t = record.get("type")
n = record.get("name")
if t is not None and n is not None:
grouped[(t, n)].append(record)
else:
# 如果 type 或 name 缺失,直接保留(不参与去重)
continue
result = []
for key, items in grouped.items():
if len(items) == 1:
result.append(items[0])
else:
# 多条重复,需要筛选
# 优先:保留 attributes 中 "名称" 存在的记录
valid_items = [item for item in items if item.get("attributes", {}).get("名称")]
if not valid_items:
# 如果都没有 "名称",则从所有中选字段最多的
candidates = items
else:
candidates = valid_items
# 从 candidates 中选择 attributes 非空字段最多的
best_record = max(candidates, key=lambda x: count_non_empty_attributes(x))
result.append(best_record)
# 添加未被分组的type/name 缺失的)原始记录
for record in records:
t = record.get("type")
n = record.get("name")
if t is None or n is None:
result.append(record)
return result
def count_non_empty_attributes(record):
attrs = record.get("attributes", {})
if not isinstance(attrs, dict):
return 0
# 统计非空、非None、非空字符串的字段数量
return sum(1 for v in attrs.values() if v not in [None, "", ""])
def deduplicate_relationships(relationships):
"""
对 relationships 列表进行去重,基于 type、from_entity 和 to_entity 三个字段。
参数:
relationships (list of dict): 包含关系记录的列表,每条记录为字典,
至少包含 'type', 'from_entity', 'to_entity' 三个键。
返回:
list of dict: 去重后的列表,保留首次出现的记录。
"""
seen = set()
unique_relationships = []
for rel in relationships:
# 构造一个不可变的元组作为唯一标识
key = (rel['type'], rel['from_entity'], rel['to_entity'])
if key not in seen:
seen.add(key)
unique_relationships.append(rel)
return unique_relationships