sorrow404null
a7412824e0
Import the pre-repair source tree as the history baseline. Runtime data (data/), virtualenvs, bytecode caches and logs are gitignored so local secrets and user state stay out of the repo.
162 lines
6.3 KiB
Python
162 lines
6.3 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
PDF 读取工具 (pdf_reader)
|
||
=========================
|
||
|
||
项目内第一个正式内置工具（builtin tool）。对外提供两项能力：
|
||
|
||
1. ``extract_pdf_text`` —— 提取 PDF 全部文本，按页结构化为提示词友好格式
|
||
2. ``extract_pdf_images`` —— 提取 PDF 内嵌的所有图片，逐张落地为 PNG
|
||
|
||
代码规范与调用模式（约定）
|
||
--------------------------
|
||
* 每个对外函数都带完整类型注解（参数 + 返回值）。
|
||
* docstring 采用 ``Args / Returns / Raises`` 分段，逐参数说明含义与单位。
|
||
* 将来的工具注册器（tools/registry.py，尚未实现）可通过
|
||
``inspect.signature`` + ``__doc__`` 自动反射出 OpenAI 格式的 Tool JSON Schema，
|
||
因此本文件的函数签名与文档即"接口契约"，请保持稳定、描述清晰。
|
||
|
||
当前阶段说明
|
||
------------
|
||
* 本工具暂**不接入** agent 的自动工具调用链路（registry 未实现），
|
||
由 UI 附件流程（``ui/views/main_window.py`` 的 PDF 处理）直接 import 调用。
|
||
* 参数与返回值已按"可被自动调用"的标准描述清楚，后续接线无需改动函数本身。
|
||
|
||
依赖与协议
|
||
----------
|
||
* 依赖 PyMuPDF（导入名 ``pymupdf``）。
|
||
* PyMuPDF 为 **AGPL-3.0 / Artifex 商业** 双许可：分发或通过网络提供服务时，
|
||
需按 AGPL 开源本项目或购买商业许可。纯自用/内部不分发则无约束。
|
||
"""
|
||
import os
|
||
import uuid
|
||
|
||
try:
|
||
import pymupdf # PyMuPDF >= 1.24 的官方导入名
|
||
except ImportError: # 兼容旧版导入名 fitz
|
||
import fitz as pymupdf # type: ignore
|
||
|
||
|
||
# 单个 PDF 的大小上限（字节）：超过则拒绝解析，避免内存与上下文爆炸
|
||
PDF_MAX_BYTES = 20 * 1024 * 1024 # 20 MB
|
||
|
||
|
||
def _open_doc(path: str, max_bytes: int) -> "pymupdf.Document":
|
||
"""内部辅助：做存在性/大小校验后打开 PDF，统一异常为 ValueError。"""
|
||
if not os.path.isfile(path):
|
||
raise ValueError("文件不存在")
|
||
size = os.path.getsize(path)
|
||
if size > max_bytes:
|
||
raise ValueError(
|
||
f"PDF 过大（{size / 1024 / 1024:.1f} MB > {max_bytes / 1024 / 1024:.0f} MB）"
|
||
)
|
||
try:
|
||
doc = pymupdf.open(path)
|
||
except Exception as e: # pymupdf 抛出的异常类型不固定，统一兜底
|
||
raise ValueError(f"PDF 打开失败（可能损坏）: {e}")
|
||
if doc.needs_pass:
|
||
doc.close()
|
||
raise ValueError("PDF 已加密，需要密码，暂不支持")
|
||
return doc
|
||
|
||
|
||
def extract_pdf_text(path: str, max_bytes: int = PDF_MAX_BYTES) -> tuple:
|
||
"""提取 PDF 的全部文本，按页结构化。
|
||
|
||
输出格式（无文本的空页会被跳过）::
|
||
|
||
[第 1 页]
|
||
<该页文本>
|
||
|
||
[第 2 页]
|
||
<该页文本>
|
||
|
||
Args:
|
||
path: PDF 文件路径（绝对路径，或相对当前工作目录的路径）。
|
||
max_bytes: 允许解析的最大文件字节数，默认 ``PDF_MAX_BYTES``（20MB）。
|
||
|
||
Returns:
|
||
``(structured_text, page_count)`` 二元组：
|
||
|
||
* ``structured_text`` (str): 按页拼接的结构化文本，各页之间以空行分隔；
|
||
若整份 PDF 无文本则为空字符串 ``""``。
|
||
* ``page_count`` (int): PDF 的总页数（含空页）。
|
||
|
||
Raises:
|
||
ValueError: 文件不存在 / 超过 ``max_bytes`` / PDF 损坏 / PDF 加密。
|
||
"""
|
||
doc = _open_doc(path, max_bytes)
|
||
try:
|
||
page_count = doc.page_count
|
||
parts = []
|
||
for i, page in enumerate(doc, 1):
|
||
text = page.get_text("text").strip()
|
||
if text:
|
||
parts.append(f"[第 {i} 页]\n{text}")
|
||
return "\n\n".join(parts), page_count
|
||
finally:
|
||
doc.close()
|
||
|
||
|
||
def extract_pdf_images(path: str, out_dir: str, max_bytes: int = PDF_MAX_BYTES) -> list:
|
||
"""提取 PDF 内嵌的所有图片，逐张保存为 PNG。
|
||
|
||
遍历每一页的内嵌图片（``page.get_images``），用 ``Pixmap`` 解码后写出；
|
||
CMYK 色彩空间会自动转换为 RGB（否则无法存为 PNG）。单张图片解码失败
|
||
会被跳过，不影响其余图片。
|
||
|
||
Args:
|
||
path: PDF 文件路径。
|
||
out_dir: 图片输出目录（建议使用绝对路径）。不存在时自动创建。
|
||
max_bytes: 允许解析的最大文件字节数，默认 ``PDF_MAX_BYTES``（20MB）。
|
||
|
||
Returns:
|
||
图片元数据列表（按 ``(page, index)`` 升序），每项为 dict::
|
||
|
||
{
|
||
"page": int, # 所在页码（从 1 开始）
|
||
"index": int, # 该页内的第几张（从 1 开始）
|
||
"abs_path": str, # 落地 PNG 的绝对路径
|
||
"mime": str, # 固定为 "image/png"
|
||
"size_kb": float, # 文件大小（KB，保留两位小数）
|
||
"width": int, # 图片像素宽
|
||
"height": int, # 图片像素高
|
||
}
|
||
|
||
PDF 中无内嵌图片时返回空列表 ``[]``。
|
||
|
||
Raises:
|
||
ValueError: 文件不存在 / 超过 ``max_bytes`` / PDF 损坏 / PDF 加密。
|
||
"""
|
||
doc = _open_doc(path, max_bytes)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
tag = uuid.uuid4().hex[:8] # 本批次文件名前缀，避免跨 PDF 撞名
|
||
results = []
|
||
try:
|
||
for page_num, page in enumerate(doc, 1):
|
||
for img_index, img_info in enumerate(page.get_images(full=True), 1):
|
||
xref = img_info[0]
|
||
try:
|
||
pix = pymupdf.Pixmap(doc, xref)
|
||
if pix.n - pix.alpha >= 4: # CMYK → RGB
|
||
pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
|
||
fname = f"pdfimg_{tag}_p{page_num}_{img_index}.png"
|
||
abs_path = os.path.join(out_dir, fname)
|
||
pix.save(abs_path)
|
||
width, height = pix.width, pix.height
|
||
pix = None # 及时释放位图内存
|
||
except Exception:
|
||
continue # 单张失败不影响整体
|
||
results.append({
|
||
"page": page_num,
|
||
"index": img_index,
|
||
"abs_path": abs_path,
|
||
"mime": "image/png",
|
||
"size_kb": round(os.path.getsize(abs_path) / 1024, 2),
|
||
"width": width,
|
||
"height": height,
|
||
})
|
||
return results
|
||
finally:
|
||
doc.close()
|