165 lines
5.7 KiB
Python
165 lines
5.7 KiB
Python
"""本地文档处理器（MVP：同步处理）。"""
|
||
|
||
from pathlib import Path
|
||
|
||
from app.core.logging import get_logger
|
||
from app.core.security import decrypt_token, encrypt_token, generate_token, hash_token
|
||
from app.models.document import Document
|
||
from app.processors.parsers.docx_parser import parse_docx_with_python_docx
|
||
from app.processors.parsers.markitdown_parser import parse_with_markitdown
|
||
from app.processors.parsers.pdf_parser import has_text_layer, parse_pdf_with_pymupdf
|
||
from app.repositories.doc_repo import DocumentRepository
|
||
from app.storage.local_storage import get_storage
|
||
from app.storage.object_keys import markdown_object_key
|
||
from sqlalchemy.orm import Session
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class LocalDocumentProcessor:
|
||
"""本地文档处理器：在同一进程中完成解析。"""
|
||
|
||
def __init__(self, session: Session) -> None:
|
||
self._session = session
|
||
self._doc_repo = DocumentRepository(session)
|
||
|
||
def process(self, document_id: str) -> None:
|
||
"""处理文档：解析 → 提取元数据 → 保存 Markdown → 更新 DB。"""
|
||
doc = self._doc_repo.get_by_id(document_id)
|
||
if doc is None:
|
||
logger.error("Document not found: %s", document_id)
|
||
return
|
||
|
||
# 更新状态为 PROCESSING
|
||
doc.status = "PROCESSING"
|
||
self._session.commit()
|
||
|
||
try:
|
||
# 读取原始文件
|
||
storage = get_storage()
|
||
file_content = storage.read(doc.storage_path)
|
||
|
||
# 解析
|
||
markdown = self._parse_file(file_content, doc.file_ext, doc.original_filename)
|
||
|
||
if markdown is None:
|
||
# 解析失败
|
||
if doc.file_ext == ".pdf" and not has_text_layer(file_content):
|
||
doc.error_code = "SCANNED_PDF_NO_TEXT_LAYER"
|
||
doc.status = "FAILED"
|
||
else:
|
||
doc.error_code = "PARSING_FAILED"
|
||
doc.status = "FAILED"
|
||
self._session.commit()
|
||
logger.warning("Document parsing failed: %s (%s)", doc.id, doc.error_code)
|
||
return
|
||
|
||
# 清洗 Markdown
|
||
markdown = self._clean_markdown(markdown)
|
||
|
||
# 保存 Markdown 文件
|
||
md_key = markdown_object_key(
|
||
user_id=doc.user_id,
|
||
knowledge_base_id=doc.knowledge_base_id,
|
||
document_id=doc.id,
|
||
)
|
||
storage.save(md_key, markdown.encode("utf-8"))
|
||
doc.markdown_path = md_key
|
||
|
||
# 提取元数据
|
||
title = self._extract_title(markdown, doc.original_filename)
|
||
summary = self._extract_summary(markdown)
|
||
keywords = self._extract_keywords(markdown)
|
||
|
||
doc.title = title
|
||
doc.content_summary = summary
|
||
doc.keywords = keywords
|
||
doc.status = "READY"
|
||
self._session.commit()
|
||
|
||
logger.info("Document processed: %s → READY", doc.id)
|
||
|
||
except Exception as e:
|
||
logger.error("Document processing error: %s - %s", doc.id, e, exc_info=True)
|
||
doc.status = "FAILED"
|
||
doc.error_code = "PARSING_FAILED"
|
||
self._session.commit()
|
||
|
||
def _parse_file(self, content: bytes, ext: str, filename: str) -> str | None:
|
||
"""按扩展名路由到对应解析器。"""
|
||
# 优先 MarkItDown
|
||
markdown = parse_with_markitdown(content, filename)
|
||
if markdown and markdown.strip():
|
||
return markdown
|
||
|
||
# Fallback
|
||
if ext == ".pdf":
|
||
return parse_pdf_with_pymupdf(content)
|
||
elif ext == ".docx":
|
||
return parse_docx_with_python_docx(content)
|
||
|
||
return None
|
||
|
||
def _clean_markdown(self, markdown: str) -> str:
|
||
"""清洗 Markdown：压缩空行、规整标题层级。"""
|
||
import re
|
||
|
||
# 压缩连续空行为最多 2 个
|
||
markdown = re.sub(r"\n{3,}", "\n\n", markdown)
|
||
# 截断超长行（> 1000 字符的行）
|
||
lines = markdown.split("\n")
|
||
cleaned = []
|
||
for line in lines:
|
||
if len(line) > 1000:
|
||
line = line[:1000] + "..."
|
||
cleaned.append(line)
|
||
return "\n".join(cleaned)
|
||
|
||
def _extract_title(self, markdown: str, fallback: str) -> str:
|
||
"""提取标题：H1 → 文件名（去扩展名）。"""
|
||
import re
|
||
|
||
# 找第一个 H1
|
||
match = re.search(r"^#\s+(.+)$", markdown, re.MULTILINE)
|
||
if match:
|
||
return match.group(1).strip()
|
||
|
||
# Fallback：文件名去扩展名
|
||
return Path(fallback).stem
|
||
|
||
def _extract_summary(self, markdown: str, max_len: int = 200) -> str:
|
||
"""抽取式摘要：取正文前 ~200 字纯文本。"""
|
||
import re
|
||
|
||
# 去掉 Markdown 标记
|
||
text = re.sub(r"[#*_`\[\]()>]", "", markdown)
|
||
text = re.sub(r"\s+", " ", text).strip()
|
||
|
||
if len(text) <= max_len:
|
||
return text
|
||
|
||
# 在 max_len 附近找句号/逗号断句
|
||
cut = text[:max_len]
|
||
for sep in ["。", "！", "？", "；", ".", "!", "?", "；"]:
|
||
idx = cut.rfind(sep)
|
||
if idx > max_len * 0.5:
|
||
return cut[: idx + 1]
|
||
return cut + "..."
|
||
|
||
def _extract_keywords(self, markdown: str, top_k: int = 10) -> str:
|
||
"""提取关键词（jieba TF-IDF）。"""
|
||
try:
|
||
import jieba.analyse
|
||
|
||
# 去掉 Markdown 标记
|
||
import re
|
||
|
||
text = re.sub(r"[#*_`\[\]()>]", "", markdown)
|
||
text = re.sub(r"\s+", " ", text).strip()
|
||
|
||
keywords = jieba.analyse.extract_tags(text, topK=top_k)
|
||
return ",".join(keywords)
|
||
except Exception as e:
|
||
logger.warning("关键词提取失败: %s", e)
|
||
return ""
|