feat: 新增多格式文件上传入库与认证体系
- 新增 JWT 认证模块，支持登录/注册/用户管理 - 新增文件上传接口，支持 .txt/.md/.html/.pdf/.docx 等格式解析入库 - 新增检索结果 AI 总结功能 - 新增文本去重缓存机制 - 新增全局认证夹具简化测试 - 新增配置项与环境变量支持 - 完善文档与测试覆盖
This commit is contained in:
@@ -0,0 +1,204 @@
|
||||
"""多格式文件文本提取：按扩展名分发到对应解析器
|
||||
|
||||
支持的扩展名：
|
||||
- .txt / .md：UTF-8 解码（errors="replace" 兜底）
|
||||
- .html / .htm：标准库 html.parser 剥离标签提取可见文本
|
||||
- .pdf：pypdf 逐页 extract_text 拼接；文本层为空（扫描件/图片型）时
|
||||
自动降级为 OCR（pypdfium2 渲染 + rapidocr-onnxruntime 识别）
|
||||
- .docx：python-docx 段落文本拼接（不含表格/页眉页脚）
|
||||
|
||||
未识别扩展名抛 ValueError("不支持的文件类型: {ext}")；
|
||||
解析异常统一包装为 ValueError("文件解析失败: {detail}")，原异常链式保留。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from collections.abc import Callable
|
||||
from html.parser import HTMLParser
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import structlog
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# 模块级懒加载 OCR 引擎单例（首次调用时初始化，避免无扫描件场景白白下载模型）
|
||||
_ocr_engine: Any | None = None
|
||||
_ocr_unavailable: bool = False # 标记 OCR 依赖不可用，后续直接跳过避免重复尝试
|
||||
|
||||
|
||||
class _VisibleTextExtractor(HTMLParser):
|
||||
"""HTMLParser 子类：累积可见文本，跳过 script/style 内容"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(convert_charrefs=True)
|
||||
self._parts: list[str] = []
|
||||
self._skip_depth = 0 # 在 script/style 标签内时 > 0
|
||||
|
||||
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
||||
if tag.lower() in {"script", "style"}:
|
||||
self._skip_depth += 1
|
||||
|
||||
def handle_endtag(self, tag: str) -> None:
|
||||
if tag.lower() in {"script", "style"} and self._skip_depth > 0:
|
||||
self._skip_depth -= 1
|
||||
|
||||
def handle_data(self, data: str) -> None:
|
||||
if self._skip_depth == 0:
|
||||
self._parts.append(data)
|
||||
|
||||
def get_text(self) -> str:
|
||||
# 块级标签间用空格连接，再折叠多余空白
|
||||
text = " ".join(self._parts)
|
||||
return " ".join(text.split())
|
||||
|
||||
|
||||
def _decode_html(content: bytes) -> str:
|
||||
"""HTML 内容解码：优先 utf-8（带 BOM），失败回退 latin-1"""
|
||||
try:
|
||||
return content.decode("utf-8-sig")
|
||||
except UnicodeDecodeError:
|
||||
return content.decode("latin-1", errors="replace")
|
||||
|
||||
|
||||
def _parse_text(content: bytes) -> str:
|
||||
"""UTF-8 解码（errors=replace 兜底），保留原字符"""
|
||||
return content.decode("utf-8", errors="replace")
|
||||
|
||||
|
||||
def _parse_html(content: bytes) -> str:
|
||||
"""HTML 剥离标签，保留可见文本"""
|
||||
parser = _VisibleTextExtractor()
|
||||
parser.feed(_decode_html(content))
|
||||
parser.close()
|
||||
return parser.get_text()
|
||||
|
||||
|
||||
def _parse_pdf(content: bytes) -> str:
|
||||
"""PDF 解析：优先 pypdf extract_text；文本层为空（扫描件）时降级 OCR
|
||||
|
||||
OCR 流程：pypdfium2 渲染每页为 PIL Image → rapidocr-onnxruntime 识别 →
|
||||
拼接每页识别出的文本。受 settings.pdf_ocr_* 控制：开关、最大页数、DPI。
|
||||
OCR 依赖未安装或运行异常时降级返回空字符串（由上游 upload 端点拒绝入库）。
|
||||
"""
|
||||
from pypdf import PdfReader
|
||||
|
||||
reader = PdfReader(io.BytesIO(content))
|
||||
parts: list[str] = []
|
||||
for page in reader.pages:
|
||||
text = page.extract_text() or ""
|
||||
if text:
|
||||
parts.append(text)
|
||||
text_layer = "\n".join(parts).strip()
|
||||
|
||||
# 文本层非空：直接返回
|
||||
if text_layer:
|
||||
return text_layer
|
||||
|
||||
# 文本层为空 → 尝试 OCR 降级
|
||||
if not settings.pdf_ocr_enabled:
|
||||
return ""
|
||||
|
||||
ocr_text = _ocr_pdf(content, max_pages=settings.pdf_ocr_max_pages, dpi=settings.pdf_ocr_dpi)
|
||||
return ocr_text
|
||||
|
||||
|
||||
def _ocr_pdf(content: bytes, max_pages: int, dpi: int) -> str:
|
||||
"""对扫描件 PDF 跑 OCR：渲染每页 → 识别 → 拼接
|
||||
|
||||
返回空字符串的场景：依赖未安装 / 渲染或识别异常 / 无识别结果。
|
||||
任何异常仅告警不抛出，由上游按"无法提取文本"处理。
|
||||
"""
|
||||
global _ocr_engine, _ocr_unavailable
|
||||
|
||||
if _ocr_unavailable:
|
||||
return ""
|
||||
|
||||
# 1. 懒加载 OCR 引擎
|
||||
if _ocr_engine is None:
|
||||
try:
|
||||
from rapidocr_onnxruntime import RapidOCR
|
||||
|
||||
_ocr_engine = RapidOCR()
|
||||
logger.info("PDF OCR 引擎已初始化", dpi=dpi)
|
||||
except Exception:
|
||||
_ocr_unavailable = True
|
||||
logger.warning("OCR 依赖不可用，扫描件 PDF 将无法提取文本", exc_info=True)
|
||||
return ""
|
||||
|
||||
# 2. 渲染并识别
|
||||
try:
|
||||
import pypdfium2 as pdfium
|
||||
|
||||
scale = max(1.0, dpi / 72.0)
|
||||
pdf = pdfium.PdfDocument(io.BytesIO(content))
|
||||
total = min(len(pdf), max(1, max_pages))
|
||||
page_texts: list[str] = []
|
||||
for i in range(total):
|
||||
page = pdf[i]
|
||||
pil_image = page.render(scale=scale).to_pil()
|
||||
result, _ = _ocr_engine(pil_image)
|
||||
if result:
|
||||
# result: [[box, text, score], ...]，按行拼接
|
||||
lines = [item[1] for item in result if item and len(item) >= 2 and item[1]]
|
||||
if lines:
|
||||
page_texts.append("\n".join(lines))
|
||||
pdf.close()
|
||||
return "\n".join(page_texts).strip()
|
||||
except Exception as exc:
|
||||
logger.warning("PDF OCR 失败，降级返回空文本", error=str(exc), exc_info=True)
|
||||
return ""
|
||||
|
||||
|
||||
def _parse_docx(content: bytes) -> str:
|
||||
"""DOCX 段落文本拼接（不含表格/页眉页脚）"""
|
||||
from docx import Document # type: ignore[import-untyped]
|
||||
|
||||
document = Document(io.BytesIO(content))
|
||||
parts = [p.text for p in document.paragraphs if p.text and p.text.strip()]
|
||||
return "\n".join(parts).strip()
|
||||
|
||||
|
||||
# 扩展名 → 解析函数映射（启动时构建，避免每次请求重复构造）
|
||||
_PARSERS: dict[str, Callable[[bytes], str]] = {
|
||||
".txt": _parse_text,
|
||||
".md": _parse_text,
|
||||
".html": _parse_html,
|
||||
".htm": _parse_html,
|
||||
".pdf": _parse_pdf,
|
||||
".docx": _parse_docx,
|
||||
}
|
||||
|
||||
|
||||
def parse_file(filename: str, content: bytes) -> str:
|
||||
"""按扩展名分发解析器提取文本
|
||||
|
||||
Args:
|
||||
filename: 文件名（用于判定扩展名）
|
||||
content: 文件二进制内容
|
||||
|
||||
Returns:
|
||||
提取的纯文本
|
||||
|
||||
Raises:
|
||||
ValueError: 未识别扩展名 → "不支持的文件类型: {ext}"
|
||||
解析失败 → "文件解析失败: {detail}"（保留原异常链）
|
||||
"""
|
||||
ext = Path(filename).suffix.lower()
|
||||
parser = _PARSERS.get(ext)
|
||||
if parser is None:
|
||||
raise ValueError(f"不支持的文件类型: {ext or '(无扩展名)'}")
|
||||
try:
|
||||
return parser(content)
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ValueError(f"文件解析失败: {exc}") from exc
|
||||
|
||||
|
||||
def supported_extensions() -> set[str]:
|
||||
"""返回当前支持的扩展名集合（含点号，全小写）"""
|
||||
return set(_PARSERS.keys())
|
||||
Reference in New Issue
Block a user