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

519 lines
23 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 os
import json
from typing import List, Tuple, Dict, Any, Optional
import re
from kg_build.extract_compose_table import get_compose_node_relation
from kg_build.extract_compose_tujie_table import get_tujiecompose_node_relation
from kg_build.extract_guzhang_table import get_guzhang_node_relation
from kg_build.extract_operation_table import get_operation_node_relation
from kg_build.extract_repair_table import get_repair_node_relation
from kg_build.extract_wxanli_table import get_wxanli_node_relation
from kg_build.get_text_node_reala import extract_text_node_relationship
from kg_build.extract_beipinbeijian_table import get_beipinbeijian_node_relation
from kg_build.extract_filtertext import get_filtertext_node_relation
import requests
from doc2pdf import Doc2PDF
import traceback
from extract_excel_node_relation import get_excel_node_relation
import logging
import asyncio # ← 新增导入
import time
# --- 日志配置 ---
logger = logging.getLogger(__name__)
if not logger.handlers:
handler = logging.StreamHandler()
fmt = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(fmt)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
converter = Doc2PDF()
def build_repair_relationships(entities: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
"""
根据维修项目编号 OR 名称在故障维修方案中的出现情况,建立“修理”关系。
匹配规则:只要编号或名称任意一个出现在维修方案文本中,即视为匹配。
"""
# 1. 数据分类:分离故障模式和维修项目
faults = [e for e in entities if e.get("type") == "故障模式"]
maintenances = [e for e in entities if e.get("type") == "维修项目"]
devices = [e for e in entities if e.get("type") == "设备"]
relationships = []
entitys = []
# 2. 遍历每一个维修项目
for mt in maintenances:
props = mt.get("properties", {})
mt_code = props.get("编号")
mt_name = props.get("名称")
# 如果连名称和编号都没有,无法进行任何匹配,直接跳过
if not mt_code and not mt_name:
continue
# 3. 遍历每一个故障模式
for fault in faults:
f_props = fault.get("properties", {})
fault_name = f_props.get("名称")
repair_text = f_props.get("维修方案", "")
if not fault_name or not repair_text:
continue
# 4. 核心匹配逻辑修改 (OR 关系)
match_by_code = mt_code in repair_text if mt_code else False
match_by_name = mt_name in repair_text if mt_name else False
# 只要满足任一条件,即建立关系
if match_by_code or match_by_name:
relationships.append({
"type": "修理",
"from_entity": mt_name, # 关系起点:维修项目名称
"to_entity": fault_name # 关系终点:故障名称
})
for mt in devices:
props = mt.get("properties", {})
mt_code = props.get("维修项目名称")
mt_name = props.get("维修项目编号")
devicename = props.get("名称")
# 如果连名称和编号都没有,无法进行任何匹配,直接跳过
if not mt_code and not mt_name:
continue
# 3. 遍历每一个故障模式
for fault in faults:
f_props = fault.get("properties", {})
fault_name = f_props.get("名称")
repair_text = f_props.get("维修方案", "")
if not fault_name or not repair_text:
continue
# 4. 核心匹配逻辑修改 (OR 关系)
match_by_code = mt_code in repair_text if mt_code else False
match_by_name = mt_name in repair_text if mt_name else False
# 只要满足任一条件,即建立关系
if match_by_code or match_by_name:
entitys.append({
"type": "维修工作",
"properties": {"名称": f"{devicename}的维修工作"}
})
relationships.append({
"type": "发生故障使用维修工作",
"from_entity": devicename,
"to_entity": f"{devicename}的维修工作"
})
relationships.append({
"type": "包含故障",
"from_entity": f"{devicename}的维修工作",
"to_entity": fault_name
})
relationships.append({
"type": "修理",
"from_entity": mt_name, # 关系起点:维修项目名称
"to_entity": fault_name
})
return entitys,relationships
async def get_md_node_relation(
content_list: str,
input_pdf_path:str,
Ontology: List[str],
Ontology_DESCRIPTION: Dict[str, List[str]],
Relationships: List[str],
RELATIONSHIP_TYPES: Dict[str, List[str]],
NODE_TYPES: Dict[str, List[str]],
PREFIX_URL : str = "http://192.168.0.111:9085",
device : str = "",
device_type : str = ""
) -> Tuple[List[Dict], List[Dict]]:
"""
核心函数:协调各模块完成知识抽取、清洗、向量化。
"""
start_time = time.time()
final_all_entity = []
final_all_relation = []
if not content_list:
logger.error("输入的 Markdown 内容为空")
return [], []
if "维修手册" in input_pdf_path or "维修说明书" in input_pdf_path or "培训手册" in input_pdf_path:
t = time.time()
all_entity_results_guzhang ,all_relation_results_guzhang =await get_guzhang_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
length_entity_guzhang = len(all_entity_results_guzhang)
length_relation_guzhang = len(all_relation_results_guzhang)
logger.info(f"故障抽取 - 实体: {length_entity_guzhang}, 关系: {length_relation_guzhang}, 耗时: {time.time()-t:.2f}s")
if length_entity_guzhang > 0 :
final_all_entity.extend(all_entity_results_guzhang)
if length_relation_guzhang > 0:
final_all_relation.extend(all_relation_results_guzhang)
t = time.time()
all_entity_results_repair ,all_relation_results_repair =await get_repair_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL)
length_entity_repair = len(all_entity_results_repair)
length_relation_repair = len(all_relation_results_repair)
logger.info(f"维修项目表抽取 - 实体: {length_entity_repair}, 关系: {length_relation_repair}, 耗时: {time.time()-t:.2f}s")
if length_entity_repair > 0 :
final_all_entity.extend(all_entity_results_repair)
if length_relation_repair > 0:
final_all_relation.extend(all_relation_results_repair)
t = time.time()
all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
length_entity_text= len(all_entity_results_text)
length_relation_text = len(all_relation_results_text)
logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
if length_entity_text > 0 :
final_all_entity.extend(all_entity_results_text)
if length_relation_text > 0:
final_all_relation.extend(all_relation_results_text)
elif "图解目录" in input_pdf_path:
# 2. 处理图解组成表
t = time.time()
all_entity_results_tujiecompose ,all_relation_results_tujiecompose =await get_tujiecompose_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device,lower_entity_type=device_type)
length_entity_tujiecompose = len(all_entity_results_tujiecompose)
length_relation_tujiecompose = len(all_relation_results_tujiecompose)
logger.info(f"图解组成表抽取 - 实体: {length_entity_tujiecompose}, 关系: {length_relation_tujiecompose}, 耗时: {time.time()-t:.2f}s")
if length_entity_tujiecompose > 0 :
final_all_entity.extend(all_entity_results_tujiecompose)
if length_relation_tujiecompose > 0:
final_all_relation.extend(all_relation_results_tujiecompose)
elif "操作使用手册" in input_pdf_path:
# 6. 处理操作项目表
t = time.time()
all_entity_results_operation ,all_relation_results_goperation =await get_operation_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
length_entity_operation= len(all_entity_results_operation)
length_relation_goperation = len(all_relation_results_goperation)
logger.info(f"操作项目表抽取 - 实体: {length_entity_operation}, 关系: {length_relation_goperation}, 耗时: {time.time()-t:.2f}s")
if length_entity_operation > 0 :
final_all_entity.extend(all_entity_results_operation)
if length_relation_goperation > 0:
final_all_relation.extend(all_relation_results_goperation)
t = time.time()
all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
length_entity_text= len(all_entity_results_text)
length_relation_text = len(all_relation_results_text)
logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
if length_entity_text > 0 :
final_all_entity.extend(all_entity_results_text)
if length_relation_text > 0:
final_all_relation.extend(all_relation_results_text)
elif "备品备件" in input_pdf_path:
# 5. 处理备品备件表
t = time.time()
all_entity_results_beipinbeijian ,all_relation_results_beipinbeijian =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
length_entity_beipinbeijian= len(all_entity_results_beipinbeijian)
length_relation_beipinbeijian = len(all_relation_results_beipinbeijian)
logger.info(f"备品备件表抽取 - 实体: {length_entity_beipinbeijian}, 关系: {length_relation_beipinbeijian}, 耗时: {time.time()-t:.2f}s")
if length_entity_beipinbeijian > 0 :
final_all_entity.extend(all_entity_results_beipinbeijian)
if length_relation_beipinbeijian > 0:
final_all_relation.extend(all_relation_results_beipinbeijian)
elif "故障案例" in input_pdf_path:
# 7. 处理维修案例
t = time.time()
all_entity_results_wxanli ,all_relation_results_wxanli =await get_wxanli_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL)
length_entity_wxanli= len(all_entity_results_wxanli)
length_relation_wxanli = len(all_relation_results_wxanli)
logger.info(f"维修案例抽取 - 实体: {length_entity_wxanli}, 关系: {length_relation_wxanli}, 耗时: {time.time()-t:.2f}s")
if length_entity_wxanli > 0 :
final_all_entity.extend(all_entity_results_wxanli)
if length_relation_wxanli > 0:
final_all_relation.extend(all_relation_results_wxanli)
else:
logger.warning("未识别的文档类型,默认进行文本抽取")
t = time.time()
all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
length_entity_text= len(all_entity_results_text)
length_relation_text = len(all_relation_results_text)
logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
if length_entity_text > 0 :
final_all_entity.extend(all_entity_results_text)
if length_relation_text > 0:
final_all_relation.extend(all_relation_results_text)
# 8. 建立修复关系
reparientitys,repair_relations = build_repair_relationships(final_all_entity)
final_all_entity.extend(reparientitys)
final_all_relation.extend(repair_relations)
logger.info(f"构建故障修复实体数量: {len(reparientitys)}")
logger.info(f"构建故障修复关系数量: {len(repair_relations)}")
logger.info(f"总实体: {len(final_all_entity)}, 总关系: {len(final_all_relation)}")
result = {
"entities": final_all_entity,
"relations": final_all_relation
}
output_json_path= "/app/wxproject.json"
# 保存为 JSON 文件
try:
with open(output_json_path, 'w', encoding='utf-8') as f:
json.dump(result, f, ensure_ascii=False, indent=4)
print(f"结果已保存至:{output_json_path}")
except Exception as e:
print(f"保存 JSON 文件时出错:{e}")
end_time = time.time()
print(f"总耗时: {end_time - start_time:.2f}")
return final_all_entity, final_all_relation
def get_content_list_json(filename, folder_path):
"""
根据传入的文件名,在指定文件夹下查找对应的 _content_list.json 文件并读取内容。
参数:
filename (str): 原始文件名,例如 "163-06A0015-B01001_发动机-操作使用手册.pdf"
folder_path (str): 目标文件夹路径
返回:
dict or None: 如果找到文件则返回 JSON 内容字典,否则返回 None
"""
# 提取文件名主体(去掉扩展名)
base_name = os.path.splitext(filename)[0]
# 构造目标 JSON 文件名
target_json_filename = f"{base_name}_content_list.json"
# 构造完整路径
target_path = os.path.join(folder_path, target_json_filename)
# 检查文件是否存在
if not os.path.exists(target_path):
print(f"未找到文件: {target_path}")
return None
# 读取并解析 JSON
try:
with open(target_path, 'r', encoding='utf-8') as f:
data = json.load(f)
return data
except json.JSONDecodeError as e:
print(f"JSON 解析错误: {e}")
return None
except Exception as e:
print(f"读取文件时发生错误: {e}")
return None
def process_file_content(
file_name: str,
file_path: str,
task_id: str,
data_dir: str,
ontology: Any,
ontology_desc: str,
relationships: List,
node_types: List,
relationship_types: List,
prefix_url: str,
api_url: str,
check_cancelled_callback: callable,
device : str,
device_type : str,
) -> Tuple[List[Dict], List[Dict]]:
"""
根据文件类型解析内容并抽取实体和关系。
Returns:
tuple: (entities, relationships)
"""
file_extension = os.path.splitext(file_name)[1].lower()
merged_output_data = {"entities": [], "relationships": []}
from fileparse_util import process_document
# --- 处理 PDF 和 MD ---
if file_extension in ['.pdf', '.md']:
async def _run_process_doc():
return await process_document(file_name)
doc_result = asyncio.run(_run_process_doc())
if doc_result is not None:
content_list, images_dict, md_content= doc_result
else:
if not os.path.exists(file_path):
raise FileNotFoundError(f"文件未找到: {file_path}")
with open(file_path, "rb") as f:
files = {"file": (os.path.basename(file_path), f, "application/pdf")}
logger.info(file_path)
response = requests.post(api_url, files=files, timeout=12000)
if response.status_code != 200:
logger.error(f"❌ API 请求失败,状态码: {response.status_code}, 响应: {response.text}")
raise Exception(f"PDF/MD 解析 API 失败: {response.status_code}")
result = response.json()
data = result.get("data", {})
content_list = data.get("content_list")
logger.info(f"[{task_id}] 开始文件 {file_name} 的实体和关系抽取")
if check_cancelled_callback():
raise InterruptedError("任务已取消")
async def run_extraction():
return await get_md_node_relation(
content_list=content_list,
input_pdf_path=file_path,
Ontology=ontology,
Ontology_DESCRIPTION=ontology_desc,
Relationships=relationships,
RELATIONSHIP_TYPES=relationship_types,
NODE_TYPES=node_types,
PREFIX_URL=prefix_url,
device = device,
device_type = device_type
)
final_all_entity, final_all_relation = asyncio.run(run_extraction())
merged_output_data = {"entities": final_all_entity, "relationships": final_all_relation}
# --- 处理 DOCX ---
elif file_extension in ['.docx']:
base_name = os.path.splitext(file_name)[0]
pdf_filename = f"{base_name}.pdf"
pdf_path = os.path.join(data_dir, pdf_filename)
logger.info(f"[{task_id}] 转换 DOCX -> PDF: {file_name}")
try:
converter.convert(file_path, output_dir=data_dir)
if not os.path.exists(pdf_path):
raise FileNotFoundError(f"PDF 转换失败: {pdf_path}")
# 与 PDF 分支保持一致:先尝试 process_document失败则上传 API
async def _run_process_doc_docx():
return await process_document(pdf_filename)
doc_result = asyncio.run(_run_process_doc_docx())
if doc_result is not None:
content_list, _ = doc_result
else:
with open(pdf_path, "rb") as f:
files = {"file": (pdf_filename, f, "application/pdf")}
logger.info(pdf_path)
response = requests.post(api_url, files=files, timeout=12000)
if response.status_code != 200:
logger.error(f"❌ API 请求失败,状态码: {response.status_code}, 响应: {response.text}")
raise Exception(f"PDF 解析 API 失败: {response.status_code}")
result = response.json()
data = result.get("data", {})
content_list = data.get("content_list")
logger.info(f"[{task_id}] 开始文件 {file_name} 的实体和关系抽取")
if check_cancelled_callback():
raise InterruptedError("任务已取消")
async def run_extraction_docx():
return await get_md_node_relation(
content_list=content_list,
input_pdf_path=pdf_path,
Ontology=ontology,
Ontology_DESCRIPTION=ontology_desc,
Relationships=relationships,
RELATIONSHIP_TYPES=relationship_types,
NODE_TYPES=node_types,
PREFIX_URL=prefix_url,
device=device,
device_type=device_type,
)
final_all_entity, final_all_relation = asyncio.run(run_extraction_docx())
merged_output_data = {"entities": final_all_entity, "relationships": final_all_relation}
except Exception as e:
logger.error(f"[{task_id}] DOCX 处理失败: {str(e)}")
logger.error(traceback.format_exc())
raise
finally:
pass
# --- 处理 XLSX ---
elif file_extension in ['.xlsx']:
base_name = os.path.splitext(file_name)[0]
pdf_filename = f"{base_name}.pdf"
logger.info(f"[{task_id}] 转换 xlsx -> PDF (仅生成索引引用): {file_name}")
# 注意:原代码中这里也调用了 converter但后续逻辑主要依赖 Excel 解析
try:
converter.convert(file_path, output_dir=os.path.join(data_dir, "uploads_v1"))
except Exception as e:
logger.warning(f"Excel 转 PDF 失败,但不影响数据提取: {e}")
all_entities, all_relationships = get_excel_node_relation(file_path, sheet_name="Sheet1")
# 补充知识源信息
for entity in all_entities:
knowledge_source = []
qiepian = entity["properties"].get("切片")
url = f"{prefix_url}/files/uploads/{pdf_filename}"
knowledge_source.append({
"filename": file_name,
"info": qiepian,
"url": url
})
entity["properties"]["knowledge_source"] = knowledge_source
merged_output_data = {"entities": all_entities, "relationships": all_relationships}
else:
raise ValueError(f"不支持的文件类型: {file_extension}")
# 统一取消检查
if check_cancelled_callback():
raise InterruptedError("任务已取消")
# 保存缓存 (可选,保持原有逻辑)
try:
safe_fname = file_name.replace(os.sep, "_").replace(" ", "_")
cache_dir = os.path.join(data_dir, "kg_cache")
os.makedirs(cache_dir, exist_ok=True)
cache_path = os.path.join(cache_dir, f"merged_output_{task_id}_{safe_fname}.json")
with open(cache_path, "w", encoding="utf-8") as cf:
json.dump(merged_output_data, cf, ensure_ascii=False, indent=2)
logger.info(f"[{task_id}] Saved merged_output_data cache: {cache_path}")
except Exception as e:
logger.warning(f"[{task_id}] Failed to save merged_output_data cache: {e}")
return merged_output_data["entities"], merged_output_data["relationships"]
def check_filename_ocr(filename: str) -> Optional[str]:
"""
根据文件名关键词判断是否需要 OCR。
Returns:
"ocr_true" — 文件名包含维修/操作类关键词
"ocr_false" — 文件名包含原理/接线/总图类关键词
None — 未命中任何关键词
"""
#ocr_true_keywords = ["维修手册", "维修说明书", "操作使用手册", "图解目录","培训手册",""]
# ocr_false_keywords = ["原理图", "接线图", "总图","随机文件目录","安装说明书","芯线表","链接图","备品备件及工具清单"]
ocr_false_keywords = ["原理图", "接线图", "总图","连接图","线图","布置图","安装图","曲线图","标志图","综合图","结构图","基座图","型线图","展开图","编号图","剖面图","吊装图","总装图","曲线图","代号图","视界图","射角图","管路图","周界图","外形图","关系图","装配图","覆盖图","尺寸图","完工图","芯线表","明细表","接线表","编制表","值表册","线路表","矩阵表","汇总表","信息表","连接表","随机文件目录", "随机文件","备品备件","图解目录","电流图","接线图和芯线表","外形图","外形图、外部接线图、安装图","设备总图", "线路图"]
# for kw in ocr_true_keywords:
# if kw in filename:
# return "ocr_true"
for kw in ocr_false_keywords:
if kw in filename:
return "ocr_false"
# else:
# return "ocr_true"
#return None
return "ocr_true"