插件

difyPlugin/pdf/tools/pdf.py

@@ -0,0 +1,61 @@
+import re
+from collections.abc import Generator
+from io import BytesIO
+from typing import Any
+
+import PyPDF2
+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
+
+        # file.blob returns bytes
+        pdf_bytes = file.blob
+        reader = PyPDF2.PdfReader(BytesIO(pdf_bytes))
+        num_pages = len(reader.pages)
+
+        toc_start = None
+        toc_end = None
+
+        toc_patterns = [
+            r'目录',
+            r'Table of Contents',
+            r'Contents',
+            r'目次'
+        ]
+
+        for page_num in range(num_pages):
+            page = reader.pages[page_num]
+            text = page.extract_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
+
+        if toc_start is None:
+            yield self.create_json_message({
+                "start": None,
+                "end": None,
+                "pages": []
+            })
+            return
+
+        toc_pages = []
+        for page_num in range(toc_start, toc_end + 1):
+            page = reader.pages[page_num]
+            toc_pages.append(page.extract_text() or "")
+
+        yield self.create_json_message({
+            "start": toc_start,
+            "end": toc_end,
+            "pages": toc_pages
+        })