From 0f8a9499fa98b604f4b4139ad62bda06eb29f5a7 Mon Sep 17 00:00:00 2001
From: sinnercode228 <sinnercode228@users.noreply.github.com>
Date: Sun, 1 Nov 2026 15:55:42 +0700
Subject: [PATCH] Keep DOCX tables in the section where they appear

load_docx read every paragraph first and only then every table, so all
table rows ended up in the last section's text under the last heading.
Walk the body with iter_inner_content() so paragraphs and tables are
handled in document order.

Fixes #10
---
 backend/src/docmind/ingestion/loaders.py | 19 +++++++++++--------
 backend/tests/test_ingestion.py          | 20 ++++++++++++++++++++
 2 files changed, 31 insertions(+), 8 deletions(-)

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"<html><head><title>Guide</title><style>p{}</style></head><body>"