1141 lines
42 KiB
Python
1141 lines
42 KiB
Python
import json
|
||
import fitz # PyMuPDF
|
||
import os
|
||
from collections import Counter
|
||
import uuid
|
||
from openai import OpenAI
|
||
import re
|
||
from typing import List, Dict, Optional, Union, Any, Tuple
|
||
# 指定 JSON 文件路径
|
||
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/122-06A0014-B01001_发动机-维修手册_content_list.json'
|
||
from filename_proceess_and_kgquery import get_entity
|
||
from pathlib import PurePath, Path as PathLib
|
||
from fastapi import FastAPI, File, Path, UploadFile, HTTPException, Form, Request, Header, Body
|
||
from fileparse_util import process_document
|
||
import threading
|
||
import aiofiles
|
||
import httpx
|
||
import logging
|
||
from config import SHIP_MODEL_NAME,API_URLS
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
_api_url_inflight = [0] * len(API_URLS)
|
||
_api_url_lock = threading.Lock()
|
||
|
||
def format_entity_text(data):
|
||
# 1. 定义实体类型与属性的映射关系
|
||
# 格式: '实体类型': [ ('前缀1', '属性1'), ('前缀2', '属性2') ]
|
||
type_mapping = {
|
||
'舰艇': [('舰艇为', '名称'), ('，舷号', '舷号')],
|
||
'系统': [('，系统：', '名称')],
|
||
'子系统': [('，子系统：', '名称')],
|
||
'设备': [('，设备：', '名称')]
|
||
}
|
||
|
||
text_parts = []
|
||
|
||
# 2. 遍历 entities 列表
|
||
for entity in data.get('entities', []):
|
||
entity_type = entity.get('type')
|
||
|
||
properties = entity.get('properties', {})
|
||
|
||
# 3. 如果当前实体类型在映射表中，则进行拼接
|
||
if entity_type in type_mapping:
|
||
entity_str = ""
|
||
# 遍历该实体需要提取的所有属性
|
||
for prefix, attr_key in type_mapping[entity_type]:
|
||
value = str(properties.get(attr_key, ''))
|
||
entity_str += f"{prefix}{value}"
|
||
text_parts.append(entity_str)
|
||
|
||
# 4. 将所有部分拼接成一段完整的文本
|
||
# 注意：因为前缀中已经包含了逗号，这里直接拼接即可
|
||
return ''.join(text_parts)
|
||
|
||
def get_api_url(task_id: str = "") -> tuple:
|
||
"""返回 (analyze_url, contentlist_url)，选当前 inflight 最少的实例。"""
|
||
with _api_url_lock:
|
||
idx = _api_url_inflight.index(min(_api_url_inflight))
|
||
_api_url_inflight[idx] += 1
|
||
base = API_URLS[idx]
|
||
return base
|
||
|
||
|
||
def release_api_url(url: str):
|
||
with _api_url_lock:
|
||
if url in API_URLS:
|
||
idx = API_URLS.index(url)
|
||
_api_url_inflight[idx] = max(0, _api_url_inflight[idx] - 1)
|
||
def extract_text_from_first_tr(html: str) -> str:
|
||
"""提取 HTML 表格中第一个 <tr> 内的文本（去除标签）"""
|
||
if not isinstance(html, str):
|
||
return ""
|
||
match = re.search(r'<tr[^>]*>(.*?)</tr>', html, re.IGNORECASE | re.DOTALL)
|
||
if not match:
|
||
return ""
|
||
first_tr_content = match.group(1)
|
||
# 去除所有 HTML 标签
|
||
text = re.sub(r'<[^>]+>', '', first_tr_content)
|
||
return text.strip()
|
||
|
||
|
||
def matches_any_pattern(text: str, patterns: List[List[str]]) -> bool:
|
||
"""检查文本是否包含任意一组关键词（全部命中）"""
|
||
if not text:
|
||
return False
|
||
text_lower = text.lower()
|
||
for pattern in patterns:
|
||
if all(keyword in text_lower for keyword in pattern):
|
||
return True
|
||
return False
|
||
|
||
|
||
def is_maintenance_table(table_body: str, patterns: List[List[str]]) -> bool:
|
||
"""判断是否为维修项目表格：仅检查第一个 <tr> 是否匹配任一模式"""
|
||
first_tr_text = extract_text_from_first_tr(table_body)
|
||
return matches_any_pattern(first_tr_text, patterns)
|
||
|
||
|
||
# ==============================
|
||
# 维修项目表格的关键词模式（可按需扩展）
|
||
# ==============================
|
||
|
||
MAINTENANCE_TABLE_PATTERNS: List[List[str]] = [
|
||
["维修项目编号", "组成编码", "名称", "维修项目", "维修间隔期"],
|
||
["维修项目编号", "维修级别", "名称"],
|
||
["项目编号", "编码", "维修内容", "周期"],
|
||
["序号", "维修项", "标准", "频次"], # 示例：可继续添加业务变体
|
||
]
|
||
|
||
|
||
# ==============================
|
||
# 主函数：提取维修项目表格及其后续内容组
|
||
# ==============================
|
||
|
||
def extract_maintenance_groups_and_remaining(
|
||
records: List[Dict[str, Any]],
|
||
patterns: List[List[str]] = MAINTENANCE_TABLE_PATTERNS,
|
||
max_following_for_last: int | None = None
|
||
) -> Tuple[List[List[Dict]], List[Dict]]:
|
||
groups = []
|
||
current_group = None
|
||
used_ids = set()
|
||
|
||
# 标记是否已经遇到第一个维修表格（用于跳过）
|
||
skipped_first = False
|
||
|
||
for record in records:
|
||
is_start = (
|
||
record.get("type") == "table" and
|
||
is_maintenance_table(record.get("table_body", ""), patterns)
|
||
)
|
||
|
||
if is_start:
|
||
# 如果是第一个维修表，跳过（不加入 groups，也不标记 used_ids）
|
||
if not skipped_first:
|
||
skipped_first = True
|
||
# 注意：这里不清空 current_group，因为前面不可能有 group（第一个 start）
|
||
current_group = None # 确保不会把之前的非 start 内容误加
|
||
continue # 跳过这个 record，不加入任何 group
|
||
|
||
# 从第二个维修表开始，才正常分组
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
current_group = [record]
|
||
else:
|
||
if current_group is not None:
|
||
current_group.append(record)
|
||
|
||
# 处理最后一个 group（仅当有有效 group 时）
|
||
if current_group is not None:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
|
||
# 限制最后一个 group 长度
|
||
if groups and max_following_for_last is not None:
|
||
last_group = groups[-1]
|
||
max_len = max_following_for_last + 1
|
||
if len(last_group) > max_len:
|
||
truncated = last_group[:max_len]
|
||
removed_ids = {r["id"] for r in last_group[max_len:]}
|
||
used_ids -= removed_ids
|
||
groups[-1] = truncated
|
||
|
||
# remaining_records 包含所有未被 used_ids 标记的记录
|
||
# 由于第一个 group 没有加入 used_ids，所以会保留在 remaining_records 中
|
||
remaining_records = [r for r in records if r["id"] not in used_ids]
|
||
return groups, remaining_records
|
||
|
||
|
||
def is_operation_table_by_first_tr(table_body: str) -> bool:
|
||
"""
|
||
判断 table_body 是否为操作项目表，依据是第一行 <tr> 中是否包含 '项目编号' 和 '操作项目'
|
||
"""
|
||
if not isinstance(table_body, str):
|
||
return False
|
||
|
||
# 非贪婪匹配第一个 <tr> ... </tr>
|
||
match = re.search(r'<tr[^>]*>(.*?)</tr>', table_body, re.IGNORECASE | re.DOTALL)
|
||
if not match:
|
||
return False
|
||
|
||
first_tr_html = match.group(1)
|
||
|
||
# 去除所有 HTML 标签，提取纯文本
|
||
text_in_first_tr = re.sub(r'<[^>]+>', '', first_tr_html)
|
||
|
||
# 转为小写（中文无所谓大小写，但统一处理更安全）
|
||
text_clean = text_in_first_tr.strip()
|
||
|
||
# 检查是否同时包含两个关键词
|
||
has_code = "项目编号" in text_clean
|
||
has_item = "操作项目" in text_clean
|
||
return has_code and has_item
|
||
|
||
def extract_operation_groups_and_remaining(records, max_following_for_last=23):
|
||
"""
|
||
通过语义判断表格是否为操作项目表（基于第一行内容），并分组。
|
||
注意：第一个识别到的操作项目表及其后续记录组成的 group 会被跳过（不加入 groups），
|
||
但这些记录仍保留在 remaining_records 中。
|
||
|
||
普通起始标记：组 = [start, ..., next_start - 1]
|
||
最后一个起始标记：组 = [start, ..., start + max_following_for_last] （若指定）
|
||
"""
|
||
groups = []
|
||
current_group = None
|
||
used_ids = set()
|
||
first_group_skipped = False # 标记是否已跳过第一个 group
|
||
|
||
for record in records:
|
||
is_start = (
|
||
record.get("type") == "table"
|
||
and is_operation_table_by_first_tr(record.get("table_body", ""))
|
||
)
|
||
|
||
if is_start:
|
||
if current_group is not None:
|
||
# 决定是否保留当前积累的 group
|
||
if not first_group_skipped:
|
||
# 跳过第一个 group：不清空 current_group，也不加入 groups 或 used_ids
|
||
first_group_skipped = True
|
||
# 注意：这里不把 current_group 加入 groups，也不更新 used_ids
|
||
else:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
# 开启新 group
|
||
current_group = [record]
|
||
else:
|
||
# 第一个起始点
|
||
current_group = [record]
|
||
else:
|
||
if current_group is not None:
|
||
current_group.append(record)
|
||
|
||
# 处理最后一个 group（如果存在）
|
||
if current_group is not None:
|
||
if not first_group_skipped:
|
||
# 整个数据中只有一个 group，且是第一个 → 跳过它
|
||
pass # 不加入 groups，也不加 id 到 used_ids
|
||
else:
|
||
groups.append(current_group)
|
||
used_ids.update(r["id"] for r in current_group)
|
||
|
||
# —————— 后处理：截断最后一个 group（如果需要） ——————
|
||
if groups and max_following_for_last is not None:
|
||
last_group = groups[-1]
|
||
if len(last_group) > max_following_for_last + 1: # +1 是起始记录本身
|
||
truncated = last_group[:max_following_for_last + 1]
|
||
removed_ids = {r["id"] for r in last_group[max_following_for_last + 1:]}
|
||
used_ids -= removed_ids
|
||
groups[-1] = truncated
|
||
|
||
remaining_records = [r for r in records if r["id"] not in used_ids]
|
||
return groups, remaining_records
|
||
|
||
|
||
def extract_wxanli(records):
|
||
keywords = [
|
||
"一、故障现象",
|
||
"单位:",
|
||
"故障名称:",
|
||
"一、基本情况",
|
||
"故障现象及原因:",
|
||
"故障名称：",
|
||
"实例故障"
|
||
]
|
||
|
||
# 存储切片
|
||
slices = []
|
||
# 标记已处理的记录索引
|
||
processed_indices = set()
|
||
|
||
# 查找所有包含关键词的记录位置
|
||
keyword_positions = []
|
||
for i, record in enumerate(records):
|
||
if record.get("type") == "text":
|
||
text = record.get("text", "")
|
||
# 检查是否包含任一关键词
|
||
if any(keyword in text for keyword in keywords):
|
||
keyword_positions.append(i)
|
||
|
||
# 为每个关键词位置创建切片
|
||
for idx, start_pos in enumerate(keyword_positions):
|
||
# 确定切片结束位置
|
||
if idx + 1 < len(keyword_positions):
|
||
end_pos = keyword_positions[idx + 1]
|
||
else:
|
||
end_pos = len(records)
|
||
|
||
# 收集当前切片
|
||
slice_records = []
|
||
|
||
# 添加前2个记录（如果存在且未超出边界）
|
||
for offset in [2, 1]:
|
||
prev_idx = start_pos - offset
|
||
if prev_idx >= 0:
|
||
slice_records.append(records[prev_idx])
|
||
processed_indices.add(prev_idx)
|
||
|
||
# 添加从关键词位置到下一个关键词位置的所有记录
|
||
for i in range(start_pos, end_pos):
|
||
slice_records.append(records[i])
|
||
processed_indices.add(i)
|
||
|
||
slices.append(slice_records)
|
||
|
||
# 收集未处理的剩余记录
|
||
remaining_records = [
|
||
records[i] for i in range(len(records))
|
||
if i not in processed_indices
|
||
]
|
||
|
||
return slices, remaining_records
|
||
|
||
def split_content_bbox(data, max_length=8000):
|
||
contents = []
|
||
highlight_lists = []
|
||
|
||
current_content = ""
|
||
current_highlights = []
|
||
|
||
for ins in data:
|
||
if "type" not in ins:
|
||
continue
|
||
if not ins.get("bbox"):
|
||
continue
|
||
# 提取通用字段
|
||
bbox = ins.get("bbox")
|
||
page_idx = ins.get("page_idx", -1)
|
||
|
||
new_text = ""
|
||
should_record_highlight = False # 标记是否需要记录 highlight（即使 new_text 为空）
|
||
|
||
if ins["type"] == "text":
|
||
text = ins.get("text", "")
|
||
level = ins.get("text_level")
|
||
if level is not None and isinstance(level, int) and level > 0:
|
||
text = "#" * min(level, 6) + " " + text # 安全限制标题级别
|
||
new_text = text + "\n"
|
||
should_record_highlight = True
|
||
|
||
elif ins["type"] == "table":
|
||
table_caption = ins.get("table_caption")
|
||
table_body = ins.get("table_body", "")
|
||
|
||
# 清理冗余属性
|
||
for pattern in [' colspan="1"', ' rowspan="1"', " colspan='1'", " rowspan='1'", " colspan=1", " rowspan=1"]:
|
||
table_body = table_body.replace(pattern, "")
|
||
|
||
caption_str = ""
|
||
if table_caption:
|
||
if isinstance(table_caption, list):
|
||
caption_str = "\n".join(str(c) for c in table_caption)
|
||
else:
|
||
caption_str = str(table_caption)
|
||
|
||
# 判断是否有实质内容
|
||
has_caption = bool(caption_str.strip())
|
||
has_body = bool(table_body.strip())
|
||
|
||
if has_caption or has_body:
|
||
new_text = caption_str + ("\n" if has_caption else "") + table_body + "\n"
|
||
else:
|
||
new_text = "" # 空内容，不拼接
|
||
|
||
# 即使内容为空，只要是 table 类型，就记录 highlight（满足你的需求）
|
||
should_record_highlight = True
|
||
|
||
elif ins["type"] == "image":
|
||
image_caption = ins.get("image_caption", "")
|
||
img_path = ins.get("img_path", "")
|
||
|
||
if isinstance(image_caption, list):
|
||
image_caption = "\n".join(str(item) for item in image_caption)
|
||
elif not isinstance(image_caption, str):
|
||
image_caption = str(image_caption)
|
||
|
||
if isinstance(img_path, list):
|
||
img_path = "\n".join(str(p) for p in img_path)
|
||
elif not isinstance(img_path, str):
|
||
img_path = str(img_path)
|
||
|
||
if image_caption.strip() or img_path.strip():
|
||
new_text = image_caption + "\n" + img_path + "\n"
|
||
else:
|
||
new_text = ""
|
||
|
||
should_record_highlight = True
|
||
|
||
elif ins["type"] == "discarded":
|
||
text = ins.get("text", "")
|
||
if text.strip():
|
||
new_text = text + "\n"
|
||
else:
|
||
new_text = ""
|
||
should_record_highlight = True
|
||
|
||
else:
|
||
continue # 不支持的 type，跳过且不记录 highlight
|
||
|
||
# --- 关键：即使 new_text 为空，只要 should_record_highlight 为 True，就要处理 chunk 切分和 highlight 记录 ---
|
||
|
||
# 构造 highlight 项（即使 new_text 为空）
|
||
highlight_item = {"page_idx": page_idx, "bbox": bbox}
|
||
|
||
if not new_text.strip() and not should_record_highlight:
|
||
continue # 完全跳过
|
||
|
||
# 检查是否需要切分（仅当 new_text 非空时才影响长度）
|
||
would_exceed = len(current_content) + len(new_text) > max_length
|
||
|
||
# 如果 new_text 非空 且 会导致超长，则先保存当前 chunk
|
||
if new_text.strip() and would_exceed:
|
||
if current_content or current_highlights:
|
||
contents.append(current_content.rstrip('\n'))
|
||
highlight_lists.append(current_highlights)
|
||
current_content = new_text
|
||
current_highlights = [highlight_item]
|
||
else:
|
||
# 追加内容（如果 new_text 非空）
|
||
if new_text.strip():
|
||
current_content += new_text
|
||
# 无论如何，只要 should_record_highlight，就加 highlight
|
||
if should_record_highlight:
|
||
current_highlights.append(highlight_item)
|
||
|
||
# 保存最后一段
|
||
if current_content or current_highlights:
|
||
contents.append(current_content.rstrip('\n'))
|
||
highlight_lists.append(current_highlights)
|
||
|
||
return contents, highlight_lists
|
||
|
||
|
||
|
||
def group_records(records):
|
||
result = []
|
||
current_group = []
|
||
|
||
for i, ins in enumerate(records):
|
||
if 'text_level' in ins and ins['text_level']:
|
||
# 当遇到新的text_level时，如果current_group非空，则先将其添加到结果中
|
||
if current_group:
|
||
result.append(current_group)
|
||
current_group = [] # 开始新一组
|
||
current_group.append(ins)
|
||
else:
|
||
# 如果当前记录没有text_level或其值为空，则继续添加到当前组
|
||
current_group.append(ins)
|
||
|
||
# 别忘了将最后一个组添加到结果中
|
||
if current_group:
|
||
result.append(current_group)
|
||
|
||
return result
|
||
|
||
def get_chunk_bbox_otherfile(data, max_length=8000):
|
||
"""
|
||
处理没有 bbox 信息的文件数据（主要面向纯表格类内容）。
|
||
每个 table 作为一个独立切片，不与其他记录拼接。
|
||
|
||
- 单个 table 内容超过 max_length 时，调用 chunk_reset 在 </tr> 处安全切分，
|
||
切分后的多个 chunk 共享同一组 positions
|
||
- 清理冗余 colspan="1" / rowspan="1" 属性
|
||
- bbox 缺失时填充空列表 []
|
||
|
||
返回结构与 get_chunk_bbox 一致：[{"content": str, "positions": [...]}]
|
||
"""
|
||
slices = []
|
||
|
||
for ins in data:
|
||
if ins.get("type") != "table":
|
||
continue
|
||
|
||
page_idx = ins.get("page_idx", -1)
|
||
bbox = ins.get("bbox", [])
|
||
|
||
table_caption = ins.get("table_caption")
|
||
table_body = ins.get("table_body", "") or ""
|
||
|
||
# 清理冗余属性
|
||
for pattern in [
|
||
' colspan="1"', ' rowspan="1"',
|
||
" colspan='1'", " rowspan='1'",
|
||
" colspan=1", " rowspan=1",
|
||
]:
|
||
table_body = table_body.replace(pattern, "")
|
||
|
||
# 拼接 caption
|
||
caption_str = ""
|
||
if table_caption:
|
||
if isinstance(table_caption, list):
|
||
caption_str = "\n".join(str(c) for c in table_caption)
|
||
else:
|
||
caption_str = str(table_caption)
|
||
|
||
has_caption = bool(caption_str.strip())
|
||
has_body = bool(table_body.strip())
|
||
|
||
if not (has_caption or has_body):
|
||
continue # 空表格跳过
|
||
|
||
content = caption_str + ("\n" if has_caption else "") + table_body
|
||
positions = [{"page_idx": page_idx, "bbox": bbox}]
|
||
|
||
# 单表超长 → 在 </tr> 处切分，多个 chunk 共享 positions
|
||
if len(content) > max_length:
|
||
sub_chunks = chunk_reset(content, max_length)
|
||
for sub in sub_chunks:
|
||
slices.append({
|
||
"content": sub.rstrip('\n'),
|
||
"positions": positions
|
||
})
|
||
else:
|
||
slices.append({
|
||
"content": content.rstrip('\n'),
|
||
"positions": positions
|
||
})
|
||
|
||
# 按 page_idx 排序，与 get_chunk_bbox 风格一致
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get("positions")
|
||
and isinstance(item["positions"], list)
|
||
and len(item["positions"]) > 0
|
||
and "page_idx" in item["positions"][0]
|
||
]
|
||
valid_slices.sort(key=lambda item: item["positions"][-1]["page_idx"])
|
||
|
||
return valid_slices
|
||
|
||
def get_chunk_bbox(data):
|
||
for idx, item in enumerate(data, start=1):
|
||
item["id"] = idx
|
||
repair_data, clean_repair_records = extract_maintenance_groups_and_remaining(data,max_following_for_last=23)
|
||
operation_data, clean_operation_records = extract_operation_groups_and_remaining(data,max_following_for_last=23)
|
||
wxanli_data, clean_wxanli_data = extract_wxanli(data)
|
||
slices = []
|
||
if len(repair_data) > 0 :
|
||
for ins in repair_data:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
groups = group_records(clean_repair_records)
|
||
for ins in groups:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
elif len(operation_data) > 0:
|
||
for ins in operation_data:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
groups = group_records(clean_operation_records)
|
||
for ins in groups:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
|
||
elif len(wxanli_data) > 0:
|
||
for ins in wxanli_data:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
for ins in clean_wxanli_data:
|
||
if not ins.get("bbox"):
|
||
continue
|
||
page_idx = ins.get("page_idx", 0)
|
||
bbox = ins["bbox"]
|
||
positions = [{"page_idx": page_idx, "bbox": bbox}]
|
||
|
||
if ins["type"] == "text":
|
||
if ins.get("text_level"):
|
||
text = ins.get("text", "")
|
||
level = ins.get("text_level")
|
||
content = "#" * level + text
|
||
else:
|
||
content = ins.get("text", "")
|
||
elif ins["type"] == "table":
|
||
content = (ins.get("table_caption") or "") + "\n" + (ins.get("table_body") or "")
|
||
elif ins["type"] == "image":
|
||
image_caption = ins.get("image_caption", [])
|
||
img_path = ins.get("img_path", [])
|
||
|
||
# 如果 image_caption 是列表，则将其转换为字符串
|
||
if isinstance(image_caption, list):
|
||
image_caption = "\n".join(str(item) for item in image_caption)
|
||
elif not isinstance(image_caption, str):
|
||
image_caption = str(image_caption)
|
||
|
||
# 如果 img_path 是列表，则将其转换为字符串
|
||
if isinstance(img_path, list):
|
||
img_path = "\n".join(str(p) for p in img_path)
|
||
elif not isinstance(img_path, str):
|
||
img_path = str(img_path)
|
||
|
||
if image_caption.strip() or img_path.strip():
|
||
content = (image_caption or "") + "\n" + (img_path or "")
|
||
else:
|
||
continue # 跳过空图像项
|
||
elif ins["type"] == "discarded":
|
||
content = ins.get("text", "")
|
||
else:
|
||
continue
|
||
|
||
slices.append({"content": content, "positions": positions})
|
||
|
||
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
else:
|
||
groups = group_records(data)
|
||
for ins in groups:
|
||
content, highlight_list = split_content_bbox(ins)
|
||
for i, (content, h_list) in enumerate(zip(content, highlight_list)):
|
||
slices.append({
|
||
"content" : content,
|
||
"positions":h_list
|
||
})
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
valid_slices = [
|
||
item for item in slices
|
||
if item.get('positions')
|
||
and isinstance(item['positions'], list)
|
||
and len(item['positions']) > 0
|
||
and 'page_idx' in item['positions'][0]
|
||
]
|
||
slices = sorted(valid_slices, key=lambda item: item['positions'][-1]['page_idx'])
|
||
return slices
|
||
|
||
def split_text_preserve_sentences(text, max_length=7500):
|
||
if len(text) <= max_length:
|
||
return [text]
|
||
|
||
# 按句子分割，保留标点
|
||
parts = re.split(r'([。！？；!?])', text)
|
||
sentences = []
|
||
for i in range(0, len(parts), 2):
|
||
s = parts[i] + (parts[i+1] if i+1 < len(parts) else '')
|
||
sentences.append(s)
|
||
|
||
chunks = []
|
||
current = ""
|
||
|
||
for sent in sentences:
|
||
if len(current) + len(sent) <= max_length:
|
||
current += sent
|
||
else:
|
||
if current:
|
||
chunks.append(current)
|
||
else:
|
||
# 单句超长，硬切（极端情况）
|
||
chunks.append(sent[:max_length])
|
||
sent = sent[max_length:]
|
||
current = sent
|
||
|
||
if current:
|
||
chunks.append(current)
|
||
|
||
return chunks
|
||
|
||
def data_replace(data, prefix):
|
||
"""
|
||
遍历 data 列表，若元素包含 'img_path' 字段，则移除前缀 'images/' 并拼接新前缀。
|
||
使用 pathlib 安全处理路径。
|
||
"""
|
||
from pathlib import PurePosixPath
|
||
for ins in data:
|
||
if isinstance(ins, dict) and "img_path" in ins:
|
||
img_path = ins["img_path"]
|
||
p = PurePosixPath(img_path)
|
||
# 如果以 images/ 开头，去掉第一级目录
|
||
if p.parts and p.parts[0] == "images":
|
||
relative_path = str(PurePosixPath(*p.parts[1:]))
|
||
ins["img_path"] = f""
|
||
else:
|
||
# 否则保留原路径（或按需处理）
|
||
ins["img_path"] = f""
|
||
return data
|
||
|
||
def filter_content_bbox(data, max_length=3000):
|
||
"""
|
||
将 data 中的 text/table/image 内容按字符长度分块（每块不超过 max_length），
|
||
返回 [{"content": str, "highlight_list": list}, ...]
|
||
"""
|
||
chunks = []
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
def _flush_chunk():
|
||
"""内部函数：将当前内容打包进 chunks，并重置"""
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list.copy()
|
||
})
|
||
|
||
for ins in data:
|
||
if "type" not in ins:
|
||
continue
|
||
|
||
added_content = ""
|
||
bbox = ins.get("bbox", [])
|
||
page_idx = ins.get("page_idx", -1)
|
||
|
||
if ins["type"] == "text":
|
||
text_content = ins.get("text", "")
|
||
added_content = text_content + "\n"
|
||
|
||
elif ins["type"] == "table":
|
||
table_caption = ins.get("table_caption")
|
||
table_body = ins.get("table_body")
|
||
|
||
caption_str = ""
|
||
if table_caption:
|
||
if isinstance(table_caption, list):
|
||
caption_str = "\n".join(str(c) for c in table_caption)
|
||
else:
|
||
caption_str = str(table_caption)
|
||
|
||
added_content = caption_str + "\n" + (table_body or "") + "\n"
|
||
|
||
elif ins["type"] == "image":
|
||
image_caption = ins.get("image_caption", "")
|
||
img_path = ins.get("img_path", "")
|
||
|
||
if isinstance(image_caption, list):
|
||
image_caption = "\n".join(str(item) for item in image_caption)
|
||
elif not isinstance(image_caption, str):
|
||
image_caption = str(image_caption)
|
||
|
||
if isinstance(img_path, list):
|
||
img_path = "\n".join(str(p) for p in img_path)
|
||
elif not isinstance(img_path, str):
|
||
img_path = str(img_path)
|
||
|
||
if image_caption or img_path:
|
||
added_content = image_caption + "\n" + img_path + "\n"
|
||
else:
|
||
continue # 跳过空图像项
|
||
else:
|
||
continue
|
||
|
||
# 检查加上 added_content 后是否超限
|
||
if len(current_content) + len(added_content) > max_length:
|
||
# 如果当前块非空，先 flush
|
||
if current_content:
|
||
_flush_chunk()
|
||
current_content = ""
|
||
current_highlight_list = []
|
||
|
||
# 如果单个 added_content 超过 max_length，仍需保留（避免丢弃）
|
||
if len(added_content) > max_length:
|
||
# 强制截断或直接放入（这里选择直接放入，避免信息丢失）
|
||
current_content = added_content[:max_length]
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
_flush_chunk()
|
||
# 剩余部分？此处简化处理：忽略超长部分（或可递归拆分，但复杂）
|
||
continue
|
||
else:
|
||
current_content = added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
else:
|
||
current_content += added_content
|
||
current_highlight_list.append({"page_idx": page_idx, "bbox": bbox})
|
||
|
||
# 处理最后一块
|
||
if current_content or current_highlight_list:
|
||
chunks.append({
|
||
"content": current_content,
|
||
"highlight_list": current_highlight_list
|
||
})
|
||
|
||
return chunks
|
||
|
||
def chunk_reset(content,MAX_length=8000):
|
||
chunks = []
|
||
start = 0
|
||
while start < len(content):
|
||
end = start + MAX_length
|
||
if end >= len(content):
|
||
chunks.append(content[start:])
|
||
break
|
||
# 尽量在 </tr> 后面切分
|
||
snippet = content[start:end]
|
||
last_tr_end = snippet.rfind("</tr>")
|
||
if last_tr_end != -1:
|
||
end = start + last_tr_end + len("</tr>")
|
||
else:
|
||
# 如果找不到 </tr>，退回到最近的换行或空白
|
||
last_newline = snippet.rfind("\n")
|
||
if last_newline != -1:
|
||
end = start + last_newline
|
||
else:
|
||
# 实在不行就硬切
|
||
pass
|
||
chunks.append(content[start:end])
|
||
start = end
|
||
return chunks
|
||
def chunk_check(data,MAX_length=8000):
|
||
slice_reset = []
|
||
for ins in data:
|
||
content = ins["content"]
|
||
positions = ins["positions"]
|
||
if len(content) <= 8000:
|
||
slice_reset.append(ins)
|
||
else:
|
||
# 使用你的 chunk_reset 函数切分 content
|
||
chunks = chunk_reset(content, MAX_length)
|
||
for chunk in chunks:
|
||
slice_reset.append({
|
||
"content": chunk,
|
||
"positions": positions # 共享原始位置信息
|
||
})
|
||
return slice_reset
|
||
|
||
def find_ship_info_by_hull(json_file_path, data):
|
||
"""
|
||
从NER结果中提取舷号，再从JSON文件中查找对应的舰船信息
|
||
:param json_file_path: JSON文件路径
|
||
:param data: get_entity() 返回的NER结果字典
|
||
:return: 包含 model_name 和 ship_name 的字典，未找到返回 None
|
||
"""
|
||
# 1. 提取舷号
|
||
hulls = [
|
||
entity['properties']['舷号']
|
||
for entity in data.get('entities', [])
|
||
if 'properties' in entity and '舷号' in entity['properties']
|
||
]
|
||
if not hulls:
|
||
return None
|
||
|
||
target_hull = hulls[0]
|
||
|
||
# 2. 根据舷号查找舰船信息
|
||
try:
|
||
with open(json_file_path, 'r', encoding='utf-8') as f:
|
||
ship_data = json.load(f)
|
||
|
||
for item in ship_data:
|
||
if str(item.get('hull_number')) == str(target_hull):
|
||
return {
|
||
"model_name": item.get('model_name'),
|
||
"ship_name": item.get('ship_name')
|
||
}
|
||
|
||
return None
|
||
|
||
except FileNotFoundError:
|
||
print(f"错误：找不到文件 {json_file_path}")
|
||
return None
|
||
except json.JSONDecodeError:
|
||
print("错误：JSON 文件格式不正确")
|
||
return None
|
||
def merge_short_slices(slices, min_length=30,filename="163-06A0014-B01001_发动机-维修手册.pdf"):
|
||
"""
|
||
合并过短的切片：
|
||
- 如果某切片 content 长度 <= min_length，则将其合并到下一个切片的开头
|
||
- 若处于末尾无下一个切片，则反向合并到上一个切片末尾
|
||
- content 用换行拼接，positions 顺序拼接
|
||
|
||
Args:
|
||
slices: [{"content": str, "positions": [...]}]
|
||
min_length: 短切片的字符长度阈值（含）
|
||
Returns:
|
||
合并后的 slices 列表
|
||
"""
|
||
if not slices:
|
||
return slices
|
||
|
||
result = []
|
||
pending_contents = [] # 缓存等待合并到"下一个"的短切片 content
|
||
pending_positions = [] # 缓存对应的 positions
|
||
|
||
for ins in slices:
|
||
content = ins.get("content", "") or ""
|
||
positions = ins.get("positions", []) or []
|
||
|
||
if len(content) <= min_length:
|
||
# 暂存，等到下一个正常长度的切片再合并
|
||
pending_contents.append(content)
|
||
pending_positions.extend(positions)
|
||
else:
|
||
# 正常切片：把暂存的短切片合并到它的开头
|
||
if pending_contents:
|
||
merged_prefix = "\n".join(pending_contents)
|
||
content = merged_prefix + ("\n" if merged_prefix else "") + content
|
||
positions = pending_positions + positions
|
||
pending_contents = []
|
||
pending_positions = []
|
||
result.append({
|
||
"content": content,
|
||
"positions": positions
|
||
})
|
||
|
||
# 收尾：如果末尾还有未合并的短切片（后面没有正常切片可合并）
|
||
# 则反向合并到上一个切片末尾
|
||
if pending_contents:
|
||
merged_suffix = "\n".join(pending_contents)
|
||
if result:
|
||
last = result[-1]
|
||
last["content"] = (last["content"] or "") + ("\n" if last["content"] else "") + merged_suffix
|
||
last["positions"] = (last["positions"] or []) + pending_positions
|
||
else:
|
||
# 极端情况：所有切片都很短，整体作为一个切片返回
|
||
result.append({
|
||
"content": merged_suffix,
|
||
"positions": pending_positions
|
||
})
|
||
try:
|
||
final_result = get_entity(filename)
|
||
print(final_result)
|
||
xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME,final_result)
|
||
if xinghao.get('model_name'):
|
||
model_name = xinghao['model_name']
|
||
else:
|
||
model_name = ""
|
||
# lower_entity = final_result.get("low_level")
|
||
otehr_data = format_entity_text(final_result)
|
||
logger.info(f"文件名实体提取成功: {otehr_data}")
|
||
except Exception as e:
|
||
logger.warning(f"文件名实体提取失败: {e}")
|
||
otehr_data = "" # 建议加个兜底,避免后面 NameError
|
||
|
||
for ins in result:
|
||
content = ins['content']
|
||
if not otehr_data:
|
||
continue
|
||
newline_idx = content.find('\n')
|
||
# 拼接成一个括号：otehr_data + 型号为xxx
|
||
info = otehr_data
|
||
if model_name:
|
||
info += f'，型号为{model_name}'
|
||
suffix = f'({info})'
|
||
if newline_idx == -1:
|
||
ins['content'] = content + suffix
|
||
else:
|
||
ins['content'] = content[:newline_idx] + suffix + content[newline_idx:]
|
||
return result
|
||
def find_and_read_content_list(directory, original_filename, encoding='utf-8'):
|
||
"""
|
||
根据原始文件名，在指定目录及其子目录下查找对应的 _content_list.json 文件并返回内容。
|
||
"""
|
||
# 1. 去除原始文件的后缀，并拼接新的后缀
|
||
# 加上 只获取不带后缀的文件名部分
|
||
base_name = os.path.splitext(original_filename)
|
||
file_name = base_name[0]
|
||
target_filename = f"{file_name}_content_list.json"
|
||
|
||
print(f"正在查找文件: {target_filename}")
|
||
|
||
# 2. 使用 os.walk 递归遍历目录
|
||
for root, dirs, files in os.walk(directory):
|
||
if target_filename in files:
|
||
file_path = os.path.join(root, target_filename)
|
||
|
||
try:
|
||
with open(file_path, 'r', encoding=encoding) as f:
|
||
json_data = json.load(f)
|
||
return json_data
|
||
except json.JSONDecodeError:
|
||
print(f"文件 {file_path} 不是有效的 JSON 格式。")
|
||
return None
|
||
except Exception as e:
|
||
print(f"读取文件 {file_path} 时发生错误: {e}")
|
||
return None
|
||
|
||
return None
|
||
|
||
async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Optional[Dict[str, Any]]:
|
||
"""异步调用 PDF 分析服务并处理响应"""
|
||
doc_result = await process_document(filename)
|
||
if doc_result is not None:
|
||
content_list, images = doc_result
|
||
if content_list and images:
|
||
replaced_content_list = data_replace(content_list, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
slices_check = merge_short_slices(slices_check, min_length=30,filename=filename)
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
api_url = get_api_url()
|
||
try:
|
||
async with aiofiles.open(pdf_path, "rb") as f:
|
||
file_content = await f.read()
|
||
files = {"file": (filename, file_content, "application/pdf")}
|
||
|
||
async with httpx.AsyncClient(timeout=60 * 20) as client:
|
||
response = await client.post(api_url, files=files)
|
||
|
||
if response.status_code != 200:
|
||
logger.error(f"❌ API 请求失败，状态码: {response.status_code}, 响应: {response.text}")
|
||
return None
|
||
|
||
result = response.json()
|
||
data = result.get("data", {})
|
||
content_list = data.get("content_list", [])
|
||
images = data.get("images", {})
|
||
replaced_content_list = data_replace(content_list, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
slices_check = merge_short_slices(slices_check, min_length=30,filename=filename)
|
||
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
except json.JSONDecodeError:
|
||
logger.error("❌ 响应不是有效的 JSON 格式")
|
||
logger.debug(response.text)
|
||
return {"slices": [], "images": {}}
|
||
except Exception as e:
|
||
logger.error(f"处理 PDF 文件时出错: {e}", exc_info=True)
|
||
return None
|
||
finally:
|
||
release_api_url(api_url)
|
||
|
||
async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) -> Optional[Dict[str, Any]]:
|
||
"""异步调用 Office 文件分析服务并处理响应"""
|
||
# 优先尝试走本地/缓存处理（保持与你原逻辑一致）
|
||
doc_result = await process_document(filename)
|
||
if doc_result is not None:
|
||
content_list, images = doc_result
|
||
if content_list and images:
|
||
replaced_content_list = data_replace(content_list, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
api_url = get_api_url()
|
||
analyze_other_url = api_url.replace("/analyze-pdf", "/analyze-otherfile")
|
||
|
||
mime_type, _ = mimetypes.guess_type(pdf_path.name)
|
||
content_type = mime_type or "application/octet-stream"
|
||
|
||
try:
|
||
async with aiofiles.open(pdf_path, "rb") as f:
|
||
file_content = await f.read()
|
||
|
||
files = {"file": (pdf_path.name, file_content, content_type)}
|
||
|
||
async with httpx.AsyncClient(timeout=60 * 20) as client:
|
||
response = await client.post(analyze_other_url, files=files)
|
||
|
||
if response.status_code != 200:
|
||
logger.error(f"❌ Office 文件 API 请求失败，状态码: {response.status_code}, 响应: {response.text}")
|
||
return None
|
||
|
||
result = response.json()
|
||
data = result.get("data", {})
|
||
content_list = data.get("content_list", [])
|
||
images = data.get("images", {})
|
||
|
||
# 后续的数据替换与切片处理逻辑
|
||
replaced_content_list = data_replace(content_list, image_prefix)
|
||
slices = get_chunk_bbox(replaced_content_list)
|
||
slices_check = chunk_check(slices, 8000)
|
||
|
||
return {"slices": slices_check, "images": images}
|
||
|
||
except json.JSONDecodeError:
|
||
error_text = response.text if 'response' in locals() else "请求未到达服务器"
|
||
logger.error(f"❌ 响应不是有效的 JSON 格式: {error_text}")
|
||
return {"slices": [], "images": {}}
|
||
except Exception as e:
|
||
logger.error(f"处理 Office 文件时出错: {e}", exc_info=True)
|
||
return None
|
||
finally:
|
||
release_api_url(api_url)
|
||
|
||
if __name__ == "__main__":
|
||
|
||
filepath = r"E:\ZKYNLP\Hjunproject\project0506\kgrag\122-06A0014-B01001_发动机-维修手册_content_list.json"
|
||
try:
|
||
with open(filepath, 'r', encoding='utf-8') as file:
|
||
data = json.load(file)
|
||
except Exception as e:
|
||
print(f"读取 JSON 文件时出错：{e}")
|
||
data = []
|
||
# print(len(data))
|
||
start_chars = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "0", "第", "十",
|
||
"一", "二", "三", "四", "五", "六", "七", "八", "九"]
|
||
|
||
data = data_replace(data=data,prefix="/api/v1/knowledge/files/images/")
|
||
slices = get_chunk_bbox(data=data)
|
||
slices_1 =chunk_check(slices, 8000)
|
||
slices_2 = merge_short_slices(slices_1, min_length=30) # ← 新增这一行
|
||
|
||
|
||
for ins in slices_2:
|
||
print(ins)
|
||
|
||
|
||
|
||
|
||
|