|
|
|
|
@@ -5,7 +5,7 @@ import logging
|
|
|
|
|
import os
|
|
|
|
|
import tempfile
|
|
|
|
|
from collections.abc import Mapping, Sequence
|
|
|
|
|
from typing import Any
|
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
|
|
|
|
|
|
import charset_normalizer
|
|
|
|
|
import docx
|
|
|
|
|
@@ -20,7 +20,6 @@ from docx.oxml.text.paragraph import CT_P
|
|
|
|
|
from docx.table import Table
|
|
|
|
|
from docx.text.paragraph import Paragraph
|
|
|
|
|
|
|
|
|
|
from configs import dify_config
|
|
|
|
|
from core.helper import ssrf_proxy
|
|
|
|
|
from core.variables import ArrayFileSegment
|
|
|
|
|
from core.variables.segments import ArrayStringSegment, FileSegment
|
|
|
|
|
@@ -29,11 +28,15 @@ from core.workflow.file import File, FileTransferMethod, file_manager
|
|
|
|
|
from core.workflow.node_events import NodeRunResult
|
|
|
|
|
from core.workflow.nodes.base.node import Node
|
|
|
|
|
|
|
|
|
|
from .entities import DocumentExtractorNodeData
|
|
|
|
|
from .entities import DocumentExtractorNodeData, UnstructuredApiConfig
|
|
|
|
|
from .exc import DocumentExtractorError, FileDownloadError, TextExtractionError, UnsupportedFileTypeError
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
|
from core.workflow.entities import GraphInitParams
|
|
|
|
|
from core.workflow.runtime import GraphRuntimeState
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class DocumentExtractorNode(Node[DocumentExtractorNodeData]):
|
|
|
|
|
"""
|
|
|
|
|
@@ -47,6 +50,23 @@ class DocumentExtractorNode(Node[DocumentExtractorNodeData]):
|
|
|
|
|
def version(cls) -> str:
|
|
|
|
|
return "1"
|
|
|
|
|
|
|
|
|
|
def __init__(
|
|
|
|
|
self,
|
|
|
|
|
id: str,
|
|
|
|
|
config: Mapping[str, Any],
|
|
|
|
|
graph_init_params: "GraphInitParams",
|
|
|
|
|
graph_runtime_state: "GraphRuntimeState",
|
|
|
|
|
*,
|
|
|
|
|
unstructured_api_config: UnstructuredApiConfig | None = None,
|
|
|
|
|
) -> None:
|
|
|
|
|
super().__init__(
|
|
|
|
|
id=id,
|
|
|
|
|
config=config,
|
|
|
|
|
graph_init_params=graph_init_params,
|
|
|
|
|
graph_runtime_state=graph_runtime_state,
|
|
|
|
|
)
|
|
|
|
|
self._unstructured_api_config = unstructured_api_config or UnstructuredApiConfig()
|
|
|
|
|
|
|
|
|
|
def _run(self):
|
|
|
|
|
variable_selector = self.node_data.variable_selector
|
|
|
|
|
variable = self.graph_runtime_state.variable_pool.get(variable_selector)
|
|
|
|
|
@@ -64,7 +84,10 @@ class DocumentExtractorNode(Node[DocumentExtractorNodeData]):
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
if isinstance(value, list):
|
|
|
|
|
extracted_text_list = list(map(_extract_text_from_file, value))
|
|
|
|
|
extracted_text_list = [
|
|
|
|
|
_extract_text_from_file(file, unstructured_api_config=self._unstructured_api_config)
|
|
|
|
|
for file in value
|
|
|
|
|
]
|
|
|
|
|
return NodeRunResult(
|
|
|
|
|
status=WorkflowNodeExecutionStatus.SUCCEEDED,
|
|
|
|
|
inputs=inputs,
|
|
|
|
|
@@ -72,7 +95,7 @@ class DocumentExtractorNode(Node[DocumentExtractorNodeData]):
|
|
|
|
|
outputs={"text": ArrayStringSegment(value=extracted_text_list)},
|
|
|
|
|
)
|
|
|
|
|
elif isinstance(value, File):
|
|
|
|
|
extracted_text = _extract_text_from_file(value)
|
|
|
|
|
extracted_text = _extract_text_from_file(value, unstructured_api_config=self._unstructured_api_config)
|
|
|
|
|
return NodeRunResult(
|
|
|
|
|
status=WorkflowNodeExecutionStatus.SUCCEEDED,
|
|
|
|
|
inputs=inputs,
|
|
|
|
|
@@ -103,7 +126,12 @@ class DocumentExtractorNode(Node[DocumentExtractorNodeData]):
|
|
|
|
|
return {node_id + ".files": typed_node_data.variable_selector}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _extract_text_by_mime_type(*, file_content: bytes, mime_type: str) -> str:
|
|
|
|
|
def _extract_text_by_mime_type(
|
|
|
|
|
*,
|
|
|
|
|
file_content: bytes,
|
|
|
|
|
mime_type: str,
|
|
|
|
|
unstructured_api_config: UnstructuredApiConfig,
|
|
|
|
|
) -> str:
|
|
|
|
|
"""Extract text from a file based on its MIME type."""
|
|
|
|
|
match mime_type:
|
|
|
|
|
case "text/plain" | "text/html" | "text/htm" | "text/markdown" | "text/xml":
|
|
|
|
|
@@ -111,7 +139,7 @@ def _extract_text_by_mime_type(*, file_content: bytes, mime_type: str) -> str:
|
|
|
|
|
case "application/pdf":
|
|
|
|
|
return _extract_text_from_pdf(file_content)
|
|
|
|
|
case "application/msword":
|
|
|
|
|
return _extract_text_from_doc(file_content)
|
|
|
|
|
return _extract_text_from_doc(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
|
|
|
|
|
return _extract_text_from_docx(file_content)
|
|
|
|
|
case "text/csv":
|
|
|
|
|
@@ -119,11 +147,11 @@ def _extract_text_by_mime_type(*, file_content: bytes, mime_type: str) -> str:
|
|
|
|
|
case "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" | "application/vnd.ms-excel":
|
|
|
|
|
return _extract_text_from_excel(file_content)
|
|
|
|
|
case "application/vnd.ms-powerpoint":
|
|
|
|
|
return _extract_text_from_ppt(file_content)
|
|
|
|
|
return _extract_text_from_ppt(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case "application/vnd.openxmlformats-officedocument.presentationml.presentation":
|
|
|
|
|
return _extract_text_from_pptx(file_content)
|
|
|
|
|
return _extract_text_from_pptx(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case "application/epub+zip":
|
|
|
|
|
return _extract_text_from_epub(file_content)
|
|
|
|
|
return _extract_text_from_epub(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case "message/rfc822":
|
|
|
|
|
return _extract_text_from_eml(file_content)
|
|
|
|
|
case "application/vnd.ms-outlook":
|
|
|
|
|
@@ -140,7 +168,12 @@ def _extract_text_by_mime_type(*, file_content: bytes, mime_type: str) -> str:
|
|
|
|
|
raise UnsupportedFileTypeError(f"Unsupported MIME type: {mime_type}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _extract_text_by_file_extension(*, file_content: bytes, file_extension: str) -> str:
|
|
|
|
|
def _extract_text_by_file_extension(
|
|
|
|
|
*,
|
|
|
|
|
file_content: bytes,
|
|
|
|
|
file_extension: str,
|
|
|
|
|
unstructured_api_config: UnstructuredApiConfig,
|
|
|
|
|
) -> str:
|
|
|
|
|
"""Extract text from a file based on its file extension."""
|
|
|
|
|
match file_extension:
|
|
|
|
|
case (
|
|
|
|
|
@@ -203,7 +236,7 @@ def _extract_text_by_file_extension(*, file_content: bytes, file_extension: str)
|
|
|
|
|
case ".pdf":
|
|
|
|
|
return _extract_text_from_pdf(file_content)
|
|
|
|
|
case ".doc":
|
|
|
|
|
return _extract_text_from_doc(file_content)
|
|
|
|
|
return _extract_text_from_doc(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case ".docx":
|
|
|
|
|
return _extract_text_from_docx(file_content)
|
|
|
|
|
case ".csv":
|
|
|
|
|
@@ -211,11 +244,11 @@ def _extract_text_by_file_extension(*, file_content: bytes, file_extension: str)
|
|
|
|
|
case ".xls" | ".xlsx":
|
|
|
|
|
return _extract_text_from_excel(file_content)
|
|
|
|
|
case ".ppt":
|
|
|
|
|
return _extract_text_from_ppt(file_content)
|
|
|
|
|
return _extract_text_from_ppt(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case ".pptx":
|
|
|
|
|
return _extract_text_from_pptx(file_content)
|
|
|
|
|
return _extract_text_from_pptx(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case ".epub":
|
|
|
|
|
return _extract_text_from_epub(file_content)
|
|
|
|
|
return _extract_text_from_epub(file_content, unstructured_api_config=unstructured_api_config)
|
|
|
|
|
case ".eml":
|
|
|
|
|
return _extract_text_from_eml(file_content)
|
|
|
|
|
case ".msg":
|
|
|
|
|
@@ -312,14 +345,15 @@ def _extract_text_from_pdf(file_content: bytes) -> str:
|
|
|
|
|
raise TextExtractionError(f"Failed to extract text from PDF: {str(e)}") from e
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _extract_text_from_doc(file_content: bytes) -> str:
|
|
|
|
|
def _extract_text_from_doc(file_content: bytes, *, unstructured_api_config: UnstructuredApiConfig) -> str:
|
|
|
|
|
"""
|
|
|
|
|
Extract text from a DOC file.
|
|
|
|
|
"""
|
|
|
|
|
from unstructured.partition.api import partition_via_api
|
|
|
|
|
|
|
|
|
|
if not dify_config.UNSTRUCTURED_API_URL:
|
|
|
|
|
raise TextExtractionError("UNSTRUCTURED_API_URL must be set")
|
|
|
|
|
if not unstructured_api_config.api_url:
|
|
|
|
|
raise TextExtractionError("Unstructured API URL is not configured for DOC file processing.")
|
|
|
|
|
api_key = unstructured_api_config.api_key or ""
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".doc", delete=False) as temp_file:
|
|
|
|
|
@@ -329,8 +363,8 @@ def _extract_text_from_doc(file_content: bytes) -> str:
|
|
|
|
|
elements = partition_via_api(
|
|
|
|
|
file=file,
|
|
|
|
|
metadata_filename=temp_file.name,
|
|
|
|
|
api_url=dify_config.UNSTRUCTURED_API_URL,
|
|
|
|
|
api_key=dify_config.UNSTRUCTURED_API_KEY, # type: ignore
|
|
|
|
|
api_url=unstructured_api_config.api_url,
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
)
|
|
|
|
|
os.unlink(temp_file.name)
|
|
|
|
|
return "\n".join([getattr(element, "text", "") for element in elements])
|
|
|
|
|
@@ -420,12 +454,20 @@ def _download_file_content(file: File) -> bytes:
|
|
|
|
|
raise FileDownloadError(f"Error downloading file: {str(e)}") from e
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _extract_text_from_file(file: File):
|
|
|
|
|
def _extract_text_from_file(file: File, *, unstructured_api_config: UnstructuredApiConfig) -> str:
|
|
|
|
|
file_content = _download_file_content(file)
|
|
|
|
|
if file.extension:
|
|
|
|
|
extracted_text = _extract_text_by_file_extension(file_content=file_content, file_extension=file.extension)
|
|
|
|
|
extracted_text = _extract_text_by_file_extension(
|
|
|
|
|
file_content=file_content,
|
|
|
|
|
file_extension=file.extension,
|
|
|
|
|
unstructured_api_config=unstructured_api_config,
|
|
|
|
|
)
|
|
|
|
|
elif file.mime_type:
|
|
|
|
|
extracted_text = _extract_text_by_mime_type(file_content=file_content, mime_type=file.mime_type)
|
|
|
|
|
extracted_text = _extract_text_by_mime_type(
|
|
|
|
|
file_content=file_content,
|
|
|
|
|
mime_type=file.mime_type,
|
|
|
|
|
unstructured_api_config=unstructured_api_config,
|
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
raise UnsupportedFileTypeError("Unable to determine file type: MIME type or file extension is missing")
|
|
|
|
|
return extracted_text
|
|
|
|
|
@@ -517,12 +559,14 @@ def _extract_text_from_excel(file_content: bytes) -> str:
|
|
|
|
|
raise TextExtractionError(f"Failed to extract text from Excel file: {str(e)}") from e
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _extract_text_from_ppt(file_content: bytes) -> str:
|
|
|
|
|
def _extract_text_from_ppt(file_content: bytes, *, unstructured_api_config: UnstructuredApiConfig) -> str:
|
|
|
|
|
from unstructured.partition.api import partition_via_api
|
|
|
|
|
from unstructured.partition.ppt import partition_ppt
|
|
|
|
|
|
|
|
|
|
api_key = unstructured_api_config.api_key or ""
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
if dify_config.UNSTRUCTURED_API_URL:
|
|
|
|
|
if unstructured_api_config.api_url:
|
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".ppt", delete=False) as temp_file:
|
|
|
|
|
temp_file.write(file_content)
|
|
|
|
|
temp_file.flush()
|
|
|
|
|
@@ -530,8 +574,8 @@ def _extract_text_from_ppt(file_content: bytes) -> str:
|
|
|
|
|
elements = partition_via_api(
|
|
|
|
|
file=file,
|
|
|
|
|
metadata_filename=temp_file.name,
|
|
|
|
|
api_url=dify_config.UNSTRUCTURED_API_URL,
|
|
|
|
|
api_key=dify_config.UNSTRUCTURED_API_KEY, # type: ignore
|
|
|
|
|
api_url=unstructured_api_config.api_url,
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
)
|
|
|
|
|
os.unlink(temp_file.name)
|
|
|
|
|
else:
|
|
|
|
|
@@ -543,12 +587,14 @@ def _extract_text_from_ppt(file_content: bytes) -> str:
|
|
|
|
|
raise TextExtractionError(f"Failed to extract text from PPTX: {str(e)}") from e
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _extract_text_from_pptx(file_content: bytes) -> str:
|
|
|
|
|
def _extract_text_from_pptx(file_content: bytes, *, unstructured_api_config: UnstructuredApiConfig) -> str:
|
|
|
|
|
from unstructured.partition.api import partition_via_api
|
|
|
|
|
from unstructured.partition.pptx import partition_pptx
|
|
|
|
|
|
|
|
|
|
api_key = unstructured_api_config.api_key or ""
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
if dify_config.UNSTRUCTURED_API_URL:
|
|
|
|
|
if unstructured_api_config.api_url:
|
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".pptx", delete=False) as temp_file:
|
|
|
|
|
temp_file.write(file_content)
|
|
|
|
|
temp_file.flush()
|
|
|
|
|
@@ -556,8 +602,8 @@ def _extract_text_from_pptx(file_content: bytes) -> str:
|
|
|
|
|
elements = partition_via_api(
|
|
|
|
|
file=file,
|
|
|
|
|
metadata_filename=temp_file.name,
|
|
|
|
|
api_url=dify_config.UNSTRUCTURED_API_URL,
|
|
|
|
|
api_key=dify_config.UNSTRUCTURED_API_KEY, # type: ignore
|
|
|
|
|
api_url=unstructured_api_config.api_url,
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
)
|
|
|
|
|
os.unlink(temp_file.name)
|
|
|
|
|
else:
|
|
|
|
|
@@ -568,12 +614,14 @@ def _extract_text_from_pptx(file_content: bytes) -> str:
|
|
|
|
|
raise TextExtractionError(f"Failed to extract text from PPTX: {str(e)}") from e
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _extract_text_from_epub(file_content: bytes) -> str:
|
|
|
|
|
def _extract_text_from_epub(file_content: bytes, *, unstructured_api_config: UnstructuredApiConfig) -> str:
|
|
|
|
|
from unstructured.partition.api import partition_via_api
|
|
|
|
|
from unstructured.partition.epub import partition_epub
|
|
|
|
|
|
|
|
|
|
api_key = unstructured_api_config.api_key or ""
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
if dify_config.UNSTRUCTURED_API_URL:
|
|
|
|
|
if unstructured_api_config.api_url:
|
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".epub", delete=False) as temp_file:
|
|
|
|
|
temp_file.write(file_content)
|
|
|
|
|
temp_file.flush()
|
|
|
|
|
@@ -581,8 +629,8 @@ def _extract_text_from_epub(file_content: bytes) -> str:
|
|
|
|
|
elements = partition_via_api(
|
|
|
|
|
file=file,
|
|
|
|
|
metadata_filename=temp_file.name,
|
|
|
|
|
api_url=dify_config.UNSTRUCTURED_API_URL,
|
|
|
|
|
api_key=dify_config.UNSTRUCTURED_API_KEY, # type: ignore
|
|
|
|
|
api_url=unstructured_api_config.api_url,
|
|
|
|
|
api_key=api_key,
|
|
|
|
|
)
|
|
|
|
|
os.unlink(temp_file.name)
|
|
|
|
|
else:
|
|
|
|
|
|