210 lines
7.3 KiB
Python
210 lines
7.3 KiB
Python
|
import json
|
|||
|
|
import re
|
|||
|
|
from collections.abc import Generator
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
import fitz
|
|||
|
|
from dify_plugin import Tool
|
|||
|
|
from dify_plugin.entities.model.llm import LLMModelConfig
|
|||
|
|
from dify_plugin.entities.model.message import SystemPromptMessage, UserPromptMessage
|
|||
|
|
from dify_plugin.entities.tool import ToolInvokeMessage
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PdfSummaryTool(Tool):
|
|||
|
|
"""Fast PDF page summary tool.
|
|||
|
|
|
|||
|
|
Default behavior is optimized for throughput in large workflows:
|
|||
|
|
- Extract plain text and lightweight table data only.
|
|||
|
|
- Skip expensive image base64 and drawing path extraction.
|
|||
|
|
- Skip LLM by default unless `use_llm=true` is explicitly passed.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
start_page = self._to_int(tool_parameters.get("pdf_start_page"), 0)
|
|||
|
|
end_page = self._to_int(tool_parameters.get("pdf_end_page"), 0)
|
|||
|
|
model_config = tool_parameters.get("model")
|
|||
|
|
use_llm = self._to_bool(tool_parameters.get("use_llm"), False)
|
|||
|
|
|
|||
|
|
max_chars_per_page = self._to_int(tool_parameters.get("max_chars_per_page"), 6000)
|
|||
|
|
max_chars_per_page = max(800, min(max_chars_per_page, 20000))
|
|||
|
|
|
|||
|
|
llm_prompt = tool_parameters.get(
|
|||
|
|
"llm_prompt",
|
|||
|
|
"请基于输入的PDF页面文本做简洁准确摘要，输出中文要点。不要输出思考过程。",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
pdf_bytes = file.blob
|
|||
|
|
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
|
|||
|
|
try:
|
|||
|
|
num_pages = len(doc)
|
|||
|
|
start_page = max(0, min(start_page, num_pages - 1))
|
|||
|
|
end_page = max(start_page, min(end_page, num_pages - 1))
|
|||
|
|
|
|||
|
|
pages_data: list[dict[str, Any]] = []
|
|||
|
|
for page_idx in range(start_page, end_page + 1):
|
|||
|
|
page = doc[page_idx]
|
|||
|
|
page_data = self._extract_page_fast(page, page_idx, max_chars_per_page)
|
|||
|
|
pages_data.append(page_data)
|
|||
|
|
|
|||
|
|
result = {
|
|||
|
|
"total_pages_extracted": len(pages_data),
|
|||
|
|
"page_range": {"start": start_page, "end": end_page},
|
|||
|
|
"pages": pages_data,
|
|||
|
|
}
|
|||
|
|
yield self.create_json_message(result)
|
|||
|
|
|
|||
|
|
# Fast local summary first (deterministic, no model latency)
|
|||
|
|
local_text = self._build_local_summary(pages_data)
|
|||
|
|
|
|||
|
|
# Optional LLM refinement, explicitly enabled only
|
|||
|
|
if use_llm and model_config:
|
|||
|
|
refined = self._summarize_with_llm(local_text, llm_prompt, model_config)
|
|||
|
|
final_text = refined if refined else local_text
|
|||
|
|
else:
|
|||
|
|
final_text = local_text
|
|||
|
|
|
|||
|
|
if final_text:
|
|||
|
|
yield self.create_text_message(final_text)
|
|||
|
|
finally:
|
|||
|
|
doc.close()
|
|||
|
|
|
|||
|
|
def _extract_page_fast(self, page: fitz.Page, page_idx: int, max_chars_per_page: int) -> dict[str, Any]:
|
|||
|
|
text = (page.get_text("text") or "").strip()
|
|||
|
|
if len(text) > max_chars_per_page:
|
|||
|
|
text = text[:max_chars_per_page] + "\n...[truncated]"
|
|||
|
|
|
|||
|
|
tables: list[dict[str, Any]] = []
|
|||
|
|
try:
|
|||
|
|
tabs = page.find_tables()
|
|||
|
|
for tab_idx, tab in enumerate(tabs.tables[:3]):
|
|||
|
|
cells = tab.extract() or []
|
|||
|
|
tables.append(
|
|||
|
|
{
|
|||
|
|
"index": tab_idx,
|
|||
|
|
"rows": tab.row_count,
|
|||
|
|
"cols": tab.col_count,
|
|||
|
|
"cells": cells[:10],
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"page_number": page_idx,
|
|||
|
|
"text": text,
|
|||
|
|
"tables": tables,
|
|||
|
|
"images": [],
|
|||
|
|
"drawings_summary": [],
|
|||
|
|
"text_blocks": [],
|
|||
|
|
"width": float(page.rect.width),
|
|||
|
|
"height": float(page.rect.height),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def _build_local_summary(self, pages_data: list[dict[str, Any]]) -> str:
|
|||
|
|
"""Output actual page content as Markdown (text + tables).
|
|||
|
|
|
|||
|
|
No LLM needed downstream — the text is already usable Markdown.
|
|||
|
|
"""
|
|||
|
|
parts: list[str] = []
|
|||
|
|
for page in pages_data:
|
|||
|
|
text = (page.get("text") or "").strip()
|
|||
|
|
tables = page.get("tables") or []
|
|||
|
|
|
|||
|
|
page_parts: list[str] = []
|
|||
|
|
if text:
|
|||
|
|
page_parts.append(text)
|
|||
|
|
|
|||
|
|
for tab in tables:
|
|||
|
|
cells = tab.get("cells") or []
|
|||
|
|
if len(cells) >= 2:
|
|||
|
|
md = self._cells_to_md_table(cells)
|
|||
|
|
if md:
|
|||
|
|
page_parts.append(md)
|
|||
|
|
|
|||
|
|
if page_parts:
|
|||
|
|
parts.append("\n\n".join(page_parts))
|
|||
|
|
|
|||
|
|
return "\n\n--- 分页 ---\n\n".join(parts)
|
|||
|
|
|
|||
|
|
@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)
|
|||
|
|
|
|||
|
|
def _summarize_with_llm(self, local_text: str, llm_prompt: str, model_config: dict[str, Any]) -> str:
|
|||
|
|
response = self.session.model.llm.invoke(
|
|||
|
|
model_config=LLMModelConfig(**model_config),
|
|||
|
|
prompt_messages=[
|
|||
|
|
SystemPromptMessage(content=llm_prompt),
|
|||
|
|
UserPromptMessage(content=local_text),
|
|||
|
|
],
|
|||
|
|
stream=False,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
llm_text = ""
|
|||
|
|
if hasattr(response, "message") and response.message:
|
|||
|
|
content = response.message.content
|
|||
|
|
if isinstance(content, str):
|
|||
|
|
llm_text = content
|
|||
|
|
elif isinstance(content, list):
|
|||
|
|
llm_text = "".join(
|
|||
|
|
item.data if hasattr(item, "data") else str(item)
|
|||
|
|
for item in content
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return self._extract_visible_answer(llm_text)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _extract_visible_answer(text: str) -> str:
|
|||
|
|
if not text:
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
box_match = re.search(r"<\|begin_of_box\|>([\s\S]*?)<\|end_of_box\|>", text)
|
|||
|
|
if box_match:
|
|||
|
|
text = box_match.group(1)
|
|||
|
|
else:
|
|||
|
|
text = re.sub(r"<think>[\s\S]*?</think>", "", text, flags=re.IGNORECASE)
|
|||
|
|
|
|||
|
|
text = re.sub(r"<\|[^>]+\|>", "", text)
|
|||
|
|
return text.strip()
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _to_int(value: Any, default: int) -> int:
|
|||
|
|
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
|