diff --git a/backend/src/docmind/ingestion/loaders.py b/backend/src/docmind/ingestion/loaders.py index 2776140..a8ce08c 100644 --- a/backend/src/docmind/ingestion/loaders.py +++ b/backend/src/docmind/ingestion/loaders.py @@ -143,6 +143,7 @@ def load_pdf(data: bytes, filename: str | None = None) -> ExtractedDocument: def load_docx(data: bytes, filename: str | None = None) -> ExtractedDocument: import docx + from docx.table import Table try: document = docx.Document(io.BytesIO(data)) @@ -160,9 +161,16 @@ def flush() -> None: sections.append(Section(text=body, heading=heading)) buffer.clear() - for paragraph in document.paragraphs: - style = (paragraph.style.name if paragraph.style is not None else "") or "" - text = paragraph.text.strip() + # Walk the body in document order so a table stays under its own heading. + for block in document.iter_inner_content(): + if isinstance(block, Table): + for row in block.rows: + cells = [cell.text.strip() for cell in row.cells if cell.text.strip()] + if cells: + buffer.append(" | ".join(cells)) + continue + style = (block.style.name if block.style is not None else "") or "" + text = block.text.strip() if not text: continue if style == "Title" and title is None: @@ -172,11 +180,6 @@ def flush() -> None: heading = text else: buffer.append(text) - for table in document.tables: - for row in table.rows: - cells = [cell.text.strip() for cell in row.cells if cell.text.strip()] - if cells: - buffer.append(" | ".join(cells)) flush() return ExtractedDocument( title=title or _title_from(filename), sections=sections, mime_type=DOCX diff --git a/backend/tests/test_ingestion.py b/backend/tests/test_ingestion.py index 79aba0d..1a2d86c 100644 --- a/backend/tests/test_ingestion.py +++ b/backend/tests/test_ingestion.py @@ -1,5 +1,6 @@ from __future__ import annotations +import io import itertools import pytest @@ -117,6 +118,25 @@ def test_docx_headings_paragraphs_tables(self) -> None: assert [s.heading for s in doc.sections][:2] == ["Eligibility", "Budget"] assert "600 EUR" in doc.sections[-1].text + def test_docx_table_stays_in_its_section(self) -> None: + import docx + + document = docx.Document() + document.add_heading("Pricing", level=1) + table = document.add_table(rows=1, cols=2) + table.rows[0].cells[0].text = "Basic" + table.rows[0].cells[1].text = "10 EUR" + document.add_heading("Contacts", level=1) + document.add_paragraph("Write to sales@example.com.") + buffer = io.BytesIO() + document.save(buffer) + + doc = extract(buffer.getvalue(), mime_type=DOCX) + assert [(s.heading, s.text) for s in doc.sections] == [ + ("Pricing", "Basic | 10 EUR"), + ("Contacts", "Write to sales@example.com."), + ] + def test_html_strips_chrome_and_scripts(self) -> None: html = ( b"Guide"