107 lines
3.2 KiB
Python
107 lines
3.2 KiB
Python
import json
|
||
import re
|
||
from collections.abc import Generator
|
||
from io import BytesIO
|
||
from typing import Any
|
||
|
||
import fitz # PyMuPDF 核心库
|
||
from dify_plugin import Tool
|
||
from dify_plugin.entities.tool import ToolInvokeMessage
|
||
|
||
|
||
class PdfTool(Tool):
|
||
def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessage]:
|
||
file = tool_parameters.get("file")
|
||
if not file:
|
||
yield self.create_text_message("Error: file is required")
|
||
return
|
||
|
||
# 从字节流加载 PDF（替换 PyPDF2）
|
||
pdf_bytes = file.blob
|
||
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
|
||
num_pages = len(doc)
|
||
|
||
toc_start = None
|
||
toc_end = None
|
||
|
||
# 目录匹配正则（与原代码一致）
|
||
toc_patterns = [
|
||
r'目录',
|
||
r'目 录',
|
||
r'目\u3000录',
|
||
r'Table of Contents',
|
||
r'Contents',
|
||
r'目次'
|
||
]
|
||
|
||
# 遍历页面识别目录页（逻辑不变，仅替换文本提取方式）
|
||
for page_num in range(num_pages):
|
||
page = doc[page_num]
|
||
text = page.get_text() or "" # PyMuPDF 提取文本
|
||
|
||
if any(re.search(pattern, text, re.IGNORECASE) for pattern in toc_patterns):
|
||
if toc_start is None:
|
||
toc_start = page_num
|
||
toc_end = page_num
|
||
elif toc_start is not None and toc_end is not None:
|
||
break
|
||
|
||
# 提取目录页文本
|
||
toc_pages = []
|
||
if toc_start is not None and toc_end is not None:
|
||
for page_num in range(toc_start, toc_end + 1):
|
||
page = doc[page_num]
|
||
toc_pages.append(page.get_text() or "")
|
||
|
||
# 关闭文档
|
||
doc.close()
|
||
|
||
result = {
|
||
"start": toc_start,
|
||
"end": toc_end,
|
||
"pages": toc_pages,
|
||
"pages_text": "\n".join(toc_pages) if toc_pages else "",
|
||
}
|
||
yield self.create_text_message(json.dumps(result, ensure_ascii=False))
|
||
yield self.create_json_message(result)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 测试代码（改用 PyMuPDF）
|
||
pdf_path = r"F:\Project\urbanLifeline\docs\AI训练资料\菱重S12R发动机说明书.pdf"
|
||
doc = fitz.open(pdf_path) # 本地文件直接打开
|
||
num_pages = len(doc)
|
||
|
||
toc_start = None
|
||
toc_end = None
|
||
|
||
toc_patterns = [
|
||
r'目录',
|
||
r'目 录',
|
||
r'目\u3000录',
|
||
r'Table of Contents',
|
||
r'Contents',
|
||
r'目次'
|
||
]
|
||
|
||
# 遍历页面找目录
|
||
for page_num in range(num_pages):
|
||
page = doc[page_num]
|
||
text = page.get_text() or ""
|
||
if any(re.search(pattern, text, re.IGNORECASE) for pattern in toc_patterns):
|
||
if toc_start is None:
|
||
toc_start = page_num
|
||
toc_end = page_num
|
||
elif toc_start is not None and toc_end is not None:
|
||
break
|
||
|
||
# 提取目录页文本
|
||
toc_pages = []
|
||
toc_start = toc_start if toc_start is not None else 18
|
||
toc_end = toc_end if toc_end is not None else toc_start + 9
|
||
for page_num in range(toc_start, toc_end):
|
||
page = doc[page_num]
|
||
toc_pages.append(page.get_text() or "")
|
||
|
||
print(toc_start, toc_end, toc_pages)
|
||
doc.close() # 关闭文档 |