Update app.py

app.py

@@ -703,11 +703,52 @@ def _row_keyword_score(row):
     return keyword_hits / len(non_empty)
 
 
+# Pattern: header cell contains balance/account info — do NOT promote away from it.
+# Matches things like "Beginning Balance:", "Ending Balance:", account numbers, "$20.48$3.46"
+_BALANCE_INFO_RE = re.compile(
+    r"(beginning|ending|opening|closing)\s+balance"
+    r"|account\s*#?\s*\d{5,}"
+    r"|\b\d{7,}\b"
+    r"|\$\d{1,3}(?:,\d{3})*\.\d{2}",
+    re.IGNORECASE,
+)
+
+
+def _header_contains_balance_info(header_row):
+    """
+    Return True if any cell in the header row contains balance or account
+    information — indicating this IS a legitimate header row even if its
+    keyword score is low (e.g. Citi's "208479667 | Beginning Balance:$20.48 | ...").
+    """
+    for cell in header_row:
+        if _BALANCE_INFO_RE.search(str(cell or "")):
+            return True
+    return False
+
+
 def _promote_misplaced_header_row(grid):
     """
     Detects and fixes the OCR artifact where real column headers land in a
     <td> data row instead of the <th> header row.
 
+    Two scenarios handled:
+
+    Scenario A — Simple misplaced header (e.g. TD Bank):
+      <th>ACCOUNT ACTIVITY</th>        ← section title, not a real header
+      <td>DATE DESCRIPTION</td> ...    ← real column headers in a data row
+      <td>01/05 ...</td>               ← data
+    → Promote the keyword row to header, discard the section title rows above.
+
+    Scenario B — Metadata header + misplaced column headers (e.g. Citi page 1):
+      <th>208479667</th> <th>Beginning Balance:$20.48 Ending Balance:$3.46</th>
+      <td>Date Description</td> <td>Debits</td> <td>Credits</td> <td>Balance</td>
+      <td>04/01 DEBIT CARD... 8.01</td> ...
+    → The existing th row has balance/account metadata (NOT a real column header).
+    → Promote the keyword data row to header.
+    → Preserve metadata row cells as a plain-text prefix OUTSIDE the table
+      by embedding them as a leading data row with colspan (kept for reference).
+    → Split any "Date Description" fused cells in data rows.
+
     Guard conditions (no-op if not met — safe for all other PDFs):
       - Current header keyword score < threshold.
       - A data row within the first 5 rows has keyword score >= threshold.
@@ -719,6 +760,7 @@ def _promote_misplaced_header_row(grid):
     if current_header_score >= _HEADER_ROW_KEYWORD_THRESHOLD:
         return grid
 
+    # Search the first few data rows for a candidate keyword header row.
     candidate_idx = None
     candidate_score = 0.0
     for i in range(1, min(len(grid), 6)):
@@ -745,6 +787,17 @@ def _promote_misplaced_header_row(grid):
     new_ncols = len(expanded_header)
     col_expansion = new_ncols - len(candidate_row)
 
+    # If the existing header row contains balance/account metadata (Scenario B),
+    # preserve it as the first data row so the information is not lost.
+    # This row will have its content in col 0 (joined) and blanks elsewhere.
+    metadata_row = None
+    if _header_contains_balance_info(grid[0]):
+        meta_cells = [str(c or "").strip() for c in grid[0]]
+        meta_text = " ".join(c for c in meta_cells if c).strip()
+        if meta_text:
+            metadata_row = [meta_text] + [""] * (new_ncols - 1)
+
+    # Re-align data rows below the candidate header
     new_data_rows = []
     for row in grid[candidate_idx + 1:]:
         if col_expansion > 0 and row:
@@ -765,6 +818,9 @@ def _promote_misplaced_header_row(grid):
         return r[:n]
 
     new_grid = [pad(expanded_header, new_ncols)]
+    # Insert metadata row first if present (preserves Beginning/Ending Balance)
+    if metadata_row is not None:
+        new_grid.append(pad(metadata_row, new_ncols))
     for row in new_data_rows:
         new_grid.append(pad(row, new_ncols))
 
@@ -1055,32 +1111,26 @@ def _reconstruct_separator_rows(grid):
 
     The artifact (Bank of America and similar):
       A section label like "Card account # XXXX XXXX XXXX 2889" sits between
-      transaction rows as a full-width label.  Because the table has columns
-      (Date, Description, Amount) and "2889" is right-aligned, OCR splits it:
-          col 0: "Card account # XXXX XXXX XXXX 2"
-          col 1: ""                                  ← description empty
-          col 2: "889"                               ← trailing digits treated as amount
-
-    How to detect a split separator row (ALL conditions must hold):
-      - The row has exactly 1 non-empty cell OR the last non-empty cell is a
-        short numeric fragment (1-4 digits, no decimal, no sign) that is NOT
-        a plausible standalone amount.
-      - The first cell does NOT match a date pattern (separators are never dates).
-      - The first cell is non-empty and looks like a label (has letters).
-      - The last cell (if non-empty and different from first) is a short digit
-        fragment that makes sense as the tail of the label in col 0.
+      transaction rows as a full-width label. OCR splits the trailing digits
+      across columns because they are right-aligned, producing variations like:
+
+        2-cell split:  ["Card account # XXXX XXXX XXXX 2", "",    "889"]
+        3-cell split:  ["Card account # XXXX XXXX XXXX",   "2",   "889"]
+        clean 1-cell:  ["Card account # XXXX XXXX XXXX 2889", "", ""]
+
+    Detection criteria (ALL must hold — guards safe for all other PDFs):
+      1. First cell is NOT a date token (transaction rows always start with date).
+      2. First cell contains letters (it is a label, not a bare number).
+      3. All cells except the first either:
+           a. are empty, OR
+           b. are a short pure-digit fragment (1-4 digits, no decimal, no sign)
+              — these are the split-off tails of the account/card number.
+      4. There must be at least one non-empty cell after the first (to detect
+         the split; single-cell rows are also handled as clean label rows).
 
     The fix:
-      - Concatenate all non-empty cells (trimming the spurious digit fragment
-        back onto the label), place the result in col 0, and blank all other cells.
-
-    Guard conditions (no-op if not met — safe for all other tables):
-      - Grid must have at least 2 columns and 2 rows.
-      - First cell must not be a date (transaction rows are left alone).
-      - Last cell, if different from first, must be a short pure-digit fragment
-        (1-4 digits, no decimal point) — this is very specific and avoids
-        touching legitimate amount cells like "-14.19".
-      - Entire row must have at most 2 non-empty cells total.
+      Concatenate all non-empty cells in order, place the result in col 0,
+      blank all other cells.
     """
     if not grid or len(grid) < 2:
         return grid
@@ -1093,56 +1143,52 @@ def _reconstruct_separator_rows(grid):
 
     for row in grid[1:]:
         cells = [str(c or "").strip() for c in row]
-
-        # Pad to ncols if needed
         while len(cells) < ncols:
             cells.append("")
 
         non_empty = [(i, c) for i, c in enumerate(cells) if c]
 
-        # Guard 1: must have 1 or 2 non-empty cells total
-        if len(non_empty) == 0 or len(non_empty) > 2:
+        # Completely empty row — leave as-is
+        if len(non_empty) == 0:
             new_grid.append(row)
             continue
 
         first_idx, first_val = non_empty[0]
 
-        # Guard 2: first cell must not be a date token
+        # Guard 1: first cell must not be a date token
         if _DATE_RE.match(first_val):
             new_grid.append(row)
             continue
 
-        # Guard 3: first cell must contain letters (it's a label, not a number)
+        # Guard 2: first cell must contain letters (label, not a bare number)
         if not re.search(r"[A-Za-z]", first_val):
             new_grid.append(row)
             continue
 
+        # Single non-empty cell — already a clean label row
         if len(non_empty) == 1:
-            # Only one cell has content — already a clean label row, ensure
-            # it is in col 0 and all other cells are blank.
             new_row = [""] * ncols
             new_row[0] = first_val
             new_grid.append(new_row)
             continue
 
-        # Exactly 2 non-empty cells
-        last_idx, last_val = non_empty[1]
-
-        # Guard 4: last cell must be a short pure-digit fragment (1-4 digits,
-        # no decimal, no sign, no $) — the hallmark of a split card/account number.
-        if not re.fullmatch(r"\d{1,4}", last_val):
-            new_grid.append(row)
-            continue
+        # Multiple non-empty cells: check that every cell AFTER the first
+        # is a short pure-digit fragment (1-4 digits, no decimal, no sign).
+        # This is the key generic guard — it allows 2, 3, or more cells
+        # as long as all the trailing cells are digit-only fragments.
+        # Real transaction rows always have amounts with decimals (-14.19)
+        # or descriptions with letters, so they never pass this check.
+        trailing = [val for _, val in non_empty[1:]]
+        all_trailing_are_digit_fragments = all(
+            re.fullmatch(r"\d{1,4}", v) for v in trailing
+        )
 
-        # Guard 5: the standalone fragment must NOT be a plausible amount on its own
-        # (amounts have decimals; a bare 2-4 digit integer alone could be a year or
-        # ref number but not a transaction amount in a USD statement)
-        if _SPLIT_AMOUNT_RE.match(last_val) and "." in last_val:
+        if not all_trailing_are_digit_fragments:
             new_grid.append(row)
             continue
 
-        # Reconstruct: join first_val + last_val (they are adjacent parts of the label)
-        reconstructed = (first_val + last_val).strip()
+        # Reconstruct: concatenate all non-empty cells in order
+        reconstructed = "".join(val for _, val in non_empty).strip()
         new_row = [""] * ncols
         new_row[0] = reconstructed
         new_grid.append(new_row)