KG_ICH/dofile/kg_project/merge_retry_results.py

# -*- coding: utf-8 -*-
"""
将重新抽取的结果合并到现有的节点和关系文件中
"""

import json
import pandas as pd
from pathlib import Path
import sys

# 添加模块路径
sys.path.append(str(Path(__file__).parent / 'src'))

from data_processing.entity_normalizer import EntityNormalizer
from data_processing.relationship_builder import RelationshipBuilder

def merge_retry_results():
    """合并重新抽取的结果"""

    # 文件路径
    retry_file = Path(__file__).parent / 'output' / 'retry_projects.json'
    nodes_file = Path(__file__).parent / 'output' / 'nodes_llm.csv'
    rels_file = Path(__file__).parent / 'output' / 'rels_llm.csv'
    ontology_file = Path(__file__).parent / 'config' / 'entity_ontology.yaml'

    print("=== 加载数据 ===")

    # 读取重新抽取的结果
    with open(retry_file, 'r', encoding='utf-8') as f:
        retry_data = json.load(f)

    print(f"重新抽取的项目数: {len(retry_data)}")

    # 读取现有的节点和关系
    existing_nodes_df = pd.read_csv(nodes_file, encoding='utf-8-sig')
    existing_rels_df = pd.read_csv(rels_file, encoding='utf-8-sig')

    print(f"现有节点数: {len(existing_nodes_df)}")
    print(f"现有关系数: {len(existing_rels_df)}")

    # 初始化规范化器和关系构建器
    normalizer = EntityNormalizer(str(ontology_file))
    builder = RelationshipBuilder(str(ontology_file))

    # 规范化重新抽取的实体
    print("\n=== 规范化实体 ===")

    # 收集所有需要规范化的实体
    all_entities_to_normalize = []
    for item in retry_data:
        result = item['extraction_result']
        entities = result.get('entities', [])
        # 添加项目ID以便跟踪
        for entity in entities:
            entity['_project_id'] = item['project_id']
        all_entities_to_normalize.extend(entities)

    # 批量规范化
    extraction_results = []
    for item in retry_data:
        extraction_results.append(item['extraction_result'])

    normalizer.normalize_batch(extraction_results)

    # 生成新的实体节点
    new_entity_nodes = normalizer.get_entity_nodes()
    print(f"新生成实体节点数: {len(new_entity_nodes)}")

    # 添加项目节点
    new_project_nodes = []
    for item in retry_data:
        new_project_nodes.append({
            'id': item['project_id'],
            'label': item['project_name'],
            'type': 'ICH_Project',
            'properties': '{}'
        })

    print(f"新增项目节点数: {len(new_project_nodes)}")

    # 构建关系
    print("\n=== 构建关系 ===")
    all_new_relationships = []

    # 获取当前所有节点的ID到标签映射（用于查找目标实体ID）
    all_nodes_for_lookup = pd.concat([
        existing_nodes_df,
        pd.DataFrame(new_entity_nodes)
    ], ignore_index=True)

    # 创建名称->ID的映射字典
    name_to_id = {}
    for idx, row in all_nodes_for_lookup.iterrows():
        if row['type'] != 'ICH_Project':  # 只映射实体节点
            name_to_id[(row['type'], row['label'])] = row['id']

    for item in retry_data:
        project_id = item['project_id']
        result = item['extraction_result']

        entities = result.get('entities', [])
        relationships = result.get('relationships', [])

        # 构建实体名称到ID的映射（仅限当前项目）
        entity_name_to_id = {}
        for entity in entities:
            # 实体名称可能在name字段或attributes.name字段
            entity_name = entity.get('name') or entity.get('attributes', {}).get('name', '')
            entity_type = entity.get('type', '')
            normalized_name = normalizer.normalize_text(entity_name)

            # 在规范化器的映射中查找（键是normalized_text，不是元组）
            entity_id = normalizer.text_to_id_map.get(normalized_name)
            if not entity_id:
                # 在所有节点中查找（可能是已存在的实体）
                entity_id = name_to_id.get((entity_type, entity_name))

            if entity_id:
                entity_name_to_id[entity_name] = entity_id

        # 手动构建关系
        for rel in relationships:
            # 关系中的字段名是source_entity和target_entity
            source_name = rel.get('source_entity') or rel.get('source')
            target_name = rel.get('target_entity') or rel.get('target')
            rel_type = rel.get('type')
            rel_props = rel.get('properties', {})

            # 查找源实体ID
            if source_name == project_id:
                source_id = project_id
            else:
                # 先尝试在当前项目实体中查找
                source_id = entity_name_to_id.get(source_name)
                if not source_id:
                    # 在规范化器的映射中查找
                    normalized_name = normalizer.normalize_text(source_name)
                    source_id = normalizer.text_to_id_map.get(normalized_name)

            # 查找目标实体ID
            # 先尝试在当前项目实体中查找
            target_id = entity_name_to_id.get(target_name)
            if not target_id:
                # 在规范化器的映射中查找
                normalized_name = normalizer.normalize_text(target_name)
                target_id = normalizer.text_to_id_map.get(normalized_name)

            # 如果都找到了，添加关系
            if source_id and target_id:
                all_new_relationships.append({
                    'source': source_id,
                    'target': target_id,
                    'type': rel_type,
                    'properties': json.dumps(rel_props, ensure_ascii=False) if rel_props else '{}'
                })

        print(f"项目 {project_id}: {len([r for r in all_new_relationships if r['source'] == project_id])} 条关系")

    print(f"新增关系总数: {len(all_new_relationships)}")

    # 合并节点
    print("\n=== 合并节点 ===")
    # 过滤掉已经存在的项目节点
    existing_project_ids = set(existing_nodes_df[existing_nodes_df['type'] == 'ICH_Project']['id'].tolist())
    new_project_nodes_filtered = [n for n in new_project_nodes if n['id'] not in existing_project_ids]

    all_nodes = pd.concat([
        existing_nodes_df,
        pd.DataFrame(new_entity_nodes),
        pd.DataFrame(new_project_nodes_filtered)
    ], ignore_index=True)

    print(f"合并后节点数: {len(all_nodes)} (新增 {len(new_entity_nodes) + len(new_project_nodes_filtered)} 个)")

    # 合并关系
    print("\n=== 合并关系 ===")
    all_rels = pd.concat([
        existing_rels_df,
        pd.DataFrame(all_new_relationships)
    ], ignore_index=True)

    print(f"合并后关系数: {len(all_rels)} (新增 {len(all_new_relationships)} 条)")

    # 保存结果
    print("\n=== 保存结果 ===")
    all_nodes.to_csv(nodes_file, index=False, encoding='utf-8-sig')
    all_rels.to_csv(rels_file, index=False, encoding='utf-8-sig')

    print(f"节点已保存: {nodes_file}")
    print(f"关系已保存: {rels_file}")

    # 验证
    print("\n=== 验证 ===")
    for item in retry_data:
        project_id = item['project_id']
        rel_count = len(all_rels[all_rels['source'] == project_id])
        print(f"{project_id}: {rel_count} 条关系")

    print("\n完成！")


if __name__ == '__main__':
    merge_retry_results()