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: fallback_summary = get_text_level_summary_fallback(content_list) if not md_content or not str(md_content).strip(): return fallback_summary try: summaries = await asyncio.wait_for( generate_summaries([str(md_content)]), timeout=SUMMARY_GENERATION_TIMEOUT_SECONDS, ) summary = summaries[0] if summaries else "" if summary and str(summary).strip(): return str(summary).strip() logger.warning( "Generated summary is empty; using text_level fallback" ) except Exception as exc: logger.error( "Summary generation failed; using text_level fallback: %s", exc, exc_info=True, ) logger.info("Fallback summary: %s", fallback_summary) return fallback_summary 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, } 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 strip_image_ocr_prefix(text: Any) -> str: text = stringify_content_field(text).strip() for prefix in IMAGE_OCR_PREFIXES: if text.startswith(prefix): return text[len(prefix):].strip() return text def normalize_image_ocr_text(value: Any) -> str: return "\n".join( line for line in (strip_image_ocr_prefix(line) for line in stringify_content_field(value).splitlines()) if line ) def split_image_caption_and_ocr(caption: Any) -> Tuple[str, str]: if isinstance(caption, list): raw_items = caption else: raw_items = [caption] caption_parts = [] ocr_parts = [] for item in raw_items: for line in stringify_content_field(item).splitlines(): line = line.strip() if not line: continue ocr_line = strip_image_ocr_prefix(line) if ocr_line != line: ocr_parts.append(ocr_line) else: caption_parts.append(line) return "\n".join(caption_parts), "\n".join(ocr_parts) 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( ins.get("table_caption"), clean_table_body(ins.get("table_body", "")), ins.get("table_footnote"), ) if record_type == "image": imagecontent = '' image_caption, caption_ocr_text = split_image_caption_and_ocr(ins.get("image_caption")) image_caption = image_caption +"如下图所示:" img_path = stringify_content_field(ins.get("img_path")).strip() if caption_ocr_text: imagecontent = image_caption + "\n" + img_path + "\n" + f"备注:图片中包含的内容为{caption_ocr_text}" else: imagecontent = image_caption + "\n" + img_path return imagecontent if record_type == "chart": return join_nonempty_parts( ins.get("chart_caption"), ins.get("content"), ins.get("img_path"), ins.get("chart_footnote"), ) if record_type == "equation": return join_nonempty_parts( ins.get("text"), ) if record_type == "code": return join_nonempty_parts( ins.get("code_caption"), ins.get("code_body"), ins.get("code_footnote"), ) if record_type == "list": return join_nonempty_parts(ins.get("list_items")) if record_type in {"discarded", "header", "footer", "page_number"}: return join_nonempty_parts(ins.get("text")) return join_nonempty_parts(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": "ocr_text", "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 ins.get("type") == "image": existing_text = join_nonempty_parts(existing_text, ins.get("image_caption")) elif 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"![图片]({prefix}{relative_path})", resolve_local_image_path(relative_path) return f"![图片]({prefix}{img_path})", 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 content_starts_with_same_lines(content: str, prefix: str) -> bool: prefix_lines = [line.strip() for line in (prefix or "").splitlines() if line.strip()] content_lines = [line.strip() for line in (content or "").splitlines() if line.strip()] return bool(prefix_lines) and content_lines[:len(prefix_lines)] == prefix_lines def content_ends_with_same_lines(content: str, suffix: str) -> bool: suffix_lines = [line.strip() for line in (suffix or "").splitlines() if line.strip()] content_lines = [line.strip() for line in (content or "").splitlines() if line.strip()] return bool(suffix_lines) and content_lines[-len(suffix_lines):] == suffix_lines def is_heading_only_content(content: str) -> bool: lines = [ line.strip() for line in (content or "").splitlines() if line.strip() ] meaningful_lines = [ line for line in lines if not line.startswith("#文件名为:") and not re.fullmatch(r"\d+", line) ] return bool(meaningful_lines) and all(line.startswith("#") for line in meaningful_lines) def merge_prefix_content(merged_prefix: str, content: str, pending_positions: list, positions: list): prefix_lines = [line.strip() for line in (merged_prefix or "").splitlines() if line.strip()] content_lines = [line.strip() for line in (content or "").splitlines() if line.strip()] duplicate_count = 0 for prefix_line, content_line in zip(prefix_lines, content_lines): if prefix_line != content_line: break duplicate_count += 1 if duplicate_count == len(prefix_lines): return content, positions content_lines_raw = (content or "").splitlines() content_without_duplicates = "\n".join(content_lines_raw[duplicate_count:]).lstrip("\n") merged_content = merged_prefix + ("\n" if merged_prefix and content_without_duplicates else "") + content_without_duplicates merged_positions = pending_positions + positions[duplicate_count:] return merged_content, merged_positions 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 or is_heading_only_content(content): # 暂存,等到下一个正常长度的切片再合并 if pending_contents: merged_prefix = "\n".join(pending_contents) if content_starts_with_same_lines(content, merged_prefix): pending_contents = [content] pending_positions = positions else: pending_contents.append(content) pending_positions.extend(positions) else: pending_contents.append(content) pending_positions.extend(positions) else: # 正常切片:把暂存的短切片合并到它的开头 if pending_contents: merged_prefix = "\n".join(pending_contents) content, positions = merge_prefix_content( merged_prefix, content, 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] if not content_ends_with_same_lines(last.get("content", ""), merged_suffix): 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: 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 logger.info(f"摘要生成:" + summary[:30]) return {"slices": slices_check, "images": images or {}, "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\163-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", "第", "十", "一", "二", "三", "四", "五", "六", "七", "八", "九"] 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[:15]: print(ins)