121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
"""图谱构建相关的纯函数工具集。
|
||
|
||
将数据变换逻辑从 MilvusGraphService 中抽离，
|
||
使 service 类专注于 I/O 和业务编排。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from typing import Any
|
||
|
||
|
||
def normalize_entity_name(text: str) -> str:
|
||
"""统一实体名称：去首尾空白、小写化、压缩内部连续空白。"""
|
||
return " ".join(text.strip().lower().split())
|
||
|
||
|
||
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.db_id = $db_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}` {{
|
||
db_id: $db_id,
|
||
normalized_name: $normalized_name,
|
||
label: $entity_label
|
||
}})
|
||
SET e.name = $name,
|
||
e.attributes = $attributes
|
||
MERGE (c)-[m:MENTIONS {{chunk_id: $chunk_id, file_id: $file_id, db_id: $db_id}}]->(e)
|
||
"""
|
||
|
||
|
||
def cypher_merge_relation(db_label: str) -> str:
|
||
"""MERGE 两个 Entity 之间的 RELATION 边。"""
|
||
return f"""
|
||
MATCH (source:Entity:MilvusKB:`{db_label}` {{
|
||
db_id: $db_id,
|
||
normalized_name: $source_name,
|
||
label: $source_label
|
||
}})
|
||
MATCH (target:Entity:MilvusKB:`{db_label}` {{
|
||
db_id: $db_id,
|
||
normalized_name: $target_name,
|
||
label: $target_label
|
||
}})
|
||
MERGE (source)-[r:RELATION {{
|
||
db_id: $db_id,
|
||
chunk_id: $chunk_id,
|
||
source_name: $source_name,
|
||
target_name: $target_name,
|
||
type: $relation_type
|
||
}}]->(target)
|
||
SET r.text = $text,
|
||
r.file_id = $file_id,
|
||
r.extractor_type = $extractor_type
|
||
""" |