diff --git a/chunk_text.py b/chunk_text.py index 8a6239a..0b5dc36 100644 --- a/chunk_text.py +++ b/chunk_text.py @@ -15,6 +15,7 @@ except ImportError: 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 @@ -44,6 +45,42 @@ 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)) @@ -355,7 +392,6 @@ def record_to_chunk_text(ins: Dict[str, Any]) -> str: return join_nonempty_parts( ocr_text, ins.get("text"), - ins.get("img_path"), ) if record_type == "code": @@ -1086,6 +1122,8 @@ async def data_replace(data, prefix): 使用 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 = [] @@ -1255,103 +1293,6 @@ def find_ship_info_by_hull(json_file_path, data): 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}' - suffix = f'({info})' - - # 将信息插入到第一行末尾 - if newline_idx == -1: - ins['content'] = content + suffix - else: - ins['content'] = content[:newline_idx] + suffix + content[newline_idx:] - - return result - def merge_short_slices(slices, min_length=30, filename="122-06A0014-B01003_雷达-使用说明书.pdf"): """ @@ -1490,14 +1431,16 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op """异步调用 PDF 分析服务并处理响应""" doc_result = await process_document(filename) if doc_result is not None: - content_list, images = doc_result + 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) - return {"slices": slices_check, "images": images} + 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: @@ -1516,19 +1459,20 @@ async def process_pdf_file(pdf_path: Path, image_prefix: str,filename:str) -> Op 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) - print(11111111111111111111111111111111111111) - print(len(images)) - return {"slices": slices_check, "images": images} + 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": {}} + return {"slices": [], "images": {}, "summary": ""} except Exception as e: logger.error(f"处理 PDF 文件时出错: {e}", exc_info=True) return None @@ -1540,13 +1484,15 @@ async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) - # 优先尝试走本地/缓存处理(保持与你原逻辑一致) doc_result = await process_document(filename) if doc_result is not None: - content_list, images = doc_result + 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) - return {"slices": slices_check, "images": images} + 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") @@ -1571,19 +1517,21 @@ async def process_other_file(pdf_path: Path, image_prefix: str, filename: str) - 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) - return {"slices": slices_check, "images": images} + 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": {}} + return {"slices": [], "images": {}, "summary": ""} except Exception as e: logger.error(f"处理 Office 文件时出错: {e}", exc_info=True) return None diff --git a/fileparse_util.py b/fileparse_util.py index 723d7cb..d14d1fa 100644 --- a/fileparse_util.py +++ b/fileparse_util.py @@ -5,7 +5,7 @@ import asyncio import logging from typing import Optional, Tuple, Dict import aiofiles # pip install aiofiles - +from config import SEARCH_DIR,IMAGE_DIR logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -44,6 +44,34 @@ async def find_and_read_content_list( return None, None +async def find_and_read_md_file( + directory: str, + original_filename: str, + encoding: str = 'utf-8' +) -> Tuple[Optional[str], Optional[str]]: + """Find and read the Markdown file matching the uploaded document basename.""" + file_name, _ = os.path.splitext(original_filename) + target_filename = f"{file_name}.md" + logger.info(f"正在查找 Markdown 文件: {target_filename}") + + def _walk_for_file(): + for root, _, files in os.walk(directory): + if target_filename in files: + return os.path.join(root, target_filename) + return None + + file_path = await asyncio.to_thread(_walk_for_file) + if file_path is None: + return None, None + + try: + async with aiofiles.open(file_path, 'r', encoding=encoding, errors='ignore') as f: + return await f.read(), file_path + except Exception as e: + logger.error(f"读取 Markdown 文件 {file_path} 时发生错误: {e}") + return None, None + + def extract_all_jpg_filenames(doc_data: list) -> set: """从文档数据中提取所有 .jpg 图片的纯文件名。纯 CPU 操作,不需要 async。""" result_set = set() @@ -103,10 +131,9 @@ async def encode_images_to_base64( async def process_document( input_file_name: str, - search_dir: Optional[str] = "/app/mineru_output", local_image_dir: Optional[str] = None, concurrency: int = 32, -) -> Optional[Tuple[list, Dict[str, str]]]: +) -> Optional[Tuple[list, Dict[str, str], Optional[str]]]: """ 异步处理文档:查找 content_list.json,提取图片引用,并发编码为 base64。 @@ -117,9 +144,10 @@ async def process_document( concurrency: 图片编码的并发数上限,默认 32 Returns: - 成功: (content_list, images_dict) + 成功: (content_list, images_dict, md_content) 失败: None """ + search_dir = SEARCH_DIR content, content_path = await find_and_read_content_list(search_dir, input_file_name) if content is None: logger.error("❌ 未找到指定的 _content_list.json 文件。") @@ -128,6 +156,13 @@ async def process_document( logger.info(f"✅ 找到文件: {content_path}") logger.info(f"内容类型: {type(content).__name__}, 长度: {len(str(content))}") + md_content, md_path = await find_and_read_md_file(search_dir, input_file_name) + if md_content is None: + logger.info("未找到对应的 Markdown 文件。") + else: + logger.info(f"找到 Markdown 文件: {md_path}") + logger.info(f"Markdown 内容长度: {len(md_content)}") + if local_image_dir is None: local_image_dir = os.path.join(os.path.dirname(content_path), "images") logger.info(f"图片目录: {local_image_dir}") @@ -140,20 +175,19 @@ async def process_document( ) logger.info(f"成功编码 {len(images_dict)} 张图片") - return content, images_dict + return content, images_dict, md_content async def main(): result = await process_document( input_file_name="163-06A0014-B01001_发动机-维修手册.pdf", - search_dir="/app/mineru_output", concurrency=32, ) if result is None: return 1 - content, images_dict = result + content, images_dict, md_content = result print(f"\n=== 处理完成 ===") print(f"文档段落数: {len(content)}") print(f"图片数量: {len(images_dict)}") diff --git a/generate_summary.py b/generate_summary.py new file mode 100644 index 0000000..5e00026 --- /dev/null +++ b/generate_summary.py @@ -0,0 +1,142 @@ +""" +文件概述生成工具 +把 wiki_engine 的文档概述能力独立出来 +输入:markdown 格式的文件内容 +输出:文件概述(SUMMARY 行 + markdown 正文) +支持并发 +""" +import asyncio +from typing import List, Optional +from config import LLM_CONFIG +from openai import AsyncOpenAI + + +# ==================== 配置 ==================== +_MAX_CONCURRENCY = 8 # 最大并发数 + +_client: Optional[AsyncOpenAI] = None +_semaphore: Optional[asyncio.Semaphore] = None + + +def _get_client() -> AsyncOpenAI: + """获取或创建 AsyncOpenAI 单例客户端""" + global _client + if _client is None: + _client = AsyncOpenAI(api_key=LLM_CONFIG['api_key'], base_url=LLM_CONFIG['base_url']) + return _client + + +def _get_semaphore() -> asyncio.Semaphore: + """获取或创建并发信号量""" + global _semaphore + if _semaphore is None: + _semaphore = asyncio.Semaphore(_MAX_CONCURRENCY) + return _semaphore + + +# ==================== Prompt(与 wiki_builder.py 的 WIKI_SUMMARY_PROMPT 一致,仅去掉 extracted_slugs 相关部分) ==================== +SUMMARY_PROMPT = """You are a wiki editor. Given the following document content, create a structured wiki summary page in Markdown format. + + + +{content} + + + + +1. The FIRST line of your output MUST be: SUMMARY: {{one sentence, 15-40 words, describing what this document is about for wiki index listing}} +2. Create a concise but useful document-level summary. +3. Include the document's main subject, scope, important procedures, standards, systems, equipment, tables, fields. +4. Do NOT invent facts. Stay grounded in the document content. +5. Write in Chinese. +6. If the content is empty or has no substantive information, output exactly: "SUMMARY: No textual content was extractable from this document." followed by a brief note. + + +Output the SUMMARY line first, then the Markdown content. Do not include any other preamble.""" + + +async def generate_summary(content: str) -> str: + """ + 从 markdown 内容生成文件概述 + + Args: + content: markdown 格式的文件内容 + + Returns: + 文件概述(SUMMARY 行 + markdown 正文) + """ + if not content or not content.strip(): + return "" + + prompt = SUMMARY_PROMPT.format(content=content) + client = _get_client() + + async with _get_semaphore(): + try: + response = await client.chat.completions.create( + model=LLM_CONFIG['model'], + messages=[ + {"role": "system", "content": "You are a grounded wiki editor. Do not invent facts."}, + {"role": "user", "content": prompt}, + ], + temperature=0.1, + max_tokens=LLM_CONFIG['max_tokens'], + stream=False, + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, + ) + return (response.choices[0].message.content or "").strip() + except Exception as exc: + print(f"[generate_summary] 概述生成失败: {exc}") + return "" + + +async def generate_summaries(contents: List[str]) -> List[str]: + """ + 并发生成多个文件的概述 + + Args: + contents: markdown 格式的文件内容列表 + + Returns: + 文件概述列表,顺序与输入一致 + """ + return await asyncio.gather(*(generate_summary(c) for c in contents)) + + +if __name__ == "__main__": +# import sys + from pathlib import Path + +# if len(sys.argv) > 1: +# with open(sys.argv[1], "r", encoding="utf-8") as f: +# test_content = f.read() +# else: +# test_content = """ +# # 船舶动力系统维护规程 + +# ## 概述 +# 本文档详细介绍了船舶动力系统的日常维护和故障处理流程。 + +# ## 发动机日常检查 +# - 润滑油位检查:每日检查发动机润滑油位,保持在标尺正常范围 +# - 冷却液位检查:确保冷却系统液位正常,无泄漏 +# - 皮带张紧度:检查传动皮带张紧度,过松或过紧均需调整 + +# ## 冷却系统维护 +# 定期清洗热交换器,检查水泵密封性,更换老化管路。 + +# ## 故障诊断流程 +# 1. 现象观察:记录故障现象和发生条件 +# 2. 数据采集:收集运行参数和报警信息 +# 3. 原因分析:对照标准参数分析故障原因 +# 4. 处理方案:制定维修方案并执行 + +# ## 安全注意事项 +# 所有维护操作必须在停机状态下进行,操作人员需佩戴防护装备。 +# """ + md_path = Path(r"E:\ZKYNLP\Hjunproject\project0506\kgrag\船舶主机燃油泵自动控制系统故障树分析.md") + with open(md_path, "r", encoding="utf-8") as f: + test_content = f.read() + + result = asyncio.run(generate_summary(test_content)) + print(result)