diff --git a/apps/worker/app/services/document_ingestion/page_estimator.py b/apps/worker/app/services/document_ingestion/page_estimator.py index f2ed89cdb..cbd66eb0d 100644 --- a/apps/worker/app/services/document_ingestion/page_estimator.py +++ b/apps/worker/app/services/document_ingestion/page_estimator.py @@ -117,9 +117,11 @@ def _count_pptx(cls, file_path: str) -> int: @classmethod def _count_docx(cls, file_path: str) -> int: """Estimate pages for DOCX using word-based counting.""" - from docx import Document + from app.services.document_parser.formats.docx.word_package import ( + open_word_document, + ) - document = Document(file_path) + document = open_word_document(file_path) total_text = "" for paragraph in document.paragraphs: diff --git a/apps/worker/app/services/document_parser/formats/docx/block_stream.py b/apps/worker/app/services/document_parser/formats/docx/block_stream.py index 3d8ddacf0..eb4a83e80 100644 --- a/apps/worker/app/services/document_parser/formats/docx/block_stream.py +++ b/apps/worker/app/services/document_parser/formats/docx/block_stream.py @@ -6,6 +6,9 @@ from app.services.document_parser.formats.docx.toc import detect_doc_tocs, detect_sdt_toc from app.services.document_parser.assets.image_size_filter import is_below_img_min_size +from app.services.document_parser.formats.docx.word_package import ( + normalize_word_package_bytes, +) from docx import Document from docx.oxml.table import CT_Tbl from docx.oxml.text.paragraph import CT_P @@ -16,6 +19,7 @@ def iter_block_items(doc_data): + doc_data = normalize_word_package_bytes(doc_data) doc_stream = io.BytesIO(doc_data) doc = Document(doc_stream) diff --git a/apps/worker/app/services/document_parser/formats/docx/word_package.py b/apps/worker/app/services/document_parser/formats/docx/word_package.py new file mode 100644 index 000000000..2c66c160a --- /dev/null +++ b/apps/worker/app/services/document_parser/formats/docx/word_package.py @@ -0,0 +1,89 @@ +"""Normalize Word OOXML packages so python-docx can open them. + +python-docx only accepts the standard WordprocessingML main document content +type. Macro-enabled (``.docm``) and template (``.dotx`` / ``.dotm``) packages +are otherwise identical ZIP+XML documents and fail with: + + file '...' is not a Word file, content type is + 'application/vnd.ms-word.document.macroEnabled.main+xml' +""" + +from __future__ import annotations + +import io +import zipfile +from typing import BinaryIO + +from docx import Document +from docx.opc.constants import CONTENT_TYPE as CT +from loguru import logger + +_CONTENT_TYPES_MEMBER = "[Content_Types].xml" +_WML_DOCUMENT_MAIN = CT.WML_DOCUMENT_MAIN +_WORD_MAIN_CONTENT_TYPES_TO_REWRITE: tuple[str, ...] = ( + "application/vnd.ms-word.document.macroEnabled.main+xml", + "application/vnd.ms-word.template.macroEnabled.main+xml", + "application/vnd.openxmlformats-officedocument.wordprocessingml.template.main+xml", +) + + +def normalize_word_package_bytes(package_bytes: bytes) -> bytes: + """Rewrite template/macro-enabled main content types to standard ``.docx``.""" + if not zipfile.is_zipfile(io.BytesIO(package_bytes)): + return package_bytes + + with zipfile.ZipFile(io.BytesIO(package_bytes), "r") as source: + try: + content_types_xml = source.read(_CONTENT_TYPES_MEMBER) + except KeyError: + return package_bytes + + rewritten_xml, original_type = _rewrite_word_main_content_type(content_types_xml) + if rewritten_xml == content_types_xml: + return package_bytes + + output = io.BytesIO() + with zipfile.ZipFile(output, "w") as destination: + for item in source.infolist(): + data = ( + rewritten_xml + if item.filename == _CONTENT_TYPES_MEMBER + else source.read(item.filename) + ) + destination.writestr( + item.filename, + data, + compress_type=item.compress_type, + ) + + logger.info( + "Normalized Word OOXML content type {!r} to standard document.main+xml", + original_type, + ) + return output.getvalue() + + +def open_word_document(source: str | bytes | BinaryIO): + """Open a Word OOXML package, including ``.docm`` / ``.dotx`` / ``.dotm``.""" + if isinstance(source, bytes): + package_bytes = source + elif isinstance(source, str): + with open(source, "rb") as handle: + package_bytes = handle.read() + else: + package_bytes = source.read() + + return Document(io.BytesIO(normalize_word_package_bytes(package_bytes))) + + +def _rewrite_word_main_content_type(content_types_xml: bytes) -> tuple[bytes, str | None]: + rewritten = content_types_xml + matched_type: str | None = None + for content_type in _WORD_MAIN_CONTENT_TYPES_TO_REWRITE: + marker = content_type.encode("ascii") + if marker not in rewritten: + continue + rewritten = rewritten.replace(marker, _WML_DOCUMENT_MAIN.encode("ascii")) + matched_type = content_type + break + return rewritten, matched_type diff --git a/apps/worker/tests/unit/test_word_package.py b/apps/worker/tests/unit/test_word_package.py new file mode 100644 index 000000000..819ac4726 --- /dev/null +++ b/apps/worker/tests/unit/test_word_package.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import io +import os +import zipfile +from pathlib import Path + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +import pytest # noqa: E402 +from docx import Document # noqa: E402 +from docx.opc.constants import CONTENT_TYPE as CT # noqa: E402 + +from app.services.document_ingestion.page_estimator import PageEstimator # noqa: E402 +from app.services.document_parser.formats.docx.block_stream import ( # noqa: E402 + iter_block_items, +) +from app.services.document_parser.formats.docx.word_package import ( # noqa: E402 + normalize_word_package_bytes, + open_word_document, +) + +_DOCM_MAIN_CONTENT_TYPE = ( + "application/vnd.ms-word.document.macroEnabled.main+xml" +) +_DOTX_MAIN_CONTENT_TYPE = ( + "application/vnd.openxmlformats-officedocument.wordprocessingml.template.main+xml" +) +_PARAGRAPH_TEXT = "Macro-enabled Word sample" + + +def test_normalize_word_package_bytes_leaves_standard_docx_unchanged() -> None: + package_bytes = _build_docx_bytes(_PARAGRAPH_TEXT) + + assert normalize_word_package_bytes(package_bytes) is package_bytes + + +def test_normalize_word_package_bytes_rewrites_docm_content_type() -> None: + package_bytes = _with_main_content_type( + _build_docx_bytes(_PARAGRAPH_TEXT), + _DOCM_MAIN_CONTENT_TYPE, + ) + + with pytest.raises(ValueError, match="is not a Word file"): + Document(io.BytesIO(package_bytes)) + + normalized = normalize_word_package_bytes(package_bytes) + document = Document(io.BytesIO(normalized)) + + assert _read_main_content_type(normalized) == CT.WML_DOCUMENT_MAIN + assert document.paragraphs[0].text == _PARAGRAPH_TEXT + + +def test_iter_block_items_opens_macro_enabled_word_package() -> None: + package_bytes = _with_main_content_type( + _build_docx_bytes(_PARAGRAPH_TEXT), + _DOCM_MAIN_CONTENT_TYPE, + ) + + blocks = list(iter_block_items(package_bytes)) + paragraph_texts = [ + block[1].text if hasattr(block[1], "text") else str(block[1]) + for block in blocks + if block[2] == "PTXT" + ] + + assert _PARAGRAPH_TEXT in paragraph_texts + + +def test_open_word_document_accepts_template_content_type(tmp_path: Path) -> None: + path = tmp_path / "sample.dotx" + path.write_bytes( + _with_main_content_type( + _build_docx_bytes(_PARAGRAPH_TEXT), + _DOTX_MAIN_CONTENT_TYPE, + ) + ) + + document = open_word_document(str(path)) + + assert document.paragraphs[0].text == _PARAGRAPH_TEXT + + +def test_page_estimator_counts_macro_enabled_docx(tmp_path: Path) -> None: + path = tmp_path / "sample.docx" + path.write_bytes( + _with_main_content_type( + _build_docx_bytes(_PARAGRAPH_TEXT), + _DOCM_MAIN_CONTENT_TYPE, + ) + ) + + estimate = PageEstimator.estimate_workload(str(path)) + + assert estimate.used_fallback is False + assert estimate.method == "docx_words" + assert estimate.page_count >= 1 + + +def _build_docx_bytes(text: str) -> bytes: + document = Document() + document.add_paragraph(text) + buffer = io.BytesIO() + document.save(buffer) + return buffer.getvalue() + + +def _with_main_content_type(package_bytes: bytes, content_type: str) -> bytes: + source = io.BytesIO(package_bytes) + output = io.BytesIO() + with zipfile.ZipFile(source, "r") as src, zipfile.ZipFile(output, "w") as dst: + for item in src.infolist(): + data = src.read(item.filename) + if item.filename == "[Content_Types].xml": + data = data.replace( + CT.WML_DOCUMENT_MAIN.encode("ascii"), + content_type.encode("ascii"), + ) + dst.writestr(item.filename, data, compress_type=item.compress_type) + return output.getvalue() + + +def _read_main_content_type(package_bytes: bytes) -> str: + with zipfile.ZipFile(io.BytesIO(package_bytes), "r") as archive: + content_types = archive.read("[Content_Types].xml").decode("utf-8") + marker = 'PartName="/word/document.xml" ContentType="' + start = content_types.index(marker) + len(marker) + end = content_types.index('"', start) + return content_types[start:end] diff --git a/packages/shared-python/shared/services/http/url_file_type.py b/packages/shared-python/shared/services/http/url_file_type.py index 14fd2a16d..bf243b409 100644 --- a/packages/shared-python/shared/services/http/url_file_type.py +++ b/packages/shared-python/shared/services/http/url_file_type.py @@ -22,6 +22,9 @@ "application/pdf": ".pdf", "application/msword": ".doc", "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", + "application/vnd.ms-word.document.macroEnabled.12": ".docx", + "application/vnd.openxmlformats-officedocument.wordprocessingml.template": ".docx", + "application/vnd.ms-word.template.macroEnabled.12": ".docx", "application/vnd.ms-excel": ".xls", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx", "application/vnd.ms-powerpoint": ".ppt",