forked from admin/PlantAI-rag-new
jerryandClaude Fable 5
4749f77544
第一批：上传去限制+进度展示、模板HTML预览+按需导出、检索TAB实时更新
第二批：edit护栏移除、导出按钮纯化+file_name落库、编辑失焦还原、
二次预览白底根治（防竞态+渲染兜底）
模板体系：7套公文模板(templates/)+templating渲染器、websearch联网
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
287 lines
9.4 KiB
Python
287 lines
9.4 KiB
Python
"""parsing.py —— 上传文档解析（docx / pdf / 文本），产出统一结构。
|
||
|
||
统一产出（dict）：
|
||
{
|
||
"outline": [{"level": int, "text": str}, ...], # 标题大纲（按出现顺序）
|
||
"title": str, # 文档标题（推断）
|
||
"paragraphs": [str, ...], # 正文段落
|
||
"tables": [[[cell, ...], ...], ...], # 表格（按出现顺序）
|
||
"style_hints": {...}, # docx 样式线索（导出对齐用）
|
||
}
|
||
|
||
文本文件 utf-8 → gb18030 回退；上传不设大小/页数/字符上限（2026-09-04）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import re
|
||
from typing import Any
|
||
|
||
import pymupdf
|
||
from docx import Document
|
||
from docx.document import Document as _DocxDoc
|
||
from docx.table import Table as _DocxTable
|
||
from docx.text.paragraph import Paragraph as _DocxParagraph
|
||
|
||
# 2026-09-04 手工测试：PDF 不再设页数上限（上传不设限）。
|
||
# 大纲嵌入切块参数（与 settings 对齐，避免循环 import）
|
||
_OUTLINE_EMBED_MAX_CHARS = 7500
|
||
_OUTLINE_EMBED_MAX_CHUNKS = 3
|
||
|
||
_HEADING_RE = re.compile(r"[Hh]eading\s*(\d+)")
|
||
_HEADING_CN_RE = re.compile(r"标题\s*(\d+)")
|
||
|
||
|
||
class ParsingError(Exception):
|
||
"""文件解析失败（格式不支持/损坏/超限）。"""
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# docx
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _heading_level(style_name: str) -> int | None:
|
||
"""从样式名推断标题级别（Heading 1 / 标题 1 → 1）。"""
|
||
m = _HEADING_RE.match(style_name or "")
|
||
if m:
|
||
return int(m.group(1))
|
||
m = _HEADING_CN_RE.match(style_name or "")
|
||
if m:
|
||
return int(m.group(1))
|
||
return None
|
||
|
||
|
||
def _docx_style_hints(doc: _DocxDoc) -> dict[str, Any]:
|
||
"""提取 docx 样式线索：正文字体/字号 + 各级标题字体/字号（尽力而为）。"""
|
||
hints: dict[str, Any] = {
|
||
"body_font": None,
|
||
"body_size": None,
|
||
"headings": {},
|
||
"used_styles": [],
|
||
}
|
||
try:
|
||
styles = doc.styles
|
||
normal = styles["Normal"] if "Normal" in [s.name for s in styles] else None
|
||
# python-docx: styles["Normal"] 可能抛 KeyError，用迭代兜底
|
||
if normal is None:
|
||
for s in styles:
|
||
if s.name == "Normal":
|
||
normal = s
|
||
break
|
||
if normal is not None and normal.font is not None:
|
||
hints["body_font"] = normal.font.name
|
||
if normal.font.size is not None:
|
||
hints["body_size"] = normal.font.size.pt
|
||
except Exception:
|
||
pass
|
||
|
||
used: set[str] = set()
|
||
for para in doc.paragraphs:
|
||
style_name = para.style.name if para.style is not None else ""
|
||
if style_name:
|
||
used.add(style_name)
|
||
level = _heading_level(style_name)
|
||
if level is None:
|
||
continue
|
||
font_name = size = None
|
||
try:
|
||
runs = para.runs
|
||
if runs:
|
||
font_name = runs[0].font.name
|
||
if runs[0].font.size is not None:
|
||
size = runs[0].font.size.pt
|
||
except Exception:
|
||
pass
|
||
if font_name or size:
|
||
entry = hints["headings"].setdefault(str(level), {})
|
||
if font_name:
|
||
entry["font"] = font_name
|
||
if size:
|
||
entry["size"] = size
|
||
hints["used_styles"] = sorted(used)
|
||
return hints
|
||
|
||
|
||
def _table_rows(table: _DocxTable) -> list[list[str]]:
|
||
"""表格 → 行×单元格文本（合并单元格重复出现，保持矩阵形状）。"""
|
||
rows: list[list[str]] = []
|
||
for row in table.rows:
|
||
rows.append([cell.text.strip() for cell in row.cells])
|
||
return rows
|
||
|
||
|
||
def extract_docx(file_bytes: bytes) -> dict[str, Any]:
|
||
"""解析 docx：大纲 + 段落 + 表格 + 样式线索。"""
|
||
try:
|
||
doc: _DocxDoc = Document(io.BytesIO(file_bytes))
|
||
except Exception as exc:
|
||
raise ParsingError(f"docx 解析失败：{exc}") from exc
|
||
|
||
outline: list[dict[str, Any]] = []
|
||
paragraphs: list[str] = []
|
||
title = ""
|
||
for para in doc.paragraphs:
|
||
text = (para.text or "").strip()
|
||
if not text:
|
||
continue
|
||
level = _heading_level(para.style.name if para.style is not None else "")
|
||
if level is not None:
|
||
outline.append({"level": level, "text": text})
|
||
continue
|
||
paragraphs.append(text)
|
||
if not title:
|
||
title = text # 首个非标题段落视作文档标题
|
||
|
||
tables: list[list[list[str]]] = []
|
||
for table in doc.tables:
|
||
rows = _table_rows(table)
|
||
if rows and any(any(cell for cell in row) for row in rows):
|
||
tables.append(rows)
|
||
|
||
return {
|
||
"outline": outline,
|
||
"title": title,
|
||
"paragraphs": paragraphs,
|
||
"tables": tables,
|
||
"style_hints": _docx_style_hints(doc),
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# pdf
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def extract_pdf(file_bytes: bytes) -> dict[str, Any]:
|
||
"""解析 PDF（pymupdf）：书签为大纲，文本块为段落。"""
|
||
try:
|
||
doc = pymupdf.open(stream=file_bytes, filetype="pdf")
|
||
except Exception as exc:
|
||
raise ParsingError(f"pdf 解析失败：{exc}") from exc
|
||
try:
|
||
outline: list[dict[str, Any]] = []
|
||
try:
|
||
for level, text, _page in doc.get_toc():
|
||
if text and text.strip():
|
||
outline.append({"level": max(1, level), "text": text.strip()})
|
||
except Exception:
|
||
outline = []
|
||
|
||
paragraphs: list[str] = []
|
||
for page in doc:
|
||
blocks = page.get_text("blocks")
|
||
for block in blocks:
|
||
try:
|
||
text = str(block[4]).strip() if len(block) > 4 else ""
|
||
except Exception:
|
||
continue
|
||
if text:
|
||
paragraphs.append(text)
|
||
|
||
title = outline[0]["text"] if outline else ""
|
||
return {
|
||
"outline": outline,
|
||
"title": title,
|
||
"paragraphs": paragraphs,
|
||
"tables": [],
|
||
"style_hints": {},
|
||
}
|
||
finally:
|
||
doc.close()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 文本文件
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def extract_text_file(file_bytes: bytes) -> dict[str, Any]:
|
||
"""解析纯文本（utf-8 → gb18030 回退）；空行分段。"""
|
||
for encoding in ("utf-8", "gb18030"):
|
||
try:
|
||
text = file_bytes.decode(encoding)
|
||
break
|
||
except UnicodeDecodeError:
|
||
continue
|
||
else:
|
||
text = file_bytes.decode("utf-8", errors="replace")
|
||
|
||
paragraphs = [line.strip() for line in text.split("\n\n") if line.strip()]
|
||
if not paragraphs:
|
||
paragraphs = [line.strip() for line in text.split("\n") if line.strip()]
|
||
title = paragraphs[0] if paragraphs else ""
|
||
return {
|
||
"outline": [],
|
||
"title": title,
|
||
"paragraphs": paragraphs,
|
||
"tables": [],
|
||
"style_hints": {},
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 大纲派生
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def outline_to_markdown(outline: list[dict[str, Any]]) -> str:
|
||
"""大纲 → markdown 标题行（level 1 → #，level 2 → ## …）。"""
|
||
lines: list[str] = []
|
||
for item in outline:
|
||
try:
|
||
level = max(1, int(item.get("level", 1)))
|
||
except (TypeError, ValueError):
|
||
level = 1
|
||
text = str(item.get("text", "")).strip()
|
||
if text:
|
||
lines.append(f"{'#' * min(level, 6)} {text}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def outline_to_embedding_chunks(
|
||
outline: list[dict[str, Any]],
|
||
max_chars: int = _OUTLINE_EMBED_MAX_CHARS,
|
||
max_chunks: int = _OUTLINE_EMBED_MAX_CHUNKS,
|
||
) -> list[str]:
|
||
"""大纲 → 嵌入文本块。
|
||
|
||
全长 ≤ max_chars 直接整块；超长按一级标题分节打包，
|
||
块数不超过 max_chunks（超出部分并入最后一块）。
|
||
"""
|
||
md = outline_to_markdown(outline)
|
||
if len(md) <= max_chars:
|
||
return [md] if md else []
|
||
|
||
sections: list[str] = []
|
||
current: list[str] = []
|
||
for line in md.split("\n"):
|
||
if line.startswith("# "): # 一级标题开新节
|
||
if current:
|
||
sections.append("\n".join(current))
|
||
current = [line]
|
||
else:
|
||
current.append(line)
|
||
if current:
|
||
sections.append("\n".join(current))
|
||
if not sections:
|
||
return [md[:max_chars]]
|
||
|
||
chunks: list[str] = []
|
||
buf: list[str] = []
|
||
for section in sections:
|
||
if buf and len("\n".join(buf)) + len(section) + 1 > max_chars and len(chunks) < max_chunks - 1:
|
||
chunks.append("\n".join(buf))
|
||
buf = [section]
|
||
else:
|
||
buf.append(section)
|
||
if buf:
|
||
chunks.append("\n".join(buf))
|
||
return chunks[:max_chunks]
|
||
|
||
|
||
def extract(file_bytes: bytes, filename: str) -> dict[str, Any]:
|
||
"""按扩展名分发解析。未知扩展按文本处理（utf-8 → gb18030）。"""
|
||
name = (filename or "").lower()
|
||
if name.endswith(".docx"):
|
||
return extract_docx(file_bytes)
|
||
if name.endswith(".pdf"):
|
||
return extract_pdf(file_bytes)
|
||
return extract_text_file(file_bytes)
|