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'
try:
from filename_proceess_and_kgquery import get_entity
except ImportError:
get_entity = None
from pathlib import PurePath, PurePosixPath, Path as PathLib
from fastapi import FastAPI, File, Path, UploadFile, HTTPException, Form, Request, Header, Body
from fileparse_util import process_document
from generate_summary import generate_summaries
from io import BytesIO
from urllib.parse import unquote, urlparse
import asyncio
import base64
import binascii
import threading
import aiofiles
import httpx
import logging
from resetlevel import reset_textlevel
from rapidocr_onnxruntime import RapidOCR
from PIL import Image, ImageOps, UnidentifiedImageError
from config import (
API_URLS,
IMAGE_DIR,
MAX_SPLIT_IMAGE_BASE64_CHARS,
MAX_SPLIT_IMAGE_BYTES,
MAX_SPLIT_IMAGE_INPUT_BYTES,
OCR_CONCURRENCY,
SPLIT_IMAGE_COMPRESS_THRESHOLD_BYTES,
SPLIT_IMAGE_JPEG_QUALITY,
SPLIT_IMAGE_MAX_DIMENSION,
)
try:
from config import SHIP_MODEL_NAME
except ImportError:
SHIP_MODEL_NAME = ""
logger = logging.getLogger(__name__)
SUMMARY_GENERATION_TIMEOUT_SECONDS = 60
def get_text_level_summary_fallback(content_list) -> str:
summary_parts = []
for record in content_list:
if not isinstance(record, dict):
continue
text_level = record.get("text_level")
if text_level is None or text_level == "":
continue
text = record.get("text")
if text:
summary_parts.append(str(text).strip())
return "\n".join(summary_parts)
async def generate_summary_from_md(md_content: Optional[str],content_list) -> str:
if not md_content or not str(md_content).strip():
return get_text_level_summary_fallback(content_list)
try:
summaries = await asyncio.wait_for(
generate_summaries([str(md_content)]),
timeout=SUMMARY_GENERATION_TIMEOUT_SECONDS,
)
summary = summaries[0] if summaries else ""
if summary:
return summary
logger.warning("生成文档摘要为空,使用 content_list text_level 兜底摘要")
except Exception as exc:
logger.error(f"生成文档摘要失败,使用 content_list text_level 兜底摘要: {exc}", exc_info=True)
return get_text_level_summary_fallback(content_list)
ocr_engine = RapidOCR()
OCR_SEMAPHORE = asyncio.Semaphore(max(1, OCR_CONCURRENCY))
_ocr_thread_lock = threading.Lock()
_image_ocr_cache: Dict[str, str] = {}
_image_ocr_cache_lock = threading.Lock()
_api_url_inflight = [0] * len(API_URLS)
_api_url_lock = threading.Lock()
IMAGE_OCR_PREFIX = "图片中文本内容为:"
IMAGE_OCR_PREFIXES = (IMAGE_OCR_PREFIX, "图片包含的文字内容为:")
def _image_mime_type(suffix: str) -> str:
return {
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".png": "image/png",
".bmp": "image/bmp",
".webp": "image/webp",
".gif": "image/gif",
".tif": "image/tiff",
".tiff": "image/tiff",
}.get((suffix or "").lower(), "image/png")
def _image_save_format(suffix: str) -> tuple[str, dict]:
suffix = (suffix or "").lower()
if suffix in {".jpg", ".jpeg"}:
return "JPEG", {"quality": SPLIT_IMAGE_JPEG_QUALITY, "optimize": True}
if suffix == ".webp":
return "WEBP", {"quality": SPLIT_IMAGE_JPEG_QUALITY, "method": 6}
if suffix == ".png":
return "PNG", {"optimize": True, "compress_level": 9}
if suffix in {".tif", ".tiff"}:
return "TIFF", {"compression": "tiff_lzw"}
if suffix == ".gif":
return "GIF", {"optimize": True}
if suffix == ".bmp":
return "BMP", {}
return "PNG", {"optimize": True, "compress_level": 9}
def _compression_suffix(suffix: str) -> str:
suffix = (suffix or "").lower()
if suffix in {".jpg", ".jpeg", ".png", ".webp"}:
return suffix
if suffix in {".bmp", ".gif", ".tif", ".tiff"}:
return ".png"
return ".png"
def _save_image_to_bytes(img: Image.Image, suffix: str) -> bytes:
image_format, save_options = _image_save_format(suffix)
if image_format in {"JPEG", "WEBP", "BMP"} and img.mode not in {"RGB", "L"}:
img = img.convert("RGB")
elif image_format in {"PNG", "TIFF", "GIF"} and img.mode not in {"RGB", "RGBA", "L", "P"}:
img = img.convert("RGBA")
output = BytesIO()
img.save(output, format=image_format, **save_options)
return output.getvalue()
def _save_image_as_jpeg(img: Image.Image) -> bytes:
if img.mode not in {"RGB", "L"}:
img = img.convert("RGB")
last_bytes = b""
for quality in (SPLIT_IMAGE_JPEG_QUALITY, 82, 74):
output = BytesIO()
img.save(output, format="JPEG", quality=quality, optimize=True)
last_bytes = output.getvalue()
if len(last_bytes) <= MAX_SPLIT_IMAGE_BYTES:
break
return last_bytes
def compress_image_bytes(image_bytes: bytes, suffix: str = "") -> tuple[bytes, str]:
if not image_bytes:
raise HTTPException(status_code=400, detail="decoded image is empty")
if len(image_bytes) > MAX_SPLIT_IMAGE_INPUT_BYTES:
raise HTTPException(status_code=413, detail="image input is too large")
suffix = (suffix or "").lower()
try:
with Image.open(BytesIO(image_bytes)) as img:
img = ImageOps.exif_transpose(img)
max_side = max(img.size)
should_resize = SPLIT_IMAGE_MAX_DIMENSION > 0 and max_side > SPLIT_IMAGE_MAX_DIMENSION
should_reencode = len(image_bytes) > SPLIT_IMAGE_COMPRESS_THRESHOLD_BYTES
if not should_resize and not should_reencode:
if len(image_bytes) > MAX_SPLIT_IMAGE_BYTES:
raise HTTPException(status_code=413, detail="image is too large")
return image_bytes, suffix
if should_resize:
scale = SPLIT_IMAGE_MAX_DIMENSION / max_side
new_size = (
max(1, int(img.width * scale)),
max(1, int(img.height * scale)),
)
resample = getattr(getattr(Image, "Resampling", Image), "LANCZOS")
img = img.resize(new_size, resample)
compressed_suffix = _compression_suffix(suffix)
compressed_bytes = _save_image_to_bytes(img, compressed_suffix)
if len(compressed_bytes) > MAX_SPLIT_IMAGE_BYTES:
compressed_bytes = _save_image_as_jpeg(img)
compressed_suffix = ".jpg"
except UnidentifiedImageError as exc:
raise HTTPException(status_code=400, detail="image content is not a supported image") from exc
except HTTPException:
raise
except Exception as exc:
logger.warning(f"Failed to compress image, using original bytes: {exc}")
compressed_bytes = image_bytes
compressed_suffix = suffix
if len(compressed_bytes) >= len(image_bytes) and len(image_bytes) <= MAX_SPLIT_IMAGE_BYTES:
return image_bytes, suffix
if len(compressed_bytes) > MAX_SPLIT_IMAGE_BYTES:
raise HTTPException(status_code=413, detail="image is too large after compression")
return compressed_bytes, compressed_suffix
def prepare_split_image(image_base64: str, suffix: str = "") -> tuple[bytes, str]:
raw = (image_base64 or "").strip()
if not raw:
raise HTTPException(status_code=400, detail="image_base64 is empty")
if "," in raw and raw.split(",", 1)[0].lower().startswith("data:image"):
raw = raw.split(",", 1)[1]
raw = "".join(raw.split())
if len(raw) > MAX_SPLIT_IMAGE_BASE64_CHARS:
raise HTTPException(status_code=413, detail="image_base64 is too large")
try:
image_bytes = base64.b64decode(raw, validate=True)
except (binascii.Error, ValueError) as exc:
raise HTTPException(status_code=400, detail="image_base64 is not valid base64") from exc
return compress_image_bytes(image_bytes, suffix)
def decode_base64_image(image_base64: str, suffix: str = "") -> bytes:
image_bytes, _ = prepare_split_image(image_base64, suffix)
return image_bytes
def safe_original_filename(filename: str) -> str:
name = os.path.basename((filename or "").strip())
if not name:
raise HTTPException(status_code=400, detail="filename is empty")
safe_name = re.sub(r"[^\w.\-\u4e00-\u9fff]+", "_", name).strip("._")
if not safe_name:
raise HTTPException(status_code=400, detail="filename is invalid")
return safe_name
def _strip_markdown_image_path(img_path: str) -> str:
value = (img_path or "").strip()
match = re.match(r"!\[[^\]]*\]\((.*?)\)", value)
return match.group(1).strip() if match else value
def resolve_local_image_path(img_path: str) -> Optional[PathLib]:
value = _strip_markdown_image_path(str(img_path or ""))
if not value:
return None
parsed_path = unquote(urlparse(value).path or value)
marker = "/images/"
if marker in parsed_path:
relative_path = parsed_path.split(marker, 1)[1]
else:
normalized = parsed_path.replace("\\", "/").lstrip("/")
if normalized.startswith("images/"):
relative_path = normalized.split("/", 1)[1]
else:
relative_path = normalized
relative_path = relative_path.replace("\\", "/").lstrip("/")
if not relative_path:
return None
image_root = PathLib(IMAGE_DIR).resolve()
local_path = (image_root / PurePosixPath(relative_path)).resolve()
try:
local_path.relative_to(image_root)
except ValueError:
logger.warning(f"Image path is outside IMAGE_DIR and will be ignored: {img_path}")
return None
return local_path if local_path.is_file() else None
def extract_image_text_sync(image_path: Union[str, PathLib]) -> dict:
resolved_path = str(PathLib(image_path).resolve())
with _image_ocr_cache_lock:
cached_text = _image_ocr_cache.get(resolved_path)
if cached_text is not None:
return {"texts": [], "full_text": cached_text, "elapse": None}
with _ocr_thread_lock:
ocr_result, elapse = ocr_engine(resolved_path)
texts = [item[1] for item in (ocr_result or [])]
full_text = IMAGE_OCR_PREFIX + ",".join(texts) if texts else ""
with _image_ocr_cache_lock:
_image_ocr_cache[resolved_path] = full_text
return {
"texts": texts,
"full_text": full_text,
"elapse": elapse,
}
def build_image_markdown_with_ocr(markdown_image: str, local_image_path: Optional[PathLib]) -> str:
if any(prefix in markdown_image for prefix in IMAGE_OCR_PREFIXES):
return markdown_image
if not local_image_path:
return markdown_image
try:
ocr_data = extract_image_text_sync(local_image_path)
except Exception as exc:
logger.warning(f"Failed to OCR local image {local_image_path}: {exc}")
return markdown_image
ocr_full_text = (ocr_data.get("full_text") or "").strip()
if not ocr_full_text:
return markdown_image
return f"{markdown_image}\n{ocr_full_text}"
async def get_image_ocr_text(local_image_path: Optional[PathLib]) -> str:
if not local_image_path:
return ""
try:
ocr_data = await extract_image_text(local_image_path)
except Exception as exc:
logger.warning(f"Failed to OCR local image {local_image_path}: {exc}")
return ""
return (ocr_data.get("full_text") or "").strip()
def stringify_content_field(value: Any) -> str:
if value is None:
return ""
if isinstance(value, list):
return "\n".join(str(item) for item in value if str(item).strip())
if isinstance(value, dict):
return json.dumps(value, ensure_ascii=False)
return str(value)
def clean_table_body(table_body: Any) -> str:
table_body = stringify_content_field(table_body)
for pattern in [' colspan="1"', ' rowspan="1"', " colspan='1'", " rowspan='1'", " colspan=1", " rowspan=1"]:
table_body = table_body.replace(pattern, "")
return table_body
def join_nonempty_parts(*parts: Any) -> str:
return "\n".join(part for part in (stringify_content_field(part).strip() for part in parts) if part)
def record_to_chunk_text(ins: Dict[str, Any]) -> str:
record_type = ins.get("type")
ocr_text = ins.get("ocr_text", "")
if record_type == "text":
text = stringify_content_field(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
return join_nonempty_parts(ocr_text, text)
if record_type == "table":
return join_nonempty_parts(
ocr_text,
ins.get("table_caption"),
clean_table_body(ins.get("table_body", "")),
ins.get("table_footnote"),
)
if record_type == "image":
return join_nonempty_parts(
ocr_text,
ins.get("image_caption"),
ins.get("img_path"),
ins.get("image_footnote"),
)
if record_type == "chart":
return join_nonempty_parts(
ocr_text,
ins.get("chart_caption"),
ins.get("content"),
ins.get("img_path"),
ins.get("chart_footnote"),
)
if record_type == "equation":
return join_nonempty_parts(
ocr_text,
ins.get("text"),
)
if record_type == "code":
return join_nonempty_parts(
ocr_text,
ins.get("code_caption"),
ins.get("code_body"),
ins.get("code_footnote"),
)
if record_type == "list":
return join_nonempty_parts(ocr_text, ins.get("list_items"))
if record_type in {"discarded", "header", "footer", "page_number"}:
return join_nonempty_parts(ocr_text, ins.get("text"))
return join_nonempty_parts(ocr_text, ins.get("text"))
def is_supported_chunk_record(ins: Dict[str, Any]) -> bool:
return ins.get("type") in {
"text",
"table",
"image",
"chart",
"equation",
"code",
"list",
"discarded",
"header",
"footer",
"page_number",
}
def append_caption_ocr(ins: Dict[str, Any], ocr_texts: List[str]) -> None:
caption_field = {
"image": "image_caption",
"table": "table_caption",
"chart": "chart_caption",
"code": "code_caption",
}.get(ins.get("type"))
if not caption_field:
caption_field = "ocr_text"
caption = ins.get(caption_field, [])
existing_text = stringify_content_field(caption)
if caption_field == "ocr_text":
existing_text = join_nonempty_parts(existing_text, ins.get("text"))
additions = []
for text in ocr_texts:
text = (text or "").strip()
text_without_prefix = text
for prefix in IMAGE_OCR_PREFIXES:
if text_without_prefix.startswith(prefix):
text_without_prefix = text_without_prefix[len(prefix):].strip()
break
if (
text
and text not in existing_text
and text_without_prefix not in existing_text
and text not in additions
):
additions.append(text)
if not additions:
return
if caption_field == "ocr_text":
ins[caption_field] = join_nonempty_parts(additions, caption)
return
if isinstance(caption, list):
ins[caption_field] = additions + caption
elif caption is None or not str(caption).strip():
ins[caption_field] = additions
else:
ins[caption_field] = "\n".join(additions) + "\n" + str(caption).lstrip()
def build_image_markdown(markdown_or_path: str, prefix: str) -> Tuple[str, Optional[PathLib]]:
img_path = str(markdown_or_path).strip()
if not img_path:
return "", None
if img_path.strip().startswith("!["):
return img_path, resolve_local_image_path(img_path)
p = PurePosixPath(img_path)
if p.parts and p.parts[0] == "images":
relative_path = str(PurePosixPath(*p.parts[1:]))
return f"", resolve_local_image_path(relative_path)
return f"", resolve_local_image_path(img_path)
def image_base64_to_data_url(image_base64: Union[str, bytes], suffix: str) -> str:
if isinstance(image_base64, bytes):
raw = base64.b64encode(image_base64).decode("ascii")
else:
raw = (image_base64 or "").strip()
if "," in raw and raw.split(",", 1)[0].lower().startswith("data:image"):
raw = raw.split(",", 1)[1]
raw = "".join(raw.split())
mime_type = _image_mime_type(suffix)
return f"data:{mime_type};base64,{raw}"
async def extract_image_text(image_path: PathLib) -> dict:
async with OCR_SEMAPHORE:
return await asyncio.to_thread(extract_image_text_sync, image_path)
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 表格中第一个
内的文本(去除标签)"""
if not isinstance(html, str):
return ""
match = re.search(r'
]*>(.*?)
', 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:
"""判断是否为维修项目表格:仅检查第一个 是否匹配任一模式"""
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 是否为操作项目表,依据是第一行
中是否包含 '项目编号' 和 '操作项目'
"""
if not isinstance(table_body, str):
return False
# 非贪婪匹配第一个
...
match = re.search(r']*>(.*?)
', 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)
if not is_supported_chunk_record(ins):
continue
new_text = record_to_chunk_text(ins)
if new_text.strip():
new_text += "\n"
should_record_highlight = True
# --- 关键:即使 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_without_parent_context(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_group_level(group):
if not group:
return None
level = group[0].get("text_level")
if isinstance(level, int) and level > 0:
return level
return None
def group_records(records):
"""
Split records by text_level, then prepend ancestor section content.
For example, a text_level=2 section will include the nearest text_level=1
group content before its own content. A text_level=3 section will include
the nearest level 1 and level 2 group content before its own content.
"""
base_groups = _group_records_without_parent_context(records)
result = []
section_stack = [] # [(level, merged_group)]
for group in base_groups:
level = _get_group_level(group)
if level is None:
result.append(group)
continue
while section_stack and section_stack[-1][0] >= level:
section_stack.pop()
parent_group = section_stack[-1][1] if section_stack else []
merged_group = parent_group + group
result.append(merged_group)
section_stack.append((level, merged_group))
return result
def get_chunk_bbox_otherfile(data, max_length=8000):
"""
处理没有 bbox 信息的文件数据。
每条支持的记录作为一个独立切片,不与其他记录拼接。
- 单个 table 内容超过 max_length 时,调用 chunk_reset 在 处安全切分,
切分后的多个 chunk 共享同一组 positions
- 清理冗余 colspan="1" / rowspan="1" 属性
- bbox 缺失时填充空列表 []
返回结构与 get_chunk_bbox 一致:[{"content": str, "positions": [...]}]
"""
slices = []
for ins in data:
if not is_supported_chunk_record(ins):
continue
page_idx = ins.get("page_idx", -1)
bbox = ins.get("bbox", [])
content = record_to_chunk_text(ins)
if not content.strip():
continue
positions = [{"page_idx": page_idx, "bbox": bbox}]
# 单表超长 → 在 处切分,多个 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 not is_supported_chunk_record(ins):
continue
content = record_to_chunk_text(ins)
if not content.strip():
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
async def data_replace(data, prefix):
"""
遍历 data 列表,若元素包含 'img_path' 字段,则移除前缀 'images/' 并拼接新前缀。
使用 pathlib 安全处理路径。
"""
for ins in data:
if isinstance(ins, dict) and ins.get("type") == "equation":
continue
if isinstance(ins, dict) and "img_path" in ins:
img_path = ins["img_path"]
local_image_paths = []
if isinstance(img_path, list):
markdown_images = []
for item in img_path:
if not str(item).strip():
continue
markdown_image, local_image_path = build_image_markdown(str(item), prefix)
markdown_images.append(markdown_image)
if local_image_path:
local_image_paths.append(local_image_path)
ins["img_path"] = "\n".join(markdown_images)
else:
markdown_image, local_image_path = build_image_markdown(str(img_path), prefix)
ins["img_path"] = markdown_image
if local_image_path:
local_image_paths.append(local_image_path)
ocr_texts = await asyncio.gather(
*(get_image_ocr_text(path) for path in local_image_paths)
)
append_caption_ocr(ins, ocr_texts)
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
bbox = ins.get("bbox", [])
page_idx = ins.get("page_idx", -1)
if not is_supported_chunk_record(ins):
continue
added_content = record_to_chunk_text(ins)
if added_content.strip():
added_content += "\n"
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
# 尽量在 后面切分
snippet = content[start:end]
last_tr_end = snippet.rfind("")
if last_tr_end != -1:
end = start + last_tr_end + len("")
else:
# 如果找不到 ,退回到最近的换行或空白
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="122-06A0014-B01003_雷达-使用说明书.pdf"):
"""
合并过短的切片:
- 如果某切片 content 长度 <= min_length,则将其合并到下一个切片的开头
- 若处于末尾无下一个切片,则反向合并到上一个切片末尾
- content 用换行拼接,positions 顺序拼接
Args:
slices: [{"content": str, "positions": [...]}]
min_length: 短切片的字符长度阈值(含)
filename: 用于实体提取的文件名
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)
xinghao = find_ship_info_by_hull(SHIP_MODEL_NAME, final_result)
if xinghao is None: # 确保 xinghao 为 None 时不会报错
xinghao = {"model_name": "", "ship_name": ""}
print("未找到匹配的舰船信息")
# 安全处理 xinghao 为 None 的情况
model_name = ""
if xinghao and 'model_name' in xinghao:
model_name = xinghao['model_name']
other_data = format_entity_text(final_result)
logger.info(f"文件名实体提取成功: {other_data}")
except Exception as e:
logger.warning(f"文件名实体提取失败: {e}")
other_data = ""
model_name = "" # 确保异常时 model_name 有定义
# 将提取的信息(如舰艇名、型号)注入到每个切片的开头
for ins in result:
content = ins['content']
# 如果没有提取到有效数据,则跳过注入
newline_idx = content.find('\n')
info = other_data
if model_name:
info += f',型号为{model_name}'
# ===== 修改部分:拼接文件名 + 实体信息 =====
# 构建前缀:#文件名为:xxx.pdf(实体信息)
filename_prefix = f"#文件名为:{filename}"
if info:
filename_prefix += f"({info})"
# 注意:这里使用全角右括号")"以匹配你的示例,如需半角请改为 ")"
# 将前缀插入到 content 的最开头
if content:
ins['content'] = filename_prefix + "\n" + content
else:
ins['content'] = filename_prefix
# ==========================================
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, md_content = doc_result
if content_list and images:
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, 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)
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
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", {})
md_content = data.get("full_content", "")
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, 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)
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
except json.JSONDecodeError:
logger.error("❌ 响应不是有效的 JSON 格式")
logger.debug(response.text)
return {"slices": [], "images": {}, "summary": ""}
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, md_content = doc_result
if content_list and images:
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
slices = get_chunk_bbox(replaced_content_list)
slices_check = chunk_check(slices, 8000)
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
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", {})
md_content = data.get("full_content", "")
# 后续的数据替换与切片处理逻辑
replaced_content_textlevel = reset_textlevel(content_list)
replaced_content_list = await data_replace(replaced_content_textlevel, image_prefix)
slices = get_chunk_bbox(replaced_content_list)
slices_check = chunk_check(slices, 8000)
summary = await generate_summary_from_md(md_content,content_list)
summary = f"文件名:{filename}\n" + summary
return {"slices": slices_check, "images": images, "summary": summary}
except json.JSONDecodeError:
error_text = response.text if 'response' in locals() else "请求未到达服务器"
logger.error(f"❌ 响应不是有效的 JSON 格式: {error_text}")
return {"slices": [], "images": {}, "summary": ""}
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\舰船抗沉损管训练仿真系统研究_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", "第", "十",
"一", "二", "三", "四", "五", "六", "七", "八", "九"]
replaced_content_textlevel = reset_textlevel(data)
data = asyncio.run(data_replace(data=replaced_content_textlevel,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[:30]:
print(ins)