422 lines
14 KiB
Python
422 lines
14 KiB
Python
"""DocumentLoader - 多格式文档解析器
|
||
|
||
支持 PDF（PyMuPDF/pdfplumber）、Word（python-docx）、Markdown（mistune）、
|
||
HTML（BeautifulSoup）、纯文本。所有格式依赖均为可选（try/except ImportError）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import logging
|
||
import uuid
|
||
from dataclasses import dataclass, field
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import TypeAlias
|
||
|
||
from agentkit.memory.base import MetadataDict
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ponytail: resource caps prevent OOM from malicious/oversized uploads.
|
||
# Ceiling: a 100MB document is ~25M tokens — beyond any useful LLM context.
|
||
# Upgrade path: stream to disk for very large files if needed.
|
||
MAX_CONTENT_SIZE = 100 * 1024 * 1024 # 100MB
|
||
MAX_ROWS_PER_SHEET = 10_000
|
||
MAX_CELL_CHARS = 10_000
|
||
|
||
# 文档元数据：source/format/parser/page_count/table_count/sheet_count/row_count/
|
||
# heading_count/created_at/title/truncated — 全部为原始标量。
|
||
DocumentMetadata: TypeAlias = MetadataDict
|
||
ParseResult: TypeAlias = tuple[str, DocumentMetadata]
|
||
|
||
|
||
@dataclass
|
||
class Document:
|
||
"""解析后的文档统一格式"""
|
||
|
||
doc_id: str
|
||
title: str
|
||
content: str
|
||
metadata: DocumentMetadata = field(default_factory=dict)
|
||
|
||
def __post_init__(self) -> None:
|
||
if "source" not in self.metadata:
|
||
self.metadata["source"] = ""
|
||
if "format" not in self.metadata:
|
||
self.metadata["format"] = "unknown"
|
||
if "page_count" not in self.metadata:
|
||
self.metadata["page_count"] = 0
|
||
if "created_at" not in self.metadata:
|
||
self.metadata["created_at"] = datetime.now(timezone.utc).isoformat()
|
||
|
||
def to_dict(self) -> dict[str, object]:
|
||
return {
|
||
"doc_id": self.doc_id,
|
||
"title": self.title,
|
||
"content": self.content,
|
||
"metadata": self.metadata,
|
||
}
|
||
|
||
|
||
def _detect_format(filename: str) -> str:
|
||
"""根据文件扩展名检测文档格式"""
|
||
ext = Path(filename).suffix.lower()
|
||
format_map = {
|
||
".pdf": "pdf",
|
||
".docx": "docx",
|
||
".doc": "docx",
|
||
".xlsx": "xlsx",
|
||
".xls": "xlsx",
|
||
".md": "markdown",
|
||
".markdown": "markdown",
|
||
".html": "html",
|
||
".htm": "html",
|
||
".txt": "text",
|
||
".csv": "text",
|
||
".json": "text",
|
||
".xml": "text",
|
||
}
|
||
return format_map.get(ext, "text")
|
||
|
||
|
||
class DocumentLoader:
|
||
"""多格式文档解析器
|
||
|
||
支持格式：
|
||
- PDF: PyMuPDF (fitz) → pdfplumber → 纯文本回退
|
||
- Word: python-docx → 纯文本回退
|
||
- Excel: openpyxl → 纯文本回退
|
||
- Markdown: mistune → 纯文本回退
|
||
- HTML: BeautifulSoup → 纯文本回退
|
||
- 纯文本: 直接读取
|
||
"""
|
||
|
||
def load(self, file_path: str | Path) -> Document:
|
||
"""从文件路径加载文档
|
||
|
||
Args:
|
||
file_path: 文件路径
|
||
|
||
Returns:
|
||
解析后的 Document 对象
|
||
|
||
Raises:
|
||
FileNotFoundError: 文件不存在
|
||
ValueError: 不支持的格式
|
||
"""
|
||
path = Path(file_path)
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"File not found: {path}")
|
||
|
||
content = path.read_bytes()
|
||
return self.load_bytes(content, path.name)
|
||
|
||
def load_bytes(self, content: bytes, filename: str) -> Document:
|
||
"""从字节内容加载文档
|
||
|
||
Args:
|
||
content: 文件字节内容
|
||
filename: 文件名（用于格式检测和元数据）
|
||
|
||
Returns:
|
||
解析后的 Document 对象
|
||
|
||
Raises:
|
||
ValueError: 内容超过 MAX_CONTENT_SIZE
|
||
"""
|
||
if len(content) > MAX_CONTENT_SIZE:
|
||
raise ValueError(
|
||
f"Content size {len(content)} bytes exceeds limit {MAX_CONTENT_SIZE} bytes"
|
||
)
|
||
doc_format = _detect_format(filename)
|
||
doc_id = str(uuid.uuid4())
|
||
|
||
parsers = {
|
||
"pdf": self._parse_pdf,
|
||
"docx": self._parse_docx,
|
||
"xlsx": self._parse_xlsx,
|
||
"markdown": self._parse_markdown,
|
||
"html": self._parse_html,
|
||
"text": self._parse_text,
|
||
}
|
||
|
||
parser = parsers.get(doc_format)
|
||
if parser is None:
|
||
logger.warning(
|
||
f"Unsupported format '{doc_format}' for {filename}, falling back to text"
|
||
)
|
||
parser = self._parse_text
|
||
|
||
text, extra_meta = parser(content, filename)
|
||
|
||
metadata: DocumentMetadata = {
|
||
"source": filename,
|
||
"format": doc_format,
|
||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||
}
|
||
metadata.update(extra_meta)
|
||
|
||
title = Path(filename).stem
|
||
if "title" in extra_meta:
|
||
title = extra_meta["title"]
|
||
|
||
return Document(
|
||
doc_id=doc_id,
|
||
title=title,
|
||
content=text,
|
||
metadata=metadata,
|
||
)
|
||
|
||
def _parse_pdf(self, content: bytes, filename: str) -> ParseResult:
|
||
"""解析 PDF 文件
|
||
|
||
优先使用 PyMuPDF (fitz)，回退到 pdfplumber，最终回退到纯文本。
|
||
"""
|
||
# 尝试 PyMuPDF
|
||
try:
|
||
import fitz # PyMuPDF
|
||
|
||
doc = fitz.open(stream=content, filetype="pdf")
|
||
pages = []
|
||
for page in doc:
|
||
pages.append(page.get_text())
|
||
text = "\n\n".join(pages)
|
||
meta = {
|
||
"page_count": len(doc),
|
||
"parser": "pymupdf",
|
||
}
|
||
# 提取 PDF 元数据中的标题
|
||
pdf_meta = doc.metadata
|
||
if pdf_meta and pdf_meta.get("title"):
|
||
meta["title"] = pdf_meta["title"]
|
||
doc.close()
|
||
return text, meta
|
||
except ImportError:
|
||
pass
|
||
except Exception as e:
|
||
logger.warning(f"PyMuPDF parsing failed for {filename}: {e}")
|
||
|
||
# 尝试 pdfplumber
|
||
try:
|
||
import pdfplumber
|
||
import io
|
||
|
||
pdf = pdfplumber.open(io.BytesIO(content))
|
||
pages = []
|
||
for page in pdf.pages:
|
||
page_text = page.extract_text()
|
||
if page_text:
|
||
pages.append(page_text)
|
||
text = "\n\n".join(pages)
|
||
meta = {
|
||
"page_count": len(pdf.pages),
|
||
"parser": "pdfplumber",
|
||
}
|
||
pdf.close()
|
||
return text, meta
|
||
except ImportError:
|
||
pass
|
||
except Exception as e:
|
||
logger.warning(f"pdfplumber parsing failed for {filename}: {e}")
|
||
|
||
# 回退到纯文本
|
||
logger.warning(f"No PDF parser available for {filename}, falling back to text extraction")
|
||
return self._parse_text(content, filename)
|
||
|
||
def _parse_docx(self, content: bytes, filename: str) -> ParseResult:
|
||
"""解析 Word 文件
|
||
|
||
使用 python-docx，回退到纯文本。
|
||
"""
|
||
try:
|
||
from docx import Document as DocxDocument
|
||
import io
|
||
|
||
doc = DocxDocument(io.BytesIO(content))
|
||
paragraphs = []
|
||
table_count = 0
|
||
|
||
# 提取段落文本
|
||
for para in doc.paragraphs:
|
||
if para.text.strip():
|
||
paragraphs.append(para.text.strip())
|
||
|
||
# 提取表格文本
|
||
for table in doc.tables:
|
||
table_count += 1
|
||
for row in table.rows:
|
||
row_text = " | ".join(cell.text.strip() for cell in row.cells)
|
||
if row_text.strip(" |"):
|
||
paragraphs.append(row_text)
|
||
|
||
text = "\n\n".join(paragraphs)
|
||
meta = {
|
||
"parser": "python-docx",
|
||
"table_count": table_count,
|
||
}
|
||
|
||
# 提取文档属性中的标题
|
||
if doc.core_properties and doc.core_properties.title:
|
||
meta["title"] = doc.core_properties.title
|
||
|
||
return text, meta
|
||
except ImportError:
|
||
logger.warning(f"python-docx not available for {filename}, falling back to text")
|
||
return self._parse_text(content, filename)
|
||
except Exception as e:
|
||
logger.warning(f"python-docx parsing failed for {filename}: {e}")
|
||
return self._parse_text(content, filename)
|
||
|
||
def _parse_xlsx(self, content: bytes, filename: str) -> ParseResult:
|
||
"""解析 Excel 文件
|
||
|
||
使用 openpyxl，回退到纯文本。每个 sheet 转为 Markdown 表格，
|
||
多个 sheet 用空行分隔，sheet 名作为 H2 标题。
|
||
|
||
注意：data_only=True 对未在 Excel 中打开过的公式返回 None（静默数据丢失）。
|
||
合并单元格仅左上角有值，其余为空。
|
||
"""
|
||
try:
|
||
from openpyxl import load_workbook
|
||
|
||
wb = load_workbook(io.BytesIO(content), data_only=True, read_only=True)
|
||
try:
|
||
sections: list[str] = []
|
||
sheet_count = 0
|
||
total_rows = 0
|
||
truncated = False
|
||
|
||
for ws in wb.worksheets:
|
||
sheet_count += 1
|
||
row_iter = ws.iter_rows(values_only=True)
|
||
rows: list[tuple] = []
|
||
for row in row_iter:
|
||
if total_rows + len(rows) >= MAX_ROWS_PER_SHEET:
|
||
truncated = True
|
||
break
|
||
rows.append(row)
|
||
if not rows:
|
||
continue
|
||
|
||
sections.append(f"## {ws.title}")
|
||
|
||
# Compute max column count for uniform Markdown table
|
||
max_cols = max(len(r) for r in rows)
|
||
|
||
for i, row in enumerate(rows):
|
||
total_rows += 1
|
||
cells = ["" if v is None else str(v)[:MAX_CELL_CHARS] for v in row]
|
||
# Pad to max_cols for valid Markdown table
|
||
cells += [""] * (max_cols - len(cells))
|
||
sections.append("| " + " | ".join(cells) + " |")
|
||
# ponytail: separator after header row for Markdown table validity
|
||
if i == 0:
|
||
sep_cells = ["---"] * max_cols
|
||
sections.append("| " + " | ".join(sep_cells) + " |")
|
||
|
||
if truncated:
|
||
sections.append(f"<!-- truncated at {MAX_ROWS_PER_SHEET} rows -->")
|
||
|
||
sections.append("") # blank line between sheets
|
||
finally:
|
||
wb.close()
|
||
text = "\n".join(sections).strip()
|
||
meta: DocumentMetadata = {
|
||
"parser": "openpyxl",
|
||
"sheet_count": sheet_count,
|
||
"row_count": total_rows,
|
||
}
|
||
if truncated:
|
||
meta["truncated"] = True
|
||
return text, meta
|
||
except ImportError:
|
||
logger.warning(f"openpyxl not available for {filename}, falling back to text")
|
||
return self._parse_text(content, filename)
|
||
except Exception as e:
|
||
logger.warning(f"openpyxl parsing failed for {filename}: {e}")
|
||
return self._parse_text(content, filename)
|
||
|
||
def _parse_markdown(self, content: bytes, filename: str) -> ParseResult:
|
||
"""解析 Markdown 文件
|
||
|
||
使用 mistune（如果可用），否则直接读取文本。
|
||
Markdown 原文保留，因为后续分块需要标题结构。
|
||
"""
|
||
try:
|
||
text = content.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
text = content.decode("utf-8", errors="replace")
|
||
|
||
# 提取第一个标题作为文档标题
|
||
title = ""
|
||
for line in text.split("\n"):
|
||
line_stripped = line.strip()
|
||
if line_stripped.startswith("#"):
|
||
title = line_stripped.lstrip("#").strip()
|
||
break
|
||
|
||
meta: DocumentMetadata = {
|
||
"parser": "markdown",
|
||
}
|
||
if title:
|
||
meta["title"] = title
|
||
|
||
# 统计标题数量（ponytail: simple string check, no mistune dependency needed）
|
||
heading_count = 0
|
||
for line in text.split("\n"):
|
||
if line.strip().startswith("#"):
|
||
heading_count += 1
|
||
meta["heading_count"] = heading_count
|
||
|
||
return text, meta
|
||
|
||
def _parse_html(self, content: bytes, filename: str) -> ParseResult:
|
||
"""解析 HTML 文件
|
||
|
||
使用 BeautifulSoup 提取文本，回退到纯文本。
|
||
"""
|
||
try:
|
||
from bs4 import BeautifulSoup
|
||
|
||
try:
|
||
html_text = content.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
html_text = content.decode("utf-8", errors="replace")
|
||
|
||
soup = BeautifulSoup(html_text, "html.parser")
|
||
|
||
# 移除 script 和 style 标签
|
||
for tag in soup(["script", "style"]):
|
||
tag.decompose()
|
||
|
||
text = soup.get_text(separator="\n", strip=True)
|
||
|
||
# 提取标题
|
||
title = ""
|
||
if soup.title and soup.title.string:
|
||
title = soup.title.string.strip()
|
||
|
||
meta: DocumentMetadata = {
|
||
"parser": "beautifulsoup",
|
||
}
|
||
if title:
|
||
meta["title"] = title
|
||
|
||
return text, meta
|
||
except ImportError:
|
||
logger.warning(f"BeautifulSoup not available for {filename}, falling back to text")
|
||
return self._parse_text(content, filename)
|
||
except Exception as e:
|
||
logger.warning(f"BeautifulSoup parsing failed for {filename}: {e}")
|
||
return self._parse_text(content, filename)
|
||
|
||
def _parse_text(self, content: bytes, filename: str) -> ParseResult:
|
||
"""解析纯文本文件"""
|
||
try:
|
||
text = content.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
text = content.decode("utf-8", errors="replace")
|
||
|
||
return text, {"parser": "text"}
|