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

155 lines
5.0 KiB
Python
Raw Permalink 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 re
from openai import OpenAI
from typing import List, Dict
import json
import time
import os
from kg_build_v2.extract_html import extract_repairproject_tables_html,extract_operation_tables_html,extract_tables
from kg_build_v2.Prompt import prompt_extract_repair_node_rela
from modelsAPI.model_api import OpenaiAPI
def read_markdown_file(file_path):
"""
读取 Markdown 文件并返回其全部内容作为字符串。
参数:
file_path (str): Markdown 文件的路径(支持相对或绝对路径)
返回:
str 或 None: 成功时返回文件内容;出错时返回 None
"""
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
return content
except FileNotFoundError:
print(f"错误:文件未找到 - {file_path}")
except PermissionError:
print(f"错误:没有权限读取文件 - {file_path}")
except UnicodeDecodeError:
print(f"错误:文件编码不是 UTF-8无法读取 - {file_path}")
except Exception as e:
print(f"读取文件时发生未知错误: {e}")
return None
def safe_json_loads(s: str):
try:
return json.loads(s)
except json.JSONDecodeError as e:
# 尝试修复常见问题
s = s.replace("\\n", "\\\n").replace("\r", "\\r").replace("\t", "\\t")
# 或者直接清理非法字符
s = re.sub(r'[\x00-\x1f\x7f]', '', s) # 移除控制字符
try:
return json.loads(s)
except:
raise e
def extract_wxindextable(chunks, skip_first=True):
"""
批量处理一个文档的多个切片chunks每个切片调用大模型结构化并返回全部结果。
:param chunks: List[str] - 同一文档的多个文本切片
:param skip_first: bool - 是否跳过第一个表格块默认True
:return: List[Dict] - 结构化结果列表
"""
all_results = []
# 确定起始索引
start_idx = 1 if skip_first else 0
for idx, block in enumerate(chunks, 1):
# 跳过第一个表格
if skip_first and idx == 1:
print(f"=== 跳过第 {idx} 个表格块(首个表格) ===\n")
continue
print(f"=== 维修项目 Block {idx} ===")
print(block[:100]) # 只打印前500字符
print("\n" + "="*60 + "\n")
final_prompt = prompt_extract_repair_node_rela.format(text=block)
raw_response = OpenaiAPI.openai_chat(final_prompt,timeout=180)
if raw_response is None:
print(f" ⚠️ 第 {idx} 切片 API 调用失败")
continue
# 尝试解析 JSON
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)
if "properties" in json_result:
json_result["properties"]["切片"] = block # 添加原始文本块
else:
print(f" ⚠️ 第 {idx} 切片的 JSON 中缺少 'properties' 字段,跳过添加切片")
# 打印美化版
all_results.append(json_result)
except Exception as e:
print(f" ⚠️ 第 {idx} 切片 JSON 解析失败: {e}")
print(f" 原始响应: {repr(raw_response[:500])}")
json_result = None
continue
return all_results
# 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
def get_repair_node_rela(mdcontent):
print("开始提取维修项目表实体和关系")
chunks, cleaned_content = extract_repairproject_tables_html(mdcontent)
print(f"找到 {len(chunks)} 个维修表格块")
if len(chunks) > 0:
result = extract_wxindextable(chunks=chunks)
return result
else:
return []
# if __name__ == "__main__":
# input_directory = "F:\zklnlp\HJ\hj_kg_code\kgrag\extract_html_json\测试wx手册_shenghna_output11_full_content.md"
# # 调用目录处理函数
# mdcontent = read_markdown_file(input_directory)
# if mdcontent is None:
# print("无法读取文件,程序退出")
# exit(1)
# result = get_repair_node_rela(mdcontent=mdcontent)
# print(len(result))