split_result接口增加文件摘要抽取功能
This commit is contained in:
parent
325fed5b2d
commit
f8c99dcfdc
166
chunk_text.py
166
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
|
||||
|
||||
@ -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()
|
||||
@ -105,7 +133,7 @@ async def process_document(
|
||||
input_file_name: str,
|
||||
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。
|
||||
|
||||
@ -116,7 +144,7 @@ async def process_document(
|
||||
concurrency: 图片编码的并发数上限,默认 32
|
||||
|
||||
Returns:
|
||||
成功: (content_list, images_dict)
|
||||
成功: (content_list, images_dict, md_content)
|
||||
失败: None
|
||||
"""
|
||||
search_dir = SEARCH_DIR
|
||||
@ -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)}")
|
||||
|
||||
142
generate_summary.py
Normal file
142
generate_summary.py
Normal file
@ -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.
|
||||
|
||||
<document>
|
||||
<content>
|
||||
{content}
|
||||
</content>
|
||||
</document>
|
||||
|
||||
<instructions>
|
||||
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.
|
||||
</instructions>
|
||||
|
||||
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)
|
||||
Loading…
x
Reference in New Issue
Block a user