wehub-resource-sync
151 lines
4.9 KiB
Python
151 lines
4.9 KiB
Python
"""图谱构建相关的纯函数工具集。
|
|
|
|
将数据变换逻辑从 MilvusGraphService 中抽离，
|
|
使 service 类专注于 I/O 和业务编排。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from yuxi.utils import hashstr
|
|
|
|
|
|
def normalize_entity_name(text: str) -> str:
|
|
"""统一实体名称：去首尾空白、小写化、压缩内部连续空白。"""
|
|
return " ".join(text.strip().lower().split())
|
|
|
|
|
|
def compute_entity_id(kb_id: str, normalized_name: str, label: str) -> str:
|
|
return hashstr(f"{kb_id}:{normalized_name}:{label}", length=32)
|
|
|
|
|
|
def compute_triple_id(
|
|
kb_id: str,
|
|
source_normalized_name: str,
|
|
source_label: str,
|
|
relation_type: str,
|
|
target_normalized_name: str,
|
|
target_label: str,
|
|
) -> str:
|
|
return hashstr(
|
|
f"{kb_id}:{source_normalized_name}:{source_label}:{relation_type}:{target_normalized_name}:{target_label}",
|
|
length=32,
|
|
)
|
|
|
|
|
|
def graph_entity_collection_name(kb_id: str) -> str:
|
|
return f"{kb_id}_entity"
|
|
|
|
|
|
def graph_triple_collection_name(kb_id: str) -> str:
|
|
return f"{kb_id}_triple"
|
|
|
|
|
|
def build_graph_payload(normalized_result: dict[str, Any]) -> dict[str, Any]:
|
|
"""将抽取器产出的标准化结果转换为 Neo4j 写入所需的图结构。
|
|
|
|
返回的 entities 已完成去重合并：同名同 label 的实体只保留一份，
|
|
属性（attributes）取并集。
|
|
"""
|
|
entities: list[dict[str, Any]] = []
|
|
entity_by_key: dict[tuple[str, str], dict[str, Any]] = {}
|
|
|
|
def add_entity(entity: dict[str, Any]) -> str:
|
|
key = (normalize_entity_name(entity["text"]), entity.get("label") or "Entity")
|
|
existing = entity_by_key.get(key)
|
|
if existing is not None:
|
|
known_attributes = {(attr["text"], attr["label"]) for attr in existing.get("attributes") or []}
|
|
for attribute in entity.get("attributes") or []:
|
|
attribute_key = (attribute["text"], attribute["label"])
|
|
if attribute_key not in known_attributes:
|
|
existing.setdefault("attributes", []).append(attribute)
|
|
known_attributes.add(attribute_key)
|
|
return existing["id"]
|
|
|
|
graph_entity = {
|
|
"id": f"e{len(entities) + 1}",
|
|
"text": entity["text"],
|
|
"label": entity.get("label") or "Entity",
|
|
"attributes": list(entity.get("attributes") or []),
|
|
}
|
|
entities.append(graph_entity)
|
|
entity_by_key[key] = graph_entity
|
|
return graph_entity["id"]
|
|
|
|
for entity in normalized_result["entities"]:
|
|
add_entity(entity)
|
|
|
|
relations = []
|
|
for relation in normalized_result["relations"]:
|
|
relations.append(
|
|
{
|
|
"source": add_entity(relation["source"]),
|
|
"target": add_entity(relation["target"]),
|
|
"text": relation["text"],
|
|
"label": relation.get("label") or "RELATED_TO",
|
|
}
|
|
)
|
|
|
|
return {"entities": entities, "relations": relations, "metadata": normalized_result["metadata"]}
|
|
|
|
|
|
# ─── Cypher 模板 ────────────────────────────────────────────────
|
|
# 将大段 Cypher 字符串集中管理，提升 write_chunk_graph 的可读性。
|
|
|
|
|
|
def cypher_merge_chunk(db_label: str) -> str:
|
|
"""MERGE Chunk 节点并写入元数据。"""
|
|
return f"""
|
|
MERGE (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
|
|
SET c.file_id = $file_id,
|
|
c.kb_id = $kb_id,
|
|
c.chunk_index = $chunk_index,
|
|
c.content_preview = $content_preview,
|
|
c.start_char_pos = $start_char_pos,
|
|
c.end_char_pos = $end_char_pos
|
|
"""
|
|
|
|
|
|
def cypher_merge_entity_mention(db_label: str) -> str:
|
|
"""MERGE Entity 节点并创建 Chunk → Entity 的 MENTIONS 关系。"""
|
|
return f"""
|
|
MATCH (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
|
|
MERGE (e:Entity:MilvusKB:`{db_label}` {{
|
|
kb_id: $kb_id,
|
|
normalized_name: $normalized_name,
|
|
label: $entity_label
|
|
}})
|
|
SET e.entity_id = $entity_id,
|
|
e.name = $name,
|
|
e.attributes = $attributes
|
|
MERGE (c)-[m:MENTIONS {{chunk_id: $chunk_id, file_id: $file_id, kb_id: $kb_id}}]->(e)
|
|
"""
|
|
|
|
|
|
def cypher_merge_relation(db_label: str) -> str:
|
|
"""MERGE 两个 Entity 之间的 RELATION 边。"""
|
|
return f"""
|
|
MATCH (source:Entity:MilvusKB:`{db_label}` {{
|
|
kb_id: $kb_id,
|
|
normalized_name: $source_name,
|
|
label: $source_label
|
|
}})
|
|
MATCH (target:Entity:MilvusKB:`{db_label}` {{
|
|
kb_id: $kb_id,
|
|
normalized_name: $target_name,
|
|
label: $target_label
|
|
}})
|
|
MERGE (source)-[r:RELATION {{
|
|
kb_id: $kb_id,
|
|
chunk_id: $chunk_id,
|
|
source_name: $source_name,
|
|
target_name: $target_name,
|
|
type: $relation_type
|
|
}}]->(target)
|
|
SET r.triple_id = $triple_id,
|
|
r.text = $text,
|
|
r.file_id = $file_id,
|
|
r.extractor_type = $extractor_type
|
|
"""
|