Update app.py

app.py

@@ -24,6 +24,11 @@ Included (data-driven, no institution names):
     grid, then slide a solitary amount token past trailing blank logical slots
     into the rightmost slot (colspan-aware). Rowspan rows are skipped for edits
     but do not disable an entire table. Stabilization runs in multiple passes.
+  - Split single cells that clearly contain transaction amount + trailing balance
+    (two money tokens, tight gap) into two cells so classifiers can see a balance column.
+  - Long runs of plain-text ledger lines (date + description + amount, optional balance)
+    are coerced into one HTML table when 5+ consecutive lines match generic patterns.
+  - thead uses th only; degenerate empty/sparse non-financial tables are dropped.
 
 Configure GLMOCR_API_KEY (environment variable). Optional: glmocr + gradio +
 pymupdf + pillow installed.
@@ -51,7 +56,7 @@ import os
 import re
 import tempfile
 from collections import Counter
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import yaml
 
@@ -73,7 +78,7 @@ logging.basicConfig(level=logging.INFO)
 # ---------------------------------------------------------------------------
 
 # Never commit secrets: Space / local runs use ZHIPU_API_KEY or GLMOCR_API_KEY.
-GLMOCR_API_KEY = "cee1d52dd91a4ab591b3f6e105f8ad89.LgbQTECuzX0zrito"
+GLMOCR_API_KEY = os.environ.get("ZHIPU_API_KEY") or os.environ.get("GLMOCR_API_KEY") or ""
 if not GLMOCR_API_KEY:
     log.warning(
         "No ZHIPU_API_KEY or GLMOCR_API_KEY in environment; GlmOcr() will fail until you set one."
@@ -304,6 +309,62 @@ _CELL = re.compile(
     re.IGNORECASE | re.DOTALL,
 )
 
+# Currency tokens inside a cell (not anchored); used to split merged amount+balance.
+_MONEY_IN_TEXT = re.compile(
+    r"(?:\$|€|£)?\s*-?\d{1,3}(?:,\d{3})*\.\d{2}\b|(?:\$|€|£)?\s*-?\d+\.\d{2}\b"
+)
+_EOL_MONEY = re.compile(
+    r"(?:\$|€|£)?\s*-?\d{1,3}(?:,\d{3})*\.\d{2}$|(?:\$|€|£)?\s*-?\d+\.\d{2}$"
+)
+
+
+def _split_cell_trailing_balance(full_cell: str) -> List[str]:
+    """
+    When OCR puts transaction amount and running balance in one <td>, split into
+    two cells so classifiers can assign separate columns. Uses only currency
+    patterns and whitespace gaps (no header names).
+    """
+    plain = _cell_plain_text(full_cell)
+    if len(plain) < 10:
+        return [full_cell]
+    spans = [(m.start(), m.end()) for m in _MONEY_IN_TEXT.finditer(plain)]
+    if len(spans) < 2:
+        return [full_cell]
+    (a0, a1), (b0, b1) = spans[-2], spans[-1]
+    if b1 < len(plain) - 16:
+        return [full_cell]
+    gap = plain[a1:b0]
+    if re.search(r"[A-Za-z]{2,}", gap):
+        return [full_cell]
+    if len(gap) > 14:
+        return [full_cell]
+    left = plain[:b0].strip()
+    right = plain[b0:].strip()
+    if not left or not right:
+        return [full_cell]
+    m = _CELL.fullmatch(full_cell.strip())
+    if not m:
+        return [full_cell]
+    tag, attrs = m.group(1), m.group(2)
+    return [
+        f"<{tag}{attrs}>{html.escape(left)}</{tag}>",
+        f"<{tag}{attrs}>{html.escape(right)}</{tag}>",
+    ]
+
+
+def _expand_tr_inner_split_merged(tr_inner: str) -> str:
+    """Insert extra td/th where a single cell clearly holds amount + trailing balance."""
+    entries = _cell_entries(tr_inner)
+    if not entries:
+        return tr_inner
+    parts: List[str] = []
+    for full, span in entries:
+        if span != 1:
+            parts.append(full)
+        else:
+            parts.extend(_split_cell_trailing_balance(full))
+    return "".join(parts)
+
 
 def _cell_entries(tr_inner: str) -> List[Tuple[str, int]]:
     """(full_cell_html, logical_width) for each td/th; 0 cells if unparseable."""
@@ -528,6 +589,29 @@ def normalize_html_table_row_widths(md: str) -> str:
         if not tr_blocks:
             return full
 
+        # Phase 1: split merged amount+balance cells so column counts match real grids.
+        phase1_parts: List[str] = []
+        last_end = 0
+        for tm in tr_blocks:
+            phase1_parts.append(body[last_end : tm.start()])
+            seg = tm.group(0)
+            op = re.search(r"<tr\b[^>]*>", seg, flags=re.IGNORECASE)
+            cl = seg.lower().rfind("</tr>")
+            if not op or cl < 0:
+                phase1_parts.append(seg)
+            else:
+                open_tr = seg[: op.end()]
+                inner = seg[op.end() : cl]
+                close_tr = seg[cl:]
+                if re.search(r"rowspan\s*=", inner, flags=re.IGNORECASE):
+                    phase1_parts.append(seg)
+                else:
+                    phase1_parts.append(open_tr + _expand_tr_inner_split_merged(inner) + close_tr)
+            last_end = tm.end()
+        phase1_parts.append(body[last_end:])
+        body = "".join(phase1_parts)
+
+        tr_blocks = list(re.finditer(r"<tr\b[^>]*>.*?</tr>", body, flags=re.IGNORECASE | re.DOTALL))
         tr_inners: List[str] = []
         for tm in tr_blocks:
             seg = tm.group(0)
@@ -649,6 +733,108 @@ def strip_degenerate_html_tables(md: str) -> str:
     return re.sub(r"\n{3,}", "\n\n", out)
 
 
+def _strip_trailing_money_tokens(rest: str, max_take: int = 2) -> Tuple[str, List[str]]:
+    """Strip 1–2 currency tokens from the right of *rest*; return (description, tokens oldest-first)."""
+    cur = rest.rstrip()
+    toks: List[str] = []
+    for _ in range(max_take):
+        cur = cur.rstrip()
+        m = _EOL_MONEY.search(cur)
+        if not m or m.end() != len(cur):
+            break
+        toks.append(m.group().strip())
+        cur = cur[: m.start()].rstrip()
+    toks.reverse()
+    return cur, toks
+
+
+def _parse_ledger_line(line: str) -> Optional[Tuple[str, str, str, str]]:
+    """
+    If *line* looks like a bank-style ledger row (date, description, amount, optional balance),
+    return (date, description, amount, balance_or_empty); else None.
+    """
+    ln = line.strip()
+    if not ln or ln[0] in "#|!<":
+        return None
+    mm = re.match(
+        r"^(?P<d>(?:\d{1,2}[/-]\d{1,2}(?:[/-]\d{2,4})?|"
+        r"(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Sept|Oct|Nov|Dec)[a-z]*\.?\s+\d{1,2},?\s+\d{4}))"
+        r"\s+",
+        ln,
+        re.I,
+    )
+    if not mm:
+        return None
+    d = mm.group("d")
+    tail = ln[mm.end() :]
+    desc, toks = _strip_trailing_money_tokens(tail, 2)
+    desc = desc.strip()
+    if not toks or len(desc) < 2 or len(desc) > 420:
+        return None
+    amt = toks[0]
+    bal = toks[1] if len(toks) > 1 else ""
+    return (d, desc, amt, bal)
+
+
+def _ledger_rows_to_html(rows: List[Tuple[str, str, str, str]]) -> str:
+    has_bal = any(r[3] for r in rows)
+    head = (
+        "<tr><th>Date</th><th>Description</th><th>Amount</th>"
+        + ("<th>Balance</th>" if has_bal else "")
+        + "</tr>"
+    )
+    body_lines = []
+    for d, desc, amt, bal in rows:
+        cells = [
+            f"<td>{html.escape(d)}</td>",
+            f"<td>{html.escape(desc)}</td>",
+            f"<td>{html.escape(amt)}</td>",
+        ]
+        if has_bal:
+            cells.append(f"<td>{html.escape(bal)}</td>" if bal else "<td></td>")
+        body_lines.append("<tr>" + "".join(cells) + "</tr>")
+    return "<table>\n" + head + "\n" + "\n".join(body_lines) + "\n</table>"
+
+
+def coalesce_loose_ledger_lines(text: str) -> str:
+    """
+    When the model emits many consecutive plain-text ledger lines (no HTML table),
+    downstream classification sees no tables and skips. Convert long runs of
+    generic date + amount lines into a single HTML table. Requires 5+ matching
+    consecutive non-empty lines; does not match institution names.
+    """
+    if not text or "<table" in text.lower():
+        return text
+    lines = text.split("\n")
+    out: List[str] = []
+    i = 0
+    while i < len(lines):
+        if not lines[i].strip():
+            out.append(lines[i])
+            i += 1
+            continue
+        run_rows: List[Tuple[str, str, str, str]] = []
+        j = i
+        while j < len(lines):
+            raw = lines[j]
+            if not raw.strip():
+                break
+            parsed = _parse_ledger_line(raw)
+            if parsed is None:
+                break
+            run_rows.append(parsed)
+            j += 1
+        if len(run_rows) >= 5:
+            if i > 0 and out and out[-1].strip():
+                out.append("")
+            out.append(_ledger_rows_to_html(run_rows))
+            i = j
+            continue
+        out.append(lines[i])
+        i += 1
+    return "\n".join(out)
+
+
 def looks_like_markdown_table(block: str) -> bool:
     lines = [ln.rstrip() for ln in block.strip().splitlines() if ln.strip()]
     if len(lines) < 2:
@@ -716,7 +902,7 @@ def light_stabilize_markdown(page_md: str) -> str:
         if looks_like_markdown_table(b):
             out_blocks.append(md_table_to_html(b))
         else:
-            out_blocks.append(b)
+            out_blocks.append(coalesce_loose_ledger_lines(b))
     merged = close_unclosed_html("\n\n".join(out_blocks))
     merged = repair_thead_cell_semantics(merged)
     merged = stabilize_table_markup(merged, rounds=4)
@@ -900,4 +1086,4 @@ def _create_gradio_demo():
 
 
 if __name__ == "__main__":
-    _create_gradio_demo().launch(share=True)
+    _create_gradio_demo().launch()