diff --git a/docs/pdf-cross-page-table-reconstruction-design.md b/docs/pdf-cross-page-table-reconstruction-design.md new file mode 100644 index 000000000..5d878c67f --- /dev/null +++ b/docs/pdf-cross-page-table-reconstruction-design.md @@ -0,0 +1,502 @@ +# Cross-page table row reconstruction during PDF indexing + +## Goal + +Improve retrieval from PDFs in which a logical table row spans multiple pages. Continuation pages may contain text in only one or two columns while visually inheriting identifying values such as the row title, reference, or category from an earlier page. + +The solution should reconstruct logical rows before chunking, preserve their inherited context, and remain independent of any one PDF parser. + +This capability must be optional and controlled by the partition's indexing preset. Existing presets must keep the current indexing behavior unless an administrator explicitly enables reconstruction. + +## Current indexing pipeline + +OpenRAG currently processes an uploaded document through the following stages: + +```text +Upload + → resolve the partition indexing preset + → dispatch to an indexing worker + → parse into page-level text blocks + → optionally caption images + → chunk the extracted text or Markdown + → optionally contextualize and tag chunks + → embed chunk text + → store chunks and metadata +``` + +Marker, PyMuPDF, and Docling ultimately expose page-level text blocks. These blocks preserve text and page numbers but do not describe columns, tables, rows, or relationships across page boundaries. + +The current chunker joins page blocks using synthetic page markers and recognizes complete Markdown tables. It can group continuation rows when an empty first column already exists inside the same valid Markdown table. It cannot restore a relationship once the parser has emitted the continuation as a separate table or as ordinary prose. + +Table splitting also happens between row groups. A single oversized logical row is therefore not subdivided safely while retaining its identity. + +Embedding and reranking operate on the final chunk text. Metadata alone cannot compensate for missing row context. + +## Findings from the sample PDF + +The supplied `LEGITEXT000006070158-1.pdf` contains 904 pages. + +### Primary regression: pages 803–805 + +The original regression is the Annex 10 table across PDF pages 803–805: + +- Page 803 contains the table header and starts row 1: `CST portant la mention "salarié"`, reference `L. 421-1`. +- Page 804 contains only the continuation of row 1's `Pièces justificatives` cell. Its first four columns are visually inherited from page 803. PyMuPDF's table detector does not identify a table on this page. +- Page 805 begins with the final continuation of row 1 in the fifth column. +- Row 2, `CST portant la mention "travailleur temporaire"`, reference `L. 421-3`, begins later on page 805. + +This is the required regression fixture for the first delivery. + +### Secondary example: pages 872–877 + +Pages 872–877 contain another instance of the same document pattern: + +- Row 53 starts on page 873 and continues at the beginning of page 874. +- Row 54 starts later on page 874 and continues throughout page 875. +- On page 875, PyMuPDF emits the continuation as ordinary prose without its table or row identity. +- Row 55 begins on page 876 and continues onto page 877. +- Page 877 begins with four empty cells containing the continuation of row 55, followed by the next complete row. + +This produces two different failure modes: + +1. The continuation remains table-shaped, but its inherited columns are missing. +2. The continuation no longer looks like a table and becomes ordinary page text. + +Improving Markdown table parsing alone cannot handle both cases. + +## Best integration point + +Add a structural-normalization stage after optional image captioning and before chunking: + +```text +Parse + → caption + → normalize document structure + → chunk + → contextualize + → embed + → store +``` + +At this point: + +- The complete parsed document is still available. +- Image placeholders have already been resolved. +- Page relationships have not yet been destroyed by chunking. +- The original PDF remains available for obtaining layout evidence. +- The implementation can work independently of the selected parser. + +Reconstruction should not live inside the parser because that would duplicate behavior across Marker, PyMuPDF, and Docling. It should not live only inside the chunker because detecting page relationships and splitting content are different responsibilities. + +The pipeline should resolve the reconstruction policy from the effective indexation configuration before running the stage. In disabled mode, the stage must be bypassed so the resulting behavior is equivalent to the current pipeline. + +## Preset configuration and operating modes + +The complete feature should provide three modes: + +| Mode | Behavior | +| --- | --- | +| Disabled | Bypass reconstruction and preserve the current parser and chunker behavior | +| Automatic | Detect likely cross-page tables and reconstruct only high-confidence cases | +| Strict | Require reliable reconstruction for detected complex tables and fail indexing when safe reconstruction is not possible | + +The initial default must be disabled for backward compatibility. Existing presets and configuration snapshots that do not contain the option should resolve to disabled. + +Automatic mode may become the recommended default after accuracy, latency, and retrieval-quality evaluation. Changing the recommendation later must not silently modify existing presets or previously stored indexation snapshots. + +The first delivery implements disabled and automatic modes. Strict mode remains part of the target design but is deferred until automatic reconstruction has been evaluated. The Admin UI must not present strict mode before the backend implements its failure contract. + +### Backend ownership + +The backend is the source of truth for this capability. It must: + +- Validate the selected mode. +- Validate confidence and detection thresholds against safe server-side ranges. +- Resolve missing values to backward-compatible defaults. +- Persist the effective values in the document's indexation configuration snapshot. +- Pass the effective configuration through the existing dispatcher and indexing worker. +- Apply the same validation regardless of whether a preset is created through the Admin UI or the API. + +Thresholds should have safe backend defaults. The UI does not need to expose every internal heuristic initially; exposing the operating mode is required, while expert threshold controls can remain optional. + +Persisting the effective configuration is important because reconstruction changes indexed content. Operators must be able to determine which policy and algorithm version produced a document's chunks. + +### Admin UI + +The indexing preset editor should expose the mode under advanced PDF parsing or table-processing settings. + +The interface should: + +- Explain that disabled preserves current behavior. +- Describe automatic as conservative and fail-open. +- Keep the option independent of the selected primary PDF parser. +- Avoid suggesting that automatic mode guarantees reconstruction for every table. + +When strict mode is implemented in a later delivery, the interface must warn that it can cause indexing jobs to fail. + +The Admin UI sends the requested preset configuration, but it must not duplicate or replace backend validation. + +## Architecture options + +| Architecture | Advantages | Limitations | +| --- | --- | --- | +| Stitch Markdown tables inside the chunker | Small and inexpensive change | Depends on valid Markdown and fails when continuation pages become ordinary prose | +| Enhance each PDF parser | Can use parser-specific structural data | Duplicates logic and produces inconsistent behavior across parsers | +| Add a parser-independent structural-normalization stage | Generic, deterministic, testable, and reusable | Requires a logical table representation and an additional layout-analysis step | +| Use a VLM or external document-analysis service | Can resolve difficult visual cases | More expensive, slower, and non-deterministic | + +The parser-independent normalization stage is the recommended foundation. A VLM or external service should only be an optional resolver for ambiguous cases. + +Parser-independent means that the reconstruction algorithm is not coupled to Marker, PyMuPDF, or Docling. It does not mean that every parser provides equal evidence. Reconstruction quality depends on the text and layout evidence available for a document and must degrade safely when that evidence is incomplete. + +## Recommended architecture + +### Structural normalizer + +Introduce a document-structure normalizer that receives both the parsed document and the original document. Its responsibility is to identify table fragments, resolve page-boundary relationships, and produce logical table-row blocks. + +The normalizer must preserve the existing parser contract. It returns another `ProcessedDocument`; it does not introduce a parallel indexing pipeline. + +Resolved rows are exposed through the existing block contract as `TextBlock` instances with `block_type="table_row"` and typed table-row data. This lets the current chunker consume normalized blocks without changing parser outputs. + +The normalizer should receive a validated policy rather than reading environment variables or UI state directly. This keeps behavior deterministic for a given indexation snapshot and makes the stage straightforward to test. + +### Exact input and output contract + +The normalizer interface receives: + +- The original `Document`, including content type and PDF bytes. +- The parser-produced `ProcessedDocument`. +- The validated table-reconstruction configuration from the effective indexing preset. + +It returns a `ProcessedDocument` with the following guarantees: + +- `raw_text_blocks` contains the untouched parser output captured before image captions can replace placeholders. It is populated only when reconstruction is enabled. +- `text_blocks` remains the existing working block view, including any caption substitutions already performed by the current pipeline. +- `normalized_text_blocks` is `None` when no safe normalization was produced. +- When normalization succeeds, `normalized_text_blocks` is a complete chunkable view of the document in source order. It contains unaffected text, residual text around table regions, and reconstructed `table_row` blocks. +- Every normalized block references its source slices through typed source fragments. +- The chunker uses `normalized_text_blocks` when present and otherwise uses the raw `text_blocks`. +- A normalization report records decisions and fallback reasons without replacing or mutating source content. + +The existing document models should be extended with the following exact contract: + +| Model | Required data | +| --- | --- | +| `SourceFragment` | Source block index, source page number, inclusive start and exclusive end character offsets, and an optional page-normalized bounding box | +| `PageBoundaryDecision` | Previous and next page numbers, separate same-table and row-continuation confidence values, decision (`merged` or `preserved`), and reason | +| `TableCellData` | Column index, optional column name, text, source fragments, and cell-assignment confidence | +| `TableRowData` | Deterministic table and row identifiers, algorithm version, inferred identity-column indexes, optional table title, section path, ordered cells, start and end pages, and the page-boundary decisions used to reconstruct it | +| `NormalizationReport` | Algorithm version, status (`unchanged`, `normalized`, or `partial_fallback`), all page-boundary decisions, reconstructed-row count, and fallback reasons | + +`TextBlock` gains optional `source_fragments` and `table_row` fields. `ProcessedDocument` gains optional `raw_text_blocks`, `normalized_text_blocks`, and `normalization_report` fields plus an `effective_text_blocks()` accessor. All new fields have backward-compatible defaults. + +Source references must point to character ranges in `raw_text_blocks`. If layout or captioned text cannot be aligned reliably with parser text, the cell-assignment confidence must remain below the automatic threshold and that region must use the existing working block without reconstruction. + +This makes normalization reversible and debuggable: an operator or test can compare the normalized row with the exact parser slices that produced it, and the algorithm version defines the deterministic whitespace and joining rules used to derive the row. + +### Evidence providers + +The normalizer should combine: + +- Page text and Markdown from the selected parser. +- Lightweight PDF layout evidence such as word positions, normalized column boundaries, table regions, and page coordinates. + +PyMuPDF can provide the first layout adapter because it is already available, but it should remain behind an interface. It supplies evidence and is not expected to reconstruct the table perfectly. + +Future adapters could consume richer Docling output, Marker structure, OCR results, or another document-analysis service without changing the reconstruction algorithm. + +Only the PyMuPDF layout-evidence adapter is included in the first delivery. Additional adapters are follow-up work. + +### Candidate detection + +Detailed layout processing should run only on likely table regions and their adjacent pages. Candidate signals include: + +- Markdown or geometric table-like structure. +- Text aligned into stable column bands. +- A table approaching the lower page boundary. +- Content beginning near the upper boundary of the following page. +- Continuation text concentrated in only a subset of the previous table's columns. + +Candidate windows should expand until the end of the logical table is found. This avoids an expensive second table-analysis pass over every page of large documents. + +### Deterministic continuation resolver + +The default resolver should operate as a page-boundary state machine. A boundary can be considered a continuation when several signals agree: + +- Pages are adjacent and have compatible dimensions and orientation. +- Column boundaries remain similar after normalization against page width. +- The preceding table reaches the lower content area. +- The next fragment begins near the upper content area. +- Headers either match or are recognized as repeated headers. +- Previously identifying columns are empty while content continues in another column. +- No heading, caption, or unrelated section occurs between the fragments. + +Repeated document headers and footers must be excluded before matching. + +The resolver must not assume that the first four columns identify a row or that only the final column can continue. Identity columns should be inferred from the table's observed row starts and value patterns. + +When the next fragment has empty identity cells, their values are inherited from the open logical row. A new logical row begins when those identity columns become populated again. + +### Confidence model + +Confidence must not be represented by one aggregate score. The resolver records and evaluates three independent values: + +- **Same-table confidence:** whether fragments on adjacent pages belong to the same logical table. +- **Row-continuation confidence:** whether the leading fragment on the next page continues the currently open logical row. +- **Cell-assignment confidence:** whether a text fragment was assigned to the correct table column and aligned to the correct parser-text range. + +Automatic reconstruction occurs only when the same-table and row-continuation thresholds pass and every cell assignment used by the merge passes its own threshold. Scores must not be averaged because a strong table match must not hide an uncertain row or cell assignment. + +The initial automatic thresholds are backend-owned and independently configurable. The first delivery uses defaults of `0.90` for all three thresholds and validates each value within the safe range `0.80–1.00`. + +The persisted nested configuration is: + +| Field | First-delivery contract | +| --- | --- | +| `mode` | `disabled` or `automatic`; default `disabled` | +| `same_table_min_confidence` | Float from `0.80` to `1.00`; default `0.90` | +| `row_continuation_min_confidence` | Float from `0.80` to `1.00`; default `0.90` | +| `cell_assignment_min_confidence` | Float from `0.80` to `1.00`; default `0.90` | +| `algorithm_version` | Backend-owned literal identifying the deterministic algorithm; initial value `adjacent-layout-v1` | + +Unknown fields are rejected inside this nested configuration. The algorithm version is persisted with the effective indexation snapshot so a document can be reproduced after heuristics evolve. + +### Logical table representation + +The normalized representation should preserve: + +- Stable table and row identifiers. +- Table title and section path. +- Column schema. +- Reconstructed cell values. +- Source fragments and their page provenance. +- Start and end pages. +- Reconstruction method, confidence, and evidence. + +Identifiers should be deterministic so reindexing the same document produces comparable metadata. + +### Row-aware chunking + +The chunker should handle normalized table rows separately from ordinary text. + +A normal-sized row should produce one table chunk containing labelled cells. + +For an oversized cell: + +- Reserve space for a compact context prefix. +- Include the section, table identity, row identity, and inherited column values. +- Split content on headings, numbered sections, list items, paragraphs, and sentences. +- Use token-based splitting only as a final fallback. +- Repeat the context prefix in every generated chunk. +- Keep each final chunk within the configured token budget. + +A self-contained chunk should convey the following information: + +```text +Section: … +Table: … +Row: Number 54; Category …; Title …; Reference L. 426-7 +Supporting documents, part 2 of 4: … +``` + +The final representation does not need to remain a Markdown table. Labelled text is generally clearer for embeddings and reranking. + +## Metadata and downstream behavior + +Each generated table chunk should carry compact metadata describing: + +- Table and row identifiers. +- Table title and section. +- Row identity. +- Content column. +- Source page start and end. +- Chunk part number and total parts. +- Reconstruction method and confidence. +- Source-fragment provenance for the content included in that chunk. + +For compatibility, the existing page value should remain the starting page. The ending page can be added as metadata and later exposed in source citations. + +The row context must also be included in the chunk text because embeddings operate on text, not metadata. The embedding and storage stages otherwise require no behavioral change. + +Optional LLM contextualization can continue to run after chunking. It may enrich the chunk, but it should not be responsible for reconstructing table relationships. + +## Ambiguity and failure handling + +False merges can be more damaging than missed merges, so reconstruction should be conservative. + +- High-confidence boundaries are reconstructed automatically. +- When row identity is certain but exact cell alignment is uncertain, preserve the fragment and attach only the confirmed context. +- When row identity is uncertain, preserve the original extraction and record the ambiguity. +- Never introduce content that cannot be traced back to an extracted source fragment. +- Never silently merge unrelated rows. + +The resolver should expose an extension point for advanced decisions. A future VLM could inspect only ambiguous page pairs, while the deterministic resolver remains the default. + +Failure behavior must follow the configured mode: + +- Disabled mode does not run detection or reconstruction. +- Automatic mode fails open. Uncertain boundaries and reconstruction errors preserve the original parser output and are reported through logs and metrics. +- Strict mode fails the indexing job when a detected complex table cannot be reconstructed with the required confidence. The error should identify the reconstruction stage and affected page boundary without exposing document content unnecessarily. + +Strict mode should not reject an ordinary document merely because no complex table was detected. It applies the reliability requirement only after the detector identifies a table that requires cross-page reconstruction. + +Strict mode is outside the first delivery. The first delivery must still define the normalizer and stage interfaces so strict failure behavior can be added without changing their input or output contracts. + +## Scalability and observability + +The normalizer should process candidate page windows incrementally and avoid loading rendered images for deterministic reconstruction. + +The raw parser-block snapshot is created only in automatic mode and released with the in-memory processed document after indexing, so disabled mode has no additional document-copy cost. + +Operational metrics should include: + +- Structural-normalization duration. +- Candidate boundaries detected. +- Logical rows reconstructed. +- Continuations merged. +- Ambiguous boundaries. +- Reconstruction fallbacks and failures. + +Reconstruction metadata and the persisted indexation snapshot should include the selected mode, effective thresholds, and algorithm version so that changes can be evaluated and documents can be selectively reindexed. + +## First delivery scope + +The first delivery is intentionally limited to: + +- The structural-normalizer interface. +- The PyMuPDF layout-evidence adapter. +- Deterministic reconstruction across adjacent pages. +- Row-aware chunking of oversized cells. +- Source-fragment provenance. +- Automatic fail-open behavior. +- Backend validation and snapshot propagation required to keep the feature optional. +- Disabled and automatic choices in the Admin UI. +- Regression and pipeline tests. + +The first delivery does not include: + +- VLM fallback. +- Strict indexing behavior. +- Layout adapters other than PyMuPDF. +- Automatic activation for existing presets. +- Extensive rollout controls or confidence-tuning controls in the Admin UI. + +## Exact implementation plan + +### Domain models and interfaces + +| File | Change | +| --- | --- | +| `openrag/core/models/document.py` | Add `SourceFragment`, `PageBoundaryDecision`, `TableCellData`, `TableRowData`, and `NormalizationReport`. Extend `TextBlock` with provenance and optional row data. Extend `ProcessedDocument` with `raw_text_blocks`, `normalized_text_blocks`, `normalization_report`, and `effective_text_blocks()`. | +| `openrag/core/indexing/structure_normalizer.py` | Add the `DocumentStructureNormalizer` and `TableLayoutEvidenceProvider` interfaces and the page/table evidence models exchanged between them. | +| `openrag/core/indexing/table_normalizer.py` | Add `DeterministicTableNormalizer`, candidate-window detection, adjacent-page state handling, independent confidence decisions, parser-text alignment, raw-block preservation, and fail-open output. | + +`DocumentStructureNormalizer.normalize` accepts a `Document`, a `ProcessedDocument`, and a validated table-reconstruction configuration. It returns a `ProcessedDocument` under the contract defined above. + +### PyMuPDF evidence + +| File | Change | +| --- | --- | +| `openrag/core/indexing/parsers/pdf/pymupdf_runtime.py` | Centralize serialized PyMuPDF execution so parsing and layout evidence cannot call the non-thread-safe library concurrently. | +| `openrag/core/indexing/parsers/pdf/pymupdf.py` | Use the shared PyMuPDF runtime without changing parser output. | +| `openrag/services/workers/layout/__init__.py` | Introduce the layout-adapter package. | +| `openrag/services/workers/layout/pymupdf_table_evidence.py` | Add `PyMuPDFTableEvidenceProvider`, extracting normalized page geometry, table fragments, cells, and text bands only for candidate page windows. | + +The adapter supplies evidence; it does not decide whether rows should be merged. + +### Pipeline integration + +| File | Change | +| --- | --- | +| `openrag/services/workers/stages/parse.py` | Add a `preserve_raw_blocks` option. When automatic reconstruction is enabled, snapshot parser-produced blocks into `raw_text_blocks` before captioning. Disabled mode keeps the current allocation and behavior. | +| `openrag/services/workers/stages/normalize_structure.py` | Add the pipeline stage. Disabled mode bypasses it. Automatic mode catches detection, alignment, and timeout failures and preserves the current processed document. | +| `openrag/services/workers/pipeline_builder.py` | Inject the normalizer, add a structure-normalization timeout to `PipelineTimeouts`, run the stage after captioning and before chunking, include its timing, and pass the validated policy from the effective indexation config. | +| `openrag/services/workers/indexer_pool.py` | Construct one PyMuPDF evidence provider and deterministic normalizer per indexing worker, inject them into the pipeline, and use the existing loader parse timeout as the initial normalization bound. | +| `openrag/services/workers/stages/chunk.py` | Continue passing a `ProcessedDocument`; no parallel chunking pipeline is introduced. | + +### Row-aware chunking files + +| File | Change | +| --- | --- | +| `openrag/core/chunking/recursive.py` | Read the processed document's effective block view. Keep the existing path for ordinary blocks and route `table_row` blocks through row-aware chunking. | +| `openrag/core/chunking/table_rows.py` | Add deterministic row serialization and oversized-cell splitting. Reserve space for repeated row context, split semantically, enforce the token limit, and attach provenance and page-range metadata to each table chunk. | + +Every emitted table chunk includes row identity in its text and the relevant source fragments in its metadata. + +### Configuration, validation, and snapshot propagation + +| File | Change | +| --- | --- | +| `openrag/core/config/table_reconstruction.py` | Add `TableReconstructionConfig`. The first delivery accepts `disabled` and `automatic`, defaults to disabled, and validates the three independent thresholds. | +| `openrag/core/config/indexation_pipeline.py` | Add the nested `table_reconstruction` field to `IndexationPipelineConfig`. | +| `openrag/api/schemas/admin/preset_schemas.py` | Add available reconstruction modes to the preset-options response. | +| `openrag/api/routers/admin/presets.py` | Return backend-supported reconstruction modes from the preset-options endpoint. | +| `openrag/services/orchestrators/preset_service.py` | Continue validating through `IndexationPipelineConfig`; add coverage that invalid modes, unknown nested fields, and unsafe thresholds are rejected. Existing seed dictionaries remain sparse and therefore resolve to disabled. | +| `openrag/services/orchestrators/indexing_service.py` | No new dispatch path. Confirm by test that the effective nested configuration is included in the existing `model_dump` passed to the dispatcher. | +| `openrag/services/workers/indexer_actor.py` | No new persistence path. Confirm by test that the existing file snapshot stores the effective nested configuration. | + +The nested table-reconstruction model should reject unknown fields so API typos cannot be silently ignored. + +### Admin UI files + +| File | Change | +| --- | --- | +| `ui/src/lib/api/presets.ts` | Add reconstruction modes to `PresetOptionsResponse`, keeping the field optional for rolling-deployment compatibility. | +| `ui/src/pages/admin/preset-config.ts` | Add pure helpers for reading and updating the nested reconstruction mode. A missing value displays as disabled. | +| `ui/src/pages/admin/presets.tsx` | Add an Advanced PDF/Table processing section with Disabled and Automatic choices and concise fail-open guidance. Do not expose thresholds or strict mode in the first delivery. | + +The backend remains authoritative even when an older or modified client submits the configuration directly. + +### Tests + +| Test file | Coverage | +| --- | --- | +| `tests/resources/cross_page_table_rows_803_805.pdf` | Minimal three-page regression fixture extracted from source pages 803–805. | +| `tests/unit/core/indexing/test_table_normalizer.py` | Candidate detection, page 803→804 and 804→805 decisions, inherited cells, row 2 beginning mid-page, three independent confidence thresholds, provenance, reversibility, and ambiguous fail-open behavior. | +| `tests/unit/services/workers/layout/test_pymupdf_table_evidence.py` | Geometry and cell evidence from the regression fixture, including page 804 where no complete table is detected. | +| `tests/unit/core/chunking/test_table_rows.py` | Oversized-cell semantic splitting, repeated row identity, token budgets, page ranges, and per-chunk source fragments. | +| `tests/unit/core/chunking/test_recursive.py` | Raw-block fallback and selection of normalized blocks when present. | +| `tests/unit/services/workers/stages/test_parse.py` | Raw parser blocks are captured before captioning only when automatic reconstruction is enabled. | +| `tests/unit/services/workers/stages/test_normalize_structure.py` | Disabled bypass, successful automatic normalization, timeout, exception, and fail-open preservation. | +| `tests/unit/services/workers/test_pipeline_builder.py` | Stage ordering, effective preset policy, PDF-only invocation, and unchanged non-PDF behavior. | +| `tests/unit/services/orchestrators/test_preset_service.py` | Default-disabled behavior and rejection of invalid modes, thresholds, and unknown nested fields. | +| `tests/unit/services/orchestrators/test_indexing_service.py` | Effective reconstruction configuration reaches the existing dispatcher. | +| `tests/unit/services/workers/test_indexer_worker.py` | Effective reconstruction configuration is retained in new-file and replacement snapshots. | +| `tests/integration/api/test_presets.py` | Preset options, automatic-mode round trip, validation errors, and backward-compatible omission. | +| `ui/src/pages/admin/preset-config.test.ts` | Missing configuration displays disabled and updates remain immutable. | +| `ui/src/pages/admin/presets.test.tsx` | Mode rendering, user guidance, and submitted automatic configuration. | + +The local 904-page PDF must not be required by CI. The committed three-page fixture is the reproducible regression input. + +## Migration impact + +No relational database migration is required. Pipeline presets and document indexation snapshots are already stored as JSONB. + +No Milvus collection migration is required. Reconstruction provenance and table identifiers use the existing dynamic chunk metadata. + +Existing presets without `table_reconstruction` resolve to disabled. Existing indexed chunks are unchanged; documents benefit only after being reindexed with automatic mode enabled. + +The processed-document changes are in-memory model extensions with backward-compatible defaults. The preset-options API receives an additive field, and the Admin UI treats it as optional for mixed-version deployments. + +Strict mode will require a later additive configuration and UI change, but no database migration. + +## Acceptance criteria + +- A logical row spanning multiple pages is reconstructed before chunking. +- Empty continuation columns inherit the values of the open row. +- A new logical row beginning partway through a later page is detected correctly. +- Oversized cells produce bounded semantic chunks. +- Every chunk remains understandable without its neighboring chunks. +- Table identity, row identity, section, and page range are preserved. +- The behavior is independent of the selected primary PDF parser. +- Non-table documents and ordinary single-page tables retain their current behavior. +- Ambiguous cases do not result in invented or silently reassigned content. +- Retrieval improves for queries combining row identity with continuation content. +- Disabled is the initial default and preserves existing indexing behavior. +- Automatic mode reconstructs only high-confidence cases and otherwise preserves the original parser output. +- The first-delivery Admin UI exposes disabled and automatic modes. +- The backend validates the mode and thresholds and remains authoritative. +- The effective policy and algorithm version are retained in the indexation configuration snapshot. +- Raw parser blocks remain unchanged and every reconstructed row is traceable to them. +- Same-table, row-continuation, and cell-assignment confidence are evaluated independently. diff --git a/docs/pdf-cross-page-table-reconstruction-implementation-plan.md b/docs/pdf-cross-page-table-reconstruction-implementation-plan.md new file mode 100644 index 000000000..9640f687b --- /dev/null +++ b/docs/pdf-cross-page-table-reconstruction-implementation-plan.md @@ -0,0 +1,235 @@ +# Cross-page table reconstruction implementation plan + +## Objective + +Add an optional structural-normalization stage that reconstructs table rows spanning adjacent PDF pages before chunking. The feature must improve retrieval without changing existing indexing unless an administrator enables it. + +The detailed architecture and model analysis are documented in [pdf-cross-page-table-reconstruction-design.md](pdf-cross-page-table-reconstruction-design.md). + +## First delivery + +The first PR includes: + +- A structural-normalizer interface. +- A PyMuPDF layout-evidence adapter. +- Deterministic reconstruction across adjacent pages. +- Row-aware chunking for oversized cells. +- Source-fragment provenance. +- Automatic fail-open behavior. +- Backend configuration and snapshot propagation. +- Disabled and automatic choices in the Admin UI. +- Regression, pipeline, and configuration tests. + +Strict indexing, VLM fallback, additional layout adapters, automatic enablement, and advanced rollout controls remain follow-up work. + +## Pipeline integration + +The new stage runs after optional image captioning and before chunking: + +1. Parse the document. +2. Preserve the raw parser blocks when automatic reconstruction is enabled. +3. Apply optional image captions to the existing working blocks. +4. Normalize cross-page table structure. +5. Chunk ordinary blocks and reconstructed rows. +6. Continue through contextualization, embedding, and storage unchanged. + +Disabled mode bypasses normalization and does not create the additional raw-block snapshot. + +## Implementation steps + +### 1. Add the backend configuration + +Introduce `TableReconstructionConfig` under `IndexationPipelineConfig`. + +The first-delivery configuration contains: + +| Field | Behavior | +| --- | --- | +| `mode` | `disabled` or `automatic`; defaults to `disabled` | +| `same_table_min_confidence` | Independent threshold; defaults to `0.90` | +| `row_continuation_min_confidence` | Independent threshold; defaults to `0.90` | +| `cell_assignment_min_confidence` | Independent threshold; defaults to `0.90` | +| `algorithm_version` | Backend-owned value; initially `adjacent-layout-v1` | + +Each threshold must be between `0.80` and `1.00`. Unknown fields in this nested configuration must be rejected. + +The effective configuration is passed through the existing dispatcher and retained in the document indexation snapshot. + +### 2. Preserve raw parser blocks + +Extend `ProcessedDocument` with an optional `raw_text_blocks` collection. + +When automatic mode is enabled, the parse stage captures the parser-produced blocks before captioning can replace image placeholders. The existing `text_blocks` collection remains the working representation used by the current pipeline. + +Raw blocks are never modified. They provide the stable source against which reconstructed content and character offsets are verified. + +### 3. Extend the processed-document contract + +Add typed models for: + +- `SourceFragment` +- `PageBoundaryDecision` +- `TableCellData` +- `TableRowData` +- `NormalizationReport` + +Extend `TextBlock` with optional provenance and table-row data. Extend `ProcessedDocument` with: + +- `raw_text_blocks` +- `normalized_text_blocks` +- `normalization_report` +- `effective_text_blocks()` + +The chunker reads `normalized_text_blocks` when normalization produced a safe complete view. Otherwise, it reads the existing `text_blocks`. + +### 4. Introduce the normalizer interfaces + +`DocumentStructureNormalizer` receives: + +- The original `Document`, including the PDF bytes. +- The current `ProcessedDocument`. +- The validated table-reconstruction configuration. + +It returns a `ProcessedDocument` and never mutates the raw parser blocks. + +`TableLayoutEvidenceProvider` supplies layout evidence independently from the selected primary parser. The reconstruction algorithm must depend on this interface rather than on PyMuPDF directly. + +### 5. Build the PyMuPDF evidence adapter + +The initial adapter extracts evidence only; it does not decide whether rows should be merged. + +For candidate pages it collects: + +- Page dimensions and orientation. +- Words and normalized bounding boxes. +- Table regions and column boundaries. +- Cell content and coordinates. +- Distance from content to page boundaries. +- Repeated headers and footers. + +PyMuPDF access must use the existing serialized execution mechanism because the library is not thread-safe. + +### 6. Implement deterministic adjacent-page reconstruction + +Process candidate pages in source order with a table state machine. + +For each open table: + +1. Record its column schema and inferred identity columns. +2. Keep an unfinished logical row open at the bottom of a page. +3. Inspect only the immediately following page. +4. Determine whether both fragments belong to the same table. +5. Determine whether the leading fragment continues the open row. +6. Assign every contributing text fragment to a column. +7. Inherit empty identity cells from the open row. +8. Close the row when populated identity columns establish a new row. + +The implementation must not assume a fixed number of columns or that only the last column can continue. + +### 7. Keep confidence decisions independent + +Record separate confidence for: + +- Whether adjacent fragments belong to the same table. +- Whether the next fragment continues the open row. +- Whether each fragment was assigned to the correct cell and raw-text range. + +Automatic reconstruction proceeds only when both boundary decisions and every involved cell assignment pass their respective thresholds. Scores are never averaged. + +Contradictory evidence, ambiguous raw-text alignment, a new heading, or incompatible columns must preserve the original content. + +### 8. Build a complete normalized block view + +Each reconstructed row becomes a `TextBlock` with `block_type="table_row"` and typed `TableRowData`. + +The normalized block view must also contain: + +- Unaffected content. +- Residual text surrounding reconstructed regions. +- Original content for uncertain regions. + +Every reconstructed value references the contributing raw block, page, character range, and optional normalized bounding box. The normalization report records decisions, confidence values, fallback reasons, and algorithm version. + +### 9. Add row-aware chunking + +A normal-sized row produces one labelled-text chunk. + +When a cell exceeds the token budget: + +- Reserve room for a compact row-context prefix. +- Split on headings, numbered items, lists, paragraphs, and sentences. +- Use token splitting only as the final fallback. +- Repeat section, table, row identity, inherited values, and content-column name in every part. +- Attach only the source fragments contributing to that part. + +This repeated context ensures that every chunk remains useful to embedding, retrieval, and reranking without depending on neighboring chunks. + +### 10. Integrate fail-open behavior + +Automatic mode catches evidence-extraction failures, timeouts, alignment failures, and unexpected normalizer errors. + +An uncertain region keeps the current parser output. A stage-level failure keeps the complete current processed document. The failure reason is recorded without inventing content or silently merging unrelated rows. + +### 11. Expose the capability in the Admin UI + +Add an advanced PDF/table-processing setting to the indexing preset editor with: + +- Disabled +- Automatic + +The interface explains that automatic mode is conservative and preserves parser output when reconstruction is uncertain. Threshold controls and strict mode are not exposed in the first PR. + +## Regression behavior + +The primary regression fixture is pages 803–805 from `LEGITEXT000006070158-1.pdf`: + +- Page 803 opens row 1 with `CST salarié` and reference `L. 421-1`. +- Page 804 continues its supporting-documents cell. +- The beginning of page 805 completes row 1. +- Row 2 begins separately later on page 805. + +Pages 872–877 are a secondary example covering both table-shaped continuations and continuations emitted as ordinary prose. + +The full 904-page local document must not be required by CI. A minimal three-page fixture should cover the primary regression. + +## Validation + +The test suite must demonstrate: + +- Disabled mode preserves current behavior. +- Automatic mode reconstructs the primary regression. +- Raw parser blocks remain unchanged. +- Every reconstructed value is traceable to source fragments. +- All three confidence decisions are enforced independently. +- Ambiguous boundaries fail open. +- Row 2 is not merged with row 1. +- Oversized cells remain within the configured token budget. +- Every split chunk repeats the row identity. +- Non-PDF and ordinary PDF indexing remain unchanged. +- Preset validation, dispatch, and configuration snapshots retain the effective policy. +- The Admin UI reads and submits the automatic mode correctly. + +## Migration impact + +No PostgreSQL migration is required because presets and indexation snapshots are JSONB. + +No Milvus migration is required because table identifiers and provenance use existing dynamic chunk metadata. + +Existing presets resolve to disabled, and existing indexed documents remain unchanged. A document must be reindexed with automatic mode enabled to benefit from reconstruction. + +## Initial local evaluation + +The implementation was evaluated on 28 July 2026. + +The three-page regression fixture reconstructed row 1 across source pages 803–805, kept row 2 separate, retained provenance from all three pages, and produced table chunks within the configured test budget. Every chunk for row 1 repeated `CST salarié` and `L. 421-1`. + +The complete 904-page source PDF produced: + +- 54 normalized logical rows. +- 71 merged adjacent-page boundaries. +- Two uncertain cases preserved through fail-open fallback. +- The expected row starting on page 803 and ending on page 805. + +On the development machine, complete PyMuPDF Markdown parsing and normalization took approximately 96.5 seconds with about 150 MB peak resident memory. This is an opt-in cost; disabled mode does not run the evidence adapter or create the raw-block snapshot. + +The automated results establish structural correctness for the primary regression. A representative manual sample of the other reconstructed rows is still required before recommending automatic mode for production presets. diff --git a/openrag/api/routers/admin/presets.py b/openrag/api/routers/admin/presets.py index c4a122a5f..73e200924 100644 --- a/openrag/api/routers/admin/presets.py +++ b/openrag/api/routers/admin/presets.py @@ -10,6 +10,7 @@ ) from core.chunking import chunking_registry from core.config.indexation_pipeline import PARSING_STRATEGIES +from core.config.table_reconstruction import TABLE_RECONSTRUCTION_MODES from core.rerankers.registry import reranker_registry from core.retrieval import retriever_registry from di.providers import get_preset_service @@ -37,6 +38,7 @@ async def get_preset_options(): return PresetOptionsResponse( chunking_strategies=chunking_registry.list_registered(), parsing_strategies=_PARSING_STRATEGIES, + table_reconstruction_modes=list(TABLE_RECONSTRUCTION_MODES), retrieval_types=retriever_registry.list_registered(), reranker_providers=_registered_or_default( reranker_registry.list_registered(), diff --git a/openrag/api/schemas/admin/preset_schemas.py b/openrag/api/schemas/admin/preset_schemas.py index ba775f3ba..136d98de7 100644 --- a/openrag/api/schemas/admin/preset_schemas.py +++ b/openrag/api/schemas/admin/preset_schemas.py @@ -88,6 +88,7 @@ class PresetOptionsResponse(BaseModel): chunking_strategies: list[str] parsing_strategies: list[str] + table_reconstruction_modes: list[str] retrieval_types: list[str] reranker_providers: list[str] diff --git a/openrag/core/chunking/recursive.py b/openrag/core/chunking/recursive.py index 5270b05ed..bf4ac2e93 100644 --- a/openrag/core/chunking/recursive.py +++ b/openrag/core/chunking/recursive.py @@ -24,8 +24,9 @@ split_md_elements, ) from core.chunking.registry import chunking_registry +from core.chunking.table_rows import chunk_table_legend, chunk_table_row from core.models.chunk import Chunk, ChunkType -from core.models.document import ProcessedDocument +from core.models.document import ProcessedDocument, TextBlock from core.utils.text import sanitize_text # Substring (case-insensitive) marking a "no useful content" image caption. @@ -130,12 +131,19 @@ def __init__( # ------------------------------------------------------------------ def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: """Split a processed document into ``Chunk`` objects.""" - content = self._content_from(document) - if not content.strip(): - return [] - metadata = self._chunk_metadata_base(document, partition) - md_chunks = self._get_chunks(content=content.strip(), metadata=metadata) + blocks = document.effective_text_blocks() + if any( + (block.block_type == "table_row" and block.table_row is not None) + or (block.block_type == "table_legend" and block.table_legend is not None) + for block in blocks + ): + md_chunks = self._get_block_aware_chunks(blocks, metadata) + else: + content = self._content_from_blocks(blocks) + if not content.strip(): + return [] + md_chunks = self._get_chunks(content=content.strip(), metadata=metadata) return [ Chunk( @@ -155,6 +163,10 @@ def chunk(self, document: ProcessedDocument, partition: str = "default") -> list # ------------------------------------------------------------------ @staticmethod def _content_from(document: ProcessedDocument) -> str: + return BaseChunker._content_from_blocks(document.effective_text_blocks()) + + @staticmethod + def _content_from_blocks(blocks: list[TextBlock]) -> str: """Reconstruct chunkable markdown from a ProcessedDocument. Single-block documents on page 1 (or with no page metadata) flow @@ -167,14 +179,14 @@ def _content_from(document: ProcessedDocument) -> str: new page begins, and we also prepend a marker for the first block if it doesn't start on page 1. """ - if not document.text_blocks: + if not blocks: return "" - if len(document.text_blocks) == 1 and document.text_blocks[0].page_number in (None, 1): - return document.text_blocks[0].text + if len(blocks) == 1 and blocks[0].page_number in (None, 1): + return blocks[0].text parts: list[str] = [] last_page: int | None = None - for index, block in enumerate(document.text_blocks): + for index, block in enumerate(blocks): if block.page_number is not None: # Emit `[PAGE_{block.page_number - 1}]` immediately *before* # this block's text so downstream resolution lands on @@ -193,6 +205,61 @@ def _content_from(document: ProcessedDocument) -> str: last_page = block.page_number return "\n\n".join(parts) + def _get_block_aware_chunks( + self, + blocks: list[TextBlock], + metadata: dict[str, Any], + ) -> list[dict[str, Any]]: + chunks: list[dict[str, Any]] = [] + ordinary: list[TextBlock] = [] + + def flush_ordinary() -> None: + if not ordinary: + return + content = self._content_from_blocks(ordinary) + if content.strip(): + chunks.extend(self._get_chunks(content.strip(), metadata)) + ordinary.clear() + + for block in blocks: + if block.block_type == "table_legend" and block.table_legend is not None: + flush_ordinary() + chunks.extend( + { + **metadata, + **legend_chunk.metadata, + "page_content": legend_chunk.text, + "page": legend_chunk.page_number, + "chunk_type": "table", + } + for legend_chunk in chunk_table_legend( + block.table_legend, + chunk_size=self.chunk_size, + length_function=self.length_function, + ) + ) + continue + if block.block_type != "table_row" or block.table_row is None: + ordinary.append(block) + continue + flush_ordinary() + chunks.extend( + { + **metadata, + **row_chunk.metadata, + "page_content": row_chunk.text, + "page": row_chunk.page_number, + "chunk_type": "table", + } + for row_chunk in chunk_table_row( + block.table_row, + chunk_size=self.chunk_size, + length_function=self.length_function, + ) + ) + flush_ordinary() + return chunks + @staticmethod def _chunk_metadata_base(document: ProcessedDocument, partition: str) -> dict[str, Any]: # Reserved identity fields must win — `chunk()` later reads diff --git a/openrag/core/chunking/table_rows.py b/openrag/core/chunking/table_rows.py new file mode 100644 index 000000000..74c187284 --- /dev/null +++ b/openrag/core/chunking/table_rows.py @@ -0,0 +1,479 @@ +"""Self-contained chunking for normalized logical table rows.""" + +from __future__ import annotations + +import re +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from core.indexing.table_text import ( + TABLE_TEXT_SERIALIZATION_VERSION, + normalize_table_text, + render_table_legend, + render_table_row, + render_table_row_context, +) +from core.models.document import ( + SourceFragment, + TableCellData, + TableLegendData, + TableRowData, +) + +_SEMANTIC_BOUNDARY_RE = re.compile(r"\n{2,}|(?=\n\s*(?:#{1,6}\s+|[-*•]\s+|\d+(?:[.)]|\.\d+[.)]?)\s+))|(?<=[.!?;:])\s+") + + +@dataclass(slots=True, frozen=True) +class TableRowChunk: + text: str + page_number: int + metadata: dict[str, Any] + + +def _identity_cells(row: TableRowData, content_columns: set[int]) -> list[TableCellData]: + return [ + cell + for cell in row.cells + if cell.column_index not in content_columns + and cell.covered_by is None + and (cell.text.strip() or cell.explicit_empty) + ] + + +def _full_row_text(row: TableRowData) -> str: + return render_table_row(row) + + +def _compose_chunk(*parts: str) -> str: + return "\n\n".join(part.strip() for part in parts if part.strip()) + + +def _select_repeated_context( + row: TableRowData, + candidates: list[TableCellData], + *, + chunk_size: int, + length_function: Callable[[str], int], +) -> tuple[str, list[TableCellData]]: + """Keep only complete identity clauses that leave room for row content.""" + reserve = "Column 999999 (999999/999999):\nx" + full_scope = render_table_row_context(row, cells=[]) + if length_function(_compose_chunk(full_scope, reserve)) > chunk_size: + minimal_scope = f"Row {row.row_index}." + return ( + minimal_scope if length_function(_compose_chunk(minimal_scope, reserve)) <= chunk_size else "", + [], + ) + + selected: list[TableCellData] = [] + prefix = full_scope + for cell in candidates: + candidate = render_table_row_context(row, cells=[*selected, cell]) + if length_function(_compose_chunk(candidate, reserve)) <= chunk_size: + selected.append(cell) + prefix = candidate + return prefix, selected + + +def _select_heading( + *, + prefix: str, + cell: TableCellData, + chunk_size: int, + length_function: Callable[[str], int], +) -> str: + label = normalize_table_text(cell.column_name or f"Column {cell.column_index + 1}") + options = ( + f"Column “{label}” (999999/999999):", + f"Column {cell.column_index + 1} (999999/999999):", + "Value (999999/999999):", + "", + ) + for heading in options: + if length_function(_compose_chunk(prefix, heading, "x")) <= chunk_size: + return heading + return "" + + +def _part_heading(template: str, part: int, total: int) -> str: + return template.replace("999999/999999", f"{part}/{total}") + + +def _hard_split( + text: str, + start: int, + budget: int, + length_function: Callable[[str], int], +) -> list[tuple[str, int, int]]: + pieces: list[tuple[str, int, int]] = [] + cursor = 0 + while cursor < len(text): + remaining = text[cursor:] + if length_function(remaining) <= budget: + pieces.append((remaining.strip(), start + cursor, start + len(text))) + break + low, high = 1, len(remaining) + while low < high: + middle = (low + high + 1) // 2 + if length_function(remaining[:middle]) <= budget: + low = middle + else: + high = middle - 1 + cut = max(1, low) + whitespace = remaining.rfind(" ", 0, cut + 1) + if whitespace > 0: + cut = whitespace + piece = remaining[:cut].strip() + if piece: + leading = len(remaining[:cut]) - len(remaining[:cut].lstrip()) + pieces.append((piece, start + cursor + leading, start + cursor + cut)) + cursor += cut + while cursor < len(text) and text[cursor].isspace(): + cursor += 1 + return pieces + + +def _semantic_split( + text: str, + budget: int, + length_function: Callable[[str], int], + *, + hard_boundaries: set[int] | None = None, +) -> list[tuple[str, int, int]]: + if not text: + return [] + boundaries = [0] + boundaries.extend(match.end() for match in _SEMANTIC_BOUNDARY_RE.finditer(text)) + boundaries.extend(hard_boundaries or ()) + boundaries.append(len(text)) + boundaries = sorted(set(boundaries)) + units = [ + (text[start:end].strip(), start, end) + for start, end in zip(boundaries, boundaries[1:], strict=False) + if text[start:end].strip() + ] + + pieces: list[tuple[str, int, int]] = [] + current_text = "" + current_start = 0 + current_end = 0 + for unit, start, end in units: + if current_text and hard_boundaries and start in hard_boundaries: + pieces.append((current_text, current_start, current_end)) + current_text = "" + candidate = f"{current_text}\n\n{unit}".strip() if current_text else unit + if current_text and length_function(candidate) > budget: + pieces.append((current_text, current_start, current_end)) + current_text = "" + if length_function(unit) > budget: + pieces.extend(_hard_split(unit, start, budget, length_function)) + continue + if not current_text: + current_text = unit + current_start = start + else: + current_text = f"{current_text}\n\n{unit}" + current_end = end + if current_text: + pieces.append((current_text, current_start, current_end)) + return pieces + + +def _overlapping_fragments( + fragments: list[SourceFragment], + start: int, + end: int, +) -> list[SourceFragment]: + return [fragment for fragment in fragments if (fragment.text_end or 0) > start and fragment.text_start < end] + + +def _dump_fragments(fragments: list[SourceFragment]) -> list[dict[str, Any]]: + return [fragment.model_dump(mode="json") for fragment in fragments] + + +def _all_row_fragments(row: TableRowData) -> list[SourceFragment]: + return [ + *row.scope_fragments, + *(fragment for cell in row.cells for fragment in cell.source_fragments), + ] + + +def _piece_page_number( + fragments: list[SourceFragment], + start: int, + end: int, + *, + fallback: int, +) -> int: + overlapping = _overlapping_fragments(fragments, start, end) + if not overlapping: + return fallback + return max( + overlapping, + key=lambda fragment: ( + min(end, fragment.text_end or end) - max(start, fragment.text_start), + -fragment.text_start, + -fragment.page_number, + ), + ).page_number + + +def _metadata( + row: TableRowData, + *, + content_cell: TableCellData, + part: int, + total: int, + source_fragments: list[dict[str, Any]], +) -> dict[str, Any]: + same_table = [decision.same_table_confidence for decision in row.boundary_decisions] + continuation = [decision.row_continuation_confidence for decision in row.boundary_decisions] + return { + "table_id": row.table_id, + "row_id": row.row_id, + "row_index": row.row_index, + "table_title": row.table_title, + "section_path": row.section_path, + "page_start": row.page_start, + "page_end": row.page_end, + "table_content_kind": "row", + "content_column": content_cell.column_name, + "table_chunk_part": part, + "table_chunk_total": total, + "table_text_serialization_version": TABLE_TEXT_SERIALIZATION_VERSION, + "reconstruction_method": ( + "deterministic_adjacent_pages" if row.boundary_decisions else "deterministic_table_structure" + ), + "reconstruction_algorithm_version": row.algorithm_version, + "same_table_confidence": min(same_table) if same_table else 1.0, + "row_continuation_confidence": min(continuation) if continuation else 1.0, + "cell_assignment_confidence": min( + (cell.assignment_confidence for cell in row.cells), + default=1.0, + ), + "source_fragments": source_fragments, + } + + +def chunk_table_row( + row: TableRowData, + *, + chunk_size: int, + length_function: Callable[[str], int], +) -> list[TableRowChunk]: + """Serialize one logical row, splitting its largest cell when necessary.""" + if not row.cells: + return [] + full_text = _full_row_text(row) + visible_cells = [ + cell for cell in row.cells if cell.covered_by is None and (cell.text.strip() or cell.explicit_empty) + ] + if not visible_cells: + return [] + fallback_content = max(visible_cells, key=lambda cell: length_function(cell.text)) + identity_columns = set(row.identity_columns) + content_cells = [ + cell + for cell in visible_cells + if cell.column_index not in identity_columns and (cell.text.strip() or cell.explicit_empty) + ] + if not content_cells: + content_cells = [fallback_content] + identity_columns = {cell.column_index for cell in visible_cells if cell is not fallback_content} + if length_function(full_text) <= chunk_size: + primary_content = max(content_cells, key=lambda cell: length_function(cell.text)) + fragments = _dump_fragments(_all_row_fragments(row)) + return [ + TableRowChunk( + text=full_text, + page_number=_piece_page_number( + primary_content.source_fragments, + 0, + len(primary_content.text), + fallback=row.page_start, + ), + metadata=_metadata( + row, + content_cell=primary_content, + part=1, + total=1, + source_fragments=fragments, + ), + ) + ] + + initial_content_columns = {cell.column_index for cell in content_cells} + identity_candidates = _identity_cells(row, initial_content_columns) + prefix, repeated_identity = _select_repeated_context( + row, + identity_candidates, + chunk_size=chunk_size, + length_function=length_function, + ) + repeated_columns = {cell.column_index for cell in repeated_identity} + content_cells = [ + *content_cells, + *(cell for cell in identity_candidates if cell.column_index not in repeated_columns), + ] + + chunks: list[TableRowChunk] = [] + identity_fragments = _dump_fragments( + [ + *row.scope_fragments, + *(fragment for cell in repeated_identity for fragment in cell.source_fragments), + ] + ) + for content_cell in content_cells: + heading_template = _select_heading( + prefix=prefix, + cell=content_cell, + chunk_size=chunk_size, + length_function=length_function, + ) + content = content_cell.text if content_cell.text.strip() else "No value is present in this column." + + def rendered_length(value: str) -> int: + return length_function( + _compose_chunk( + prefix, + heading_template, + normalize_table_text(value), + ) + ) + + fragment_boundaries = { + boundary + for fragment in content_cell.source_fragments + for boundary in (fragment.text_start, fragment.text_end or len(content_cell.text)) + if 0 < boundary < len(content_cell.text) + } + pieces = _semantic_split( + content, + chunk_size, + rendered_length, + hard_boundaries=fragment_boundaries, + ) + total = len(pieces) + for part, (piece, start, end) in enumerate(pieces, start=1): + heading = _part_heading(heading_template, part, total) + text = _compose_chunk(prefix, heading, normalize_table_text(piece)) + content_fragments = _overlapping_fragments(content_cell.source_fragments, start, end) + metadata = _metadata( + row, + content_cell=content_cell, + part=part, + total=total, + source_fragments=[ + *identity_fragments, + *(fragment.model_dump(mode="json") for fragment in content_fragments), + ], + ) + metadata.update( + { + "context_columns": [ + cell.column_name or f"Column {cell.column_index + 1}" for cell in repeated_identity + ], + "deferred_context_columns": [ + cell.column_name or f"Column {cell.column_index + 1}" + for cell in identity_candidates + if cell.column_index not in repeated_columns + ], + } + ) + chunks.append( + TableRowChunk( + text=text, + page_number=_piece_page_number( + content_cell.source_fragments, + start, + end, + fallback=row.page_start, + ), + metadata=metadata, + ) + ) + return chunks + + +def chunk_table_legend( + legend: TableLegendData, + *, + chunk_size: int, + length_function: Callable[[str], int], +) -> list[TableRowChunk]: + """Serialize a table legend independently from the table's data rows.""" + rendered = render_table_legend(legend) + if not rendered: + return [] + + if length_function(rendered) <= chunk_size: + groups = [(legend, rendered)] + else: + groups: list[tuple[TableLegendData, str]] = [] + for entry in legend.entries: + + def full_render(meaning: str) -> str: + partial = legend.model_copy(update={"entries": [entry.model_copy(update={"meaning": meaning})]}) + return render_table_legend(partial) + + render_meaning: Callable[[str], str] + if length_function(full_render("x")) <= chunk_size: + render_meaning = full_render + elif length_function(f"{entry.abbreviation} means “x”.") <= chunk_size: + + def render_meaning(meaning: str) -> str: + return f"{entry.abbreviation} means “{normalize_table_text(meaning)}”." + elif length_function(f"{entry.abbreviation}: x") <= chunk_size: + + def render_meaning(meaning: str) -> str: + return f"{entry.abbreviation}: {normalize_table_text(meaning)}" + else: + render_meaning = normalize_table_text + + pieces = _semantic_split( + entry.meaning, + chunk_size, + lambda meaning: length_function(render_meaning(meaning)), + ) + groups.extend( + ( + legend.model_copy(update={"entries": [entry.model_copy(update={"meaning": piece})]}), + render_meaning(piece), + ) + for piece, _, _ in pieces + ) + + total = len(groups) + chunks: list[TableRowChunk] = [] + for part, (group, text) in enumerate(groups, start=1): + fragments = _dump_fragments( + [ + *group.scope_fragments, + *(fragment for entry in group.entries for fragment in entry.source_fragments), + ] + ) + chunks.append( + TableRowChunk( + text=text, + page_number=legend.page_number, + metadata={ + "table_id": legend.table_id, + "table_title": legend.table_title, + "section_path": legend.section_path, + "page_start": legend.page_number, + "page_end": legend.page_number, + "table_content_kind": "legend", + "table_chunk_part": part, + "table_chunk_total": total, + "table_text_serialization_version": TABLE_TEXT_SERIALIZATION_VERSION, + "reconstruction_algorithm_version": legend.algorithm_version, + "legend_abbreviations": [entry.abbreviation for entry in group.entries], + "source_fragments": fragments, + }, + ) + ) + return chunks + + +__all__ = ["TableRowChunk", "chunk_table_legend", "chunk_table_row"] diff --git a/openrag/core/config/indexation_pipeline.py b/openrag/core/config/indexation_pipeline.py index 5b53c3908..bda99c68d 100644 --- a/openrag/core/config/indexation_pipeline.py +++ b/openrag/core/config/indexation_pipeline.py @@ -10,6 +10,7 @@ from typing import Literal from core.config.chunking import ChunkerConfig +from core.config.table_reconstruction import TableReconstructionConfig from pydantic import BaseModel, ConfigDict, Field # PDF parsing backends a preset may explicitly select. ``None`` (the default) @@ -28,6 +29,7 @@ class IndexationPipelineConfig(BaseModel): chunking: ChunkerConfig = Field(default_factory=ChunkerConfig) # None => inherit the global PDFLOADER (see PARSING_STRATEGIES above). parsing_strategy: Literal["pymupdf", "marker", "docling"] | None = None + table_reconstruction: TableReconstructionConfig = Field(default_factory=TableReconstructionConfig) # VLM / image captioning vlm: str | None = None # endpoint name; None = use global default diff --git a/openrag/core/config/table_reconstruction.py b/openrag/core/config/table_reconstruction.py new file mode 100644 index 000000000..1af37760f --- /dev/null +++ b/openrag/core/config/table_reconstruction.py @@ -0,0 +1,29 @@ +"""Configuration for optional cross-page PDF table reconstruction.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +TABLE_RECONSTRUCTION_MODES: tuple[str, ...] = ("disabled", "automatic") +TABLE_RECONSTRUCTION_ALGORITHM_VERSION = "adjacent-layout-v1" + + +class TableReconstructionConfig(BaseModel): + """Conservative policy for reconstructing logical rows across PDF pages.""" + + model_config = ConfigDict(extra="forbid") + + mode: Literal["disabled", "automatic"] = "disabled" + same_table_min_confidence: float = Field(default=0.90, ge=0.80, le=1.00) + row_continuation_min_confidence: float = Field(default=0.90, ge=0.80, le=1.00) + cell_assignment_min_confidence: float = Field(default=0.90, ge=0.80, le=1.00) + algorithm_version: Literal["adjacent-layout-v1"] = TABLE_RECONSTRUCTION_ALGORITHM_VERSION + + +__all__ = [ + "TABLE_RECONSTRUCTION_ALGORITHM_VERSION", + "TABLE_RECONSTRUCTION_MODES", + "TableReconstructionConfig", +] diff --git a/openrag/core/indexing/contextualize.py b/openrag/core/indexing/contextualize.py index 897d528c8..e60c7c7cd 100644 --- a/openrag/core/indexing/contextualize.py +++ b/openrag/core/indexing/contextualize.py @@ -22,13 +22,22 @@ from tqdm.asyncio import tqdm from ..llm import LLM -from ..models.chunk import Chunk +from ..models.chunk import Chunk, ChunkType from ..prompts.contextualization_builder import build_messages, wrap_chunk_with_context logger = logging.getLogger(__name__) DEFAULT_TIMEOUT_SECONDS = 30.0 DEFAULT_BATCH_SIZE = 4 +_STRUCTURED_TABLE_CONTENT_KINDS = frozenset({"row", "legend"}) + + +def _is_structured_table_chunk(chunk: Chunk) -> bool: + """Return whether a chunk contains deterministic structured-table text.""" + return ( + chunk.chunk_type == ChunkType.TABLE + and chunk.metadata.get("table_content_kind") in _STRUCTURED_TABLE_CONTENT_KINDS + ) class ChunkContextualizer: @@ -88,10 +97,12 @@ async def contextualize( ) -> list[Chunk]: """Return new chunks with context prepended to ``text``. - Each returned chunk preserves the input's id, metadata, and other - fields; ``text`` is rewritten to the formatted (context + content) - string used for embedding, ``context`` holds the generated context, - and ``content`` holds the original chunk text. + Ordinary chunks preserve the input's id, metadata, and other fields; + ``text`` is rewritten to the formatted (context + content) string used + for embedding, ``context`` holds the generated context, and ``content`` + holds the original chunk text. Structured table rows and legends are + returned unchanged so LLM-generated context cannot mix their distinct + retrieval semantics. Falls back to returning the input chunks unchanged on any unrecoverable error. @@ -102,11 +113,13 @@ async def contextualize( try: first_chunks = chunks[:2] - contexts: list[str] = [] + eligible_indices = [index for index, chunk in enumerate(chunks) if not _is_structured_table_chunk(chunk)] + contexts: list[str | None] = [None] * len(chunks) # Schedule one batch at a time so prompt strings + coroutine # objects don't all sit in memory upfront on large documents. - for start in range(0, len(chunks), self._batch_size): - end = min(start + self._batch_size, len(chunks)) + for start in range(0, len(eligible_indices), self._batch_size): + end = min(start + self._batch_size, len(eligible_indices)) + batch_indices = eligible_indices[start:end] batch = [ self._generate_context( first_chunks=first_chunks, @@ -115,17 +128,19 @@ async def contextualize( filename=filename, lang=lang, ) - for i in range(start, end) + for i in batch_indices ] - contexts.extend( - await tqdm.gather( - *batch, - desc=f"Contextualizing chunks of *{filename}* [{start + 1}-{end}/{len(chunks)}]", - ) + generated = await tqdm.gather( + *batch, + desc=(f"Contextualizing chunks of *{filename}* [{start + 1}-{end}/{len(eligible_indices)}]"), ) + for index, context in zip(batch_indices, generated, strict=True): + contexts[index] = context return [ - chunk.model_copy( + chunk + if context is None + else chunk.model_copy( update={ "text": wrap_chunk_with_context( content=chunk.text, diff --git a/openrag/core/indexing/parsers/pdf/pymupdf.py b/openrag/core/indexing/parsers/pdf/pymupdf.py index 4915e79f5..1edde8b49 100644 --- a/openrag/core/indexing/parsers/pdf/pymupdf.py +++ b/openrag/core/indexing/parsers/pdf/pymupdf.py @@ -15,16 +15,14 @@ ``page.get_text`` / ``pymupdf4llm.to_markdown`` from different threads can raise ``ValueError: not a textpage of this page`` (upstream maintainer position: documented limitation, won't fix). We therefore -serialize all pymupdf work onto a single dedicated worker thread via -``_PYMUPDF_EXECUTOR``. The async ``parse`` method stays concurrent — -multiple callers will queue on the executor, but only one pymupdf -operation runs at a time. +serialize all PyMuPDF parsing and layout-evidence work onto the shared +executor in ``pymupdf_runtime``. The async ``parse`` method stays concurrent: +multiple callers queue on that executor, but only one PyMuPDF operation runs +at a time. """ from __future__ import annotations -import asyncio -from concurrent.futures import ThreadPoolExecutor from typing import Literal import pymupdf @@ -34,14 +32,12 @@ from ....models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock from ..document_parser import DocumentParser from ..registry import parser_registry +from .pymupdf_runtime import run_pymupdf ParseMode = Literal["markdown", "text"] logger = get_logger() -# Single dedicated worker for pymupdf — see "Threading note" in module docstring. -_PYMUPDF_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pymupdf") - def _extract_text(raw: bytes, filename: str) -> tuple[list[str], list[ImageBlock]]: """Return one stripped plain-text string per page; no images.""" @@ -111,9 +107,7 @@ async def parse(self, document: Document) -> ProcessedDocument: metadata=dict(document.metadata), ) - pages, images = await asyncio.get_running_loop().run_in_executor( - _PYMUPDF_EXECUTOR, self._extract, document.raw_bytes, document.filename - ) + pages, images = await run_pymupdf(self._extract, document.raw_bytes, document.filename) # Keep one TextBlock per source page (including empties) so callers # can preserve a 1-to-1 mapping with the original PDF's pagination. text_blocks = [TextBlock(text=text, page_number=i) for i, text in enumerate(pages, start=1)] diff --git a/openrag/core/indexing/parsers/pdf/pymupdf_runtime.py b/openrag/core/indexing/parsers/pdf/pymupdf_runtime.py new file mode 100644 index 000000000..f8eafd9e5 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/pymupdf_runtime.py @@ -0,0 +1,22 @@ +"""Serialized execution for all in-process PyMuPDF work.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from functools import partial +from typing import Any + +# PyMuPDF is not thread-safe. Parsing and structural evidence extraction must +# share this executor rather than each serializing only its own calls. +_PYMUPDF_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pymupdf") + + +async def run_pymupdf[T](function: Callable[..., T], /, *args: Any, **kwargs: Any) -> T: + """Run one callable on the process-wide serialized PyMuPDF executor.""" + call = partial(function, *args, **kwargs) + return await asyncio.get_running_loop().run_in_executor(_PYMUPDF_EXECUTOR, call) + + +__all__ = ["run_pymupdf"] diff --git a/openrag/core/indexing/structure_normalizer.py b/openrag/core/indexing/structure_normalizer.py new file mode 100644 index 000000000..0ddaf2b0b --- /dev/null +++ b/openrag/core/indexing/structure_normalizer.py @@ -0,0 +1,101 @@ +"""Contracts and evidence models for parser-independent structure normalization.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Literal + +from core.config.table_reconstruction import TableReconstructionConfig +from core.models.document import Document, ProcessedDocument + +NormalizedBBox = tuple[float, float, float, float] + + +@dataclass(slots=True, frozen=True) +class LayoutWord: + text: str + bbox: NormalizedBBox + block_number: int + line_number: int + word_number: int + + +@dataclass(slots=True, frozen=True) +class LayoutCellEvidence: + column_index: int + text: str + bbox: NormalizedBBox | None + slot_state: Literal["value", "explicit_empty", "covered", "unknown"] = "value" + column_span: int = 1 + row_span: int = 1 + covered_by: tuple[int, int] | None = None + + +@dataclass(slots=True, frozen=True) +class LayoutRowEvidence: + cells: tuple[LayoutCellEvidence, ...] + bbox: NormalizedBBox + + +@dataclass(slots=True, frozen=True) +class LayoutTableEvidence: + page_number: int + bbox: NormalizedBBox + column_bounds: tuple[tuple[float, float], ...] + rows: tuple[LayoutRowEvidence, ...] + + +@dataclass(slots=True, frozen=True) +class PageLayoutEvidence: + page_number: int + width: float + height: float + words: tuple[LayoutWord, ...] = field(default_factory=tuple) + tables: tuple[LayoutTableEvidence, ...] = field(default_factory=tuple) + + +class TableLayoutEvidenceProvider(ABC): + """Provide PDF geometry without deciding table or row relationships.""" + + provider_id = "layout" + + async def discover(self, document: Document) -> set[int]: + """Return one-based pages that may contain tables. + + Providers that cannot discover candidates cheaply may keep the empty + default and continue serving explicitly requested pages through + :meth:`collect`. + """ + return set() + + @abstractmethod + async def collect(self, document: Document, page_numbers: set[int]) -> list[PageLayoutEvidence]: + """Collect evidence for the requested one-based page numbers.""" + ... + + +class DocumentStructureNormalizer(ABC): + """Normalize parsed structure while retaining the existing document model.""" + + @abstractmethod + async def normalize( + self, + document: Document, + processed_document: ProcessedDocument, + config: TableReconstructionConfig, + ) -> ProcessedDocument: + """Return a processed document with an optional normalized block view.""" + ... + + +__all__ = [ + "DocumentStructureNormalizer", + "LayoutCellEvidence", + "LayoutRowEvidence", + "LayoutTableEvidence", + "LayoutWord", + "NormalizedBBox", + "PageLayoutEvidence", + "TableLayoutEvidenceProvider", +] diff --git a/openrag/core/indexing/table_normalizer.py b/openrag/core/indexing/table_normalizer.py new file mode 100644 index 000000000..627dec41e --- /dev/null +++ b/openrag/core/indexing/table_normalizer.py @@ -0,0 +1,2306 @@ +"""Deterministic reconstruction of logical table rows across adjacent PDF pages.""" + +from __future__ import annotations + +import hashlib +import html +import re +import unicodedata +from collections import Counter +from dataclasses import dataclass, field +from difflib import SequenceMatcher + +from core.config.table_reconstruction import TableReconstructionConfig +from core.indexing.structure_normalizer import ( + DocumentStructureNormalizer, + LayoutCellEvidence, + LayoutRowEvidence, + LayoutTableEvidence, + LayoutWord, + PageLayoutEvidence, + TableLayoutEvidenceProvider, +) +from core.indexing.table_text import render_table_legend, render_table_row +from core.models.document import ( + Document, + DocumentType, + NormalizationReport, + PageBoundaryDecision, + ProcessedDocument, + SourceFragment, + TableCellData, + TableLegendData, + TableLegendEntry, + TableRowData, + TextBlock, +) +from core.prompts.vlm_prompt_builder import wrap_caption + +_MARKDOWN_SEPARATOR_RE = re.compile(r"^:?-{3,}:?$") +_HEADING_RE = re.compile(r"(?m)^\s*(?P#{1,6})\s+(?P.+?)\s*$") +_SYNTHETIC_COLUMN_RE = re.compile(r"^col(?:umn)?\s*\d+$", re.IGNORECASE) +_ABBREVIATION_RE = re.compile(r"^([A-Z][A-Z0-9.-]{1,11})\s*[:=]\s*(.+)$") +_TABLE_TITLE_RE = re.compile(r"^(?:table(?:au)?\b|annexe\b|appendix\b)", re.IGNORECASE) +_HEADER_TERMS = { + "area", + "category", + "catégorie", + "city", + "description", + "document", + "documents", + "id", + "intitulé", + "libellé", + "name", + "number", + "owner", + "permit", + "pièce", + "pièces", + "reference", + "référence", + "region", + "status", + "title", + "titre", + "type", + "value", + "valeur", +} + + +@dataclass(slots=True, frozen=True) +class _MarkdownCell: + text: str + start: int + end: int + + +@dataclass(slots=True, frozen=True) +class _MarkdownRow: + cells: tuple[_MarkdownCell, ...] + start: int + end: int + separator: bool = False + + +@dataclass(slots=True, frozen=True) +class _MarkdownTable: + rows: tuple[_MarkdownRow, ...] + start: int + end: int + + +@dataclass(slots=True) +class _TableAlignment: + region: _AlignedRegion + rows: list[list[tuple[SourceFragment | None, float]]] + column_names: tuple[str, ...] | None = None + parser_header: bool = False + + +@dataclass(slots=True, frozen=True) +class _AlignedRegion: + source_block_index: int + raw_start: int + raw_end: int + working_start: int + working_end: int + page_number: int + kind: str + source_ref: str | None = None + confidence: float = 1.0 + + +@dataclass(slots=True) +class _RowAlignment: + cells: list[tuple[SourceFragment | None, float]] + regions: list[_AlignedRegion] + + +@dataclass(slots=True) +class _MutableCell: + column_index: int + column_name: str | None + parts: list[str] = field(default_factory=list) + source_fragments: list[SourceFragment] = field(default_factory=list) + assignment_confidence: float = 1.0 + column_span: int = 1 + row_span: int = 1 + inherited: bool = False + inherited_from: tuple[int, int] | None = None + explicit_empty: bool = False + covered_by: tuple[int, int] | None = None + + def append(self, text: str, fragments: list[SourceFragment], confidence: float) -> None: + clean = _clean_cell_text(text) + if not clean: + return + prefix_length = sum(len(part) for part in self.parts) + (2 * len(self.parts)) + self.parts.append(clean) + for fragment in fragments: + fragment_length = fragment.text_end - fragment.text_start if fragment.text_end is not None else len(clean) + self.source_fragments.append( + fragment.model_copy( + update={ + "text_start": prefix_length, + "text_end": prefix_length + max(fragment_length, len(clean)), + } + ) + ) + self.assignment_confidence = min(self.assignment_confidence, confidence) + + @property + def text(self) -> str: + return "\n\n".join(self.parts) + + def freeze(self) -> TableCellData: + return TableCellData( + column_index=self.column_index, + column_name=self.column_name, + text=self.text, + source_fragments=self.source_fragments, + assignment_confidence=self.assignment_confidence, + column_span=self.column_span, + row_span=self.row_span, + inherited=self.inherited, + inherited_from=self.inherited_from, + explicit_empty=self.explicit_empty, + covered_by=self.covered_by, + ) + + +@dataclass(slots=True) +class _MutableRow: + table_id: str + algorithm_version: str + table_title: str | None + section_path: list[str] + scope_fragments: list[SourceFragment] + cells: list[_MutableCell] + page_start: int + page_end: int + insertion_block_index: int + insertion_offset: int + row_index: int + boundary_decisions: list[PageBoundaryDecision] = field(default_factory=list) + + def merge( + self, + row: LayoutRowEvidence, + aligned_cells: list[tuple[SourceFragment | None, float]], + decision: PageBoundaryDecision, + ) -> None: + for cell_evidence, (fragment, confidence), target in zip(row.cells, aligned_cells, self.cells, strict=True): + fragments = [fragment] if fragment is not None else [] + target.append(cell_evidence.text, fragments, confidence) + self.page_end = max(self.page_end, decision.next_page) + self.boundary_decisions.append(decision) + + def freeze(self, identity_columns: tuple[int, ...]) -> TableRowData: + identity = "\x1f".join(self.cells[index].text for index in identity_columns) + row_hash = hashlib.sha256( + ( + f"{self.table_id}\x1e{identity}\x1e{self.page_start}\x1e" + f"{self.row_index}\x1e{self.insertion_block_index}\x1e{self.insertion_offset}" + ).encode() + ).hexdigest()[:20] + return TableRowData( + table_id=self.table_id, + row_id=f"row-{row_hash}", + algorithm_version=self.algorithm_version, + table_title=self.table_title, + section_path=self.section_path, + scope_fragments=self.scope_fragments, + cells=[cell.freeze() for cell in self.cells], + identity_columns=list(identity_columns), + row_index=self.row_index, + page_start=self.page_start, + page_end=self.page_end, + boundary_decisions=self.boundary_decisions, + ) + + +@dataclass(slots=True) +class _NormalizedChain: + rows: list[_MutableRow] + identity_columns: tuple[int, ...] + consumed_regions: list[_AlignedRegion] + decisions: list[PageBoundaryDecision] + used_tables: set[tuple[int, int]] + fallback_reasons: list[str] + legends: list[_LegendInsertion] = field(default_factory=list) + + +@dataclass(slots=True, frozen=True) +class _LegendInsertion: + legend: TableLegendData + insertion_block_index: int + insertion_offset: int + + +@dataclass(slots=True, frozen=True) +class _ScopeContext: + section_path: tuple[str, ...] + table_title: str | None + source_fragments: tuple[SourceFragment, ...] + title_fragment: SourceFragment | None = None + + +def _clean_cell_text(text: str) -> str: + text = html.unescape(text or "") + text = re.sub(r"", "\n", text, flags=re.IGNORECASE) + text = re.sub(r"<[^>]+>", " ", text) + text = re.sub(r"[ \t]+\n", "\n", text) + text = re.sub(r"\n[ \t]+", "\n", text) + text = re.sub(r"[ \t]{2,}", " ", text) + return text.strip() + + +def _normalised_alnum_with_offsets(text: str) -> tuple[str, list[int]]: + ignored = [False] * len(text) + for match in re.finditer(r"<[^>]*>", text): + for index in range(match.start(), match.end()): + ignored[index] = True + + chars: list[str] = [] + offsets: list[int] = [] + for index, character in enumerate(text): + if ignored[index]: + continue + for normalized in unicodedata.normalize("NFKC", character).casefold(): + if normalized.isalnum(): + chars.append(normalized) + offsets.append(index) + return "".join(chars), offsets + + +def _normalised_alnum(text: str) -> str: + return _normalised_alnum_with_offsets(text)[0] + + +def _similarity(expected: str, actual: str) -> float: + left = _normalised_alnum(expected) + right = _normalised_alnum(actual) + if not left and not right: + return 1.0 + if not left or not right: + return 0.0 + if left == right: + return 1.0 + return SequenceMatcher(None, left, right, autojunk=False).ratio() + + +def _find_normalised_range(needle: str, haystack: str) -> tuple[int, int, float] | None: + normalized_needle = _normalised_alnum(needle) + normalized_haystack, offsets = _normalised_alnum_with_offsets(haystack) + if not normalized_needle or not normalized_haystack: + return None + positions: list[int] = [] + start = 0 + while True: + found = normalized_haystack.find(normalized_needle, start) + if found < 0: + break + positions.append(found) + start = found + 1 + if len(positions) > 1: + break + if len(positions) != 1: + return None + normalized_start = positions[0] + normalized_end = normalized_start + len(normalized_needle) - 1 + return offsets[normalized_start], offsets[normalized_end] + 1, 1.0 + + +def _split_markdown_row(line: str, line_start: int) -> _MarkdownRow | None: + stripped_start = len(line) - len(line.lstrip()) + stripped = line.strip() + if not stripped.startswith("|") or not stripped.endswith("|"): + return None + + content_start = line_start + stripped_start + pipes: list[int] = [] + escaped = False + for index, character in enumerate(stripped): + if character == "\\" and not escaped: + escaped = True + continue + if character == "|" and not escaped: + pipes.append(index) + escaped = False + if len(pipes) < 2: + return None + + cells: list[_MarkdownCell] = [] + for left, right in zip(pipes, pipes[1:], strict=False): + raw = stripped[left + 1 : right] + leading = len(raw) - len(raw.lstrip()) + trailing = len(raw.rstrip()) + start = content_start + left + 1 + leading + end = content_start + left + 1 + trailing + cells.append(_MarkdownCell(text=raw.strip(), start=start, end=max(start, end))) + + separator = bool(cells) and all(_MARKDOWN_SEPARATOR_RE.fullmatch(cell.text.strip()) for cell in cells) + return _MarkdownRow( + cells=tuple(cells), + start=content_start, + end=content_start + len(stripped), + separator=separator, + ) + + +def _markdown_tables(text: str) -> list[_MarkdownTable]: + tables: list[_MarkdownTable] = [] + current: list[_MarkdownRow] = [] + offset = 0 + for line in text.splitlines(keepends=True): + row = _split_markdown_row(line.rstrip("\r\n"), offset) + if row is None: + if current: + if any(candidate.separator for candidate in current): + tables.append(_MarkdownTable(tuple(current), current[0].start, current[-1].end)) + current = [] + else: + current.append(row) + offset += len(line) + if current and any(candidate.separator for candidate in current): + tables.append(_MarkdownTable(tuple(current), current[0].start, current[-1].end)) + return tables + + +def _reliable_parser_header(row: _MarkdownRow) -> bool: + """Reject synthetic continuation rows that only satisfy Markdown syntax.""" + real_labels = [ + cell.text.strip() + for cell in row.cells + if cell.text.strip() + and not _SYNTHETIC_COLUMN_RE.fullmatch(cell.text.strip()) + and len(_normalised_alnum(cell.text)) <= 80 + ] + return len(real_labels) >= max(2, (len(row.cells) + 1) // 2) and not any( + re.search(r"\d", label) for label in real_labels + ) + + +def _parser_header_row(table: _MarkdownTable) -> _MarkdownRow | None: + for row, next_row in zip(table.rows, table.rows[1:], strict=False): + if not row.separator and next_row.separator and _reliable_parser_header(row): + return row + return None + + +def _unmatched_parser_rows_are_redundant( + *, + block_text: str, + table: _MarkdownTable, + content_rows: list[_MarkdownRow], + matched_rows: list[_MarkdownRow], +) -> bool: + unmatched = list(content_rows) + for matched in matched_rows: + for index, candidate in enumerate(unmatched): + if candidate is matched: + unmatched.pop(index) + break + + outside = f"{block_text[: table.start]}\n{block_text[table.end :]}" + for row in unmatched: + text = " ".join(cell.text for cell in row.cells) + if len(_normalised_alnum(text)) < 40: + return False + duplicated = _find_normalised_range(text, outside) is not None or any( + _similarity(text, paragraph) >= 0.90 for paragraph in re.split(r"\n{2,}", outside) if paragraph.strip() + ) + if not duplicated: + return False + return True + + +def _unique_exact_range(needle: str, haystack: str) -> tuple[int, int] | None: + if not needle: + return None + start = haystack.find(needle) + if start < 0 or haystack.find(needle, start + 1) >= 0: + return None + return start, start + len(needle) + + +def _image_regions( + processed_document: ProcessedDocument, + page_number: int, +) -> list[tuple[_AlignedRegion, str]]: + raw_blocks = processed_document.raw_text_blocks or [] + regions: list[tuple[_AlignedRegion, str]] = [] + for image in processed_document.images: + if image.page_number != page_number or not image.caption: + continue + markdown_ref = image.metadata.get("markdown_ref") + if not isinstance(markdown_ref, str) or not markdown_ref: + continue + + raw_matches: list[tuple[int, int, int]] = [] + for block_index, block in enumerate(raw_blocks): + if block.page_number != page_number: + continue + found = _unique_exact_range(markdown_ref, block.text) + if found is not None: + raw_matches.append((block_index, found[0], found[1])) + if len(raw_matches) != 1: + continue + + block_index, raw_start, raw_end = raw_matches[0] + if block_index >= len(processed_document.text_blocks): + continue + working = processed_document.text_blocks[block_index].text + wrapped_caption = wrap_caption(image.caption) + working_range = _unique_exact_range(wrapped_caption, working) + if working_range is None: + working_range = _unique_exact_range(markdown_ref, working) + if working_range is None: + continue + + source_ref = image.metadata.get("marker_key") + regions.append( + ( + _AlignedRegion( + source_block_index=block_index, + raw_start=raw_start, + raw_end=raw_end, + working_start=working_range[0], + working_end=working_range[1], + page_number=page_number, + kind="image_caption", + source_ref=str(source_ref or markdown_ref), + ), + image.caption, + ) + ) + return regions + + +def _map_parser_region( + processed_document: ProcessedDocument, + image_regions: dict[int, list[tuple[_AlignedRegion, str]]], + block_index: int, + start: int, + end: int, +) -> tuple[int, int] | None: + raw_blocks = processed_document.raw_text_blocks or [] + if block_index >= len(raw_blocks) or block_index >= len(processed_document.text_blocks): + return None + raw_block = raw_blocks[block_index] + working_block = processed_document.text_blocks[block_index] + direct = _map_interval(raw_block.text, working_block.text, start, end) + if direct is not None: + return direct + + replacements = sorted( + ( + region + for region, _ in image_regions.get(raw_block.page_number or 0, []) + if region.source_block_index == block_index + ), + key=lambda region: region.raw_start, + ) + if not replacements: + return None + + cursor = 0 + rebuilt: list[str] = [] + for region in replacements: + if region.raw_start < cursor: + return None + rebuilt.append(raw_block.text[cursor : region.raw_start]) + rebuilt.append(working_block.text[region.working_start : region.working_end]) + cursor = region.raw_end + rebuilt.append(raw_block.text[cursor:]) + if "".join(rebuilt) != working_block.text: + return None + + def map_offset(offset: int) -> int | None: + delta = 0 + for region in replacements: + if region.raw_start < offset < region.raw_end: + return None + if region.raw_end <= offset: + delta += (region.working_end - region.working_start) - (region.raw_end - region.raw_start) + return offset + delta + + mapped_start = map_offset(start) + mapped_end = map_offset(end) + if mapped_start is None or mapped_end is None: + return None + return mapped_start, mapped_end + + +def _evidence_overlap(expected: str, actual: str) -> float: + normalized_expected = _normalised_alnum(expected) + normalized_actual = _normalised_alnum(actual) + if not normalized_expected or not normalized_actual: + return 0.0 + width = 5 + + def shingles(value: str) -> Counter[str]: + if len(value) <= width: + return Counter({value: 1}) + return Counter(value[index : index + width] for index in range(len(value) - width + 1)) + + expected_shingles = shingles(normalized_expected) + actual_shingles = shingles(normalized_actual) + matched = sum((expected_shingles & actual_shingles).values()) + return min( + matched / expected_shingles.total(), + matched / actual_shingles.total(), + ) + + +def _match_image_region( + expected: str, + page_number: int, + image_regions: dict[int, list[tuple[_AlignedRegion, str]]], + threshold: float, +) -> _AlignedRegion | None: + # Short labels are too easy to associate with an unrelated page image. + if len(_normalised_alnum(expected)) < 80: + return None + candidates = [ + (score, region) + for region, caption in image_regions.get(page_number, []) + if (score := _evidence_overlap(expected, caption)) >= threshold + ] + if not candidates: + return None + candidates.sort(key=lambda candidate: candidate[0], reverse=True) + if len(candidates) > 1 and candidates[0][0] - candidates[1][0] < 0.05: + return None + score, region = candidates[0] + return _AlignedRegion( + source_block_index=region.source_block_index, + raw_start=region.raw_start, + raw_end=region.raw_end, + working_start=region.working_start, + working_end=region.working_end, + page_number=region.page_number, + kind=region.kind, + source_ref=region.source_ref, + confidence=score, + ) + + +def _layout_fragment( + cell: LayoutCellEvidence, + region: _AlignedRegion, + evidence_provider: str, +) -> SourceFragment | None: + if cell.bbox is None: + return None + return SourceFragment( + source_block_index=region.source_block_index, + page_number=region.page_number, + char_start=region.raw_start, + char_end=region.raw_end, + source_kind="pdf_layout", + evidence_provider=evidence_provider, + source_ref=region.source_ref, + bbox=cell.bbox, + text_start=0, + text_end=max(1, len(_clean_cell_text(cell.text))), + ) + + +def _markdown_cell_alignment( + evidence_cell: LayoutCellEvidence, + markdown_cell: _MarkdownCell, + *, + page_number: int, + block_index: int, + processed_document: ProcessedDocument, + image_regions: dict[int, list[tuple[_AlignedRegion, str]]], + threshold: float, +) -> tuple[float, _AlignedRegion | None]: + if not evidence_cell.text.strip(): + confidence = ( + 1.0 if not markdown_cell.text.strip() or _SYNTHETIC_COLUMN_RE.fullmatch(markdown_cell.text.strip()) else 0.0 + ) + return confidence, None + + # Table normalization is allowed to repair spacing and punctuation, but it + # must never replace parser-only content with a merely similar layout + # value. Require complete alphanumeric agreement for parser text. Image + # captions remain a separately evidenced replacement path below. + confidence = 1.0 if _normalised_alnum(evidence_cell.text) == _normalised_alnum(markdown_cell.text) else 0.0 + if confidence >= threshold and markdown_cell.end > markdown_cell.start: + return confidence, None + + image_region = _match_image_region( + evidence_cell.text, + page_number, + image_regions, + threshold, + ) + if ( + image_region is None + or image_region.source_block_index != block_index + or image_region.raw_start < markdown_cell.start + or image_region.raw_end > markdown_cell.end + ): + return confidence, None + return image_region.confidence, image_region + + +def _align_plain_table( + table: LayoutTableEvidence, + processed_document: ProcessedDocument, + image_regions: dict[int, list[tuple[_AlignedRegion, str]]], +) -> _TableAlignment | None: + """Align a layout table to parser text that contains no Markdown grid. + + This path is deliberately strict: every non-empty layout cell must occur + in order, and the complete consumed parser range must contain exactly + those cell values after alphanumeric normalization. + """ + raw_blocks = processed_document.raw_text_blocks or [] + values = [cell.text for row in table.rows for cell in row.cells if cell.text.strip()] + normalized_values = [_normalised_alnum(value) for value in values] + if not values or any(not value for value in normalized_values): + return None + normalized_table = "".join(normalized_values) + + candidates: list[tuple[int, int, int, list[tuple[int, int]]]] = [] + for block_index, block in enumerate(raw_blocks): + if block.page_number != table.page_number: + continue + markdown_ranges = [(candidate.start, candidate.end) for candidate in _markdown_tables(block.text)] + segment_starts = [0, *(end for _, end in markdown_ranges)] + segment_ends = [*(start for start, _ in markdown_ranges), len(block.text)] + for segment_start, segment_end in zip( + segment_starts, + segment_ends, + strict=True, + ): + segment = block.text[segment_start:segment_end] + normalized_block, offsets = _normalised_alnum_with_offsets(segment) + starts: list[int] = [] + cursor = 0 + while True: + found = normalized_block.find(normalized_table, cursor) + if found < 0: + break + starts.append(found) + cursor = found + 1 + if len(starts) > 1: + break + if len(starts) != 1: + continue + normalized_ranges: list[tuple[int, int]] = [] + cursor = starts[0] + for value in normalized_values: + normalized_ranges.append((cursor, cursor + len(value))) + cursor += len(value) + char_ranges = [ + ( + segment_start + offsets[start], + segment_start + offsets[end - 1] + 1, + ) + for start, end in normalized_ranges + ] + raw_start = char_ranges[0][0] + raw_end = char_ranges[-1][1] + mapped = _map_parser_region( + processed_document, + image_regions, + block_index, + raw_start, + raw_end, + ) + if mapped is not None: + candidates.append((block_index, raw_start, raw_end, char_ranges)) + + if len(candidates) != 1: + return None + block_index, raw_start, raw_end, char_ranges = candidates[0] + mapped = _map_parser_region( + processed_document, + image_regions, + block_index, + raw_start, + raw_end, + ) + if mapped is None: + return None + + range_cursor = 0 + aligned_rows: list[list[tuple[SourceFragment | None, float]]] = [] + for row in table.rows: + aligned_cells: list[tuple[SourceFragment | None, float]] = [] + for cell in row.cells: + if not cell.text.strip(): + aligned_cells.append((None, 1.0)) + continue + start, end = char_ranges[range_cursor] + range_cursor += 1 + aligned_cells.append( + ( + SourceFragment( + source_block_index=block_index, + page_number=table.page_number, + char_start=start, + char_end=end, + source_kind="parser_text", + bbox=cell.bbox, + text_start=0, + text_end=max(1, len(_clean_cell_text(cell.text))), + ), + 1.0, + ) + ) + aligned_rows.append(aligned_cells) + + return _TableAlignment( + region=_AlignedRegion( + source_block_index=block_index, + raw_start=raw_start, + raw_end=raw_end, + working_start=mapped[0], + working_end=mapped[1], + page_number=table.page_number, + kind="plain_table", + ), + rows=aligned_rows, + column_names=_column_names(table.rows[0]), + ) + + +def _align_table( + table: LayoutTableEvidence, + processed_document: ProcessedDocument, + image_regions: dict[int, list[tuple[_AlignedRegion, str]]], + evidence_provider: str, + threshold: float, +) -> _TableAlignment | None: + raw_blocks = processed_document.raw_text_blocks or [] + candidates: list[tuple[float, int, _MarkdownTable, list[_MarkdownRow], bool]] = [] + for block_index, block in enumerate(raw_blocks): + if block.page_number != table.page_number: + continue + markdown_tables = _markdown_tables(block.text) + for markdown_table in markdown_tables: + content_rows = [row for row in markdown_table.rows if not row.separator] + parser_header_row = _parser_header_row(markdown_table) + matched: list[_MarkdownRow] = [] + cursor = 0 + scores: list[float] = [] + for evidence_row in table.rows: + best_score = 0.0 + best_index: int | None = None + for row_index in range(cursor, len(content_rows)): + markdown_row = content_rows[row_index] + if len(evidence_row.cells) != len(markdown_row.cells): + continue + cell_scores = [ + _markdown_cell_alignment( + evidence_cell, + markdown_cell, + page_number=table.page_number, + block_index=block_index, + processed_document=processed_document, + image_regions=image_regions, + threshold=threshold, + )[0] + for evidence_cell, markdown_cell in zip( + evidence_row.cells, + markdown_row.cells, + strict=True, + ) + ] + if any(cell_score < threshold for cell_score in cell_scores): + continue + score = sum(cell_scores) / len(cell_scores) if cell_scores else 0.0 + if score > best_score: + best_score = score + best_index = row_index + if best_index is None: + break + matched.append(content_rows[best_index]) + scores.append(best_score) + cursor = best_index + 1 + if len(matched) == len(table.rows) and _unmatched_parser_rows_are_redundant( + block_text=block.text, + table=markdown_table, + content_rows=content_rows, + matched_rows=matched, + ): + candidates.append( + ( + sum(scores) / len(scores), + block_index, + markdown_table, + matched, + bool(matched and matched[0] is parser_header_row), + ) + ) + + if not candidates: + table_text = "\n".join(cell.text for row in table.rows for cell in row.cells if cell.text.strip()) + image_region = _match_image_region(table_text, table.page_number, image_regions, threshold) + if image_region is None: + return _align_plain_table( + table, + processed_document, + image_regions, + ) + return _TableAlignment( + region=image_region, + rows=[ + [ + (None, 1.0) + if not cell.text.strip() + else ( + _layout_fragment(cell, image_region, evidence_provider), + image_region.confidence, + ) + for cell in row.cells + ] + for row in table.rows + ], + ) + + score, block_index, markdown_table, matched_rows, parser_header = max( + candidates, + key=lambda candidate: candidate[0], + ) + if score < threshold: + return None + if block_index >= len(processed_document.text_blocks): + return None + mapped_table = _map_parser_region( + processed_document, + image_regions, + block_index, + markdown_table.start, + markdown_table.end, + ) + if mapped_table is None: + return None + table_region = _AlignedRegion( + source_block_index=block_index, + raw_start=markdown_table.start, + raw_end=markdown_table.end, + working_start=mapped_table[0], + working_end=mapped_table[1], + page_number=table.page_number, + kind="markdown_table", + confidence=score, + ) + + aligned_rows: list[list[tuple[SourceFragment | None, float]]] = [] + for evidence_row, markdown_row in zip(table.rows, matched_rows, strict=True): + aligned_cells: list[tuple[SourceFragment | None, float]] = [] + for evidence_cell, markdown_cell in zip(evidence_row.cells, markdown_row.cells, strict=True): + if not evidence_cell.text.strip(): + confidence, _ = _markdown_cell_alignment( + evidence_cell, + markdown_cell, + page_number=table.page_number, + block_index=block_index, + processed_document=processed_document, + image_regions=image_regions, + threshold=threshold, + ) + aligned_cells.append((None, confidence)) + continue + confidence, image_region = _markdown_cell_alignment( + evidence_cell, + markdown_cell, + page_number=table.page_number, + block_index=block_index, + processed_document=processed_document, + image_regions=image_regions, + threshold=threshold, + ) + if image_region is not None: + aligned_cells.append( + ( + _layout_fragment(evidence_cell, image_region, evidence_provider), + image_region.confidence, + ) + ) + continue + if confidence < threshold or markdown_cell.end <= markdown_cell.start: + aligned_cells.append((None, confidence)) + continue + uses_layout_text = evidence_cell.bbox is not None and _clean_cell_text( + evidence_cell.text + ) != _clean_cell_text(markdown_cell.text) + aligned_cells.append( + ( + SourceFragment( + source_block_index=block_index, + page_number=table.page_number, + char_start=markdown_cell.start, + char_end=markdown_cell.end, + source_kind=("pdf_layout" if uses_layout_text else "parser_text"), + evidence_provider=(evidence_provider if uses_layout_text else None), + bbox=evidence_cell.bbox, + text_start=0, + text_end=max(1, len(_clean_cell_text(evidence_cell.text))), + ), + confidence, + ) + ) + aligned_rows.append(aligned_cells) + + column_names: list[str] = [] + for column_index, evidence_cell in enumerate(table.rows[0].cells): + # Layout words preserve the visual boundary between a concise header and + # any legend printed beneath it. Parser Markdown may collapse those + # lines into one key even when its alphanumeric content still aligns. + column_names.append(_concise_column_name(evidence_cell.text, column_index)) + + return _TableAlignment( + region=table_region, + rows=aligned_rows, + column_names=tuple(column_names), + parser_header=parser_header, + ) + + +def _looks_like_header( + row: LayoutRowEvidence, + following_row: LayoutRowEvidence | None = None, +) -> bool: + nonempty = sum(bool(cell.text.strip()) for cell in row.cells) + if nonempty < max(2, (len(row.cells) + 1) // 2): + return False + matched_labels = sum( + bool(set(re.findall(r"\w+", cell.text.casefold())) & _HEADER_TERMS) and len(_normalised_alnum(cell.text)) <= 200 + for cell in row.cells + ) + if matched_labels >= 2: + return True + + # Generic column labels such as ``aa | bb | cc`` have no semantic header + # vocabulary. Treat them as a header only when the next row supplies a + # strong data-type contrast. This preserves the first row of headerless + # tables such as ``22 | Paris | Active``. + if following_row is None or len(row.cells) != len(following_row.cells): + return False + labels = [_clean_cell_text(cell.text) for cell in row.cells] + if not all( + label and len(label) <= 40 and re.fullmatch(r"[\wÀ-ÖØ-öø-ÿ .()/+-]+", label) and not re.search(r"\d", label) + for label in labels + ): + return False + return any( + bool(re.search(r"\d", following.text)) and not bool(re.search(r"\d", header.text)) + for header, following in zip(row.cells, following_row.cells, strict=True) + ) + + +def _content_column(rows: tuple[LayoutRowEvidence, ...]) -> int: + column_count = len(rows[0].cells) + lengths = [0] * column_count + for row in rows: + for cell in row.cells: + lengths[cell.column_index] += len(_normalised_alnum(cell.text)) + return max(range(column_count), key=lambda index: lengths[index]) + + +def _concise_column_name(text: str, column_index: int) -> str: + clean = _clean_cell_text(text) + lines = [line.strip() for line in clean.splitlines() if line.strip()] + concise = lines[0] if lines else "" + for continuation in lines[1:]: + candidate = f"{concise} {continuation}".strip() + if ":" in continuation or len(candidate) > 80: + break + concise = candidate + return concise or f"Column {column_index + 1}" + + +def _column_names(header: LayoutRowEvidence) -> tuple[str, ...]: + return tuple(_concise_column_name(cell.text, cell.column_index) for cell in header.cells) + + +def _is_repeated_header( + row: LayoutRowEvidence, + column_names: tuple[str, ...], +) -> bool: + """Recognize only complete, exact repetitions of the established header.""" + if len(row.cells) != len(column_names): + return False + return all( + bool(_normalised_alnum(column_name)) + and _normalised_alnum(_concise_column_name(cell.text, cell.column_index)) == _normalised_alnum(column_name) + for cell, column_name in zip(row.cells, column_names, strict=True) + ) + + +def _legend_entries( + header: LayoutRowEvidence, + aligned_header: list[tuple[SourceFragment | None, float]], +) -> list[TableLegendEntry]: + entries: list[TableLegendEntry] = [] + seen: set[str] = set() + for cell, (fragment, _) in zip(header.cells, aligned_header, strict=True): + lines = [line.strip() for line in _clean_cell_text(cell.text).splitlines() if line.strip()] + current_abbreviation: str | None = None + current_meaning: list[str] = [] + + def flush() -> None: + nonlocal current_abbreviation, current_meaning + if current_abbreviation is None: + return + meaning = " ".join(current_meaning).strip() + if meaning and current_abbreviation not in seen: + fragments = [fragment.model_copy(deep=True)] if fragment is not None else [] + entries.append( + TableLegendEntry( + abbreviation=current_abbreviation, + meaning=meaning, + source_fragments=fragments, + ) + ) + seen.add(current_abbreviation) + current_abbreviation = None + current_meaning = [] + + for line in lines: + match = _ABBREVIATION_RE.fullmatch(line) + if match is not None: + flush() + current_abbreviation = match.group(1) + current_meaning = [match.group(2)] + elif current_abbreviation is not None: + current_meaning.append(line) + flush() + return entries + + +def _spans_are_unambiguous(table: LayoutTableEvidence, *, header_rows: int) -> bool: + if any(cell.slot_state == "unknown" for row in table.rows for cell in row.cells): + return False + return not any( + cell.column_span > 1 or cell.row_span > 1 or cell.slot_state == "covered" + for row in table.rows[:header_rows] + for cell in row.cells + ) + + +def _apply_merged_cell_inheritance( + rows: list[_MutableRow], + evidence_rows: tuple[LayoutRowEvidence, ...], + *, + evidence_row_offset: int, +) -> bool: + """Resolve page-local covered slots without guessing across boundaries.""" + for logical_index, (row, evidence) in enumerate(zip(rows, evidence_rows, strict=True)): + for cell_evidence, target in zip(evidence.cells, row.cells, strict=True): + if cell_evidence.slot_state != "covered": + continue + if cell_evidence.covered_by is None: + return False + anchor_evidence_row, anchor_column = cell_evidence.covered_by + anchor_logical_index = anchor_evidence_row - evidence_row_offset + if anchor_logical_index == logical_index: + target.covered_by = (row.row_index, anchor_column) + continue + if not 0 <= anchor_logical_index < logical_index: + return False + anchor = rows[anchor_logical_index].cells[anchor_column] + if not anchor.text or anchor.row_span <= logical_index - anchor_logical_index: + return False + target.parts = list(anchor.parts) + target.source_fragments = [fragment.model_copy(deep=True) for fragment in anchor.source_fragments] + target.assignment_confidence = anchor.assignment_confidence + target.column_span = anchor.column_span + target.row_span = 1 + target.inherited = True + target.inherited_from = ( + rows[anchor_logical_index].row_index, + anchor_column, + ) + target.explicit_empty = False + target.covered_by = None + return True + + +def _table_geometry_confidence( + expected: tuple[tuple[float, float], ...], + actual: tuple[tuple[float, float], ...], +) -> float: + if len(expected) != len(actual): + return 0.0 + delta = max( + abs(left - right) + for pair_a, pair_b in zip(expected, actual, strict=True) + for left, right in zip(pair_a, pair_b, strict=True) + ) + return max(0.0, 1.0 - (delta / 0.10)) + + +def _compatible_pages(previous: PageLayoutEvidence, following: PageLayoutEvidence) -> bool: + previous_landscape = previous.width > previous.height + following_landscape = following.width > following.height + if previous_landscape != following_landscape: + return False + previous_ratio = previous.width / previous.height + following_ratio = following.width / following.height + return abs(previous_ratio - following_ratio) / max(previous_ratio, following_ratio) <= 0.05 + + +def _page_starts_with_heading(raw_blocks: list[TextBlock], page_number: int) -> bool: + for block in raw_blocks: + if block.page_number != page_number: + continue + first_content = block.text.lstrip() + if first_content: + return bool(re.match(r"^#{1,6}\s+", first_content)) + return False + + +def _has_parser_content_before_table( + raw_blocks: list[TextBlock], + page_number: int, +) -> bool: + preceding_content = False + for block in raw_blocks: + if block.page_number != page_number: + continue + tables = _markdown_tables(block.text) + if tables: + return preceding_content or bool(block.text[: tables[0].start].strip()) + preceding_content = preceding_content or bool(block.text.strip()) + return False + + +def _page_edge_confidences( + previous_bottom: float, + next_top: float, +) -> tuple[float, float]: + """Score how strongly the two regions touch their respective page edges.""" + bottom = 0.90 + 0.10 * min(1.0, max(0.0, (previous_bottom - 0.82) / 0.18)) + top = 0.90 + 0.10 * min(1.0, max(0.0, (0.15 - next_top) / 0.15)) + return bottom, top + + +def _has_content_before_regions( + raw_blocks: list[TextBlock], + regions: list[_AlignedRegion], +) -> bool: + """Reject continuation evidence preceded by unaccounted page content.""" + if not regions: + return True + first = min( + regions, + key=lambda region: (region.source_block_index, region.raw_start), + ) + for block_index, block in enumerate(raw_blocks): + if block.page_number != first.page_number: + continue + if block_index < first.source_block_index and block.text.strip(): + return True + if block_index == first.source_block_index: + return bool(block.text[: first.raw_start].strip()) + return True + + +def _words_to_text(words: list[LayoutWord]) -> str: + lines: list[str] = [] + current_key: tuple[int, int] | None = None + current_words: list[str] = [] + previous_block: int | None = None + for word in sorted(words, key=lambda item: (item.block_number, item.line_number, item.word_number)): + key = (word.block_number, word.line_number) + if key != current_key: + if current_words: + lines.append(" ".join(current_words)) + if previous_block is not None and word.block_number != previous_block: + lines.append("") + current_words = [] + current_key = key + previous_block = word.block_number + current_words.append(word.text) + if current_words: + lines.append(" ".join(current_words)) + return "\n".join(lines).strip() + + +def _sparse_continuation( + page: PageLayoutEvidence, + column_bounds: tuple[tuple[float, float], ...], + identity_columns: tuple[int, ...], +) -> LayoutRowEvidence | None: + grouped: list[list[LayoutWord]] = [[] for _ in column_bounds] + for word in page.words: + x0, y0, x1, y1 = word.bbox + if y1 >= 0.90: + continue + center = (x0 + x1) / 2 + for column_index, (left, right) in enumerate(column_bounds): + if left <= center <= right: + grouped[column_index].append(word) + break + + if any(grouped[index] for index in identity_columns): + return None + populated = [index for index, words in enumerate(grouped) if words] + if not populated: + return None + all_words = [word for words in grouped for word in words] + if min(word.bbox[1] for word in all_words) > 0.12: + return None + + cells: list[LayoutCellEvidence] = [] + for column_index, words in enumerate(grouped): + if words: + bbox = ( + min(word.bbox[0] for word in words), + min(word.bbox[1] for word in words), + max(word.bbox[2] for word in words), + max(word.bbox[3] for word in words), + ) + text = _words_to_text(words) + else: + bbox = None + text = "" + cells.append(LayoutCellEvidence(column_index=column_index, text=text, bbox=bbox)) + + return LayoutRowEvidence( + cells=tuple(cells), + bbox=( + min(word.bbox[0] for word in all_words), + min(word.bbox[1] for word in all_words), + max(word.bbox[2] for word in all_words), + max(word.bbox[3] for word in all_words), + ), + ) + + +def _align_sparse_row( + row: LayoutRowEvidence, + page_number: int, + processed_document: ProcessedDocument, + image_regions: dict[int, list[tuple[_AlignedRegion, str]]], + evidence_provider: str, + threshold: float, +) -> _RowAlignment | None: + raw_blocks = processed_document.raw_text_blocks or [] + aligned: list[tuple[SourceFragment | None, float]] = [] + regions: list[_AlignedRegion] = [] + direct_alignment_failed = False + for cell in row.cells: + if not cell.text: + aligned.append((None, 1.0)) + continue + matches: list[tuple[int, int, int, float]] = [] + for block_index, block in enumerate(raw_blocks): + if block.page_number != page_number: + continue + found = _find_normalised_range(cell.text, block.text) + if found is not None: + matches.append((block_index, *found)) + if len(matches) != 1: + direct_alignment_failed = True + break + block_index, start, end, confidence = matches[0] + source_text = raw_blocks[block_index].text + while start > 0 and not source_text[start - 1].isalnum(): + start -= 1 + while end < len(source_text) and not source_text[end].isalnum(): + end += 1 + if block_index >= len(processed_document.text_blocks): + direct_alignment_failed = True + break + mapped = _map_parser_region( + processed_document, + image_regions, + block_index, + start, + end, + ) + if mapped is None: + direct_alignment_failed = True + break + aligned.append( + ( + SourceFragment( + source_block_index=block_index, + page_number=page_number, + char_start=start, + char_end=end, + bbox=cell.bbox, + text_start=0, + text_end=max(1, len(_clean_cell_text(cell.text))), + ), + confidence, + ) + ) + regions.append( + _AlignedRegion( + source_block_index=block_index, + raw_start=start, + raw_end=end, + working_start=mapped[0], + working_end=mapped[1], + page_number=page_number, + kind="parser_text", + confidence=confidence, + ) + ) + if not direct_alignment_failed: + return _RowAlignment(cells=aligned, regions=regions) + + expected = "\n".join(cell.text for cell in row.cells if cell.text.strip()) + image_region = _match_image_region( + expected, + page_number, + image_regions, + threshold, + ) + if image_region is None: + return None + return _RowAlignment( + cells=[ + (None, 1.0) + if not cell.text.strip() + else ( + _layout_fragment(cell, image_region, evidence_provider), + image_region.confidence, + ) + for cell in row.cells + ], + regions=[image_region], + ) + + +def _scope_fragment( + *, + block: TextBlock, + block_index: int, + start: int, + end: int, + fallback_page: int, + text: str, +) -> SourceFragment: + return SourceFragment( + source_block_index=block_index, + page_number=block.page_number or fallback_page, + char_start=start, + char_end=end, + text_start=0, + text_end=max(1, len(text)), + ) + + +def _scope_context( + raw_blocks: list[TextBlock], + *, + anchor_block_index: int, + anchor_offset: int, + fallback_page: int, +) -> _ScopeContext: + """Resolve heading ancestry and a separately evidenced table caption.""" + headings: list[tuple[int, str, SourceFragment]] = [] + title_candidate: tuple[str, SourceFragment, int, int] | None = None + last_heading_end = 0 + + for block_index, block in enumerate(raw_blocks[: anchor_block_index + 1]): + limit = anchor_offset if block_index == anchor_block_index else len(block.text) + for match in _HEADING_RE.finditer(block.text[:limit]): + level = len(match.group("marks")) + raw_text = match.group("text") + clean = _clean_cell_text(raw_text).strip("*_ ") + if not clean: + continue + while headings and headings[-1][0] >= level: + headings.pop() + fragment = _scope_fragment( + block=block, + block_index=block_index, + start=match.start("text"), + end=match.end("text"), + fallback_page=fallback_page, + text=clean, + ) + if _TABLE_TITLE_RE.match(clean): + title_candidate = (clean, fragment, block_index, match.end()) + else: + headings.append((level, clean, fragment)) + title_candidate = None + if block_index == anchor_block_index: + last_heading_end = match.end() + + anchor_block = raw_blocks[anchor_block_index] + cursor = last_heading_end + for line in anchor_block.text[last_heading_end:anchor_offset].splitlines(keepends=True): + raw_line = line.rstrip("\r\n") + stripped = raw_line.strip() + clean = _clean_cell_text(stripped).strip("*_ ") + if clean and len(clean) <= 120 and _TABLE_TITLE_RE.match(clean): + leading = len(raw_line) - len(raw_line.lstrip()) + start = cursor + leading + title_candidate = ( + clean, + _scope_fragment( + block=anchor_block, + block_index=anchor_block_index, + start=start, + end=start + len(stripped), + fallback_page=fallback_page, + text=clean, + ), + anchor_block_index, + cursor + len(line), + ) + cursor += len(line) + + title: tuple[str, SourceFragment] | None = None + if title_candidate is not None: + clean, fragment, block_index, end = title_candidate + intervening = [raw_blocks[block_index].text[end:]] + intervening.extend(block.text for block in raw_blocks[block_index + 1 : anchor_block_index]) + if block_index != anchor_block_index: + intervening.append(anchor_block.text[:anchor_offset]) + else: + intervening[0] = anchor_block.text[end:anchor_offset] + # Introductory prose between a caption and its first table is common. + # The caption stops applying once another Markdown table has appeared, + # preventing a later untitled table from inheriting a stale name. + if not _markdown_tables("\n".join(intervening)): + title = (clean, fragment) + + fragments = [fragment for _, _, fragment in headings] + if title is not None: + fragments.append(title[1]) + return _ScopeContext( + section_path=tuple(text for _, text, _ in headings), + table_title=title[0] if title is not None else None, + source_fragments=tuple(fragments), + title_fragment=title[1] if title is not None else None, + ) + + +def _row_from_evidence( + *, + table_id: str, + algorithm_version: str, + table_title: str | None, + section_path: list[str], + scope_fragments: list[SourceFragment], + column_names: tuple[str, ...], + evidence: LayoutRowEvidence, + aligned: list[tuple[SourceFragment | None, float]], + page_number: int, + insertion_block_index: int, + insertion_offset: int, + row_index: int, +) -> _MutableRow: + cells: list[_MutableCell] = [] + for name, cell_evidence, (fragment, confidence) in zip(column_names, evidence.cells, aligned, strict=True): + cell = _MutableCell( + column_index=cell_evidence.column_index, + column_name=name, + column_span=cell_evidence.column_span, + row_span=cell_evidence.row_span, + explicit_empty=cell_evidence.slot_state == "explicit_empty", + covered_by=cell_evidence.covered_by, + ) + cell.append(cell_evidence.text, [fragment] if fragment is not None else [], confidence) + cells.append(cell) + return _MutableRow( + table_id=table_id, + algorithm_version=algorithm_version, + table_title=table_title, + section_path=section_path, + scope_fragments=scope_fragments, + cells=cells, + page_start=page_number, + page_end=page_number, + insertion_block_index=insertion_block_index, + insertion_offset=insertion_offset, + row_index=row_index, + ) + + +def _all_assignments_pass( + row: LayoutRowEvidence, + aligned: list[tuple[SourceFragment | None, float]], + threshold: float, +) -> bool: + return all( + confidence >= threshold and (not cell.text.strip() or fragment is not None) + for cell, (fragment, confidence) in zip(row.cells, aligned, strict=True) + ) + + +class DeterministicTableNormalizer(DocumentStructureNormalizer): + """Reconstruct high-confidence adjacent-page row continuations.""" + + def __init__(self, evidence_provider: TableLayoutEvidenceProvider) -> None: + self._evidence_provider = evidence_provider + + async def normalize( + self, + document: Document, + processed_document: ProcessedDocument, + config: TableReconstructionConfig, + ) -> ProcessedDocument: + if config.mode == "disabled" or document.content_type is not DocumentType.PDF: + return processed_document + raw_blocks = processed_document.raw_text_blocks + if not raw_blocks or not document.raw_bytes: + return self._unchanged(processed_document, config, "raw parser blocks or PDF bytes are unavailable") + + table_pages = { + block.page_number for block in raw_blocks if block.page_number is not None and _markdown_tables(block.text) + } + table_pages.update( + page + for page in await self._evidence_provider.discover(document) + if 1 <= page <= processed_document.page_count + ) + candidate_image_pages = { + image.page_number + for image in processed_document.images + if image.page_number is not None + and image.caption + and len(_normalised_alnum(image.caption)) >= 80 + and image.metadata.get("markdown_ref") + } + image_regions = { + page_number: regions + for page_number in candidate_image_pages + if (regions := _image_regions(processed_document, page_number)) + } + table_pages.update(image_regions) + if not table_pages: + return self._unchanged(processed_document, config, "no parser-level table candidates were found") + + candidate_pages = { + page + for table_page in table_pages + for page in (table_page - 1, table_page, table_page + 1) + if 1 <= page <= processed_document.page_count + } + evidence = await self._evidence_provider.collect(document, candidate_pages) + pages = {page.page_number: page for page in evidence} + + chains: list[_NormalizedChain] = [] + used_tables: set[tuple[int, int]] = set() + claimed_captions: set[tuple[int, int, int]] = set() + preserved_candidates: list[str] = [] + for page_number in sorted(pages): + page = pages[page_number] + for table_index, table in enumerate(page.tables): + key = (page_number, table_index) + if key in used_tables: + continue + chain = await self._build_chain( + document, + table, + table_index, + pages, + processed_document, + image_regions, + config, + processed_document.page_count, + claimed_captions, + ) + if chain is None: + preserved_candidates.append( + f"preserved table candidate {page_number}:{table_index} because a confidence or alignment gate failed" + ) + continue + chains.append(chain) + used_tables.update(chain.used_tables) + + if not chains: + return self._unchanged(processed_document, config, "no table candidate passed all confidence gates") + + rows: list[tuple[_MutableRow, tuple[int, ...]]] = [] + legends: list[_LegendInsertion] = [] + consumed: list[_AlignedRegion] = [] + decisions: list[PageBoundaryDecision] = [] + fallback_reasons: list[str] = list(preserved_candidates) + for chain in chains: + rows.extend((row, chain.identity_columns) for row in chain.rows) + consumed.extend(chain.consumed_regions) + decisions.extend(chain.decisions) + fallback_reasons.extend(chain.fallback_reasons) + legends.extend(chain.legends) + + if not _regions_are_disjoint(consumed): + return self._unchanged(processed_document, config, "candidate table overlays overlap") + + normalized_blocks = _build_normalized_blocks( + processed_document, + rows, + legends, + consumed, + ) + if normalized_blocks is None: + return self._unchanged(processed_document, config, "raw and captioned block ranges could not be aligned") + + status = "partial_fallback" if fallback_reasons else "normalized" + return processed_document.model_copy( + update={ + "normalized_text_blocks": normalized_blocks, + "normalization_report": NormalizationReport( + algorithm_version=config.algorithm_version, + status=status, + boundary_decisions=decisions, + reconstructed_row_count=len(rows), + fallback_reasons=fallback_reasons, + ), + } + ) + + async def _build_chain( + self, + document: Document, + anchor: LayoutTableEvidence, + anchor_index: int, + pages: dict[int, PageLayoutEvidence], + processed_document: ProcessedDocument, + image_regions: dict[int, list[tuple[_AlignedRegion, str]]], + config: TableReconstructionConfig, + page_count: int, + claimed_captions: set[tuple[int, int, int]], + ) -> _NormalizedChain | None: + raw_blocks = processed_document.raw_text_blocks or [] + if not anchor.rows: + return None + + alignment = _align_table( + anchor, + processed_document, + image_regions, + self._evidence_provider.provider_id, + config.cell_assignment_min_confidence, + ) + if alignment is None: + return None + has_header = alignment.parser_header or ( + len(anchor.rows) >= 2 + and _looks_like_header( + anchor.rows[0], + anchor.rows[1], + ) + ) + if not has_header: + first_nonempty = sum(bool(cell.text.strip()) for cell in anchor.rows[0].cells) + if first_nonempty < max(2, (len(anchor.rows[0].cells) + 1) // 2): + # A sparse leading row at the top of a page is likely a + # continuation, not a safe headerless table anchor. + return None + header_rows = 1 if has_header else 0 + if not _spans_are_unambiguous(anchor, header_rows=header_rows): + return None + data_rows = anchor.rows[header_rows:] + aligned_data = alignment.rows[header_rows:] + if not data_rows: + return None + if any( + not _all_assignments_pass(row, aligned, config.cell_assignment_min_confidence) + for row, aligned in zip(data_rows, aligned_data, strict=True) + ): + return None + + column_names = ( + (alignment.column_names or _column_names(anchor.rows[0])) + if has_header + else tuple(f"Column {index + 1}" for index in range(len(anchor.rows[0].cells))) + ) + content_column = _content_column(data_rows) + identity_columns = tuple(index for index in range(len(column_names)) if index != content_column) + scope = _scope_context( + raw_blocks, + anchor_block_index=alignment.region.source_block_index, + anchor_offset=alignment.region.raw_start, + fallback_page=anchor.page_number, + ) + sections = list(scope.section_path) + table_title = scope.table_title + scope_fragments = [fragment.model_copy(deep=True) for fragment in scope.source_fragments] + caption_key = ( + ( + scope.title_fragment.source_block_index, + scope.title_fragment.char_start, + scope.title_fragment.char_end, + ) + if scope.title_fragment is not None + else None + ) + if caption_key is not None and caption_key in claimed_captions: + table_title = None + scope_fragments = [ + fragment + for fragment in scope_fragments + if ( + fragment.source_block_index, + fragment.char_start, + fragment.char_end, + ) + != caption_key + ] + table_hash = hashlib.sha256( + ( + f"{document.id}\x1e{anchor.page_number}\x1e" + f"{alignment.region.source_block_index}\x1e{alignment.region.raw_start}\x1e" + f"{anchor.bbox}\x1e{column_names}" + ).encode() + ).hexdigest()[:20] + table_id = f"table-{table_hash}" + + rows = [ + _row_from_evidence( + table_id=table_id, + algorithm_version=config.algorithm_version, + table_title=table_title, + section_path=sections, + scope_fragments=scope_fragments, + column_names=column_names, + evidence=row, + aligned=aligned, + page_number=anchor.page_number, + insertion_block_index=alignment.region.source_block_index, + insertion_offset=alignment.region.raw_start, + row_index=row_index, + ) + for row_index, (row, aligned) in enumerate( + zip(data_rows, aligned_data, strict=True), + start=1, + ) + ] + if not rows: + return None + if not _apply_merged_cell_inheritance( + rows, + data_rows, + evidence_row_offset=header_rows, + ): + return None + + legends: list[_LegendInsertion] = [] + if has_header: + legend_entries = _legend_entries(anchor.rows[0], alignment.rows[0]) + if legend_entries: + legends.append( + _LegendInsertion( + legend=TableLegendData( + table_id=table_id, + algorithm_version=config.algorithm_version, + table_title=table_title, + section_path=sections, + scope_fragments=[fragment.model_copy(deep=True) for fragment in scope_fragments], + entries=legend_entries, + page_number=anchor.page_number, + ), + insertion_block_index=alignment.region.source_block_index, + insertion_offset=alignment.region.raw_start, + ) + ) + + consumed = [alignment.region] + decisions: list[PageBoundaryDecision] = [] + used_tables = {(anchor.page_number, anchor_index)} + fallback_reasons: list[str] = [] + open_row = rows[-1] + current_page = anchor.page_number + reaches_bottom = anchor.bbox[3] >= 0.82 + previous_bottom = anchor.bbox[3] + + while reaches_bottom: + next_page_number = current_page + 1 + next_page = pages.get(next_page_number) + if next_page is None and next_page_number <= page_count: + additional = await self._evidence_provider.collect(document, {next_page_number}) + pages.update((page.page_number, page) for page in additional) + next_page = pages.get(next_page_number) + if next_page is None: + break + current_evidence = pages.get(current_page) + if current_evidence is None or not _compatible_pages(current_evidence, next_page): + fallback_reasons.append(f"preserved incompatible page boundary {current_page}->{next_page_number}") + break + + if _page_starts_with_heading(raw_blocks, next_page_number): + fallback_reasons.append(f"preserved section boundary on page {next_page_number}") + break + + compatible = self._compatible_top_table(anchor, next_page) + if compatible is not None: + if _has_parser_content_before_table(raw_blocks, next_page_number): + fallback_reasons.append(f"preserved content before a table on page {next_page_number}") + break + next_index, next_table, geometry_confidence = compatible + repeated_header = bool(next_table.rows) and _is_repeated_header( + next_table.rows[0], + column_names, + ) + data_start = 1 if repeated_header else 0 + if len(next_table.rows) <= data_start: + fallback_reasons.append(f"preserved a repeated header without data on page {next_page_number}") + break + if not _spans_are_unambiguous(next_table, header_rows=data_start): + fallback_reasons.append(f"preserved ambiguous merged cells on page {next_page_number}") + return None + leading = next_table.rows[data_start] + identity_empty = all(not leading.cells[index].text.strip() for index in identity_columns) + content_present = any( + cell.text.strip() for cell in leading.cells if cell.column_index not in identity_columns + ) + bottom_confidence, top_confidence = _page_edge_confidences( + previous_bottom, + next_table.bbox[1], + ) + same_table_confidence = min( + geometry_confidence, + bottom_confidence, + top_confidence, + ) + row_confidence = top_confidence if identity_empty and content_present else 0.0 + if not identity_empty or not content_present: + decisions.append( + PageBoundaryDecision( + previous_page=current_page, + next_page=next_page_number, + same_table_confidence=same_table_confidence, + row_continuation_confidence=row_confidence, + decision="preserved", + reason="the leading row does not safely continue the open row", + ) + ) + fallback_reasons.append(f"preserved distinct table boundary {current_page}->{next_page_number}") + break + + next_alignment = _align_table( + next_table, + processed_document, + image_regions, + self._evidence_provider.provider_id, + config.cell_assignment_min_confidence, + ) + if next_alignment is None: + fallback_reasons.append(f"preserved an unaligned table on page {next_page_number}") + return None + if not _all_assignments_pass( + leading, + next_alignment.rows[data_start], + config.cell_assignment_min_confidence, + ): + fallback_reasons.append(f"could not align table cells on page {next_page_number}") + return None + if _has_content_before_regions( + raw_blocks, + [next_alignment.region], + ): + fallback_reasons.append(f"preserved unaccounted content before a table on page {next_page_number}") + break + aligned_leading = next_alignment.rows[data_start] + assignment_confidence = min( + ( + confidence + for cell, (_, confidence) in zip( + leading.cells, + aligned_leading, + strict=True, + ) + if cell.text.strip() + ), + default=0.0, + ) + row_confidence = min(row_confidence, assignment_confidence) + if ( + same_table_confidence < config.same_table_min_confidence + or row_confidence < config.row_continuation_min_confidence + ): + return None + + decision = PageBoundaryDecision( + previous_page=current_page, + next_page=next_page_number, + same_table_confidence=same_table_confidence, + row_continuation_confidence=row_confidence, + decision="merged", + reason="compatible columns and empty identity cells continue the open row", + ) + open_row.merge(leading, aligned_leading, decision) + decisions.append(decision) + consumed.append(next_alignment.region) + used_tables.add((next_page_number, next_index)) + + next_rows: list[_MutableRow] = [] + following_start = data_start + 1 + for evidence_index, row in enumerate( + next_table.rows[following_start:], + start=following_start, + ): + aligned = next_alignment.rows[evidence_index] + if not _all_assignments_pass(row, aligned, config.cell_assignment_min_confidence): + fallback_reasons.append(f"preserved an unaligned row on page {next_page_number}") + return None + new_row = _row_from_evidence( + table_id=table_id, + algorithm_version=config.algorithm_version, + table_title=table_title, + section_path=sections, + scope_fragments=[fragment.model_copy(deep=True) for fragment in scope_fragments], + column_names=column_names, + evidence=row, + aligned=aligned, + page_number=next_page_number, + insertion_block_index=next_alignment.region.source_block_index, + insertion_offset=next_alignment.region.raw_start, + row_index=len(rows) + len(next_rows) + 1, + ) + next_rows.append(new_row) + if not _apply_merged_cell_inheritance( + next_rows, + next_table.rows[following_start:], + evidence_row_offset=following_start, + ): + fallback_reasons.append(f"preserved ambiguous merged-cell inheritance on page {next_page_number}") + return None + rows.extend(next_rows) + if next_rows: + open_row = next_rows[-1] + + current_page = next_page_number + reaches_bottom = next_table.bbox[3] >= 0.82 + previous_bottom = next_table.bbox[3] + continue + + sparse = _sparse_continuation(next_page, anchor.column_bounds, identity_columns) + if sparse is None: + break + aligned_sparse = _align_sparse_row( + sparse, + next_page_number, + processed_document, + image_regions, + self._evidence_provider.provider_id, + config.cell_assignment_min_confidence, + ) + if aligned_sparse is None or not _all_assignments_pass( + sparse, + aligned_sparse.cells, + config.cell_assignment_min_confidence, + ): + fallback_reasons.append(f"could not align sparse continuation on page {next_page_number}") + return None + if _has_content_before_regions(raw_blocks, aligned_sparse.regions): + fallback_reasons.append( + f"preserved unaccounted content before a continuation on page {next_page_number}" + ) + break + + bottom_confidence, top_confidence = _page_edge_confidences( + previous_bottom, + sparse.bbox[1], + ) + assignment_confidence = min( + ( + confidence + for cell, (_, confidence) in zip( + sparse.cells, + aligned_sparse.cells, + strict=True, + ) + if cell.text.strip() + ), + default=0.0, + ) + same_table_confidence = min(bottom_confidence, top_confidence) + row_confidence = min(top_confidence, assignment_confidence) + if ( + same_table_confidence < config.same_table_min_confidence + or row_confidence < config.row_continuation_min_confidence + ): + return None + decision = PageBoundaryDecision( + previous_page=current_page, + next_page=next_page_number, + same_table_confidence=same_table_confidence, + row_continuation_confidence=row_confidence, + decision="merged", + reason="content remains inside continuation columns at the next page boundary", + ) + open_row.merge(sparse, aligned_sparse.cells, decision) + decisions.append(decision) + consumed.extend(aligned_sparse.regions) + current_page = next_page_number + reaches_bottom = sparse.bbox[3] >= 0.82 + previous_bottom = sparse.bbox[3] + + chain = _NormalizedChain( + rows=rows, + identity_columns=identity_columns, + consumed_regions=consumed, + decisions=decisions, + used_tables=used_tables, + fallback_reasons=fallback_reasons, + legends=legends, + ) + if caption_key is not None: + claimed_captions.add(caption_key) + return chain + + @staticmethod + def _compatible_top_table( + anchor: LayoutTableEvidence, + page: PageLayoutEvidence, + ) -> tuple[int, LayoutTableEvidence, float] | None: + for table_index, table in sorted(enumerate(page.tables), key=lambda item: item[1].bbox[1]): + if table.bbox[1] > 0.15: + continue + confidence = _table_geometry_confidence(anchor.column_bounds, table.column_bounds) + if confidence > 0.0: + return table_index, table, confidence + return None + + @staticmethod + def _unchanged( + processed_document: ProcessedDocument, + config: TableReconstructionConfig, + reason: str, + ) -> ProcessedDocument: + return processed_document.model_copy( + update={ + "normalization_report": NormalizationReport( + algorithm_version=config.algorithm_version, + status="unchanged", + fallback_reasons=[reason], + ) + } + ) + + +def _row_text(row: TableRowData) -> str: + return render_table_row(row) + + +def _legend_text(legend: TableLegendData) -> str: + return render_table_legend(legend) + + +def _merge_intervals(intervals: list[tuple[int, int]]) -> list[tuple[int, int]]: + merged: list[list[int]] = [] + for start, end in sorted(intervals): + if not merged or start > merged[-1][1]: + merged.append([start, end]) + else: + merged[-1][1] = max(merged[-1][1], end) + return [(start, end) for start, end in merged] + + +def _regions_are_disjoint(regions: list[_AlignedRegion]) -> bool: + by_block: dict[int, list[_AlignedRegion]] = {} + for region in regions: + by_block.setdefault(region.source_block_index, []).append(region) + for block_regions in by_block.values(): + ordered = sorted(block_regions, key=lambda region: (region.working_start, region.working_end)) + if any(current.working_start < previous.working_end for previous, current in zip(ordered, ordered[1:])): + return False + return True + + +def _map_interval(raw: str, working: str, start: int, end: int) -> tuple[int, int] | None: + if raw == working: + return start, end + found = _find_normalised_range(raw[start:end], working) + if found is None: + return None + mapped_start, mapped_end, _ = found + return mapped_start, mapped_end + + +def _build_normalized_blocks( + processed_document: ProcessedDocument, + rows: list[tuple[_MutableRow, tuple[int, ...]]], + legends: list[_LegendInsertion], + consumed: list[_AlignedRegion], +) -> list[TextBlock] | None: + raw_blocks = processed_document.raw_text_blocks + if raw_blocks is None: + return None + + consumed_by_block: dict[int, list[_AlignedRegion]] = {} + for region in consumed: + consumed_by_block.setdefault(region.source_block_index, []).append(region) + rows_by_block: dict[int, list[tuple[_MutableRow, tuple[int, ...]]]] = {} + for row in rows: + rows_by_block.setdefault(row[0].insertion_block_index, []).append(row) + legends_by_block: dict[int, list[_LegendInsertion]] = {} + for legend in legends: + legends_by_block.setdefault(legend.insertion_block_index, []).append(legend) + + normalized: list[TextBlock] = [] + for block_index, working_block in enumerate(processed_document.text_blocks): + if block_index >= len(raw_blocks): + normalized.append(working_block) + continue + raw_block = raw_blocks[block_index] + regions = consumed_by_block.get(block_index, []) + if not regions: + normalized.append(working_block) + continue + + mapped_regions: list[tuple[int, int, int, int]] = [] + for region in regions: + if ( + region.raw_end > len(raw_block.text) + or region.working_end > len(working_block.text) + or region.raw_start >= region.raw_end + or region.working_start >= region.working_end + ): + return None + mapped_regions.append( + ( + region.raw_start, + region.raw_end, + region.working_start, + region.working_end, + ) + ) + merged_working = _merge_intervals([(start, end) for _, _, start, end in mapped_regions]) + + insertions = sorted(rows_by_block.get(block_index, []), key=lambda item: item[0].insertion_offset) + legend_insertions = sorted( + legends_by_block.get(block_index, []), + key=lambda item: item.insertion_offset, + ) + inserted: set[int] = set() + inserted_legends: set[int] = set() + cursor = 0 + for working_start, working_end in merged_working: + residual = working_block.text[cursor:working_start].strip() + if residual: + normalized.append( + working_block.model_copy( + update={ + "text": residual, + "source_fragments": [], + "table_row": None, + "table_legend": None, + } + ) + ) + raw_ranges = [ + (raw_start, raw_end) + for raw_start, raw_end, mapped_start, mapped_end in mapped_regions + if mapped_start <= working_end and mapped_end >= working_start + ] + raw_limit = max((end for _, end in raw_ranges), default=len(raw_block.text)) + for insertion_index, insertion in enumerate(legend_insertions): + if insertion_index in inserted_legends or insertion.insertion_offset > raw_limit: + continue + legend = insertion.legend + fragments = [ + *legend.scope_fragments, + *(fragment for entry in legend.entries for fragment in entry.source_fragments), + ] + normalized.append( + TextBlock( + text=_legend_text(legend), + page_number=legend.page_number, + block_type="table_legend", + metadata={ + "table_id": legend.table_id, + "table_content_kind": "legend", + }, + source_fragments=fragments, + table_legend=legend, + ) + ) + inserted_legends.add(insertion_index) + for insertion_index, (mutable_row, identity_columns) in enumerate(insertions): + if insertion_index in inserted or mutable_row.insertion_offset > raw_limit: + continue + row = mutable_row.freeze(identity_columns) + fragments = [ + *row.scope_fragments, + *(fragment for cell in row.cells for fragment in cell.source_fragments), + ] + normalized.append( + TextBlock( + text=_row_text(row), + page_number=row.page_start, + block_type="table_row", + metadata={ + "table_id": row.table_id, + "row_id": row.row_id, + "row_index": row.row_index, + "page_end": row.page_end, + }, + source_fragments=fragments, + table_row=row, + ) + ) + inserted.add(insertion_index) + cursor = working_end + + residual = working_block.text[cursor:].strip() + if residual: + normalized.append( + working_block.model_copy( + update={ + "text": residual, + "source_fragments": [], + "table_row": None, + "table_legend": None, + } + ) + ) + for insertion_index, insertion in enumerate(legend_insertions): + if insertion_index in inserted_legends: + continue + legend = insertion.legend + fragments = [ + *legend.scope_fragments, + *(fragment for entry in legend.entries for fragment in entry.source_fragments), + ] + normalized.append( + TextBlock( + text=_legend_text(legend), + page_number=legend.page_number, + block_type="table_legend", + metadata={ + "table_id": legend.table_id, + "table_content_kind": "legend", + }, + source_fragments=fragments, + table_legend=legend, + ) + ) + for insertion_index, (mutable_row, identity_columns) in enumerate(insertions): + if insertion_index in inserted: + continue + row = mutable_row.freeze(identity_columns) + fragments = [ + *row.scope_fragments, + *(fragment for cell in row.cells for fragment in cell.source_fragments), + ] + normalized.append( + TextBlock( + text=_row_text(row), + page_number=row.page_start, + block_type="table_row", + metadata={ + "table_id": row.table_id, + "row_id": row.row_id, + "row_index": row.row_index, + "page_end": row.page_end, + }, + source_fragments=fragments, + table_row=row, + ) + ) + + return normalized + + +__all__ = ["DeterministicTableNormalizer"] diff --git a/openrag/core/indexing/table_text.py b/openrag/core/indexing/table_text.py new file mode 100644 index 000000000..eb5e3f49b --- /dev/null +++ b/openrag/core/indexing/table_text.py @@ -0,0 +1,223 @@ +"""Deterministic, readable text rendering for structured table content.""" + +from __future__ import annotations + +import html +import re +import unicodedata +from collections.abc import Iterable + +from core.models.document import TableCellData, TableLegendData, TableRowData + +TABLE_TEXT_SERIALIZATION_VERSION = "natural-language-v1" + +_HTML_BREAK_RE = re.compile(r"", re.IGNORECASE) +_HTML_TAG_RE = re.compile(r"<[^>]+>") +_LIST_ITEM_RE = re.compile(r"^(?:[-*•]\s*|\d+(?:\.\d+)*[.)]\s+)") +_ORDINALS = { + 1: "first", + 2: "second", + 3: "third", + 4: "fourth", + 5: "fifth", + 6: "sixth", + 7: "seventh", + 8: "eighth", + 9: "ninth", + 10: "tenth", +} + + +def normalize_table_text(text: str) -> str: + """Reflow visual PDF lines while preserving semantic list boundaries.""" + value = unicodedata.normalize("NFC", html.unescape(text or "")).replace("\u00ad", "") + value = _HTML_BREAK_RE.sub("\n", value) + value = _HTML_TAG_RE.sub(" ", value) + value = re.sub(r"\b([cdjlmnst])\s+'(?=\w)", r"\1'", value, flags=re.IGNORECASE) + value = re.sub(r"\b(le|la|les|du|des|au|aux)(?=\d)", r"\1 ", value, flags=re.IGNORECASE) + value = re.sub(r'"\s*([^"\n]*?\S)\s*"', r'"\1"', value) + value = re.sub(r"[ \t]+\n", "\n", value) + value = re.sub(r"\n[ \t]+", "\n", value) + value = re.sub(r"[ \t]{2,}", " ", value) + + output: list[str] = [] + paragraph: list[str] = [] + + def flush_paragraph() -> None: + if paragraph: + output.append(" ".join(paragraph)) + paragraph.clear() + + for raw_line in value.splitlines(): + line = raw_line.strip() + if not line: + flush_paragraph() + if output and output[-1] != "": + output.append("") + continue + if _LIST_ITEM_RE.match(line): + flush_paragraph() + line = re.sub(r"^([-*•])(?=\S)", r"\1 ", line) + paragraph.append(line) + else: + paragraph.append(line) + flush_paragraph() + + normalized = "\n".join(output) + normalized = re.sub(r"\n{3,}", "\n\n", normalized) + return normalized.strip() + + +def _quoted(value: str) -> str: + return f"“{normalize_table_text(value)}”" + + +def _table_scope( + *, + table_title: str | None, + section_path: list[str], +) -> str: + section = " > ".join(part.strip() for part in section_path if part.strip()) + if table_title and section_path and section_path[-1].strip() == table_title.strip(): + parent = " > ".join(part.strip() for part in section_path[:-1] if part.strip()) + return ( + f"In section {_quoted(parent)}, table {_quoted(table_title)}" + if parent + else f"In table {_quoted(table_title)}" + ) + if section and table_title: + return f"In section {_quoted(section)}, table {_quoted(table_title)}" + if section: + return f"In section {_quoted(section)}, the table" + if table_title: + return f"In table {_quoted(table_title)}" + return "In the table" + + +def _compact_table_scope( + *, + table_title: str | None, + section_path: list[str], +) -> str: + section = " > ".join(part.strip() for part in section_path if part.strip()) + if table_title and section_path and section_path[-1].strip() == table_title.strip(): + parent = " > ".join(part.strip() for part in section_path[:-1] if part.strip()) + return f"Section {_quoted(parent)}, table {_quoted(table_title)}" if parent else f"Table {_quoted(table_title)}" + if section and table_title: + return f"Section {_quoted(section)}, table {_quoted(table_title)}" + if section: + return f"Section {_quoted(section)}, table" + if table_title: + return f"Table {_quoted(table_title)}" + return "Table" + + +def _column_names(row: TableRowData, cell: TableCellData) -> list[str]: + end = cell.column_index + max(cell.column_span, 1) + names = [ + candidate.column_name or f"Column {candidate.column_index + 1}" + for candidate in row.cells + if cell.column_index <= candidate.column_index < end + ] + return names or [cell.column_name or f"Column {cell.column_index + 1}"] + + +def _joined_quoted(values: Iterable[str]) -> str: + quoted = [_quoted(value) for value in values] + if len(quoted) <= 1: + return quoted[0] if quoted else "" + if len(quoted) == 2: + return f"{quoted[0]} and {quoted[1]}" + return f"{', '.join(quoted[:-1])}, and {quoted[-1]}" + + +def _cell_clause(row: TableRowData, cell: TableCellData) -> str | None: + if cell.covered_by is not None: + return None + names = _column_names(row, cell) + column_label = f"column {_quoted(names[0])}" if len(names) == 1 else f"columns {_joined_quoted(names)}" + value = normalize_table_text(cell.text) + if not value: + return f"no value in {column_label}" if cell.explicit_empty else None + inherited = "inherited " if cell.inherited else "" + return f"the {inherited}value {_quoted(value)} in {column_label}" + + +def _join_clauses(clauses: list[str]) -> str: + if not clauses: + return "" + if len(clauses) == 1: + return clauses[0] + if len(clauses) == 2: + return f"{clauses[0]} and {clauses[1]}" + return f"{', '.join(clauses[:-1])}, and {clauses[-1]}" + + +def row_reference(row: TableRowData) -> str: + """Return a human-readable logical row reference with a stable number.""" + ordinal = _ORDINALS.get(row.row_index) + return f"the {ordinal} row (row {row.row_index})" if ordinal else f"row {row.row_index}" + + +def _compact_row_reference(row: TableRowData) -> str: + ordinal = _ORDINALS.get(row.row_index) + return f"{ordinal} row ({row.row_index})" if ordinal else f"row {row.row_index}" + + +def render_table_row( + row: TableRowData, + *, + cells: Iterable[TableCellData] | None = None, +) -> str: + """Render a complete or partial row as deterministic natural language.""" + selected = list(row.cells if cells is None else cells) + clauses = [clause for cell in selected if (clause := _cell_clause(row, cell)) is not None] + scope = _table_scope(table_title=row.table_title, section_path=row.section_path) + if not clauses: + return f"{scope} contains {row_reference(row)}." + return f"{scope}, {row_reference(row)} has {_join_clauses(clauses)}." + + +def render_table_row_context( + row: TableRowData, + *, + cells: Iterable[TableCellData], +) -> str: + """Render compact, self-contained identity context for split row parts.""" + selected = [ + cell for cell in cells if cell.covered_by is None and (normalize_table_text(cell.text) or cell.explicit_empty) + ] + scope = _compact_table_scope(table_title=row.table_title, section_path=row.section_path) + clauses = [ + ( + f"{_quoted(cell.column_name or f'Column {cell.column_index + 1}')} = " + f"{_quoted(cell.text) if normalize_table_text(cell.text) else 'empty'}" + ) + for cell in selected + ] + if not clauses: + return f"{scope}; {_compact_row_reference(row)}." + return f"{scope}; {_compact_row_reference(row)}; {'; '.join(clauses)}." + + +def render_table_legend(legend: TableLegendData) -> str: + """Render abbreviation definitions independently from table rows.""" + scope = _table_scope(table_title=legend.table_title, section_path=legend.section_path) + definitions = [ + f"{entry.abbreviation} means {_quoted(entry.meaning)}" + for entry in legend.entries + if entry.abbreviation.strip() and entry.meaning.strip() + ] + if not definitions: + return "" + return f"{scope}, the abbreviation legend defines the following terms: {_join_clauses(definitions)}." + + +__all__ = [ + "TABLE_TEXT_SERIALIZATION_VERSION", + "normalize_table_text", + "render_table_legend", + "render_table_row", + "render_table_row_context", + "row_reference", +] diff --git a/openrag/core/models/document.py b/openrag/core/models/document.py index f0d0e21bb..bf474f284 100644 --- a/openrag/core/models/document.py +++ b/openrag/core/models/document.py @@ -12,9 +12,9 @@ from datetime import UTC, datetime from enum import Enum from pathlib import Path -from typing import Any +from typing import Any, Literal -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator class DocumentType(str, Enum): @@ -31,6 +31,121 @@ class DocumentType(str, Enum): EML = "eml" +class SourceFragment(BaseModel): + """A trace from normalized content back to its original evidence. + + ``parser_text`` fragments identify text copied from an untouched parser + block. ``pdf_layout`` fragments identify text recovered from PDF geometry; + their character range is a visual anchor in the parser block + (for example, Marker's original image placeholder), while ``bbox`` and + ``evidence_provider`` identify the canonical evidence. + """ + + source_block_index: int = Field(ge=0) + page_number: int = Field(ge=1) + char_start: int = Field(ge=0) + char_end: int = Field(gt=0) + source_kind: Literal["parser_text", "pdf_layout"] = "parser_text" + evidence_provider: str | None = None + source_ref: str | None = None + bbox: tuple[float, float, float, float] | None = None + text_start: int = Field(default=0, ge=0) + text_end: int | None = Field(default=None, gt=0) + + @model_validator(mode="after") + def validate_ranges(self) -> SourceFragment: + if self.char_end <= self.char_start: + raise ValueError("char_end must be greater than char_start") + if self.text_end is not None and self.text_end <= self.text_start: + raise ValueError("text_end must be greater than text_start") + if self.source_kind == "pdf_layout": + if self.bbox is None: + raise ValueError("pdf_layout fragments require a bbox") + if not self.evidence_provider: + raise ValueError("pdf_layout fragments require an evidence_provider") + return self + + +class PageBoundaryDecision(BaseModel): + """Independent evidence recorded for one adjacent-page decision.""" + + previous_page: int = Field(ge=1) + next_page: int = Field(ge=1) + same_table_confidence: float = Field(ge=0.0, le=1.0) + row_continuation_confidence: float = Field(ge=0.0, le=1.0) + decision: Literal["merged", "preserved"] + reason: str + + +class TableCellData(BaseModel): + """One logical cell assembled from source fragments.""" + + column_index: int = Field(ge=0) + column_name: str | None = None + text: str = "" + source_fragments: list[SourceFragment] = Field(default_factory=list) + assignment_confidence: float = Field(default=1.0, ge=0.0, le=1.0) + column_span: int = Field(default=1, ge=1) + row_span: int = Field(default=1, ge=1) + inherited: bool = False + inherited_from: tuple[int, int] | None = None + explicit_empty: bool = False + covered_by: tuple[int, int] | None = None + + +class TableRowData(BaseModel): + """A logical table row, potentially reconstructed across several pages.""" + + table_id: str + row_id: str + algorithm_version: str + table_title: str | None = None + section_path: list[str] = Field(default_factory=list) + scope_fragments: list[SourceFragment] = Field(default_factory=list) + cells: list[TableCellData] = Field(default_factory=list) + identity_columns: list[int] = Field(default_factory=list) + row_index: int = Field(default=1, ge=1) + page_start: int = Field(ge=1) + page_end: int = Field(ge=1) + boundary_decisions: list[PageBoundaryDecision] = Field(default_factory=list) + + @model_validator(mode="after") + def validate_page_range(self) -> TableRowData: + if self.page_end < self.page_start: + raise ValueError("page_end must be greater than or equal to page_start") + return self + + +class TableLegendEntry(BaseModel): + """One abbreviation definition extracted from a table header.""" + + abbreviation: str + meaning: str + source_fragments: list[SourceFragment] = Field(default_factory=list) + + +class TableLegendData(BaseModel): + """A table legend kept separate from its data rows.""" + + table_id: str + algorithm_version: str + table_title: str | None = None + section_path: list[str] = Field(default_factory=list) + scope_fragments: list[SourceFragment] = Field(default_factory=list) + entries: list[TableLegendEntry] = Field(default_factory=list) + page_number: int = Field(ge=1) + + +class NormalizationReport(BaseModel): + """Debug record for structural normalization decisions.""" + + algorithm_version: str + status: Literal["unchanged", "normalized", "partial_fallback"] + boundary_decisions: list[PageBoundaryDecision] = Field(default_factory=list) + reconstructed_row_count: int = Field(default=0, ge=0) + fallback_reasons: list[str] = Field(default_factory=list) + + class TextBlock(BaseModel): """A block of text extracted from a document.""" @@ -38,6 +153,9 @@ class TextBlock(BaseModel): page_number: int | None = None block_type: str = "paragraph" metadata: dict[str, Any] = Field(default_factory=dict) + source_fragments: list[SourceFragment] = Field(default_factory=list) + table_row: TableRowData | None = None + table_legend: TableLegendData | None = None class ImageBlock(BaseModel): @@ -225,6 +343,13 @@ class ProcessedDocument(BaseModel): document_id: str = "" text_blocks: list[TextBlock] = Field(default_factory=list) + raw_text_blocks: list[TextBlock] | None = None + normalized_text_blocks: list[TextBlock] | None = None + normalization_report: NormalizationReport | None = None images: list[ImageBlock] = Field(default_factory=list) metadata: dict[str, Any] = Field(default_factory=dict) page_count: int = 0 + + def effective_text_blocks(self) -> list[TextBlock]: + """Return the complete block view downstream chunkers should consume.""" + return self.normalized_text_blocks if self.normalized_text_blocks is not None else self.text_blocks diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index c8fedb012..4806d30c2 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -22,7 +22,7 @@ # and workers on the same protocol generation whenever their remote contract or # cross-process indexing semantics change, so new replicas cannot attach to a # partially compatible actor fleet left by the previous release. -_INDEXER_ACTOR_PROTOCOL_VERSION = "v2" +_INDEXER_ACTOR_PROTOCOL_VERSION = "v3" _INDEXER_POOL_DISPATCHER_ACTOR_NAME = f"IndexerPoolDispatcher-{_INDEXER_ACTOR_PROTOCOL_VERSION}" @@ -52,9 +52,11 @@ def __init__(self) -> None: import services.inference.vllm_client # noqa: F401 from core.config import load_config from core.embeddings import embedder_registry + from core.indexing.table_normalizer import DeterministicTableNormalizer from core.utils.logging import get_logger from services.storage.milvus_store import MilvusVectorStore from services.storage.postgres_store import PostgresStore + from services.workers.layout import PyMuPDFTableEvidenceProvider from services.workers.parsers.parser_dispatcher import ( build_caption_vlm, build_parser_dispatcher, @@ -80,6 +82,7 @@ def __init__(self) -> None: vlm_factory = _build_vlm_factory(cfg) contextualizer_factory = _build_contextualizer_factory(cfg) topic_tagger_factory = _build_topic_tagger_factory(cfg) + structure_normalizer = DeterministicTableNormalizer(PyMuPDFTableEvidenceProvider()) embed_cfg = cfg.embedder embedder = embedder_registry.create( @@ -108,6 +111,7 @@ def __init__(self) -> None: vlm_factory=vlm_factory, contextualizer_factory=contextualizer_factory, topic_tagger_factory=topic_tagger_factory, + structure_normalizer=structure_normalizer, defer_replace_cleanup=True, ) self._catalog_store = PostgresStore(_catalog_rdb_config(cfg), run_migrations=False) @@ -634,7 +638,10 @@ def _build_pipeline_timeouts(cfg: Settings) -> Any: """ from services.workers.pipeline_builder import PipelineTimeouts - return PipelineTimeouts(parse=cfg.loader.parse_timeout) + return PipelineTimeouts( + parse=cfg.loader.parse_timeout, + normalize_structure=cfg.loader.parse_timeout, + ) def _build_embedder_factory(cfg: Settings) -> Any: diff --git a/openrag/services/workers/layout/__init__.py b/openrag/services/workers/layout/__init__.py new file mode 100644 index 000000000..579c6585f --- /dev/null +++ b/openrag/services/workers/layout/__init__.py @@ -0,0 +1,5 @@ +"""Layout evidence adapters used by indexing workers.""" + +from .pymupdf_table_evidence import PyMuPDFTableEvidenceProvider + +__all__ = ["PyMuPDFTableEvidenceProvider"] diff --git a/openrag/services/workers/layout/pymupdf_table_evidence.py b/openrag/services/workers/layout/pymupdf_table_evidence.py new file mode 100644 index 000000000..04a52eae9 --- /dev/null +++ b/openrag/services/workers/layout/pymupdf_table_evidence.py @@ -0,0 +1,343 @@ +"""PyMuPDF adapter for table and word-position evidence.""" + +from __future__ import annotations + +import unicodedata +from statistics import median +from typing import Any + +import pymupdf +from core.indexing.parsers.pdf.pymupdf_runtime import run_pymupdf +from core.indexing.structure_normalizer import ( + LayoutCellEvidence, + LayoutRowEvidence, + LayoutTableEvidence, + LayoutWord, + NormalizedBBox, + PageLayoutEvidence, + TableLayoutEvidenceProvider, +) +from core.models.document import Document, DocumentType + + +def _normalize_bbox(bbox: tuple[float, float, float, float], width: float, height: float) -> NormalizedBBox: + x0, y0, x1, y1 = bbox + return (x0 / width, y0 / height, x1 / width, y1 / height) + + +def _row_bbox(cells: list[tuple[float, float, float, float] | None]) -> tuple[float, float, float, float] | None: + present = [cell for cell in cells if cell is not None] + if not present: + return None + return ( + min(cell[0] for cell in present), + min(cell[1] for cell in present), + max(cell[2] for cell in present), + max(cell[3] for cell in present), + ) + + +def _column_bounds( + rows: list[Any], + table_bbox: tuple[float, float, float, float], + column_count: int, + width: float, +) -> tuple[tuple[float, float], ...]: + for row in rows: + if len(row.cells) != column_count or any(cell is None for cell in row.cells): + continue + return tuple((cell[0] / width, cell[2] / width) for cell in row.cells) + + x0, _, x1, _ = table_bbox + column_width = (x1 - x0) / max(column_count, 1) + return tuple( + ((x0 + index * column_width) / width, (x0 + (index + 1) * column_width) / width) + for index in range(column_count) + ) + + +def _normalized_alnum(text: str) -> str: + return "".join(character for character in unicodedata.normalize("NFKC", text).casefold() if character.isalnum()) + + +def _word_is_inside_cell(word: LayoutWord, bbox: NormalizedBBox) -> bool: + x0, y0, x1, y1 = word.bbox + center_x = (x0 + x1) / 2 + center_y = (y0 + y1) / 2 + cell_x0, cell_y0, cell_x1, cell_y1 = bbox + return cell_x0 <= center_x <= cell_x1 and cell_y0 <= center_y <= cell_y1 + + +def _text_from_positioned_words( + words: tuple[LayoutWord, ...], + bbox: NormalizedBBox | None, + extracted_text: str, +) -> str: + """Rebuild cell text from PDF words when they match the extracted value. + + ``Table.extract()`` occasionally inserts spaces inside words or removes + spaces between adjacent words. PyMuPDF's positioned words retain the + intended token boundaries. The alphanumeric equality gate ensures geometry + is used only when it accounts for the complete extracted value; otherwise + the parser value remains the safe fallback. + """ + + if bbox is None: + return extracted_text + + cell_words = sorted( + (word for word in words if word.text.strip() and _word_is_inside_cell(word, bbox)), + key=lambda word: ( + word.block_number, + word.line_number, + word.word_number, + word.bbox[1], + word.bbox[0], + ), + ) + if not cell_words: + return extracted_text + + parts: list[str] = [] + current_key: tuple[int, int] | None = None + current_words: list[str] = [] + previous_key: tuple[int, int] | None = None + + def flush_line() -> None: + nonlocal previous_key + if current_key is None or not current_words: + return + if previous_key is not None: + previous_block, previous_line = previous_key + current_block, current_line = current_key + parts.append("\n\n" if current_block != previous_block or current_line > previous_line + 1 else "\n") + parts.append(" ".join(current_words)) + previous_key = current_key + + for word in cell_words: + key = (word.block_number, word.line_number) + if key != current_key: + flush_line() + current_key = key + current_words = [] + current_words.append(word.text.strip()) + flush_line() + + rebuilt = "".join(parts).strip() + if not rebuilt: + return extracted_text + if extracted_text.strip() and _normalized_alnum(rebuilt) != _normalized_alnum(extracted_text): + return extracted_text + return rebuilt + + +def _cell_grid_semantics( + table_rows: list[Any], + column_bounds: tuple[tuple[float, float], ...], + table_bbox: tuple[float, float, float, float], + width: float, + height: float, +) -> dict[tuple[int, int], tuple[str, int, int, tuple[int, int] | None]]: + """Resolve explicit empty and merged slots from PyMuPDF's cell geometry.""" + table_top = table_bbox[1] + table_bottom = table_bbox[3] + equal_height = (table_bottom - table_top) / max(len(table_rows), 1) + row_tops: list[float | None] = [] + for row in table_rows: + cell_tops = [cell[1] for cell in row.cells if cell is not None] + row_tops.append(float(median(cell_tops)) if cell_tops else None) + + known = [(index, top) for index, top in enumerate(row_tops) if top is not None] + for index, top in enumerate(row_tops): + if top is not None: + continue + previous = next( + ((candidate_index, candidate) for candidate_index, candidate in reversed(known) if candidate_index < index), + None, + ) + following = next( + ((candidate_index, candidate) for candidate_index, candidate in known if candidate_index > index), + None, + ) + if previous is not None and following is not None: + left_index, left = previous + right_index, right = following + ratio = (index - left_index) / (right_index - left_index) + row_tops[index] = left + (right - left) * ratio + else: + row_tops[index] = table_top + index * equal_height + + resolved_tops = [float(top) for top in row_tops] + if any(right <= left for left, right in zip(resolved_tops, resolved_tops[1:], strict=False)): + resolved_tops = [table_top + index * equal_height for index in range(len(table_rows))] + row_centers = [ + (top + (resolved_tops[index + 1] if index + 1 < len(resolved_tops) else table_bottom)) / (2 * height) + for index, top in enumerate(resolved_tops) + ] + anchors: dict[tuple[int, int], tuple[int, int]] = {} + coverage: dict[tuple[int, int], list[tuple[int, int]]] = {} + + for row_index, row in enumerate(table_rows): + for column_index, raw_bbox in enumerate(row.cells): + if raw_bbox is None: + continue + bbox = _normalize_bbox(raw_bbox, width, height) + covered_columns = [ + index for index, (left, right) in enumerate(column_bounds) if bbox[0] <= (left + right) / 2 <= bbox[2] + ] + covered_rows = [index for index, center in enumerate(row_centers) if bbox[1] <= center <= bbox[3]] + if not covered_columns: + covered_columns = [column_index] + if not covered_rows: + covered_rows = [row_index] + column_span = max(covered_columns) - min(covered_columns) + 1 + row_span = max(covered_rows) - min(covered_rows) + 1 + anchors[(row_index, column_index)] = (column_span, row_span) + for covered_row in covered_rows: + for covered_column in covered_columns: + if (covered_row, covered_column) != (row_index, column_index): + coverage.setdefault((covered_row, covered_column), []).append((row_index, column_index)) + + semantics: dict[tuple[int, int], tuple[str, int, int, tuple[int, int] | None]] = {} + for row_index, row in enumerate(table_rows): + for column_index, raw_bbox in enumerate(row.cells): + key = (row_index, column_index) + if raw_bbox is not None: + column_span, row_span = anchors[key] + semantics[key] = ("value", column_span, row_span, None) + continue + covering = coverage.get(key, []) + if len(covering) == 1: + semantics[key] = ("covered", 1, 1, covering[0]) + else: + semantics[key] = ("unknown", 1, 1, None) + return semantics + + +def _collect_evidence(raw_bytes: bytes, page_numbers: tuple[int, ...]) -> list[PageLayoutEvidence]: + collected: list[PageLayoutEvidence] = [] + with pymupdf.open(stream=raw_bytes, filetype="pdf") as pdf: + for page_number in page_numbers: + if page_number < 1 or page_number > pdf.page_count: + continue + + page = pdf[page_number - 1] + width = float(page.rect.width) + height = float(page.rect.height) + words = tuple( + LayoutWord( + text=str(word[4]), + bbox=_normalize_bbox((word[0], word[1], word[2], word[3]), width, height), + block_number=int(word[5]), + line_number=int(word[6]), + word_number=int(word[7]), + ) + for word in page.get_text("words", sort=True) + ) + + tables: list[LayoutTableEvidence] = [] + for table in page.find_tables().tables: + extracted = table.extract() + table_rows = list(table.rows) + column_bounds = _column_bounds( + table_rows, + table.bbox, + table.col_count, + width, + ) + grid_semantics = _cell_grid_semantics( + table_rows, + column_bounds, + table.bbox, + width, + height, + ) + rows: list[LayoutRowEvidence] = [] + for row_index, row in enumerate(table_rows): + bbox = getattr(row, "bbox", None) or _row_bbox(row.cells) + if bbox is None: + continue + values = extracted[row_index] if row_index < len(extracted) else [] + cells: list[LayoutCellEvidence] = [] + for column_index, cell_bbox in enumerate(row.cells): + normalized_bbox = _normalize_bbox(cell_bbox, width, height) if cell_bbox is not None else None + extracted_value = values[column_index] if column_index < len(values) else None + extracted_text = str(extracted_value or "") + slot_state, column_span, row_span, covered_by = grid_semantics[(row_index, column_index)] + if slot_state == "value" and not extracted_text: + slot_state = "explicit_empty" + cells.append( + LayoutCellEvidence( + column_index=column_index, + text=_text_from_positioned_words(words, normalized_bbox, extracted_text), + bbox=normalized_bbox, + slot_state=slot_state, + column_span=column_span, + row_span=row_span, + covered_by=covered_by, + ) + ) + rows.append( + LayoutRowEvidence( + cells=tuple(cells), + bbox=_normalize_bbox(bbox, width, height), + ) + ) + + if not rows: + continue + tables.append( + LayoutTableEvidence( + page_number=page_number, + bbox=_normalize_bbox(table.bbox, width, height), + column_bounds=column_bounds, + rows=tuple(rows), + ) + ) + + collected.append( + PageLayoutEvidence( + page_number=page_number, + width=width, + height=height, + words=words, + tables=tuple(tables), + ) + ) + return collected + + +def _discover_table_pages(raw_bytes: bytes) -> set[int]: + """Find table-bearing pages without extracting their text or geometry.""" + discovered: set[int] = set() + with pymupdf.open(stream=raw_bytes, filetype="pdf") as pdf: + for page_number, page in enumerate(pdf, start=1): + if page.find_tables().tables: + discovered.add(page_number) + return discovered + + +class PyMuPDFTableEvidenceProvider(TableLayoutEvidenceProvider): + """Collect layout evidence from the original PDF on the shared executor.""" + + provider_id = "pymupdf" + + async def discover(self, document: Document) -> set[int]: + if document.content_type is not DocumentType.PDF or not document.raw_bytes: + return set() + return await run_pymupdf( + _discover_table_pages, + document.raw_bytes, + ) + + async def collect(self, document: Document, page_numbers: set[int]) -> list[PageLayoutEvidence]: + if document.content_type is not DocumentType.PDF or not document.raw_bytes or not page_numbers: + return [] + return await run_pymupdf( + _collect_evidence, + document.raw_bytes, + tuple(sorted(page_numbers)), + ) + + +__all__ = ["PyMuPDFTableEvidenceProvider"] diff --git a/openrag/services/workers/parsers/marker_workers.py b/openrag/services/workers/parsers/marker_workers.py index ca5fd1109..5f03f0d47 100644 --- a/openrag/services/workers/parsers/marker_workers.py +++ b/openrag/services/workers/parsers/marker_workers.py @@ -403,6 +403,7 @@ async def process_pdf(self, file_path: str): _MARKER_KEY_PAGE_RE = re.compile(r"_page_(\d+)_") +_HTML_BREAK_RE = re.compile(r"", re.IGNORECASE) def _marker_key_to_page(key: str) -> int | None: @@ -422,6 +423,17 @@ def _marker_key_to_page(key: str) -> int | None: return None +def _clean_marker_breaks(markdown: str) -> str: + """Keep cell-internal breaks in pipe tables and remove them elsewhere.""" + cleaned_lines: list[str] = [] + for line in markdown.splitlines(keepends=True): + content = line.rstrip("\r\n") + stripped = content.strip() + replacement = "
" if stripped.startswith("|") and stripped.endswith("|") else "" + cleaned_lines.append(_HTML_BREAK_RE.sub(replacement, line)) + return "".join(cleaned_lines) + + class MarkerLoader(BasePooledParser): """Public ``BasePooledParser`` facade for the Marker Ray pool. @@ -531,8 +543,9 @@ def _split_pages(cls, markdown: str) -> list[tuple[int, str]]: Marker emits ``{1}[PAGE_SEP]{2}[PAGE_SEP]…``. We drop the leading ``[PAGE_SEP]`` segment (Marker prefixes one), - strip ``
``, then split on each ``{N}[PAGE_SEP]`` marker — - the captured ``N`` is the 1-indexed page that just ended. + preserve canonical ``
`` tags inside Markdown pipe-table rows, + remove them elsewhere, then split on each ``{N}[PAGE_SEP]`` marker. + The captured ``N`` is the 1-indexed page that just ended. Blank pages are preserved (text=``""``) so ``page_number`` and ``page_count`` reflect the source document, not just the @@ -544,7 +557,7 @@ def _split_pages(cls, markdown: str) -> list[tuple[int, str]]: return [] if cls.PAGE_SEP in markdown: markdown = markdown.split(cls.PAGE_SEP, 1)[1] - markdown = markdown.replace("
", "") + markdown = _clean_marker_breaks(markdown) pairs: list[tuple[int, str]] = [] cursor = 0 diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index d1988d011..10a8e8be9 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -10,6 +10,7 @@ from core.embeddings.embedder import Embedder from core.indexing.contextualize import ChunkContextualizer from core.indexing.parsers.document_parser import DocumentParser +from core.indexing.structure_normalizer import DocumentStructureNormalizer from core.indexing.topic_tags import TopicTagger from core.models.document import Document, DocumentType from core.utils.logging import get_logger @@ -20,6 +21,7 @@ from services.workers.stages.chunk import chunk_stage from services.workers.stages.contextualize import contextualize_stage from services.workers.stages.embed import embed_stage +from services.workers.stages.normalize_structure import normalize_structure_stage from services.workers.stages.parse import parse_stage from services.workers.stages.store import store_stage from services.workers.stages.topic_tag import topic_tag_stage @@ -35,6 +37,7 @@ class PipelineTimeouts: """Per-stage timeout configuration for an indexing pipeline row.""" parse: float | None = None + normalize_structure: float | None = None caption: float | None = None caption_per_image: float = 0.0 chunk: float | None = None @@ -59,6 +62,7 @@ class IndexingPipeline: caption_prompt: str | None = None contextualizer: ChunkContextualizer | None = None topic_tagger: TopicTagger | None = None + structure_normalizer: DocumentStructureNormalizer | None = None timeouts: PipelineTimeouts = PipelineTimeouts() indexation_config: IndexationPipelineConfig | None = None parser_factory: Callable[[str], DocumentParser] | None = None @@ -105,7 +109,23 @@ async def _timed(name: str, coro: Any) -> None: timings[name] = (time.perf_counter() - start) * 1000.0 try: - await _timed("parse", parse_stage(row, parser, timeout=self.timeouts.parse)) + reconstruction = config.table_reconstruction if config is not None else None + should_normalize = ( + reconstruction is not None + and reconstruction.mode == "automatic" + and self.structure_normalizer is not None + and isinstance(row.get("document"), Document) + and row["document"].content_type is DocumentType.PDF + ) + await _timed( + "parse", + parse_stage( + row, + parser, + timeout=self.timeouts.parse, + preserve_raw_blocks=should_normalize, + ), + ) # The caption decision needs the parsed document (standalone images # always caption), so the VLM is resolved after parse. vlm, vlm_name = self._select_vlm(config) if self._should_caption(row, config) else (None, None) @@ -140,6 +160,16 @@ async def _timed(name: str, coro: Any) -> None: per_image_timeout=self.timeouts.caption_per_image, ), ) + if should_normalize and self.structure_normalizer is not None and reconstruction is not None: + await _timed( + "normalize_structure", + normalize_structure_stage( + row, + self.structure_normalizer, + reconstruction, + timeout=self.timeouts.normalize_structure, + ), + ) await _timed("chunk", chunk_stage(row, chunker, timeout=self.timeouts.chunk)) if contextualizer is not None: await _timed( @@ -387,6 +417,7 @@ def build_indexing_pipeline( caption_prompt: str | None = None, contextualizer: ChunkContextualizer | None = None, topic_tagger: TopicTagger | None = None, + structure_normalizer: DocumentStructureNormalizer | None = None, timeouts: PipelineTimeouts | None = None, indexation_config: IndexationPipelineConfig | None = None, parser_factory: Callable[[str], DocumentParser] | None = None, @@ -408,6 +439,7 @@ def build_indexing_pipeline( caption_prompt=caption_prompt, contextualizer=contextualizer, topic_tagger=topic_tagger, + structure_normalizer=structure_normalizer, timeouts=timeouts or PipelineTimeouts(), indexation_config=indexation_config, parser_factory=parser_factory, diff --git a/openrag/services/workers/stages/normalize_structure.py b/openrag/services/workers/stages/normalize_structure.py new file mode 100644 index 000000000..c419193da --- /dev/null +++ b/openrag/services/workers/stages/normalize_structure.py @@ -0,0 +1,67 @@ +"""Fail-open structural normalization stage for PDF indexing.""" + +from __future__ import annotations + +from collections.abc import MutableMapping +from typing import Any + +from core.config.table_reconstruction import TableReconstructionConfig +from core.indexing.structure_normalizer import DocumentStructureNormalizer +from core.models.document import Document, NormalizationReport, ProcessedDocument +from core.utils.logging import get_logger +from services.workers.stages._common import run_with_optional_timeout, scrub_credentials + +logger = get_logger() + + +async def normalize_structure_stage( + row: MutableMapping[str, Any], + normalizer: DocumentStructureNormalizer, + config: TableReconstructionConfig, + *, + timeout: float | None = None, +) -> MutableMapping[str, Any]: + """Normalize document structure while preserving usable output on failure.""" + try: + document = row.get("document") + processed_document = row.get("processed_document") + if not isinstance(document, Document): + raise ValueError("normalize_structure_stage row must contain a Document under 'document'") + if not isinstance(processed_document, ProcessedDocument): + raise ValueError( + "normalize_structure_stage row must contain a ProcessedDocument under 'processed_document'" + ) + + async def run() -> ProcessedDocument: + return await normalizer.normalize(document, processed_document, config) + + row["processed_document"] = await run_with_optional_timeout(run, timeout) + row["stage"] = "structure_normalized" + row.pop("error", None) + return row + except Exception as exc: # noqa: BLE001 - automatic mode is explicitly fail-open + processed_document = row.get("processed_document") + if isinstance(processed_document, ProcessedDocument): + row["processed_document"] = processed_document.model_copy( + update={ + "normalized_text_blocks": None, + "normalization_report": NormalizationReport( + algorithm_version=config.algorithm_version, + status="partial_fallback", + fallback_reasons=[f"{type(exc).__name__}: {exc}"], + ), + } + ) + logger.bind( + task_id=row.get("task_id"), + filename=row.get("filename", ""), + error_type=type(exc).__name__, + ).warning("Structural normalization failed open; preserving parser output") + row["stage"] = "structure_normalization_fallback" + row.pop("error", None) + return row + finally: + scrub_credentials(row) + + +__all__ = ["normalize_structure_stage"] diff --git a/openrag/services/workers/stages/parse.py b/openrag/services/workers/stages/parse.py index 22a43979a..65a73268e 100644 --- a/openrag/services/workers/stages/parse.py +++ b/openrag/services/workers/stages/parse.py @@ -13,6 +13,7 @@ async def parse_stage( parser: DocumentParser, *, timeout: float | None = None, + preserve_raw_blocks: bool = False, ) -> MutableMapping[str, Any]: """Parse ``row["document"]`` and mutate the row with the stage result.""" @@ -22,6 +23,12 @@ async def parse_stage( raise ValueError("parse_stage row must contain a Document under 'document'") processed = await _parse_with_timeout(parser, document, timeout) + if preserve_raw_blocks: + processed = processed.model_copy( + update={ + "raw_text_blocks": [block.model_copy(deep=True) for block in processed.text_blocks], + } + ) row["processed_document"] = processed row["stage"] = "parsed" row.pop("error", None) diff --git a/tests/integration/api/test_presets.py b/tests/integration/api/test_presets.py index 59a81bdbf..4c540c0d5 100644 --- a/tests/integration/api/test_presets.py +++ b/tests/integration/api/test_presets.py @@ -22,6 +22,7 @@ def _indexation_config(chunk_size: int) -> dict: "chunk_overlap_rate": 0.1, }, "parsing_strategy": "pymupdf", + "table_reconstruction": {"mode": "automatic"}, "enable_image_captioning": False, "enable_contextualization": False, "enable_metadata_extraction": False, @@ -63,6 +64,7 @@ def test_preset_options_crud_and_rename(api_client): assert "single" in option_data["retrieval_types"] # Parsing strategies are derived from IndexationPipelineConfig's Literal. assert set(option_data["parsing_strategies"]) == {"pymupdf", "marker", "docling"} + assert option_data["table_reconstruction_modes"] == ["disabled", "automatic"] create_indexation = api_client.post( "/presets/", @@ -80,6 +82,7 @@ def test_preset_options_crud_and_rename(api_client): ) _assert_success(update_indexation, context="update indexation preset") assert update_indexation.json()["config"]["chunking"]["chunk_size"] == 96 + assert update_indexation.json()["config"]["table_reconstruction"]["mode"] == "automatic" create_retrieval = api_client.post( "/presets/", diff --git a/tests/resources/cross_page_table_rows_803_805.pdf b/tests/resources/cross_page_table_rows_803_805.pdf new file mode 100644 index 000000000..27bea32c3 Binary files /dev/null and b/tests/resources/cross_page_table_rows_803_805.pdf differ diff --git a/tests/unit/api/routers/admin/test_phase14_admin_routers.py b/tests/unit/api/routers/admin/test_phase14_admin_routers.py index 348145a81..a90959039 100644 --- a/tests/unit/api/routers/admin/test_phase14_admin_routers.py +++ b/tests/unit/api/routers/admin/test_phase14_admin_routers.py @@ -571,6 +571,7 @@ async def test_preset_options_return_registered_choices(async_client_factory): assert response.status_code == 200 body = response.json() assert body["chunking_strategies"] == ["recursive_splitter"] + assert body["table_reconstruction_modes"] == ["disabled", "automatic"] assert set(body["retrieval_types"]) == {"single", "multiQuery", "hyde"} assert body["reranker_providers"] == ["infinity", "openai", "tei"] diff --git a/tests/unit/core/chunking/test_recursive.py b/tests/unit/core/chunking/test_recursive.py index c69b2e647..bbe5363b4 100644 --- a/tests/unit/core/chunking/test_recursive.py +++ b/tests/unit/core/chunking/test_recursive.py @@ -76,6 +76,19 @@ def test_recursive_splitter_empty_document_returns_empty(): assert splitter.chunk(doc, partition="p1") == [] +def test_recursive_splitter_prefers_the_complete_normalized_block_view(): + splitter = RecursiveSplitter(chunk_size=20, chunk_overlap_rate=0.0, length_function=_word_tokens) + doc = ProcessedDocument( + document_id="d1", + text_blocks=[TextBlock(text="raw parser output", page_number=1)], + normalized_text_blocks=[TextBlock(text="normalized logical row", page_number=1)], + ) + + chunks = splitter.chunk(doc, partition="p1") + + assert [chunk.text for chunk in chunks] == ["normalized logical row"] + + def test_recursive_splitter_requires_length_function(): import pytest diff --git a/tests/unit/core/chunking/test_table_rows.py b/tests/unit/core/chunking/test_table_rows.py new file mode 100644 index 000000000..5f578b0d0 --- /dev/null +++ b/tests/unit/core/chunking/test_table_rows.py @@ -0,0 +1,407 @@ +import re + +from core.chunking.table_rows import chunk_table_legend, chunk_table_row +from core.models.document import ( + PageBoundaryDecision, + SourceFragment, + TableCellData, + TableLegendData, + TableLegendEntry, + TableRowData, +) + + +def _words(text: str) -> int: + return len(text.split()) + + +def test_oversized_cell_chunks_repeat_row_context_and_preserve_provenance(): + supporting_documents = " ".join(f"requirement-{index}." for index in range(100)) + row = TableRowData( + table_id="table-1", + row_id="row-1", + algorithm_version="adjacent-layout-v1", + table_title="Supporting documents", + section_path=["Annex 10"], + identity_columns=[0, 1, 2], + page_start=803, + page_end=805, + cells=[ + TableCellData(column_index=0, column_name="Number", text="1"), + TableCellData(column_index=1, column_name="Title", text="CST salarié"), + TableCellData(column_index=2, column_name="Reference", text="L. 421-1"), + TableCellData( + column_index=3, + column_name="Supporting documents", + text=supporting_documents, + source_fragments=[ + SourceFragment( + source_block_index=1, + page_number=804, + char_start=0, + char_end=len(supporting_documents), + text_start=0, + text_end=len(supporting_documents), + ) + ], + ), + ], + boundary_decisions=[ + PageBoundaryDecision( + previous_page=803, + next_page=804, + same_table_confidence=0.97, + row_continuation_confidence=0.98, + decision="merged", + reason="test", + ) + ], + ) + + chunks = chunk_table_row(row, chunk_size=50, length_function=_words) + + assert len(chunks) > 1 + assert all(_words(chunk.text) <= 50 for chunk in chunks) + assert all("CST salarié" in chunk.text for chunk in chunks) + assert all("L. 421-1" in chunk.text for chunk in chunks) + assert all("row-1" not in chunk.text for chunk in chunks) + assert all("table-1" not in chunk.text for chunk in chunks) + assert all(chunk.page_number == 804 for chunk in chunks) + assert all(chunk.metadata["page_start"] == 803 for chunk in chunks) + assert all(chunk.metadata["page_end"] == 805 for chunk in chunks) + assert all(chunk.metadata["source_fragments"] for chunk in chunks) + + +def test_unsplit_row_keeps_internal_ids_in_metadata_only(): + row = TableRowData( + table_id="table-opaque-hash", + row_id="row-opaque-hash", + algorithm_version="adjacent-layout-v1", + table_title="Residence permits", + section_path=["Annex"], + identity_columns=[0], + page_start=2, + page_end=3, + cells=[ + TableCellData(column_index=0, column_name="Reference", text="L. 421-1"), + TableCellData( + column_index=1, + column_name="Requirement", + text="Proof of employment", + source_fragments=[ + SourceFragment( + source_block_index=2, + page_number=3, + char_start=0, + char_end=len("Proof of employment"), + text_start=0, + text_end=len("Proof of employment"), + ) + ], + ), + ], + ) + + [chunk] = chunk_table_row(row, chunk_size=100, length_function=_words) + + assert "the first row (row 1)" in chunk.text + assert "the value “L. 421-1” in column “Reference”" in chunk.text + assert "the value “Proof of employment” in column “Requirement”" in chunk.text + assert "table-opaque-hash" not in chunk.text + assert "row-opaque-hash" not in chunk.text + assert chunk.metadata["table_id"] == "table-opaque-hash" + assert chunk.metadata["row_id"] == "row-opaque-hash" + assert chunk.metadata["page_start"] == 2 + assert chunk.metadata["page_end"] == 3 + assert chunk.page_number == 3 + + +def test_split_content_uses_page_from_overlapping_source_fragment(): + page_two = " ".join(f"page-two-{index}." for index in range(24)) + page_three = " ".join(f"page-three-{index}." for index in range(24)) + supporting_documents = f"{page_two}\n\n{page_three}" + page_three_start = len(page_two) + 2 + row = TableRowData( + table_id="table-1", + row_id="row-1", + algorithm_version="adjacent-layout-v1", + table_title="Supporting documents", + identity_columns=[0], + page_start=1, + page_end=3, + cells=[ + TableCellData(column_index=0, column_name="Reference", text="L. 421-1"), + TableCellData( + column_index=1, + column_name="Supporting documents", + text=supporting_documents, + source_fragments=[ + SourceFragment( + source_block_index=1, + page_number=2, + char_start=0, + char_end=len(page_two), + text_start=0, + text_end=len(page_two), + ), + SourceFragment( + source_block_index=2, + page_number=3, + char_start=0, + char_end=len(page_three), + text_start=page_three_start, + text_end=len(supporting_documents), + ), + ], + ), + ], + ) + + chunks = chunk_table_row(row, chunk_size=15, length_function=_words) + page_three_chunks = [chunk for chunk in chunks if "page-three" in chunk.text and "page-two" not in chunk.text] + + assert page_three_chunks + assert all(chunk.page_number == 3 for chunk in page_three_chunks) + assert all( + any(fragment["page_number"] == 3 for fragment in chunk.metadata["source_fragments"]) + for chunk in page_three_chunks + ) + + +def test_normal_table_rows_render_as_independent_natural_language_evidence(): + first = TableRowData( + table_id="table-a", + row_id="row-a-1", + algorithm_version="adjacent-layout-v1", + table_title="Table A", + row_index=1, + page_start=1, + page_end=1, + cells=[ + TableCellData(column_index=0, column_name="aa", text="22"), + TableCellData(column_index=1, column_name="bb", text="Paris"), + TableCellData(column_index=2, column_name="cc", text="Active"), + ], + ) + second = first.model_copy( + update={ + "row_id": "row-a-2", + "row_index": 2, + "cells": [ + TableCellData(column_index=0, column_name="aa", text="35"), + TableCellData(column_index=1, column_name="bb", text="Lyon"), + TableCellData(column_index=2, column_name="cc", text="Inactive"), + ], + } + ) + + first_chunk = chunk_table_row(first, chunk_size=100, length_function=_words)[0] + second_chunk = chunk_table_row(second, chunk_size=100, length_function=_words)[0] + + assert first_chunk.text == ( + "In table “Table A”, the first row (row 1) has the value “22” in column “aa”, " + "the value “Paris” in column “bb”, and the value “Active” in column “cc”." + ) + assert "the second row (row 2)" in second_chunk.text + assert "the value “35” in column “aa”" in second_chunk.text + assert first_chunk.metadata["row_index"] == 1 + assert first_chunk.metadata["table_content_kind"] == "row" + + query_terms = set(re.findall(r"\w+", "Table A first row aa 22".casefold())) + scores = [ + len(query_terms & set(re.findall(r"\w+", chunk.text.casefold()))) for chunk in (first_chunk, second_chunk) + ] + assert scores[0] > scores[1] + + +def test_legend_is_a_separate_searchable_chunk(): + legend = TableLegendData( + table_id="permits", + algorithm_version="adjacent-layout-v1", + table_title="ANNEXE", + section_path=["Article Annexe 10", "ANNEXE"], + page_number=1, + entries=[ + TableLegendEntry(abbreviation="CST", meaning="carte de séjour temporaire"), + TableLegendEntry(abbreviation="CSP", meaning="carte de séjour pluriannuelle"), + ], + ) + + [chunk] = chunk_table_legend(legend, chunk_size=100, length_function=_words) + + assert "CST means “carte de séjour temporaire”" in chunk.text + assert "CSP means “carte de séjour pluriannuelle”" in chunk.text + assert chunk.metadata["table_content_kind"] == "legend" + assert chunk.metadata["legend_abbreviations"] == ["CST", "CSP"] + + +def test_merged_and_empty_cells_do_not_shift_or_duplicate_values(): + row = TableRowData( + table_id="merged", + row_id="merged-row", + algorithm_version="adjacent-layout-v1", + row_index=1, + page_start=1, + page_end=1, + cells=[ + TableCellData( + column_index=0, + column_name="Region", + text="North", + column_span=2, + ), + TableCellData( + column_index=1, + column_name="Area", + covered_by=(1, 0), + ), + TableCellData( + column_index=2, + column_name="Owner", + text="", + explicit_empty=True, + ), + ], + ) + + [chunk] = chunk_table_row(row, chunk_size=100, length_function=_words) + + assert "the value “North” in columns “Region” and “Area”" in chunk.text + assert chunk.text.count("North") == 1 + assert "no value in column “Owner”" in chunk.text + + +def test_oversized_identity_is_preserved_as_complete_text_instead_of_truncated(): + row = TableRowData( + table_id="long-identity", + row_id="long-identity-row", + algorithm_version="adjacent-layout-v1", + identity_columns=[0, 1, 2], + page_start=1, + page_end=1, + cells=[ + TableCellData( + column_index=0, + column_name="Category", + text="A very long professional residence permit category", + ), + TableCellData(column_index=1, column_name="Permit", text="CST salarié"), + TableCellData(column_index=2, column_name="Reference", text="L. 421-1"), + TableCellData( + column_index=3, + column_name="Supporting documents", + text=" ".join(f"requirement-{index}" for index in range(40)), + ), + ], + ) + + chunks = chunk_table_row(row, chunk_size=30, length_function=_words) + + assert len(chunks) > 1 + assert all(_words(chunk.text) <= 30 for chunk in chunks) + assert all("A very long professional residence permit category" in chunk.text for chunk in chunks) + assert all("CST salarié" in chunk.text for chunk in chunks) + assert all("L. 421-1" in chunk.text for chunk in chunks) + assert all("has “professional" not in chunk.text for chunk in chunks) + + +def test_single_oversized_legend_definition_is_split_without_losing_its_meaning(): + meaning = " ".join(f"definition-{index}" for index in range(40)) + legend = TableLegendData( + table_id="long-legend", + algorithm_version="adjacent-layout-v1", + table_title="Table A", + page_number=1, + entries=[ + TableLegendEntry(abbreviation="ABC", meaning=meaning), + ], + ) + + chunks = chunk_table_legend(legend, chunk_size=16, length_function=_words) + + assert len(chunks) > 1 + assert all(_words(chunk.text) <= 16 for chunk in chunks) + assert all("ABC means" in chunk.text for chunk in chunks) + recovered = " ".join(re.search(r"ABC means “(.+)”\.", chunk.text).group(1) for chunk in chunks) + assert recovered == meaning + + +def test_context_that_exceeds_chunk_size_is_emitted_as_bounded_row_parts(): + identity = " ".join(f"identity-{index}" for index in range(20)) + content = " ".join(f"content-{index}" for index in range(10)) + row = TableRowData( + table_id="extreme", + row_id="extreme-row", + algorithm_version="adjacent-layout-v1", + identity_columns=[0], + page_start=1, + page_end=1, + cells=[ + TableCellData(column_index=0, column_name="Identity", text=identity), + TableCellData(column_index=1, column_name="Content", text=content), + ], + ) + + chunks = chunk_table_row(row, chunk_size=12, length_function=_words) + + assert len(chunks) > 1 + assert all(_words(chunk.text) <= 12 for chunk in chunks) + recovered_content = " ".join( + chunk.text.rsplit("\n\n", 1)[-1] for chunk in chunks if chunk.metadata["content_column"] == "Content" + ) + recovered_identity = " ".join( + chunk.text.rsplit("\n\n", 1)[-1] for chunk in chunks if chunk.metadata["content_column"] == "Identity" + ) + assert recovered_content == content + assert recovered_identity == identity + assert all(chunk.metadata["row_id"] == "extreme-row" for chunk in chunks) + + +def test_legend_scope_that_exceeds_chunk_size_still_preserves_every_word(): + meaning = " ".join(f"meaning-{index}" for index in range(10)) + legend = TableLegendData( + table_id="extreme-legend", + algorithm_version="adjacent-layout-v1", + table_title="A deliberately long table title", + page_number=1, + entries=[TableLegendEntry(abbreviation="XYZ", meaning=meaning)], + ) + + chunks = chunk_table_legend(legend, chunk_size=6, length_function=_words) + + assert len(chunks) > 1 + assert all("XYZ means" in chunk.text for chunk in chunks) + assert all(_words(chunk.text) <= 6 for chunk in chunks) + recovered = " ".join(re.search(r"XYZ means “(.+)”\.", chunk.text).group(1) for chunk in chunks) + assert recovered == meaning + + +def test_oversized_row_repeats_explicit_empty_identity_context(): + row = TableRowData( + table_id="empty-context", + row_id="empty-context-row", + algorithm_version="adjacent-layout-v1", + identity_columns=[0, 1], + page_start=1, + page_end=1, + cells=[ + TableCellData(column_index=0, column_name="Reference", text="A-1"), + TableCellData( + column_index=1, + column_name="Owner", + text="", + explicit_empty=True, + ), + TableCellData( + column_index=2, + column_name="Details", + text=" ".join(f"detail-{index}" for index in range(30)), + ), + ], + ) + + chunks = chunk_table_row(row, chunk_size=20, length_function=_words) + + assert len(chunks) > 1 + assert all(_words(chunk.text) <= 20 for chunk in chunks) + assert all("“Owner” = empty" in chunk.text for chunk in chunks) diff --git a/tests/unit/core/config/test_pipeline_configs.py b/tests/unit/core/config/test_pipeline_configs.py index cef085fd6..1f3ecb501 100644 --- a/tests/unit/core/config/test_pipeline_configs.py +++ b/tests/unit/core/config/test_pipeline_configs.py @@ -46,6 +46,27 @@ def test_indexation_pipeline_topic_tagging_defaults_off(): assert IndexationPipelineConfig().enable_topic_tagging is False +def test_table_reconstruction_defaults_to_disabled(): + config = IndexationPipelineConfig() + + assert config.table_reconstruction.mode == "disabled" + assert config.table_reconstruction.algorithm_version == "adjacent-layout-v1" + + +@pytest.mark.parametrize( + "payload", + [ + {"mode": "strict"}, + {"mode": "automatic", "same_table_min_confidence": 0.79}, + {"mode": "automatic", "row_continuation_min_confidence": 1.01}, + {"mode": "automatic", "unknown_threshold": 0.9}, + ], +) +def test_table_reconstruction_rejects_unsupported_or_unsafe_config(payload: dict): + with pytest.raises(ValidationError): + IndexationPipelineConfig(table_reconstruction=payload) + + def test_retrieval_pipeline_rejects_unknown_type(): """Retrieval presets reject unsupported retrieval modes.""" with pytest.raises(ValidationError): diff --git a/tests/unit/core/indexing/test_contextualize.py b/tests/unit/core/indexing/test_contextualize.py index 1385bbad7..05420a746 100644 --- a/tests/unit/core/indexing/test_contextualize.py +++ b/tests/unit/core/indexing/test_contextualize.py @@ -2,7 +2,7 @@ import pytest from core.indexing.contextualize import ChunkContextualizer -from core.models.chunk import Chunk +from core.models.chunk import Chunk, ChunkType class DictLLM: @@ -55,3 +55,63 @@ async def test_contextualizer_holds_llm_semaphore_around_chat(): # The injected gate was entered for the chat call and released afterwards. assert gate_held == [True] assert _TrackingGate.depth == 0 + + +class _RecordingLLM: + def __init__(self) -> None: + self.calls: list[list[dict[str, str]]] = [] + + async def chat(self, messages, **kwargs): + self.calls.append(messages) + return {"choices": [{"message": {"content": "ordinary context"}}]} + + +@pytest.mark.asyncio +async def test_contextualizer_leaves_structured_table_rows_and_legends_unchanged(): + llm = _RecordingLLM() + contextualizer = ChunkContextualizer(llm, "System prompt") + row = Chunk( + id="row", + text="deterministic row text", + chunk_type=ChunkType.TABLE, + metadata={"table_content_kind": "row", "table_id": "table-a"}, + ) + legend = Chunk( + id="legend", + text="deterministic legend text", + chunk_type=ChunkType.TABLE, + metadata={"table_content_kind": "legend", "table_id": "table-a"}, + ) + + result = await contextualizer.contextualize([row, legend], filename="table.pdf") + + assert llm.calls == [] + assert result[0] is row + assert result[1] is legend + assert result[0].model_dump() == row.model_dump() + assert result[1].model_dump() == legend.model_dump() + + +@pytest.mark.asyncio +async def test_contextualizer_still_processes_ordinary_chunks_in_a_mixed_batch(): + llm = _RecordingLLM() + contextualizer = ChunkContextualizer(llm, "System prompt") + table_row = Chunk( + id="row", + text="deterministic row text", + chunk_type=ChunkType.TABLE, + metadata={"table_content_kind": "row", "table_id": "table-a"}, + ) + ordinary = Chunk(id="text", text="ordinary paragraph", chunk_type=ChunkType.TEXT) + + result = await contextualizer.contextualize([table_row, ordinary], filename="mixed.pdf") + + assert len(llm.calls) == 1 + assert result[0] is table_row + assert result[0].text == "deterministic row text" + assert result[0].context is None + assert result[0].content is None + assert result[1].context == "ordinary context" + assert result[1].content == "ordinary paragraph" + assert "[CONTEXT]" in result[1].text + assert "ordinary paragraph" in result[1].text diff --git a/tests/unit/core/indexing/test_table_normalizer.py b/tests/unit/core/indexing/test_table_normalizer.py new file mode 100644 index 000000000..96b467b64 --- /dev/null +++ b/tests/unit/core/indexing/test_table_normalizer.py @@ -0,0 +1,1350 @@ +from pathlib import Path + +import pytest +from core.config.table_reconstruction import TableReconstructionConfig +from core.indexing.parsers.pdf.pymupdf import PyMuPDFParser +from core.indexing.structure_normalizer import ( + LayoutCellEvidence, + LayoutRowEvidence, + LayoutTableEvidence, + LayoutWord, + PageLayoutEvidence, + TableLayoutEvidenceProvider, +) +from core.indexing.table_normalizer import DeterministicTableNormalizer +from core.models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from core.prompts.vlm_prompt_builder import wrap_caption +from services.workers.layout import PyMuPDFTableEvidenceProvider + +FIXTURE = Path(__file__).parents[3] / "resources" / "cross_page_table_rows_803_805.pdf" + + +async def _normalize(config: TableReconstructionConfig | None = None, *, corrupt_identity: bool = False): + document = Document( + filename=FIXTURE.name, + content_type=DocumentType.PDF, + raw_bytes=FIXTURE.read_bytes(), + ) + parsed = await PyMuPDFParser().parse(document) + raw_blocks = [block.model_copy(deep=True) for block in parsed.text_blocks] + if corrupt_identity: + raw_blocks[0].text = raw_blocks[0].text.replace('CST portant la mention " salarié "', "unrelated identity") + parsed = parsed.model_copy(update={"raw_text_blocks": raw_blocks}) + result = await DeterministicTableNormalizer(PyMuPDFTableEvidenceProvider()).normalize( + document, + parsed, + config or TableReconstructionConfig(mode="automatic"), + ) + return parsed, result + + +@pytest.mark.asyncio +async def test_reconstructs_primary_regression_and_keeps_the_next_row_separate(): + parsed, result = await _normalize() + + # This is the production Marker shape observed for the fixture: a Markdown + # anchor, a plain-text middle continuation, then a synthetic-column table. + assert "|Col1|Col2|Col3|Col4|" not in parsed.raw_text_blocks[0].text + assert "2.1. Si vous occupez toujours" in parsed.raw_text_blocks[1].text + assert "|" not in parsed.raw_text_blocks[1].text + assert parsed.raw_text_blocks[2].text.startswith("|Col1|Col2|Col3|Col4|") + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + legends = [block.table_legend for block in result.effective_text_blocks() if block.table_legend is not None] + assert result.normalization_report.status == "normalized" + assert result.normalization_report.reconstructed_row_count == 2 + assert len(rows) == 2 + assert len(legends) == 1 + + first, second = rows + first_text = "\n".join(cell.text for cell in first.cells) + second_text = "\n".join(cell.text for cell in second.cells) + assert first.page_start == 1 + assert first.page_end == 3 + assert first.table_title == "ANNEXE" + assert first.section_path == ["Article Annexe 10"] + assert len(first.scope_fragments) == 2 + assert "CST portant la mention" in first_text + assert "L. 421-1" in first_text + assert "4.2. Si vous n'occupez plus" in first_text + assert "travailleur temporaire" not in first_text + assert "travailleur temporaire" in second_text + assert "L. 421-3" in second_text + assert [cell.column_name for cell in first.cells][1:5] == [ + "Catégorie de titre de séjour", + "Libellé", + "Référence du CESEDA", + "Pièces justificatives", + ] + assert [(entry.abbreviation, entry.meaning) for entry in legends[0].entries] == [ + ("APS", "autorisation provisoire de séjour"), + ("CST", "carte de séjour temporaire"), + ("CSP", "carte de séjour pluriannuelle"), + ("CR", "carte de résident"), + ] + assert parsed.raw_text_blocks == result.raw_text_blocks + normalized_text = "\n".join(block.text for block in result.effective_text_blocks()) + assert normalized_text.count("2. Pièces à fournir lorsque") == 1 + assert "Row: row-" not in normalized_text + assert "LibelléAPS" not in normalized_text + assert "CST means “carte de séjour temporaire”" in normalized_text + row_text = next(block.text for block in result.effective_text_blocks() if block.table_row is first) + assert "APS means" not in row_text + assert "the value “L. 421-1” in column “Référence du CESEDA”" in row_text + assert '"salarié"' in row_text + assert '" salarié "' not in row_text + assert "le 5 juin" in row_text + for spacing_artifact in ("c as", "s éjour", "d 'emploi"): + assert spacing_artifact not in row_text + + for row in rows: + for cell in row.cells: + if cell.text: + assert cell.source_fragments + for fragment in cell.source_fragments: + raw = result.raw_text_blocks[fragment.source_block_index].text + assert raw[fragment.char_start : fragment.char_end] + + +@pytest.mark.asyncio +async def test_reconstructs_when_marker_caption_replaces_the_sparse_continuation(): + document = Document( + filename=FIXTURE.name, + content_type=DocumentType.PDF, + raw_bytes=FIXTURE.read_bytes(), + ) + parsed = await PyMuPDFParser().parse(document) + caption = parsed.text_blocks[1].text + markdown_ref = "![](_page_1_Picture_0.jpeg)" + unrelated_ref = "![](_page_1_Picture_1.jpeg)" + unrelated_caption = "A decorative seal that is unrelated to the legal table." + raw_blocks = [block.model_copy(deep=True) for block in parsed.text_blocks] + raw_blocks[1] = raw_blocks[1].model_copy(update={"text": f"{markdown_ref}\n\n{unrelated_ref}"}) + working_blocks = [block.model_copy(deep=True) for block in parsed.text_blocks] + working_blocks[1] = working_blocks[1].model_copy( + update={"text": f"{wrap_caption(caption)}\n\n{wrap_caption(unrelated_caption)}"} + ) + marker_output = parsed.model_copy( + update={ + "raw_text_blocks": raw_blocks, + "text_blocks": working_blocks, + "images": [ + ImageBlock( + page_number=2, + caption=caption, + metadata={ + "markdown_ref": markdown_ref, + "marker_key": "_page_1_Picture_0.jpeg", + }, + ), + ImageBlock( + page_number=2, + caption=unrelated_caption, + metadata={ + "markdown_ref": unrelated_ref, + "marker_key": "_page_1_Picture_1.jpeg", + }, + ), + ], + } + ) + + result = await DeterministicTableNormalizer(PyMuPDFTableEvidenceProvider()).normalize( + document, + marker_output, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert result.normalization_report.status == "normalized" + assert len(rows) == 2 + assert rows[0].page_start == 1 + assert rows[0].page_end == 3 + assert "L. 421-1" in "\n".join(cell.text for cell in rows[0].cells) + assert "4.2. Si vous n'occupez plus" in "\n".join(cell.text for cell in rows[0].cells) + assert "L. 421-3" in "\n".join(cell.text for cell in rows[1].cells) + assert result.raw_text_blocks == raw_blocks + normalized_text = "\n".join(block.text for block in result.effective_text_blocks()) + assert markdown_ref not in normalized_text + assert caption not in normalized_text + assert unrelated_caption in normalized_text + assert normalized_text.count("4.2. Si vous n'occupez plus") == 1 + + image_fragments = [ + fragment + for cell in rows[0].cells + for fragment in cell.source_fragments + if fragment.source_kind == "pdf_layout" and fragment.source_ref is not None + ] + assert image_fragments + assert all(fragment.evidence_provider == "pymupdf" for fragment in image_fragments) + assert all(fragment.page_number == 2 for fragment in image_fragments) + assert all( + raw_blocks[fragment.source_block_index].text[fragment.char_start : fragment.char_end] == markdown_ref + for fragment in image_fragments + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ambiguous", [False, True]) +async def test_marker_caption_alignment_fails_open_when_evidence_is_unsafe(ambiguous): + document = Document( + filename=FIXTURE.name, + content_type=DocumentType.PDF, + raw_bytes=FIXTURE.read_bytes(), + ) + parsed = await PyMuPDFParser().parse(document) + evidence_text = parsed.text_blocks[1].text + captions = [evidence_text, evidence_text] if ambiguous else ["An unrelated photograph of a building."] + refs = [f"![](_page_1_Picture_{index}.jpeg)" for index in range(len(captions))] + raw_blocks = [block.model_copy(deep=True) for block in parsed.text_blocks] + raw_blocks[1] = raw_blocks[1].model_copy(update={"text": "\n\n".join(refs)}) + working_blocks = [block.model_copy(deep=True) for block in parsed.text_blocks] + working_blocks[1] = working_blocks[1].model_copy( + update={"text": "\n\n".join(wrap_caption(caption) for caption in captions)} + ) + marker_output = parsed.model_copy( + update={ + "raw_text_blocks": raw_blocks, + "text_blocks": working_blocks, + "images": [ + ImageBlock( + page_number=2, + caption=caption, + metadata={ + "markdown_ref": ref, + "marker_key": f"_page_1_Picture_{index}.jpeg", + }, + ) + for index, (ref, caption) in enumerate(zip(refs, captions, strict=True)) + ], + } + ) + + result = await DeterministicTableNormalizer(PyMuPDFTableEvidenceProvider()).normalize( + document, + marker_output, + TableReconstructionConfig(mode="automatic"), + ) + + assert result.normalized_text_blocks is None + assert result.normalization_report.status == "unchanged" + assert result.raw_text_blocks == raw_blocks + assert result.effective_text_blocks() == working_blocks + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "config", + [ + TableReconstructionConfig(mode="automatic", same_table_min_confidence=0.98), + TableReconstructionConfig(mode="automatic", row_continuation_min_confidence=0.99), + ], +) +async def test_independent_boundary_thresholds_can_reject_the_merge(config): + _, result = await _normalize(config) + + assert result.normalized_text_blocks is None + assert result.normalization_report.status == "unchanged" + + +@pytest.mark.asyncio +async def test_uncertain_cell_alignment_preserves_the_parser_output(): + _, result = await _normalize(corrupt_identity=True) + + assert result.normalized_text_blocks is None + assert result.normalization_report.status == "unchanged" + + +class FakeEvidenceProvider(TableLayoutEvidenceProvider): + def __init__(self, pages, *, discovered=None): + self.pages = pages + self.discovered = set(discovered or ()) + self.calls = [] + + async def discover(self, document): + return set(self.discovered) + + async def collect(self, document, page_numbers): + self.calls.append(set(page_numbers)) + return [self.pages[page] for page in sorted(page_numbers) if page in self.pages] + + +@pytest.mark.asyncio +async def test_short_decorative_caption_does_not_trigger_layout_scanning(): + markdown_ref = "![](_page_0_Picture_0.jpeg)" + raw = TextBlock(text=markdown_ref, page_number=1) + processed = ProcessedDocument( + text_blocks=[raw.model_copy(update={"text": wrap_caption("Company logo")})], + raw_text_blocks=[raw], + images=[ + ImageBlock( + page_number=1, + caption="Company logo", + metadata={"markdown_ref": markdown_ref}, + ) + ], + page_count=1, + ) + provider = FakeEvidenceProvider({}) + + result = await DeterministicTableNormalizer(provider).normalize( + Document(filename="decorative.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf"), + processed, + TableReconstructionConfig(mode="automatic"), + ) + + assert provider.calls == [] + assert result.normalized_text_blocks is None + assert result.effective_text_blocks() == processed.text_blocks + + +@pytest.mark.asyncio +async def test_untrusted_parser_header_cannot_rename_reconstructed_columns(): + header = LayoutRowEvidence( + cells=( + LayoutCellEvidence(0, "Reference", (0.0, 0.70, 0.30, 0.72)), + LayoutCellEvidence(1, "Documents", (0.30, 0.70, 1.0, 0.72)), + ), + bbox=(0.0, 0.70, 1.0, 0.72), + ) + data_rows = tuple( + LayoutRowEvidence( + cells=( + LayoutCellEvidence(0, f"A{index}", (0.0, 0.72, 0.30, 0.95)), + LayoutCellEvidence(1, f"description-{index}", (0.30, 0.72, 1.0, 0.95)), + ), + bbox=(0.0, 0.72, 1.0, 0.95), + ) + for index in range(15) + ) + table = LayoutTableEvidence( + page_number=1, + bbox=(0.0, 0.70, 1.0, 0.95), + column_bounds=((0.0, 0.30), (0.30, 1.0)), + rows=(header, *data_rows), + ) + pages = { + 1: PageLayoutEvidence(page_number=1, width=100, height=100, tables=(table,)), + 2: PageLayoutEvidence( + page_number=2, + width=100, + height=100, + words=(LayoutWord("continued", (0.40, 0.05, 0.65, 0.10), 0, 0, 0),), + ), + } + markdown_rows = "\n".join(f"| A{index} | description-{index} |" for index in range(15)) + raw_blocks = [ + TextBlock( + text=f"| Reference wrong | Documents wrong |\n|---|---|\n{markdown_rows}", + page_number=1, + ), + TextBlock(text="continued", page_number=2), + ] + processed = ProcessedDocument( + text_blocks=[block.model_copy(deep=True) for block in raw_blocks], + raw_text_blocks=raw_blocks, + page_count=2, + ) + document = Document(filename="synthetic.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf") + + result = await DeterministicTableNormalizer(FakeEvidenceProvider(pages)).normalize( + document, + processed, + TableReconstructionConfig(mode="automatic"), + ) + + assert result.normalized_text_blocks is None + assert result.effective_text_blocks() == processed.text_blocks + assert "Reference wrong" in result.effective_text_blocks()[0].text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("whole_table_image", [False, True]) +async def test_reconstructs_when_marker_caption_represents_table_content(whole_table_image): + caption = " ".join(f"supporting-document-{index}" for index in range(20)) + markdown_ref = "![](_page_0_Picture_0.jpeg)" + header = LayoutRowEvidence( + cells=( + LayoutCellEvidence(0, "Type C visa", (0.0, 0.80, 0.30, 0.85)), + LayoutCellEvidence(1, "Description", (0.30, 0.80, 1.0, 0.85)), + ), + bbox=(0.0, 0.80, 1.0, 0.85), + ) + data = LayoutRowEvidence( + cells=( + LayoutCellEvidence(0, "A", (0.0, 0.85, 0.30, 0.96)), + LayoutCellEvidence(1, caption, (0.30, 0.85, 1.0, 0.96)), + ), + bbox=(0.0, 0.85, 1.0, 0.96), + ) + pages = { + 1: PageLayoutEvidence( + page_number=1, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=1, + bbox=(0.0, 0.80, 1.0, 0.96), + column_bounds=((0.0, 0.30), (0.30, 1.0)), + rows=(header, data), + ), + ), + ), + 2: PageLayoutEvidence( + page_number=2, + width=100, + height=100, + words=(LayoutWord("continued", (0.40, 0.05, 0.65, 0.10), 0, 0, 0),), + ), + } + anchor_text = ( + markdown_ref if whole_table_image else f"| Type C visa | Description |\n|---|---|\n| A | {markdown_ref} |" + ) + raw_blocks = [TextBlock(text=anchor_text, page_number=1), TextBlock(text="continued", page_number=2)] + working_blocks = [ + raw_blocks[0].model_copy(update={"text": raw_blocks[0].text.replace(markdown_ref, wrap_caption(caption))}), + raw_blocks[1].model_copy(deep=True), + ] + processed = ProcessedDocument( + text_blocks=working_blocks, + raw_text_blocks=raw_blocks, + images=[ + ImageBlock( + page_number=1, + caption=caption, + metadata={"markdown_ref": markdown_ref, "marker_key": "_page_0_Picture_0.jpeg"}, + ) + ], + page_count=2, + ) + document = Document(filename="synthetic.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf") + + result = await DeterministicTableNormalizer(FakeEvidenceProvider(pages)).normalize( + document, + processed, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert len(rows) == 1 + assert rows[0].page_end == 2 + assert rows[0].cells[0].column_name == "Type C visa" + assert caption in rows[0].cells[1].text + assert "continued" in rows[0].cells[1].text + assert any(fragment.source_kind == "pdf_layout" for fragment in rows[0].cells[1].source_fragments) + normalized = "\n".join(block.text for block in result.effective_text_blocks()) + assert normalized.count(caption) == 1 + assert markdown_ref not in normalized + assert "" not in normalized + + +@pytest.mark.asyncio +async def test_candidate_window_expands_across_two_sparse_continuation_pages(): + header = LayoutRowEvidence( + cells=( + LayoutCellEvidence(0, "ID", (0.0, 0.80, 0.30, 0.85)), + LayoutCellEvidence(1, "Description", (0.30, 0.80, 1.0, 0.85)), + ), + bbox=(0.0, 0.80, 1.0, 0.85), + ) + data = LayoutRowEvidence( + cells=( + LayoutCellEvidence(0, "A", (0.0, 0.85, 0.30, 0.96)), + LayoutCellEvidence(1, "beginning", (0.30, 0.85, 1.0, 0.96)), + ), + bbox=(0.0, 0.85, 1.0, 0.96), + ) + pages = { + 1: PageLayoutEvidence( + page_number=1, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=1, + bbox=(0.0, 0.80, 1.0, 0.96), + column_bounds=((0.0, 0.30), (0.30, 1.0)), + rows=(header, data), + ), + ), + ), + 2: PageLayoutEvidence( + page_number=2, + width=100, + height=100, + words=( + LayoutWord("middle", (0.40, 0.05, 0.60, 0.10), 0, 0, 0), + LayoutWord("continues", (0.40, 0.84, 0.70, 0.88), 1, 0, 0), + ), + ), + 3: PageLayoutEvidence( + page_number=3, + width=100, + height=100, + words=(LayoutWord("end", (0.40, 0.05, 0.55, 0.10), 0, 0, 0),), + ), + } + provider = FakeEvidenceProvider(pages) + document = Document(filename="synthetic.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf") + raw_blocks = [ + TextBlock(text="| ID | Description |\n|---|---|\n| A | beginning |", page_number=1), + TextBlock(text="middle\n\ncontinues", page_number=2), + TextBlock(text="end", page_number=3), + ] + processed = ProcessedDocument( + text_blocks=[block.model_copy(deep=True) for block in raw_blocks], + raw_text_blocks=raw_blocks, + page_count=3, + ) + + result = await DeterministicTableNormalizer(provider).normalize( + document, + processed, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert provider.calls[0] == {1, 2} + assert provider.calls[-1] == {3} + assert len(rows) == 1 + assert rows[0].page_end == 3 + assert "middle" in rows[0].cells[1].text + assert "end" in rows[0].cells[1].text + + +_THREE_COLUMN_BOUNDS = ((0.10, 0.36), (0.36, 0.63), (0.63, 0.90)) + + +def _three_column_row( + values: tuple[str, str, str], + y0: float, + y1: float, + *, + cell_options: dict[int, dict] | None = None, +) -> LayoutRowEvidence: + options = cell_options or {} + cells = tuple( + LayoutCellEvidence( + column_index=column_index, + text=value, + bbox=None if options.get(column_index, {}).get("slot_state") == "covered" else (left, y0, right, y1), + **options.get(column_index, {}), + ) + for column_index, (value, (left, right)) in enumerate(zip(values, _THREE_COLUMN_BOUNDS, strict=True)) + ) + return LayoutRowEvidence(cells=cells, bbox=(0.10, y0, 0.90, y1)) + + +def _one_page_table( + rows: tuple[LayoutRowEvidence, ...], + *, + bbox: tuple[float, float, float, float] = (0.10, 0.25, 0.90, 0.65), +) -> PageLayoutEvidence: + return PageLayoutEvidence( + page_number=1, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=1, + bbox=bbox, + column_bounds=_THREE_COLUMN_BOUNDS, + rows=rows, + ), + ), + ) + + +async def _normalize_one_page_table( + markdown: str, + evidence: PageLayoutEvidence, + *, + filename: str, +) -> tuple[ProcessedDocument, ProcessedDocument]: + raw = TextBlock(text=markdown, page_number=1) + processed = ProcessedDocument( + text_blocks=[raw.model_copy(deep=True)], + raw_text_blocks=[raw], + page_count=1, + ) + result = await DeterministicTableNormalizer(FakeEvidenceProvider({1: evidence})).normalize( + Document( + filename=filename, + content_type=DocumentType.PDF, + raw_bytes=b"pdf", + ), + processed, + TableReconstructionConfig(mode="automatic"), + ) + return processed, result + + +@pytest.mark.asyncio +async def test_normalizes_an_ordinary_table_and_preserves_surrounding_paragraphs(): + evidence = _one_page_table( + ( + _three_column_row(("aa", "bb", "cc"), 0.25, 0.35), + _three_column_row(("22", "Paris", "Active"), 0.35, 0.45), + _three_column_row(("35", "Lyon", "Inactive"), 0.45, 0.55), + ) + ) + markdown = ( + "Paragraph before the table.\n\n" + "# Table A\n\n" + "| aa | bb | cc |\n" + "|---|---|---|\n" + "| 22 | Paris | Active |\n" + "| 35 | Lyon | Inactive |\n\n" + "Paragraph after the table." + ) + + parsed, result = await _normalize_one_page_table( + markdown, + evidence, + filename="table-a.pdf", + ) + + blocks = result.effective_text_blocks() + rows = [block.table_row for block in blocks if block.table_row is not None] + assert result.normalization_report.status == "normalized" + assert result.normalization_report.reconstructed_row_count == 2 + assert len(rows) == 2 + assert [row.row_index for row in rows] == [1, 2] + assert all(row.table_title == "Table A" for row in rows) + assert all(row.section_path == [] for row in rows) + assert [[cell.column_name for cell in row.cells] for row in rows] == [ + ["aa", "bb", "cc"], + ["aa", "bb", "cc"], + ] + assert [[cell.text for cell in row.cells] for row in rows] == [ + ["22", "Paris", "Active"], + ["35", "Lyon", "Inactive"], + ] + ordinary_text = "\n".join(block.text for block in blocks if block.table_row is None and block.table_legend is None) + assert "Paragraph before the table." in ordinary_text + assert "Paragraph after the table." in ordinary_text + assert parsed.raw_text_blocks == result.raw_text_blocks + + +@pytest.mark.asyncio +async def test_untitled_table_does_not_invent_the_filename_as_its_title(): + evidence = _one_page_table( + ( + _three_column_row(("aa", "bb", "cc"), 0.25, 0.35), + _three_column_row(("22", "Paris", "Active"), 0.35, 0.45), + _three_column_row(("35", "Lyon", "Inactive"), 0.45, 0.55), + ) + ) + markdown = ( + "Paragraph before.\n\n" + "| aa | bb | cc |\n" + "|---|---|---|\n" + "| 22 | Paris | Active |\n" + "| 35 | Lyon | Inactive |\n\n" + "Paragraph after." + ) + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="must-not-become-the-table-title.pdf", + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert len(rows) == 2 + assert all(row.table_title is None for row in rows) + assert all(row.section_path == [] for row in rows) + assert "must-not-become-the-table-title.pdf" not in "\n".join( + block.text for block in result.effective_text_blocks() + ) + + +@pytest.mark.asyncio +async def test_untitled_table_uses_heading_ancestry_without_inventing_a_title(): + evidence = _one_page_table( + ( + _three_column_row(("aa", "bb", "cc"), 0.25, 0.35), + _three_column_row(("22", "Paris", "Active"), 0.35, 0.45), + ) + ) + markdown = ( + "# Report\n\n" + "## Previous section\n\n" + "Previous prose.\n\n" + "## Results\n\n" + "| aa | bb | cc |\n" + "|---|---|---|\n" + "| 22 | Paris | Active |" + ) + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="report.pdf", + ) + + [row] = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert row.table_title is None + assert row.section_path == ["Report", "Results"] + assert "Previous section" not in row.section_path + assert len(row.scope_fragments) == 2 + + +@pytest.mark.asyncio +async def test_vertical_merge_is_inherited_while_explicit_empty_stays_empty(): + evidence = _one_page_table( + ( + _three_column_row(("Category", "City", "Status"), 0.25, 0.35), + _three_column_row( + ("France", "Paris", "Active"), + 0.35, + 0.45, + cell_options={0: {"row_span": 2}}, + ), + _three_column_row( + ("", "", "Inactive"), + 0.45, + 0.55, + cell_options={ + 0: { + "slot_state": "covered", + "covered_by": (1, 0), + }, + 1: {"slot_state": "explicit_empty"}, + }, + ), + ) + ) + markdown = "| Category | City | Status |\n|---|---|---|\n| France | Paris | Active |\n| | | Inactive |" + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="merged.pdf", + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert len(rows) == 2 + first, second = rows + assert first.cells[0].text == "France" + assert first.cells[0].row_span == 2 + assert second.cells[0].text == "France" + assert second.cells[0].inherited is True + assert second.cells[0].inherited_from == (1, 0) + assert second.cells[0].explicit_empty is False + assert second.cells[1].text == "" + assert second.cells[1].explicit_empty is True + assert second.cells[1].inherited is False + assert second.cells[2].text == "Inactive" + + +@pytest.mark.asyncio +async def test_unknown_merged_slot_fails_open_without_consuming_the_table(): + evidence = _one_page_table( + ( + _three_column_row(("Category", "City", "Status"), 0.25, 0.35), + _three_column_row(("France", "Paris", "Active"), 0.35, 0.45), + _three_column_row( + ("", "Lyon", "Inactive"), + 0.45, + 0.55, + cell_options={0: {"slot_state": "unknown"}}, + ), + ) + ) + markdown = ( + "Usable paragraph before the table.\n\n" + "| Category | City | Status |\n" + "|---|---|---|\n" + "| France | Paris | Active |\n" + "| | Lyon | Inactive |\n\n" + "Usable paragraph after the table." + ) + + parsed, result = await _normalize_one_page_table( + markdown, + evidence, + filename="ambiguous.pdf", + ) + + assert result.normalized_text_blocks is None + assert result.normalization_report.status == "unchanged" + assert result.effective_text_blocks() == parsed.text_blocks + assert result.raw_text_blocks == parsed.raw_text_blocks + + +@pytest.mark.asyncio +async def test_rows_after_a_continuation_receive_distinct_logical_numbers(): + page_one_rows = ( + _three_column_row(("ID", "Type", "Description"), 0.75, 0.82), + _three_column_row(("A", "Alpha", "Beginning"), 0.82, 0.96), + ) + page_two_rows = ( + _three_column_row(("", "", "Continuation"), 0.03, 0.12), + _three_column_row(("B", "Beta", "Second row"), 0.12, 0.24), + _three_column_row(("C", "Gamma", "Third row"), 0.24, 0.36), + ) + pages = { + 1: _one_page_table(page_one_rows, bbox=(0.10, 0.75, 0.90, 0.96)), + 2: PageLayoutEvidence( + page_number=2, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=2, + bbox=(0.10, 0.03, 0.90, 0.36), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=page_two_rows, + ), + ), + ), + } + raw_blocks = [ + TextBlock( + text="| ID | Type | Description |\n|---|---|---|\n| A | Alpha | Beginning |", + page_number=1, + ), + TextBlock( + text=( + "| Col1 | Col2 | Continuation |\n|---|---|---|\n| B | Beta | Second row |\n| C | Gamma | Third row |" + ), + page_number=2, + ), + ] + processed = ProcessedDocument( + text_blocks=[block.model_copy(deep=True) for block in raw_blocks], + raw_text_blocks=raw_blocks, + page_count=2, + ) + + result = await DeterministicTableNormalizer(FakeEvidenceProvider(pages)).normalize( + Document(filename="continued.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf"), + processed, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert [row.row_index for row in rows] == [1, 2, 3] + assert len({row.row_id for row in rows}) == 3 + assert "Continuation" in rows[0].cells[2].text + assert [row.cells[0].text for row in rows[1:]] == ["B", "C"] + + +@pytest.mark.asyncio +async def test_distinct_table_on_the_next_page_does_not_discard_the_first_table(): + page_one_rows = ( + _three_column_row(("ID", "Type", "Description"), 0.75, 0.82), + _three_column_row(("A", "Alpha", "First table"), 0.82, 0.96), + ) + page_two_rows = ( + _three_column_row(("ID", "Type", "Description"), 0.03, 0.10), + _three_column_row(("B", "Beta", "Second table"), 0.10, 0.24), + ) + pages = { + 1: _one_page_table(page_one_rows, bbox=(0.10, 0.75, 0.90, 0.96)), + 2: PageLayoutEvidence( + page_number=2, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=2, + bbox=(0.10, 0.03, 0.90, 0.24), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=page_two_rows, + ), + ), + ), + } + raw_blocks = [ + TextBlock( + text="| ID | Type | Description |\n|---|---|---|\n| A | Alpha | First table |", + page_number=1, + ), + TextBlock( + text="| ID | Type | Description |\n|---|---|---|\n| B | Beta | Second table |", + page_number=2, + ), + ] + processed = ProcessedDocument( + text_blocks=[block.model_copy(deep=True) for block in raw_blocks], + raw_text_blocks=raw_blocks, + page_count=2, + ) + + result = await DeterministicTableNormalizer(FakeEvidenceProvider(pages)).normalize( + Document(filename="two-tables.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf"), + processed, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert result.normalization_report.status == "partial_fallback" + assert [row.cells[0].text for row in rows] == ["A", "B"] + assert len({row.table_id for row in rows}) == 2 + + +@pytest.mark.asyncio +async def test_unmatched_parser_row_fails_open_without_losing_its_content(): + evidence = _one_page_table( + ( + _three_column_row(("Name", "City", "Status"), 0.25, 0.35), + _three_column_row(("one", "Paris", "Active"), 0.35, 0.45), + # PyMuPDF missed the middle row but still found the row below it. + _three_column_row(("three", "Lyon", "Inactive"), 0.55, 0.65), + ) + ) + markdown = ( + "| Name | City | Status |\n" + "|---|---|---|\n" + "| one | Paris | Active |\n" + "| two | Marseille | Pending |\n" + "| three | Lyon | Inactive |" + ) + + parsed, result = await _normalize_one_page_table( + markdown, + evidence, + filename="parser-row-mismatch.pdf", + ) + + assert result.normalized_text_blocks is None + assert result.normalization_report.status == "unchanged" + assert result.effective_text_blocks() == parsed.text_blocks + assert "two | Marseille | Pending" in result.effective_text_blocks()[0].text + + +@pytest.mark.asyncio +async def test_parser_value_in_layout_empty_cell_fails_open(): + evidence = _one_page_table( + ( + _three_column_row(("Name", "City", "Status"), 0.15, 0.25), + _three_column_row(("one", "Paris", "Active"), 0.25, 0.35), + _three_column_row(("two", "Lyon", "Inactive"), 0.35, 0.45), + _three_column_row(("three", "Nice", "Pending"), 0.45, 0.55), + _three_column_row( + ("", "Toulouse", "Active"), + 0.55, + 0.65, + cell_options={0: {"slot_state": "explicit_empty"}}, + ), + ), + bbox=(0.10, 0.15, 0.90, 0.65), + ) + markdown = ( + "| Name | City | Status |\n" + "|---|---|---|\n" + "| one | Paris | Active |\n" + "| two | Lyon | Inactive |\n" + "| three | Nice | Pending |\n" + "| SECRET | Toulouse | Active |" + ) + + parsed, result = await _normalize_one_page_table( + markdown, + evidence, + filename="empty-cell-disagreement.pdf", + ) + + assert result.normalized_text_blocks is None + assert result.normalization_report.status == "unchanged" + assert result.effective_text_blocks() == parsed.text_blocks + assert "SECRET" in result.effective_text_blocks()[0].text + + +@pytest.mark.asyncio +async def test_same_page_tables_with_the_same_headers_have_distinct_table_ids(): + first_table = LayoutTableEvidence( + page_number=1, + bbox=(0.10, 0.10, 0.90, 0.35), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("aa", "bb", "cc"), 0.10, 0.20), + _three_column_row(("22", "Paris", "Active"), 0.20, 0.30), + ), + ) + second_table = LayoutTableEvidence( + page_number=1, + bbox=(0.10, 0.50, 0.90, 0.75), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("aa", "bb", "cc"), 0.50, 0.60), + _three_column_row(("35", "Lyon", "Inactive"), 0.60, 0.70), + ), + ) + evidence = PageLayoutEvidence( + page_number=1, + width=100, + height=100, + tables=(first_table, second_table), + ) + markdown = ( + "# Table A\n\n" + "| aa | bb | cc |\n" + "|---|---|---|\n" + "| 22 | Paris | Active |\n\n" + "# Table B\n\n" + "| aa | bb | cc |\n" + "|---|---|---|\n" + "| 35 | Lyon | Inactive |" + ) + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="same-headers.pdf", + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert len(rows) == 2 + assert len({row.table_id for row in rows}) == 2 + + +@pytest.mark.asyncio +async def test_layout_value_that_omits_parser_suffix_fails_open(): + evidence = _one_page_table( + ( + _three_column_row(("Name", "City", "Description"), 0.25, 0.35), + _three_column_row( + ("one", "Paris", "A sufficiently long description without its suffix"), + 0.35, + 0.45, + ), + ) + ) + markdown = ( + "| Name | City | Description |\n" + "|---|---|---|\n" + "| one | Paris | A sufficiently long description without its suffix SECRET-42 |" + ) + + parsed, result = await _normalize_one_page_table( + markdown, + evidence, + filename="parser-suffix.pdf", + ) + + assert result.normalized_text_blocks is None + assert result.effective_text_blocks() == parsed.text_blocks + assert "SECRET-42" in result.effective_text_blocks()[0].text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("page_prefix", ["# New section", "A separate table follows"]) +async def test_unaccounted_page_prefix_prevents_a_compatible_cross_page_merge( + page_prefix, +): + page_one = _one_page_table( + ( + _three_column_row(("ID", "Type", "Description"), 0.75, 0.82), + _three_column_row(("A", "Alpha", "Beginning"), 0.82, 0.96), + ), + bbox=(0.10, 0.75, 0.90, 0.96), + ) + page_two_table = LayoutTableEvidence( + page_number=2, + bbox=(0.10, 0.03, 0.90, 0.20), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("", "", "Unrelated text"), 0.03, 0.10), + _three_column_row(("B", "Beta", "New section row"), 0.10, 0.20), + ), + ) + pages = { + 1: page_one, + 2: PageLayoutEvidence( + page_number=2, + width=100, + height=100, + tables=(page_two_table,), + ), + } + raw_blocks = [ + TextBlock( + text="| ID | Type | Description |\n|---|---|---|\n| A | Alpha | Beginning |", + page_number=1, + ), + TextBlock( + text=( + f"{page_prefix}\n\n" + "| Col1 | Col2 | Description |\n" + "|---|---|---|\n" + "| | | Unrelated text |\n" + "| B | Beta | New section row |" + ), + page_number=2, + ), + ] + processed = ProcessedDocument( + text_blocks=[block.model_copy(deep=True) for block in raw_blocks], + raw_text_blocks=raw_blocks, + page_count=2, + ) + + result = await DeterministicTableNormalizer(FakeEvidenceProvider(pages)).normalize( + Document(filename="sections.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf"), + processed, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + first = next(row for row in rows if row.cells[0].text == "A") + assert first.page_end == 1 + assert "Unrelated text" not in first.cells[2].text + + +@pytest.mark.asyncio +async def test_repeated_page_header_is_not_emitted_as_data_before_continuation(): + pages = { + 1: _one_page_table( + ( + _three_column_row(("ID", "Type", "Description"), 0.75, 0.82), + _three_column_row(("A", "Alpha", "Beginning"), 0.82, 0.96), + ), + bbox=(0.10, 0.75, 0.90, 0.96), + ), + 2: PageLayoutEvidence( + page_number=2, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=2, + bbox=(0.10, 0.03, 0.90, 0.30), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("ID", "Type", "Description"), 0.03, 0.09), + _three_column_row(("", "", "Continuation"), 0.09, 0.16), + _three_column_row(("B", "Beta", "Second row"), 0.16, 0.30), + ), + ), + ), + ), + } + raw_blocks = [ + TextBlock( + text="| ID | Type | Description |\n|---|---|---|\n| A | Alpha | Beginning |", + page_number=1, + ), + TextBlock( + text=("| ID | Type | Description |\n|---|---|---|\n| | | Continuation |\n| B | Beta | Second row |"), + page_number=2, + ), + ] + processed = ProcessedDocument( + text_blocks=[block.model_copy(deep=True) for block in raw_blocks], + raw_text_blocks=raw_blocks, + page_count=2, + ) + + result = await DeterministicTableNormalizer(FakeEvidenceProvider(pages)).normalize( + Document(filename="repeated-header.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf"), + processed, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert [row.cells[0].text for row in rows] == ["A", "B"] + assert "Continuation" in rows[0].cells[2].text + assert all(row.cells[0].text != "ID" for row in rows) + + +@pytest.mark.asyncio +async def test_headerless_table_preserves_its_first_data_row(): + evidence = _one_page_table( + ( + _three_column_row(("22", "Paris", "Active"), 0.25, 0.35), + _three_column_row(("35", "Lyon", "Inactive"), 0.35, 0.45), + ) + ) + markdown = "| 22 | Paris | Active |\n|---|---|---|\n| 35 | Lyon | Inactive |" + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="headerless.pdf", + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert [[cell.text for cell in row.cells] for row in rows] == [ + ["22", "Paris", "Active"], + ["35", "Lyon", "Inactive"], + ] + assert [cell.column_name for cell in rows[0].cells] == [ + "Column 1", + "Column 2", + "Column 3", + ] + + +@pytest.mark.asyncio +async def test_explicit_all_text_markdown_header_supplies_column_names(): + evidence = _one_page_table( + ( + _three_column_row(("Field", "Place", "Condition"), 0.25, 0.35), + _three_column_row(("Alice", "Paris", "Active"), 0.35, 0.45), + ) + ) + markdown = "| Field | Place | Condition |\n|---|---|---|\n| Alice | Paris | Active |" + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="text-header.pdf", + ) + + [row] = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert [cell.column_name for cell in row.cells] == [ + "Field", + "Place", + "Condition", + ] + assert [cell.text for cell in row.cells] == ["Alice", "Paris", "Active"] + + +@pytest.mark.asyncio +async def test_later_untitled_table_does_not_inherit_an_earlier_title(): + evidence = PageLayoutEvidence( + page_number=1, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=1, + bbox=(0.10, 0.10, 0.90, 0.35), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("aa", "bb", "cc"), 0.10, 0.20), + _three_column_row(("22", "Paris", "Active"), 0.20, 0.30), + ), + ), + LayoutTableEvidence( + page_number=1, + bbox=(0.10, 0.55, 0.90, 0.80), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("aa", "bb", "cc"), 0.55, 0.65), + _three_column_row(("35", "Lyon", "Inactive"), 0.65, 0.75), + ), + ), + ), + ) + markdown = ( + "# Table A\n\n" + "| aa | bb | cc |\n" + "|---|---|---|\n" + "| 22 | Paris | Active |\n\n" + "The next results are separate.\n\n" + "aa bb cc 35 Lyon Inactive" + ) + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="local-caption.pdf", + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert [row.table_title for row in rows] == ["Table A", None] + + +@pytest.mark.asyncio +async def test_plain_text_tables_claim_a_caption_only_once(): + evidence = PageLayoutEvidence( + page_number=1, + width=100, + height=100, + tables=( + LayoutTableEvidence( + page_number=1, + bbox=(0.10, 0.10, 0.90, 0.35), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("aa", "bb", "cc"), 0.10, 0.20), + _three_column_row(("22", "Paris", "Active"), 0.20, 0.30), + ), + ), + LayoutTableEvidence( + page_number=1, + bbox=(0.10, 0.55, 0.90, 0.80), + column_bounds=_THREE_COLUMN_BOUNDS, + rows=( + _three_column_row(("aa", "bb", "cc"), 0.55, 0.65), + _three_column_row(("35", "Lyon", "Inactive"), 0.65, 0.75), + ), + ), + ), + ) + raw = TextBlock( + text=("Table A\n\naa bb cc 22 Paris Active\n\nSeparate results follow.\n\naa bb cc 35 Lyon Inactive"), + page_number=1, + ) + processed = ProcessedDocument( + text_blocks=[raw.model_copy(deep=True)], + raw_text_blocks=[raw], + page_count=1, + ) + + result = await DeterministicTableNormalizer(FakeEvidenceProvider({1: evidence}, discovered={1})).normalize( + Document(filename="plain-tables.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf"), + processed, + TableReconstructionConfig(mode="automatic"), + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert [row.table_title for row in rows] == ["Table A", None] + + +@pytest.mark.asyncio +async def test_repeated_identity_values_still_receive_distinct_row_ids(): + evidence = _one_page_table( + ( + _three_column_row(("ID", "Type", "Description"), 0.20, 0.30), + _three_column_row(("A", "Alpha", "First detail"), 0.30, 0.40), + _three_column_row(("A", "Alpha", "Second detail"), 0.40, 0.50), + ) + ) + markdown = "| ID | Type | Description |\n|---|---|---|\n| A | Alpha | First detail |\n| A | Alpha | Second detail |" + + _, result = await _normalize_one_page_table( + markdown, + evidence, + filename="duplicate-identities.pdf", + ) + + rows = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert len(rows) == 2 + assert [cell.text for cell in rows[0].cells[:2]] == [cell.text for cell in rows[1].cells[:2]] + assert rows[0].row_id != rows[1].row_id + + +@pytest.mark.asyncio +async def test_layout_discovery_normalizes_table_without_parser_markdown(): + evidence = _one_page_table( + ( + _three_column_row(("aa", "bb", "cc"), 0.25, 0.35), + _three_column_row(("22", "Paris", "Active"), 0.35, 0.45), + ) + ) + raw = TextBlock(text="aa bb cc 22 Paris Active", page_number=1) + processed = ProcessedDocument( + text_blocks=[raw.model_copy(deep=True)], + raw_text_blocks=[raw], + page_count=1, + ) + provider = FakeEvidenceProvider({1: evidence}, discovered={1}) + + result = await DeterministicTableNormalizer(provider).normalize( + Document(filename="plain-table.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf"), + processed, + TableReconstructionConfig(mode="automatic"), + ) + + [row] = [block.table_row for block in result.effective_text_blocks() if block.table_row is not None] + assert [cell.text for cell in row.cells] == ["22", "Paris", "Active"] + assert "the value “22” in column “aa”" in next( + block.text for block in result.effective_text_blocks() if block.table_row is row + ) diff --git a/tests/unit/services/orchestrators/test_indexing_service.py b/tests/unit/services/orchestrators/test_indexing_service.py index d3280b2d8..91c1a8ee7 100644 --- a/tests/unit/services/orchestrators/test_indexing_service.py +++ b/tests/unit/services/orchestrators/test_indexing_service.py @@ -106,6 +106,7 @@ def _config_with_partition(partition: str = "tenant-a"): embedder="embed-fast", indexation=IndexationPipelineConfig( parsing_strategy="pymupdf", + table_reconstruction={"mode": "automatic"}, enable_image_captioning=False, enable_contextualization=True, contextualization_llm="llm-context", @@ -416,6 +417,13 @@ async def test_add_file_dispatches_partition_indexation_config_and_embedder(tmp_ assert sent["indexation_config"]["enable_image_captioning"] is False assert sent["indexation_config"]["enable_contextualization"] is True assert sent["indexation_config"]["contextualization_llm"] == "llm-context" + assert sent["indexation_config"]["table_reconstruction"] == { + "mode": "automatic", + "same_table_min_confidence": 0.9, + "row_continuation_min_confidence": 0.9, + "cell_assignment_min_confidence": 0.9, + "algorithm_version": "adjacent-layout-v1", + } assert sent["require_existing_partition"] is True assert sent["allow_legacy_require_existing_partition_retry"] is True diff --git a/tests/unit/services/orchestrators/test_preset_service.py b/tests/unit/services/orchestrators/test_preset_service.py index 90ff89986..3866dd42b 100644 --- a/tests/unit/services/orchestrators/test_preset_service.py +++ b/tests/unit/services/orchestrators/test_preset_service.py @@ -356,6 +356,19 @@ async def test_create_preset_rejects_invalid_config(): await svc.create_preset("bad", "retrieval", bad_config) +@pytest.mark.asyncio +async def test_create_preset_rejects_invalid_table_reconstruction_config(): + from core.utils.exceptions import ValidationError + + svc = _make_service() + with pytest.raises(ValidationError, match="Invalid indexation preset config"): + await svc.create_preset( + "bad-table-mode", + "indexation", + {"table_reconstruction": {"mode": "automatic", "unknown_threshold": 0.9}}, + ) + + @pytest.mark.asyncio async def test_create_preset_inserts_and_returns_row(): repo = _FakePresetRepo() diff --git a/tests/unit/services/storage/test_milvus_store.py b/tests/unit/services/storage/test_milvus_store.py index 49ce9d219..5c0916204 100644 --- a/tests/unit/services/storage/test_milvus_store.py +++ b/tests/unit/services/storage/test_milvus_store.py @@ -456,6 +456,39 @@ def test_id_is_not_in_entity(self) -> None: assert "_id" not in entity +@pytest.mark.asyncio +async def test_upsert_preserves_table_text_and_metadata_at_insert_boundary( + store: MilvusVectorStore, +) -> None: + text = ( + "In table “Table A”, the first row (row 1) has the value “22” in column “aa”, " + "the value “Paris” in column “bb”, and the value “Active” in column “cc”." + ) + metadata = { + "table_id": "table-a", + "row_id": "row-a-1", + "row_index": 1, + "table_title": "Table A", + "table_content_kind": "row", + "table_text_serialization_version": "natural-language-v1", + } + chunk = _make_chunk( + text=text, + chunk_type=ChunkType.TABLE, + metadata=metadata, + ) + store._async_client.insert = AsyncMock(return_value={"insert_count": 1}) + + assert await store.upsert([chunk]) == 1 + + insert_call = store._async_client.insert.await_args + assert insert_call.kwargs["collection_name"] == "test_collection" + [entity] = insert_call.kwargs["data"] + assert entity["text"] == text + assert entity["chunk_type"] == "table" + assert {key: entity[key] for key in metadata} == metadata + + # --------------------------------------------------------------------------- # Surface-level ABC-vs-bound-collection enforcement # --------------------------------------------------------------------------- diff --git a/tests/unit/services/workers/layout/test_pymupdf_table_evidence.py b/tests/unit/services/workers/layout/test_pymupdf_table_evidence.py new file mode 100644 index 000000000..9552055e1 --- /dev/null +++ b/tests/unit/services/workers/layout/test_pymupdf_table_evidence.py @@ -0,0 +1,199 @@ +from pathlib import Path +from types import SimpleNamespace + +import pymupdf +import pytest +from core.indexing.structure_normalizer import ( + LayoutWord, + PageLayoutEvidence, + TableLayoutEvidenceProvider, +) +from core.models.document import Document, DocumentType +from services.workers.layout import PyMuPDFTableEvidenceProvider +from services.workers.layout.pymupdf_table_evidence import ( + _cell_grid_semantics, + _collect_evidence, + _text_from_positioned_words, +) + +FIXTURE = Path(__file__).parents[4] / "resources" / "cross_page_table_rows_803_805.pdf" + + +class CollectOnlyEvidenceProvider(TableLayoutEvidenceProvider): + async def collect(self, document: Document, page_numbers: set[int]) -> list[PageLayoutEvidence]: + return [] + + +@pytest.mark.asyncio +async def test_provider_discovery_defaults_to_no_candidates(): + document = Document( + filename="document.pdf", + content_type=DocumentType.PDF, + raw_bytes=b"pdf", + ) + + assert await CollectOnlyEvidenceProvider().discover(document) == set() + + +@pytest.mark.asyncio +async def test_adapter_discovers_table_pages_without_collecting_evidence(): + document = Document( + filename=FIXTURE.name, + content_type=DocumentType.PDF, + raw_bytes=FIXTURE.read_bytes(), + ) + + assert await PyMuPDFTableEvidenceProvider().discover(document) == {1, 3} + + +@pytest.mark.asyncio +async def test_adapter_exposes_tables_and_sparse_continuation_evidence(): + document = Document( + filename=FIXTURE.name, + content_type=DocumentType.PDF, + raw_bytes=FIXTURE.read_bytes(), + ) + + pages = await PyMuPDFTableEvidenceProvider().collect(document, {1, 2, 3}) + + assert [len(page.tables) for page in pages] == [1, 0, 1] + assert len(pages[0].tables[0].column_bounds) == 5 + page_two_body = [word for word in pages[1].words if word.bbox[3] < 0.90] + assert page_two_body + assert min(word.bbox[0] for word in page_two_body) > pages[0].tables[0].column_bounds[-1][0] + + +@pytest.mark.asyncio +async def test_adapter_repairs_cell_spacing_from_positioned_words(): + document = Document( + filename=FIXTURE.name, + content_type=DocumentType.PDF, + raw_bytes=FIXTURE.read_bytes(), + ) + + pages = await PyMuPDFTableEvidenceProvider().collect(document, {1, 3}) + table_text = "\n".join( + cell.text for page in pages for table in page.tables for row in table.rows for cell in row.cells + ) + reflowed = " ".join(table_text.split()) + + assert "cas :" in reflowed + assert "séjour en cours de validité" in reflowed + assert "d'emploi" in reflowed + assert "cas :\n\n-visa" in table_text + for spacing_artifact in ("c as", "s éjour", "d 'emploi"): + assert spacing_artifact not in table_text + + +def test_positioned_word_rebuild_requires_complete_matching_evidence(): + bbox = (0.0, 0.0, 1.0, 1.0) + words = ( + LayoutWord("Alpha", (0.1, 0.1, 0.2, 0.2), 0, 0, 0), + LayoutWord("Beta", (0.1, 0.3, 0.2, 0.4), 0, 2, 0), + ) + + assert _text_from_positioned_words(words, bbox, "AlphaBeta") == "Alpha\n\nBeta" + assert _text_from_positioned_words(words, bbox, "different value") == "different value" + assert _text_from_positioned_words((), bbox, "fallback") == "fallback" + + +def test_grid_geometry_distinguishes_horizontal_merge_from_an_empty_cell(): + rows = [ + SimpleNamespace( + bbox=(0.0, 0.0, 300.0, 50.0), + cells=[ + (0.0, 0.0, 200.0, 50.0), + None, + (200.0, 0.0, 300.0, 50.0), + ], + ), + SimpleNamespace( + bbox=(0.0, 50.0, 300.0, 100.0), + cells=[ + (0.0, 50.0, 100.0, 100.0), + (100.0, 50.0, 200.0, 100.0), + (200.0, 50.0, 300.0, 100.0), + ], + ), + ] + + semantics = _cell_grid_semantics( + rows, + ((0.0, 1 / 3), (1 / 3, 2 / 3), (2 / 3, 1.0)), + (0.0, 0.0, 300.0, 100.0), + 300.0, + 100.0, + ) + + assert semantics[(0, 0)] == ("value", 2, 1, None) + assert semantics[(0, 1)] == ("covered", 1, 1, (0, 0)) + # This slot has its own geometry. An empty extracted value is therefore + # an explicit empty cell, not a merged-cell continuation. + assert semantics[(0, 2)] == ("value", 1, 1, None) + + +def test_grid_geometry_identifies_vertical_merged_cell_coverage(): + rows = [ + SimpleNamespace( + bbox=(0.0, 0.0, 200.0, 50.0), + cells=[ + (0.0, 0.0, 100.0, 100.0), + (100.0, 0.0, 200.0, 50.0), + ], + ), + SimpleNamespace( + bbox=(0.0, 50.0, 200.0, 100.0), + cells=[ + None, + (100.0, 50.0, 200.0, 100.0), + ], + ), + ] + + semantics = _cell_grid_semantics( + rows, + ((0.0, 0.5), (0.5, 1.0)), + (0.0, 0.0, 200.0, 100.0), + 200.0, + 100.0, + ) + + assert semantics[(0, 0)] == ("value", 1, 2, None) + assert semantics[(1, 0)] == ("covered", 1, 1, (0, 0)) + + +def test_real_pymupdf_rows_do_not_inherit_rowspan_from_a_neighboring_anchor(): + pdf = pymupdf.open() + page = pdf.new_page(width=300, height=240) + x_positions = (30, 110, 190, 270) + y_positions = (30, 80, 130, 180) + for x_position in x_positions: + page.draw_line((x_position, y_positions[0]), (x_position, y_positions[-1])) + for y_position in (y_positions[0], y_positions[1], y_positions[-1]): + page.draw_line((x_positions[0], y_position), (x_positions[-1], y_position)) + page.draw_line( + (x_positions[1], y_positions[2]), + (x_positions[-1], y_positions[2]), + ) + for x_position, y_position, text in ( + (40, 60, "Category"), + (120, 60, "City"), + (200, 60, "Status"), + (40, 110, "France"), + (120, 110, "Paris"), + (200, 110, "Active"), + (120, 160, "Lyon"), + (200, 160, "Inactive"), + ): + page.insert_text((x_position, y_position), text, fontsize=9) + raw_bytes = pdf.tobytes() + pdf.close() + + [evidence] = _collect_evidence(raw_bytes, (1,)) + table = evidence.tables[0] + + assert table.rows[1].cells[0].row_span == 2 + assert table.rows[2].cells[0].slot_state == "covered" + assert table.rows[2].cells[0].covered_by == (1, 0) + assert table.rows[2].cells[1].row_span == 1 + assert table.rows[2].cells[2].row_span == 1 diff --git a/tests/unit/services/workers/parsers/test_marker_workers.py b/tests/unit/services/workers/parsers/test_marker_workers.py index 055396f4b..d18d61812 100644 --- a/tests/unit/services/workers/parsers/test_marker_workers.py +++ b/tests/unit/services/workers/parsers/test_marker_workers.py @@ -19,6 +19,23 @@ def test_marker_num_gpus_uses_ray_cluster_resources_when_cuda_is_hidden(monkeypa assert marker_workers._marker_num_gpus(_config()) == 0.25 +def test_split_pages_preserves_canonical_breaks_inside_pipe_tables(): + markdown = "[PAGE_SEP]\n| Header
Legend | Value
Unit |\n|---|---|\n| Alpha
Beta | 22 |\n{1}[PAGE_SEP]" + + assert marker_workers.MarkerLoader._split_pages(markdown) == [ + ( + 1, + "| Header
Legend | Value
Unit |\n|---|---|\n| Alpha
Beta | 22 |", + ) + ] + + +def test_split_pages_removes_breaks_outside_pipe_tables(): + markdown = "[PAGE_SEP]\nFirst
paragraph.
\nSecond
paragraph.\n{1}[PAGE_SEP]" + + assert marker_workers.MarkerLoader._split_pages(markdown) == [(1, "Firstparagraph.\nSecondparagraph.")] + + # --------------------------------------------------------------------------- # _force_kill_executor — reclaiming a wedged Marker worker (#659) # --------------------------------------------------------------------------- diff --git a/tests/unit/services/workers/stages/test_normalize_structure.py b/tests/unit/services/workers/stages/test_normalize_structure.py new file mode 100644 index 000000000..b20bb3a21 --- /dev/null +++ b/tests/unit/services/workers/stages/test_normalize_structure.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import pytest +from core.config.table_reconstruction import TableReconstructionConfig +from core.indexing.structure_normalizer import DocumentStructureNormalizer +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from services.workers.stages.normalize_structure import normalize_structure_stage + + +class FakeNormalizer(DocumentStructureNormalizer): + def __init__(self, output: ProcessedDocument | None = None, error: Exception | None = None) -> None: + self.output = output + self.error = error + self.calls = [] + + async def normalize(self, document, processed_document, config): + self.calls.append((document, processed_document, config)) + if self.error is not None: + raise self.error + assert self.output is not None + return self.output + + +@pytest.mark.asyncio +async def test_normalize_structure_stage_replaces_the_processed_document(): + document = Document(filename="sample.pdf", content_type=DocumentType.PDF) + parsed = ProcessedDocument(text_blocks=[TextBlock(text="raw")]) + normalized = ProcessedDocument(text_blocks=[TextBlock(text="raw")], normalized_text_blocks=[TextBlock(text="row")]) + normalizer = FakeNormalizer(output=normalized) + row = {"document": document, "processed_document": parsed} + config = TableReconstructionConfig(mode="automatic") + + result = await normalize_structure_stage(row, normalizer, config) + + assert result["processed_document"] is normalized + assert result["stage"] == "structure_normalized" + assert normalizer.calls == [(document, parsed, config)] + + +@pytest.mark.asyncio +async def test_normalize_structure_stage_fails_open_on_an_unexpected_error(): + document = Document(filename="sample.pdf", content_type=DocumentType.PDF) + parsed = ProcessedDocument(text_blocks=[TextBlock(text="usable parser output")]) + row = {"document": document, "processed_document": parsed, "token": "secret"} + + result = await normalize_structure_stage( + row, + FakeNormalizer(error=RuntimeError("layout failed")), + TableReconstructionConfig(mode="automatic"), + ) + + assert result["processed_document"].effective_text_blocks()[0].text == "usable parser output" + assert result["processed_document"].normalization_report.status == "partial_fallback" + assert result["stage"] == "structure_normalization_fallback" + assert "error" not in result + assert "token" not in result diff --git a/tests/unit/services/workers/stages/test_parse.py b/tests/unit/services/workers/stages/test_parse.py index 48c607804..d9f5f02a0 100644 --- a/tests/unit/services/workers/stages/test_parse.py +++ b/tests/unit/services/workers/stages/test_parse.py @@ -48,6 +48,24 @@ async def test_parse_stage_mutates_row_with_processed_document_and_scrubs_creden assert "token" not in row +@pytest.mark.asyncio +async def test_parse_stage_can_preserve_an_independent_raw_block_snapshot(): + document = Document(id="doc-1", filename="sample.pdf", content_type=DocumentType.PDF) + processed = ProcessedDocument( + document_id="doc-1", + text_blocks=[TextBlock(text="parser output", page_number=1)], + ) + row = {"document": document} + + await parse_stage(row, FakeParser(output=processed), preserve_raw_blocks=True) + + result = row["processed_document"] + assert result.raw_text_blocks == processed.text_blocks + assert result.raw_text_blocks is not processed.text_blocks + result.text_blocks[0].text = "working output changed" + assert result.raw_text_blocks[0].text == "parser output" + + @pytest.mark.asyncio async def test_parse_stage_marks_error_and_scrubs_credentials_when_parser_fails(): document = Document(id="doc-1", filename="note.txt", content_type=DocumentType.TEXT, text="hello") diff --git a/tests/unit/services/workers/stages/test_pipeline_stages.py b/tests/unit/services/workers/stages/test_pipeline_stages.py index 5a789c97a..04e767329 100644 --- a/tests/unit/services/workers/stages/test_pipeline_stages.py +++ b/tests/unit/services/workers/stages/test_pipeline_stages.py @@ -6,7 +6,7 @@ from core.chunking.chunking_strategy import ChunkingStrategy from core.embeddings.embedder import Embedder from core.indexing.contextualize import ChunkContextualizer -from core.models.chunk import Chunk +from core.models.chunk import Chunk, ChunkType from core.models.document import ImageBlock, ProcessedDocument, TextBlock from core.prompts.vlm_prompt_builder import wrap_caption from core.vector_stores.vector_store import VectorStore @@ -308,6 +308,39 @@ async def test_embed_stage_attaches_vectors_by_chunk_order(): assert "secret" not in row +@pytest.mark.asyncio +async def test_embed_stage_preserves_table_text_and_row_metadata(): + text = ( + "In table “Table A”, the first row (row 1) has the value “22” in column “aa”, " + "the value “Paris” in column “bb”, and the value “Active” in column “cc”." + ) + metadata = { + "table_id": "table-a", + "row_id": "row-a-1", + "row_index": 1, + "table_title": "Table A", + "table_content_kind": "row", + "table_text_serialization_version": "natural-language-v1", + } + chunk = Chunk( + id="table-chunk", + text=text, + chunk_type=ChunkType.TABLE, + metadata=metadata, + ) + embedder = FakeEmbedder([[0.25, 0.75]]) + row = {"chunks": [chunk]} + + await embed_stage(row, embedder) + + [embedded] = row["chunks"] + assert embedder.text_batches == [[text]] + assert embedded.text == text + assert embedded.metadata == metadata + assert embedded.chunk_type is ChunkType.TABLE + assert embedded.embedding == [0.25, 0.75] + + @pytest.mark.asyncio async def test_store_stage_upserts_to_default_collection_with_chunk_partitions(): chunks = [Chunk(id="c1", text="alpha", embedding=[1.0])] diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index b728a448a..6a6cdf43b 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -123,7 +123,7 @@ def fake_options(**kwargs): opts = options_calls[0] # A protocol-specific name prevents a rolling deployment from attaching to # a detached actor that still runs the previous claim implementation. - assert opts["name"] == "IndexerPoolDispatcher-v2" + assert opts["name"] == "IndexerPoolDispatcher-v3" assert opts["namespace"] == "openrag" assert opts["get_if_exists"] is True assert opts["lifetime"] == "detached" @@ -159,9 +159,9 @@ def fake_options(**kwargs): # One detached worker actor per pool_size slot, each capped at max_tasks_per_worker. assert len(pool._workers) == 3 assert {c["name"] for c in calls} == { - "IndexerWorker-v2-0", - "IndexerWorker-v2-1", - "IndexerWorker-v2-2", + "IndexerWorker-v3-0", + "IndexerWorker-v3-1", + "IndexerWorker-v3-2", } for c in calls: assert c["lifetime"] == "detached" @@ -1093,7 +1093,7 @@ async def test_pool_drain_rejects_new_work_and_reports_accepted_work() -> None: await pool.submit(task_id="accepted-before-drain") assert await pool.begin_drain() == { - "protocol_version": "v2", + "protocol_version": "v3", "accepting_tasks": False, "inflight_jobs": 1, "worker_names": ["test-worker-0"], @@ -1104,7 +1104,7 @@ async def test_pool_drain_rejects_new_work_and_reports_accepted_work() -> None: await _settle_pool_release_tasks(pool, worker.futures[0]) assert await pool.status() == { - "protocol_version": "v2", + "protocol_version": "v3", "accepting_tasks": False, "inflight_jobs": 0, "worker_names": ["test-worker-0"], @@ -1121,7 +1121,7 @@ async def test_pool_abort_drain_restores_acceptance() -> None: await pool.submit(task_id="rejected-while-draining") assert await pool.abort_drain() == { - "protocol_version": "v2", + "protocol_version": "v3", "accepting_tasks": True, "inflight_jobs": 0, "worker_names": ["test-worker-0"], @@ -1136,7 +1136,7 @@ async def test_pool_abort_drain_restores_acceptance() -> None: async def test_pool_reports_current_protocol_version() -> None: pool = _bare_pool([_FakeWorker()]) - assert await pool.protocol_version() == "v2" + assert await pool.protocol_version() == "v3" @pytest.mark.asyncio diff --git a/tests/unit/services/workers/test_indexer_worker.py b/tests/unit/services/workers/test_indexer_worker.py index 8c171a30a..c32c0f173 100644 --- a/tests/unit/services/workers/test_indexer_worker.py +++ b/tests/unit/services/workers/test_indexer_worker.py @@ -406,7 +406,17 @@ async def test_process_file_stores_indexation_config_snapshot_on_new_file(tmp_pa task_state_manager=_fake_tsm(), document_repo=repo, ) - indexation_config = {"parsing_strategy": "pymupdf", "enable_image_captioning": False} + indexation_config = { + "parsing_strategy": "pymupdf", + "enable_image_captioning": False, + "table_reconstruction": { + "mode": "automatic", + "same_table_min_confidence": 0.9, + "row_continuation_min_confidence": 0.9, + "cell_assignment_min_confidence": 0.9, + "algorithm_version": "adjacent-layout-v1", + }, + } await worker.process_file( task_id="t-new", @@ -571,7 +581,11 @@ async def test_process_file_stores_indexation_config_snapshot_on_replace(tmp_pat task_state_manager=_fake_tsm(), document_repo=repo, ) - indexation_config = {"parsing_strategy": "marker", "enable_contextualization": True} + indexation_config = { + "parsing_strategy": "marker", + "enable_contextualization": True, + "table_reconstruction": {"mode": "automatic"}, + } await worker.process_file( task_id="t-replace", diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index 6ce4e218f..8063b8b3c 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -2,6 +2,7 @@ import pytest from core.config.indexation_pipeline import IndexationPipelineConfig +from core.indexing.structure_normalizer import DocumentStructureNormalizer from core.models.chunk import Chunk from core.models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock from services.workers.pipeline_builder import build_indexing_pipeline @@ -105,6 +106,17 @@ async def tag( return self.tags +class FakeStructureNormalizer(DocumentStructureNormalizer): + def __init__(self) -> None: + self.calls = [] + + async def normalize(self, document, processed_document, config): + self.calls.append((document, processed_document, config)) + return processed_document.model_copy( + update={"normalized_text_blocks": [TextBlock(text="normalized row", page_number=1)]} + ) + + @pytest.mark.asyncio async def test_pipeline_runs_required_stages_in_order_and_keeps_row_object(): document = Document(filename="note.txt", text="hello", partition="tenant-a") @@ -191,6 +203,61 @@ async def test_pipeline_indexation_config_disables_caption_and_contextualization assert row["stage"] == "stored" +@pytest.mark.asyncio +async def test_pipeline_normalizes_automatic_pdf_after_preserving_raw_blocks(): + document = Document( + filename="table.pdf", + content_type=DocumentType.PDF, + raw_bytes=b"pdf", + partition="tenant-a", + ) + processed = ProcessedDocument( + document_id=document.id, + text_blocks=[TextBlock(text="parser row", page_number=1)], + page_count=1, + ) + normalizer = FakeStructureNormalizer() + chunker = FakeChunker([Chunk(id="c1", text="normalized row", partition="tenant-a")]) + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=chunker, + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + structure_normalizer=normalizer, + indexation_config=IndexationPipelineConfig( + table_reconstruction={"mode": "automatic"}, + enable_image_captioning=False, + ), + ) + + await pipeline.run({"document": document, "partition": "tenant-a"}) + + assert len(normalizer.calls) == 1 + normalized_input = normalizer.calls[0][1] + assert normalized_input.raw_text_blocks[0].text == "parser row" + assert chunker.calls[0][0].effective_text_blocks()[0].text == "normalized row" + + +@pytest.mark.asyncio +async def test_pipeline_does_not_allocate_or_normalize_when_feature_is_disabled(): + document = Document(filename="table.pdf", content_type=DocumentType.PDF, raw_bytes=b"pdf") + processed = ProcessedDocument(text_blocks=[TextBlock(text="parser row", page_number=1)]) + normalizer = FakeStructureNormalizer() + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker([Chunk(id="c1", text="parser row")]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + structure_normalizer=normalizer, + indexation_config=IndexationPipelineConfig(enable_image_captioning=False), + ) + + await pipeline.run({"document": document}) + + assert normalizer.calls == [] + assert pipeline.parser.processed.raw_text_blocks is None + + def _caption_pipeline(vlm: FakeVLM, caption_prompt: str | None): document = Document(filename="note.txt", text="hello", partition="tenant-a") processed = ProcessedDocument( diff --git a/ui/src/lib/api/presets.ts b/ui/src/lib/api/presets.ts index fe10b8dc6..aab440d54 100644 --- a/ui/src/lib/api/presets.ts +++ b/ui/src/lib/api/presets.ts @@ -37,6 +37,8 @@ export interface PresetOptionsResponse { chunking_strategies: string[]; // Optional: older backends don't return this (UI falls back to a default list). parsing_strategies?: string[]; + // Optional during rolling upgrades; missing means the UI offers safe defaults. + table_reconstruction_modes?: string[]; retrieval_types: string[]; reranker_providers: string[]; } diff --git a/ui/src/pages/admin/preset-config.test.ts b/ui/src/pages/admin/preset-config.test.ts index 34f06d267..2abff42d1 100644 --- a/ui/src/pages/admin/preset-config.test.ts +++ b/ui/src/pages/admin/preset-config.test.ts @@ -1,5 +1,10 @@ import { describe, it, expect } from "vitest"; -import { applyParsingStrategyChange, PARSING_STRATEGY_INHERIT } from "./preset-config"; +import { + applyParsingStrategyChange, + applyTableReconstructionMode, + getTableReconstructionMode, + PARSING_STRATEGY_INHERIT, +} from "./preset-config"; // Regression guard for the bug where the parsing Strategy dropdown displayed a // fabricated "marker" default that was never written to config unless the user @@ -68,3 +73,28 @@ describe("applyParsingStrategyChange", () => { expect(input).toEqual({ parsing_strategy: "marker" }); }); }); + +describe("table reconstruction config", () => { + it("shows missing configuration as disabled", () => { + expect(getTableReconstructionMode({})).toBe("disabled"); + }); + + it("updates the nested mode without discarding backend thresholds", () => { + const input = { + parsing_strategy: "marker", + table_reconstruction: { + mode: "disabled", + same_table_min_confidence: 0.95, + }, + }; + + expect(applyTableReconstructionMode(input, "automatic")).toEqual({ + parsing_strategy: "marker", + table_reconstruction: { + mode: "automatic", + same_table_min_confidence: 0.95, + }, + }); + expect(getTableReconstructionMode(input)).toBe("disabled"); + }); +}); diff --git a/ui/src/pages/admin/preset-config.ts b/ui/src/pages/admin/preset-config.ts index fa8fac3c6..d7019be63 100644 --- a/ui/src/pages/admin/preset-config.ts +++ b/ui/src/pages/admin/preset-config.ts @@ -28,6 +28,7 @@ export function configUnset(prev: Config, key: string): Config { // Mirrors the retrieval form's "__none__" convention so the shown value always // matches what gets persisted (WYSIWYG). export const PARSING_STRATEGY_INHERIT = "__inherit__"; +export const TABLE_RECONSTRUCTION_DISABLED = "disabled"; // Map a parsing-strategy selection to the next config. Kept pure so the // persistence behavior is unit-testable without driving the Radix Select. @@ -58,3 +59,33 @@ export function applyParsingStrategyChange(config: Config, value: string): Confi return next; } + +export function getTableReconstructionMode(config: Config): string { + const tableReconstruction = config.table_reconstruction; + if (!tableReconstruction || typeof tableReconstruction !== "object") { + return TABLE_RECONSTRUCTION_DISABLED; + } + return configGet( + tableReconstruction as Config, + "mode", + TABLE_RECONSTRUCTION_DISABLED, + ); +} + +export function applyTableReconstructionMode( + config: Config, + mode: string, +): Config { + const current = + config.table_reconstruction && + typeof config.table_reconstruction === "object" + ? (config.table_reconstruction as Config) + : {}; + return { + ...config, + table_reconstruction: { + ...current, + mode, + }, + }; +} diff --git a/ui/src/pages/admin/presets.test.tsx b/ui/src/pages/admin/presets.test.tsx index 75f1a7260..370dca8a9 100644 --- a/ui/src/pages/admin/presets.test.tsx +++ b/ui/src/pages/admin/presets.test.tsx @@ -19,6 +19,7 @@ vi.mock("@/lib/api/presets", async () => { getPresetOptions: vi.fn().mockResolvedValue({ chunking_strategies: [], parsing_strategies: [], + table_reconstruction_modes: ["disabled", "automatic"], retrieval_types: [], reranker_providers: [], }), @@ -66,6 +67,14 @@ function renderPage() { describe("PresetsPage usage badge", () => { beforeEach(() => { + vi.stubGlobal( + "ResizeObserver", + class { + observe() {} + unobserve() {} + disconnect() {} + }, + ); listPresetsMock.mockReset(); deletePresetMock.mockReset(); vi.mocked(toast.error).mockClear(); @@ -117,4 +126,19 @@ describe("PresetsPage usage badge", () => { ), ); }); + + it("exposes conservative cross-page table reconstruction in indexation presets", async () => { + listPresetsMock.mockResolvedValue([]); + const user = userEvent.setup(); + + renderPage(); + + await user.click(await screen.findByRole("button", { name: /add preset/i })); + expect( + await screen.findByText("Cross-page table reconstruction"), + ).toBeTruthy(); + expect( + screen.getByText(/Uncertain cases keep the parser output/i), + ).toBeTruthy(); + }); }); diff --git a/ui/src/pages/admin/presets.tsx b/ui/src/pages/admin/presets.tsx index f6c5de3cd..9b8a453d4 100644 --- a/ui/src/pages/admin/presets.tsx +++ b/ui/src/pages/admin/presets.tsx @@ -41,9 +41,11 @@ import { Skeleton } from "@/components/ui/skeleton"; import { formatDate, intOr, numOr } from "@/lib/utils"; import { type Config, + applyTableReconstructionMode, configGet, configSet, applyParsingStrategyChange, + getTableReconstructionMode, PARSING_STRATEGY_INHERIT, } from "./preset-config"; @@ -249,6 +251,7 @@ function IndexationPresetForm({ onChange, chunkingStrategies, parsingStrategies, + tableReconstructionModes, vlms, llms, prompts, @@ -259,6 +262,7 @@ function IndexationPresetForm({ onChange: (c: Config) => void; chunkingStrategies: string[]; parsingStrategies: string[]; + tableReconstructionModes: string[]; vlms: string[]; llms: string[]; prompts: PromptResponse[]; @@ -365,6 +369,30 @@ function IndexationPresetForm({ +
+ + +

+ Automatic mode conservatively rebuilds table rows that continue + across PDF pages. Uncertain cases keep the parser output. +

+
@@ -888,6 +916,9 @@ function PresetDialog({ onChange={setConfig} chunkingStrategies={options?.chunking_strategies ?? []} parsingStrategies={options?.parsing_strategies ?? ["marker", "pymupdf"]} + tableReconstructionModes={ + options?.table_reconstruction_modes ?? ["disabled", "automatic"] + } vlms={vlms} llms={llms} prompts={allPrompts}