336 lines
12 KiB
Python
336 lines
12 KiB
Python
|
import base64
|
|||
|
|
import re
|
|||
|
|
from collections import OrderedDict
|
|||
|
|
from collections.abc import Generator
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import fitz
|
|||
|
|
from dify_plugin import Tool
|
|||
|
|
from dify_plugin.entities.tool import ToolInvokeMessage
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PdfToMarkdownTool(Tool):
|
|||
|
|
"""Convert PDF to a single Markdown file. No LLM needed.
|
|||
|
|
|
|||
|
|
- Auto-detect TOC and organize content by chapters.
|
|||
|
|
- Extract text and tables as Markdown.
|
|||
|
|
- Embed raster images as base64.
|
|||
|
|
- Render vector drawings as base64 PNG.
|
|||
|
|
- Output one .md file via create_blob_message.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
_TOC_PATTERNS = [
|
|||
|
|
r"目录", r"目 录", r"目\u3000录",
|
|||
|
|
r"Table of Contents", r"Contents", r"目次",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# ── entry point ──────────────────────────────────────────
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
include_images = self._to_bool(tool_parameters.get("include_images"), True)
|
|||
|
|
image_dpi = self._to_int(tool_parameters.get("image_dpi"), 150)
|
|||
|
|
image_dpi = max(72, min(image_dpi, 300))
|
|||
|
|
max_image_bytes = 2 * 1024 * 1024 # skip images > 2 MB raw
|
|||
|
|
|
|||
|
|
doc = fitz.open(stream=file.blob, filetype="pdf")
|
|||
|
|
try:
|
|||
|
|
num_pages = len(doc)
|
|||
|
|
|
|||
|
|
# 1) Build chapter map (metadata TOC → printed TOC → none)
|
|||
|
|
chapters, content_offset = self._build_chapter_map(doc, num_pages)
|
|||
|
|
|
|||
|
|
# 2) Convert every page
|
|||
|
|
page_mds: list[str] = []
|
|||
|
|
for idx in range(num_pages):
|
|||
|
|
md = self._page_to_markdown(
|
|||
|
|
doc, doc[idx], idx,
|
|||
|
|
include_images, image_dpi, max_image_bytes,
|
|||
|
|
)
|
|||
|
|
page_mds.append(md)
|
|||
|
|
|
|||
|
|
# 3) Assemble
|
|||
|
|
if chapters:
|
|||
|
|
final_md = self._assemble_by_chapters(
|
|||
|
|
chapters, page_mds, content_offset, num_pages,
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
final_md = "\n\n---\n\n".join(m for m in page_mds if m.strip())
|
|||
|
|
|
|||
|
|
# 4) Output: text (for variable aggregation) + blob (.md file)
|
|||
|
|
yield self.create_text_message(final_md)
|
|||
|
|
md_bytes = final_md.encode("utf-8")
|
|||
|
|
yield self.create_blob_message(
|
|||
|
|
blob=md_bytes,
|
|||
|
|
meta={"mime_type": "text/markdown"},
|
|||
|
|
)
|
|||
|
|
finally:
|
|||
|
|
doc.close()
|
|||
|
|
|
|||
|
|
# ── chapter detection ────────────────────────────────────
|
|||
|
|
|
|||
|
|
def _build_chapter_map(
|
|||
|
|
self, doc: fitz.Document, num_pages: int,
|
|||
|
|
) -> tuple[dict, int]:
|
|||
|
|
"""Return (chapters_dict, content_offset).
|
|||
|
|
|
|||
|
|
Try embedded PDF TOC metadata first (reliable page mapping).
|
|||
|
|
Fall back to scanning printed TOC pages.
|
|||
|
|
"""
|
|||
|
|
toc = doc.get_toc()
|
|||
|
|
if toc:
|
|||
|
|
chapters = self._chapters_from_metadata(toc, num_pages)
|
|||
|
|
if chapters:
|
|||
|
|
return chapters, 0
|
|||
|
|
|
|||
|
|
toc_start, toc_end = self._find_toc_pages(doc, num_pages)
|
|||
|
|
if toc_start is not None and toc_end is not None:
|
|||
|
|
toc_text = "\n".join(
|
|||
|
|
doc[i].get_text() or "" for i in range(toc_start, toc_end + 1)
|
|||
|
|
)
|
|||
|
|
chapters = self._parse_toc_lines(toc_text)
|
|||
|
|
if chapters:
|
|||
|
|
offset = self._guess_offset(chapters, toc_end)
|
|||
|
|
return chapters, offset
|
|||
|
|
|
|||
|
|
return {}, 0
|
|||
|
|
|
|||
|
|
def _chapters_from_metadata(
|
|||
|
|
self, toc: list, num_pages: int,
|
|||
|
|
) -> dict[str, dict[str, int]]:
|
|||
|
|
top = [(t, max(0, p - 1)) for lvl, t, p in toc if lvl <= 2 and p >= 1]
|
|||
|
|
if not top:
|
|||
|
|
return {}
|
|||
|
|
chapters: dict[str, dict[str, int]] = OrderedDict()
|
|||
|
|
for i, (title, start) in enumerate(top):
|
|||
|
|
end = top[i + 1][1] - 1 if i + 1 < len(top) else num_pages - 1
|
|||
|
|
chapters[title] = {"start": start, "end": max(start, end)}
|
|||
|
|
return chapters
|
|||
|
|
|
|||
|
|
def _find_toc_pages(self, doc, num_pages):
|
|||
|
|
toc_start = toc_end = None
|
|||
|
|
for pn in range(min(num_pages, 30)):
|
|||
|
|
text = doc[pn].get_text() or ""
|
|||
|
|
if any(re.search(p, text, re.IGNORECASE) for p in self._TOC_PATTERNS):
|
|||
|
|
if toc_start is None:
|
|||
|
|
toc_start = pn
|
|||
|
|
toc_end = pn
|
|||
|
|
elif toc_start is not None:
|
|||
|
|
break
|
|||
|
|
return toc_start, toc_end
|
|||
|
|
|
|||
|
|
def _parse_toc_lines(self, text: str) -> dict[str, dict[str, int]]:
|
|||
|
|
m = re.search(
|
|||
|
|
r"^(List\s+of\s+Figures|List\s+of\s+Tables|图目录|表目录)",
|
|||
|
|
text, re.IGNORECASE | re.MULTILINE,
|
|||
|
|
)
|
|||
|
|
if m:
|
|||
|
|
text = text[: m.start()]
|
|||
|
|
|
|||
|
|
pat = re.compile(
|
|||
|
|
r"^\s*(?P<title>.+?)\s*(?:\.{2,}|\s)\s*(?P<page>\d{1,5})\s*$"
|
|||
|
|
)
|
|||
|
|
entries: list[tuple[str, int]] = []
|
|||
|
|
for raw in text.splitlines():
|
|||
|
|
line = raw.strip()
|
|||
|
|
if not line or len(line) < 3 or re.fullmatch(r"\d+", line):
|
|||
|
|
continue
|
|||
|
|
m2 = pat.match(line)
|
|||
|
|
if not m2:
|
|||
|
|
continue
|
|||
|
|
title = re.sub(r"\s+", " ", m2.group("title")).strip("-_:： ")
|
|||
|
|
page = self._to_int(m2.group("page"), None)
|
|||
|
|
if not title or page is None or len(title) <= 1:
|
|||
|
|
continue
|
|||
|
|
if title.lower() in {"page", "pages", "目录", "contents"}:
|
|||
|
|
continue
|
|||
|
|
entries.append((title, page))
|
|||
|
|
|
|||
|
|
if not entries:
|
|||
|
|
return {}
|
|||
|
|
|
|||
|
|
dedup: OrderedDict[str, int] = OrderedDict()
|
|||
|
|
for t, p in entries:
|
|||
|
|
dedup.setdefault(t, p)
|
|||
|
|
|
|||
|
|
titles = list(dedup.keys())
|
|||
|
|
pages = [dedup[t] for t in titles]
|
|||
|
|
catalog: dict[str, dict[str, int]] = OrderedDict()
|
|||
|
|
for i, t in enumerate(titles):
|
|||
|
|
s = pages[i]
|
|||
|
|
e = max(s, pages[i + 1] - 1) if i + 1 < len(pages) else s
|
|||
|
|
catalog[t] = {"start": s, "end": e}
|
|||
|
|
return catalog
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _guess_offset(chapters: dict, toc_end: int) -> int:
|
|||
|
|
first_page = None
|
|||
|
|
for info in chapters.values():
|
|||
|
|
s = info["start"]
|
|||
|
|
if first_page is None or s < first_page:
|
|||
|
|
first_page = s
|
|||
|
|
if first_page is None:
|
|||
|
|
return 0
|
|||
|
|
return (toc_end + 1) - first_page
|
|||
|
|
|
|||
|
|
# ── per-page conversion ──────────────────────────────────
|
|||
|
|
|
|||
|
|
def _page_to_markdown(
|
|||
|
|
self,
|
|||
|
|
doc: fitz.Document,
|
|||
|
|
page: fitz.Page,
|
|||
|
|
page_idx: int,
|
|||
|
|
include_images: bool,
|
|||
|
|
image_dpi: int,
|
|||
|
|
max_image_bytes: int,
|
|||
|
|
) -> str:
|
|||
|
|
parts: list[str] = []
|
|||
|
|
|
|||
|
|
# ── text ──
|
|||
|
|
text = (page.get_text("text", sort=True) or "").strip()
|
|||
|
|
if text:
|
|||
|
|
parts.append(text)
|
|||
|
|
|
|||
|
|
# ── tables → Markdown ──
|
|||
|
|
try:
|
|||
|
|
for tab in (page.find_tables().tables or [])[:5]:
|
|||
|
|
cells = tab.extract() or []
|
|||
|
|
if len(cells) >= 2:
|
|||
|
|
md = self._cells_to_md_table(cells)
|
|||
|
|
if md:
|
|||
|
|
parts.append(md)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
if not include_images:
|
|||
|
|
return "\n\n".join(parts)
|
|||
|
|
|
|||
|
|
# ── embedded raster images ──
|
|||
|
|
try:
|
|||
|
|
for img_idx, img_info in enumerate(page.get_images(full=True)):
|
|||
|
|
xref = img_info[0]
|
|||
|
|
try:
|
|||
|
|
data = doc.extract_image(xref)
|
|||
|
|
if not data or not data.get("image"):
|
|||
|
|
continue
|
|||
|
|
raw = data["image"]
|
|||
|
|
if len(raw) > max_image_bytes:
|
|||
|
|
continue
|
|||
|
|
# skip tiny icons (< 20x20)
|
|||
|
|
w = data.get("width", 0)
|
|||
|
|
h = data.get("height", 0)
|
|||
|
|
if w < 20 and h < 20:
|
|||
|
|
continue
|
|||
|
|
ext = data.get("ext", "png")
|
|||
|
|
mime = "image/jpeg" if ext in ("jpg", "jpeg") else f"image/{ext}"
|
|||
|
|
b64 = base64.b64encode(raw).decode("ascii")
|
|||
|
|
parts.append(
|
|||
|
|
f""
|
|||
|
|
)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
# ── vector drawings → render as PNG ──
|
|||
|
|
try:
|
|||
|
|
drawings = page.get_drawings()
|
|||
|
|
if len(drawings) >= 3:
|
|||
|
|
valid_rects: list[fitz.Rect] = []
|
|||
|
|
for d in drawings:
|
|||
|
|
r = d.get("rect")
|
|||
|
|
if r:
|
|||
|
|
try:
|
|||
|
|
rect = fitz.Rect(r)
|
|||
|
|
if rect.is_valid and not rect.is_empty:
|
|||
|
|
valid_rects.append(rect)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
if valid_rects:
|
|||
|
|
bbox = valid_rects[0]
|
|||
|
|
for r in valid_rects[1:]:
|
|||
|
|
bbox |= r
|
|||
|
|
bbox &= page.rect
|
|||
|
|
if bbox.width > 30 and bbox.height > 30:
|
|||
|
|
scale = image_dpi / 72
|
|||
|
|
mat = fitz.Matrix(scale, scale)
|
|||
|
|
pix = page.get_pixmap(matrix=mat, clip=bbox)
|
|||
|
|
png = pix.tobytes("png")
|
|||
|
|
if len(png) <= max_image_bytes:
|
|||
|
|
b64 = base64.b64encode(png).decode("ascii")
|
|||
|
|
parts.append(
|
|||
|
|
f""
|
|||
|
|
)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
return "\n\n".join(parts)
|
|||
|
|
|
|||
|
|
# ── assembly ─────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def _assemble_by_chapters(
|
|||
|
|
self,
|
|||
|
|
chapters: dict[str, dict[str, int]],
|
|||
|
|
page_mds: list[str],
|
|||
|
|
offset: int,
|
|||
|
|
num_pages: int,
|
|||
|
|
) -> str:
|
|||
|
|
parts: list[str] = []
|
|||
|
|
for name, info in chapters.items():
|
|||
|
|
s = info["start"] + offset
|
|||
|
|
e = info["end"] + offset
|
|||
|
|
s = max(0, min(s, num_pages - 1))
|
|||
|
|
e = max(s, min(e, num_pages - 1))
|
|||
|
|
ch: list[str] = [f"# {name}\n"]
|
|||
|
|
for idx in range(s, e + 1):
|
|||
|
|
if idx < len(page_mds) and page_mds[idx].strip():
|
|||
|
|
ch.append(page_mds[idx])
|
|||
|
|
parts.append("\n\n".join(ch))
|
|||
|
|
return "\n\n---\n\n".join(parts)
|
|||
|
|
|
|||
|
|
# ── helpers ──────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _cells_to_md_table(cells: list) -> str:
|
|||
|
|
if not cells:
|
|||
|
|
return ""
|
|||
|
|
header = cells[0]
|
|||
|
|
ncols = len(header)
|
|||
|
|
if ncols == 0:
|
|||
|
|
return ""
|
|||
|
|
clean = lambda c: str(c or "").replace("|", "\\|").replace("\n", " ").strip()
|
|||
|
|
lines = [
|
|||
|
|
"| " + " | ".join(clean(c) for c in header) + " |",
|
|||
|
|
"| " + " | ".join("---" for _ in range(ncols)) + " |",
|
|||
|
|
]
|
|||
|
|
for row in cells[1:]:
|
|||
|
|
padded = list(row) + [""] * max(0, ncols - len(row))
|
|||
|
|
lines.append("| " + " | ".join(clean(c) for c in padded[:ncols]) + " |")
|
|||
|
|
return "\n".join(lines)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _to_int(value: Any, default: int | None) -> int | None:
|
|||
|
|
try:
|
|||
|
|
if value is None or value == "":
|
|||
|
|
return default
|
|||
|
|
return int(value)
|
|||
|
|
except Exception:
|
|||
|
|
return default
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _to_bool(value: Any, default: bool) -> bool:
|
|||
|
|
if value is None:
|
|||
|
|
return default
|
|||
|
|
if isinstance(value, bool):
|
|||
|
|
return value
|
|||
|
|
s = str(value).strip().lower()
|
|||
|
|
if s in {"1", "true", "yes", "on"}:
|
|||
|
|
return True
|
|||
|
|
if s in {"0", "false", "no", "off"}:
|
|||
|
|
return False
|
|||
|
|
return default
|