264 lines
9.3 KiB
Python
264 lines
9.3 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
import fitz
|
||
from collections import Counter
|
||
import uuid
|
||
from openai import OpenAI
|
||
from kg_build.Prompt import Prompt_anli
|
||
from extract_util import safe_json_loads,highlight_pdf_file,get_content_bbox
|
||
from modelsAPI.model_api import OpenaiAPI
|
||
|
||
# 指定 JSON 文件路径
|
||
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/890DAA发动机装备故障案例_content_list.json'
|
||
|
||
# try:
|
||
# # 打开并读取 JSON 文件
|
||
# with open(file_path, 'r', encoding='utf-8') as file:
|
||
# data = json.load(file)
|
||
|
||
# except Exception as e:
|
||
# print(f"发生未知错误：{e}")
|
||
|
||
|
||
# for idx, item in enumerate(data, start=1):
|
||
# item["id"] = idx
|
||
def extract_chunks_by_start_keywords(data, start_keywords):
|
||
"""
|
||
通用版：根据多个起始关键词提取段落块。
|
||
|
||
参数：
|
||
- data: List[dict]，每项含 "id", "type", "text"（若 type=="text"）
|
||
- start_keywords: List[str]，如 ["一、故障现象", "二、原因分析"]
|
||
|
||
规则：
|
||
- 起始点：type == "text" 且 text.strip() in start_keywords
|
||
- 每个段落包含：
|
||
a) 起始点的前一条记录（id-1，如果存在）
|
||
b) 起始点本身
|
||
c) 起始点之后的所有连续记录（任意 type），直到（但不包含）下一个起始关键词出现
|
||
- 所有被选中的记录将从原始 data 中移除
|
||
|
||
返回：
|
||
- groups: List[List[record]]，每个 group 是一个完整段落（含前一条）
|
||
- filtered_data: 剩余未被移除的记录列表
|
||
"""
|
||
if not data or not start_keywords:
|
||
return [], data
|
||
|
||
# 构建 id -> record 映射
|
||
record_map = {item["id"]: item for item in data}
|
||
max_id = max(record_map.keys())
|
||
|
||
# 第一步：找出所有起始关键词的位置（按 id 顺序）
|
||
start_keyword_set = set(kw.strip() for kw in start_keywords)
|
||
start_records = []
|
||
|
||
for item in data:
|
||
if item.get("type") == "text":
|
||
text_content = item.get("text", "").strip()
|
||
if text_content in start_keyword_set:
|
||
start_records.append(item)
|
||
|
||
if not start_records:
|
||
return [], data
|
||
|
||
# 按 id 排序（虽然 data 通常有序，但保险起见）
|
||
start_records.sort(key=lambda x: x["id"])
|
||
|
||
groups = []
|
||
ids_to_remove = set()
|
||
|
||
# 第二步：遍历每个起始点，构建段落
|
||
for record in start_records:
|
||
start_id = record["id"]
|
||
group = []
|
||
|
||
# 1. 添加前一条记录（id-1）
|
||
prev_id = start_id - 1
|
||
if prev_id >= 1 and prev_id in record_map:
|
||
group.append(record_map[prev_id])
|
||
ids_to_remove.add(prev_id)
|
||
|
||
# 2. 添加当前起始记录
|
||
group.append(record)
|
||
ids_to_remove.add(start_id)
|
||
|
||
# 3. 向后收集，直到遇到下一个起始关键词（不含）
|
||
next_id = start_id + 1
|
||
while next_id <= max_id:
|
||
next_record = record_map.get(next_id)
|
||
if next_record is None:
|
||
break
|
||
|
||
# 检查是否是任意一个起始关键词
|
||
if (next_record.get("type") == "text" and
|
||
next_record.get("text", "").strip() in start_keyword_set):
|
||
break # 遇到下一个起始点，停止
|
||
|
||
group.append(next_record)
|
||
ids_to_remove.add(next_id)
|
||
next_id += 1
|
||
|
||
groups.append(group)
|
||
|
||
# 第三步：过滤原始 data
|
||
filtered_data = [item for item in data if item["id"] not in ids_to_remove]
|
||
|
||
return groups, filtered_data
|
||
|
||
# def openai_chat(query: str, model="Qwen3-32B"):
|
||
# """调用OpenAI API"""
|
||
# try:
|
||
# client = OpenAI(
|
||
# api_key="none",
|
||
# base_url="http://192.168.0.46:59800/v1"
|
||
# )
|
||
|
||
# response = client.chat.completions.create(
|
||
# model=model,
|
||
# messages=[
|
||
# {"role": "user", "content": query},
|
||
# ],
|
||
# temperature=0.1,
|
||
# stream=False,
|
||
# response_format={"type": "json_object"},
|
||
# extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||
# )
|
||
# return response.choices[0].message.content
|
||
# except Exception as e:
|
||
# print(f"调用OpenAI API时出错: {e}")
|
||
# return None
|
||
|
||
async def extract_wxanli(chunk):
|
||
"""
|
||
成功时返回 {'entities': [...], 'relationships': [...]},
|
||
任何失败（超时、解析错、格式错）都返回 []。
|
||
"""
|
||
final_prompt = Prompt_anli.replace("text", chunk)
|
||
raw_response =await OpenaiAPI.openai_chat_aysnc(final_prompt,timeout=600)
|
||
|
||
# 如果模型没返回内容（如超时），直接返回 []
|
||
if raw_response is None:
|
||
return []
|
||
|
||
try:
|
||
cleaned = raw_response.strip()
|
||
if cleaned.startswith("```json"):
|
||
cleaned = cleaned[7:].lstrip()
|
||
if cleaned.endswith("```"):
|
||
cleaned = cleaned[:-3].rstrip()
|
||
|
||
json_result = safe_json_loads(cleaned)
|
||
|
||
# 必须是 dict，且包含 entities 字段（且为 list）
|
||
if not isinstance(json_result, dict):
|
||
return []
|
||
|
||
entities = json_result.get("entities", [])
|
||
relationships = json_result.get("relationships", [])
|
||
|
||
if not isinstance(entities, list) or not isinstance(relationships, list):
|
||
return []
|
||
|
||
# 返回标准结构
|
||
return {
|
||
"entities": entities,
|
||
"relationships": relationships
|
||
}
|
||
|
||
except Exception:
|
||
# 任何解析异常都返回 []
|
||
return []
|
||
|
||
async def get_wxanli_node_relation(input_pdf_path,data,prefix_url = "http://192.168.0.46:59085"):
|
||
|
||
filter = [
|
||
"一、故障现象",
|
||
"单位:",
|
||
"故障名称:",
|
||
"一、基本情况",
|
||
"故障现象及原因:",
|
||
"故障名称：",
|
||
"实例故障"
|
||
]
|
||
results, filtered_data = extract_chunks_by_start_keywords(data=data,start_keywords=filter)
|
||
if not results:
|
||
print("未找到符合条件的图解组成表格，跳过处理。")
|
||
all_entity_results,all_relation_results = [],[]
|
||
else:
|
||
all_entity_results = []
|
||
all_relation_results = []
|
||
|
||
for i, compose in enumerate(results):
|
||
try:
|
||
print(f"开始处理第{i+1}个维修案例")
|
||
content, highlight_list = get_content_bbox(compose)
|
||
print("维修案例切片内容为：", content[:200])
|
||
|
||
result =await extract_wxanli(content)
|
||
|
||
# 关键判断：只有 result 是 dict 才继续
|
||
if not isinstance(result, dict):
|
||
print(f" ⚠️ 第 {i+1} 组：提取结果无效，跳过")
|
||
continue
|
||
|
||
entities = result.get("entities", [])
|
||
relationships = result.get("relationships", [])
|
||
|
||
# 添加 knowledge_source
|
||
highlight_filename, most_common_page,filename = highlight_pdf_file(input_pdf_path,highlight_list)
|
||
url = f"{prefix_url}/upload/{highlight_filename}#page={most_common_page}"
|
||
# for entity in entities:
|
||
# knowledge_source = [{
|
||
# "filename": filename,
|
||
# "info": content,
|
||
# "url": url
|
||
# }]
|
||
# if "properties" not in entity or not isinstance(entity["properties"], dict):
|
||
# entity["properties"] = {}
|
||
# entity["properties"]["knowledge_source"] = knowledge_source
|
||
for entity in entities:
|
||
# 构造原始数据对象 (List of Dict)
|
||
source_data_list = [{
|
||
"filename": filename,
|
||
"info": content,
|
||
"url": url
|
||
}]
|
||
|
||
# 2. 【关键修改】将列表序列化为 JSON 字符串
|
||
# ensure_ascii=False 确保中文内容不会被转义成 \uXXXX
|
||
knowledge_source_str = json.dumps(source_data_list, ensure_ascii=False)
|
||
|
||
# 确保 properties 存在且是字典
|
||
if "properties" not in entity or not isinstance(entity["properties"], dict):
|
||
entity["properties"] = {}
|
||
|
||
# 3. 赋值字符串而不是对象列表
|
||
entity["properties"]["knowledge_source"] = knowledge_source_str
|
||
# 安全追加
|
||
all_entity_results.extend(entities)
|
||
all_relation_results.extend(relationships)
|
||
|
||
except Exception as e:
|
||
print(f" ❌ 处理第 {i+1} 组时发生未预期错误: {e}")
|
||
# 不 re-raise，继续下一轮
|
||
continue
|
||
return all_entity_results,all_relation_results
|
||
# if __name__ == "__main__":
|
||
# file_path = '/app/files/890DAA发动机装备故障案例_content_list.json'
|
||
# input_pdf_path = "/app/files/890DAA发动机装备故障案例.pdf"
|
||
|
||
# try:
|
||
# with open(file_path, 'r', encoding='utf-8') as file:
|
||
# data = json.load(file)
|
||
# except Exception as e:
|
||
# print(f"读取 JSON 文件时出错：{e}")
|
||
# data = []
|
||
# all_entity_results,all_relation_results = get_wxanli_node_relation(input_pdf_path,data=data)
|
||
# length_entity = len(all_entity_results)
|
||
# length_relation = len(all_relation_results)
|
||
|
||
# print(f"抽取{length_entity}个实体，抽取{length_relation}个关系")
|
||
|