Update app.py

app.py

@@ -164,7 +164,7 @@ if CONFIG_PATH:
         with open(CONFIG_PATH, "r", encoding="utf-8") as f:
             config = yaml.safe_load(f)
         config.setdefault("pipeline", {}).setdefault("maas", {})
-        config["pipeline"]["maas"]["enabled"] = True
+        config["pipeline"]["maas"]["enabled"] = TrueCONFIG_PATH
         config["pipeline"]["maas"]["api_key"] = GLMOCR_API_KEY
         with open(CONFIG_PATH, "w", encoding="utf-8") as f:
             yaml.dump(config, f, default_flow_style=False, sort_keys=False)
@@ -910,6 +910,288 @@ def light_stabilize_markdown(page_md: str) -> str:
     return merged
 
 
+def _parse_amount_or_none(s: str):
+    t = (s or "").strip()
+    if not t:
+        return None
+    t = t.replace("$", "").replace(",", "").replace("(", "-").replace(")", "")
+    t = t.replace("−", "-").replace("–", "-").replace("—", "-")
+    if not re.search(r"\d", t):
+        return None
+    try:
+        return float(t)
+    except Exception:
+        return None
+
+
+def _extract_rows_plain_from_table(full_table: str) -> List[List[str]]:
+    rows: List[List[str]] = []
+    for tr in re.finditer(r"<tr\b[^>]*>.*?</tr>", full_table, flags=re.IGNORECASE | re.DOTALL):
+        inner_m = re.search(r"<tr\b[^>]*>(.*)</tr>", tr.group(0), flags=re.IGNORECASE | re.DOTALL)
+        if not inner_m:
+            continue
+        inner = inner_m.group(1)
+        cells = [_cell_plain_text(m.group(0)) for m in _CELL.finditer(inner)]
+        if cells:
+            rows.append(cells)
+    return rows
+
+
+def _extract_daily_balance_by_date(md: str):
+    """
+    Build a date->balance map from daily-balance style tables.
+    Supports compact statements with repeated Date/Balance column pairs.
+    """
+    out = {}
+    for tm in re.finditer(r"<table\b[^>]*>.*?</table>", md, flags=re.IGNORECASE | re.DOTALL):
+        t = tm.group(0)
+        rows = _extract_rows_plain_from_table(t)
+        if len(rows) < 2:
+            continue
+        hdr = [c.strip().lower() for c in rows[0]]
+        if not hdr:
+            continue
+        date_cols = [i for i, c in enumerate(hdr) if "date" in c]
+        bal_cols = [i for i, c in enumerate(hdr) if "balance" in c]
+        if not date_cols or not bal_cols:
+            continue
+        pairs = []
+        for di in date_cols:
+            bi = next((b for b in bal_cols if b > di), None)
+            if bi is not None:
+                pairs.append((di, bi))
+        if not pairs:
+            continue
+        for r in rows[1:]:
+            for di, bi in pairs:
+                if di >= len(r) or bi >= len(r):
+                    continue
+                d = (r[di] or "").strip()
+                if not re.search(r"\b\d{1,2}/\d{1,2}(?:/\d{2,4})?\b", d):
+                    continue
+                m = re.search(r"\d{1,2}/\d{1,2}(?:/\d{2,4})?", d)
+                if not m:
+                    continue
+                key = m.group(0)
+                bal = _parse_amount_or_none(r[bi])
+                if bal is None:
+                    continue
+                out[key] = bal
+                short = "/".join(key.split("/")[:2])
+                out[short] = bal
+    return out
+
+
+def enrich_transaction_tables_with_daily_balances(md: str) -> str:
+    """
+    If transaction tables have Date+Amount but no Balance column, append a Balance
+    column using date-matched daily ending balances from the same document.
+    """
+    if not md or "<table" not in md.lower():
+        return md
+    dmap = _extract_daily_balance_by_date(md)
+    if not dmap:
+        return md
+
+    def repl(m: re.Match) -> str:
+        table = m.group(0)
+        if re.search(r"(?:colspan|rowspan)\s*=", table, flags=re.IGNORECASE):
+            return table
+        tr_blocks = list(re.finditer(r"<tr\b[^>]*>.*?</tr>", table, flags=re.IGNORECASE | re.DOTALL))
+        if len(tr_blocks) < 2:
+            return table
+        first = tr_blocks[0].group(0)
+        h_inner_m = re.search(r"<tr\b[^>]*>(.*)</tr>", first, flags=re.IGNORECASE | re.DOTALL)
+        if not h_inner_m:
+            return table
+        h_inner = h_inner_m.group(1)
+        h_cells = list(_CELL.finditer(h_inner))
+        if not h_cells:
+            return table
+        hdr_txt = [_cell_plain_text(c.group(0)).strip().lower() for c in h_cells]
+        if any("balance" in c for c in hdr_txt):
+            return table
+        if not any("date" in c for c in hdr_txt):
+            return table
+        if not any(("amount" in c) or ("debit" in c) or ("credit" in c) for c in hdr_txt):
+            return table
+        date_idx = next((i for i, c in enumerate(hdr_txt) if "date" in c), -1)
+        if date_idx < 0:
+            return table
+
+        new_parts: List[str] = []
+        cursor = 0
+        added = 0
+        for i, trm in enumerate(tr_blocks):
+            new_parts.append(table[cursor : trm.start()])
+            tr_seg = trm.group(0)
+            inner_m = re.search(r"<tr\b[^>]*>(.*)</tr>", tr_seg, flags=re.IGNORECASE | re.DOTALL)
+            if not inner_m:
+                new_parts.append(tr_seg)
+                cursor = trm.end()
+                continue
+            tr_open_m = re.search(r"<tr\b[^>]*>", tr_seg, flags=re.IGNORECASE)
+            tr_open = tr_open_m.group(0) if tr_open_m else "<tr>"
+            tr_close = "</tr>"
+            inner = inner_m.group(1)
+            cells = [c.group(0) for c in _CELL.finditer(inner)]
+            if i == 0:
+                new_inner = inner + "<th>Balance</th>"
+                new_parts.append(tr_open + new_inner + tr_close)
+            else:
+                if date_idx >= len(cells):
+                    new_parts.append(tr_seg)
+                    cursor = trm.end()
+                    continue
+                d_txt = _cell_plain_text(cells[date_idx])
+                dm = re.search(r"\d{1,2}/\d{1,2}(?:/\d{2,4})?", d_txt or "")
+                bal = dmap.get(dm.group(0)) if dm else None
+                if bal is None and dm:
+                    bal = dmap.get("/".join(dm.group(0).split("/")[:2]))
+                if bal is None:
+                    new_parts.append(tr_seg)
+                    cursor = trm.end()
+                    continue
+                new_inner = inner + f"<td>{bal:,.2f}</td>"
+                new_parts.append(tr_open + new_inner + tr_close)
+                added += 1
+            cursor = trm.end()
+        new_parts.append(table[cursor:])
+        if added < 3:
+            return table
+        return "".join(new_parts)
+
+    return re.sub(r"<table\b[^>]*>.*?</table>", repl, md, flags=re.IGNORECASE | re.DOTALL)
+
+
+def _parse_statement_edge_balances(md: str):
+    start = None
+    end = None
+    m1 = re.search(
+        r"(?:Beginning|Starting)\s+balance[^$\n]{0,80}\$?\s*(-?\d{1,3}(?:,\d{3})*(?:\.\d{2})?)",
+        md,
+        flags=re.IGNORECASE,
+    )
+    if m1:
+        start = _parse_amount_or_none(m1.group(1))
+    m2 = re.search(
+        r"Ending\s+balance[^$\n]{0,80}\$?\s*(-?\d{1,3}(?:,\d{3})*(?:\.\d{2})?)",
+        md,
+        flags=re.IGNORECASE,
+    )
+    if m2:
+        end = _parse_amount_or_none(m2.group(1))
+    return start, end
+
+
+def synthesize_running_balance_columns(md: str) -> str:
+    """
+    Fallback when no explicit balance column exists: infer running balances from
+    statement beginning/ending balance and signed transaction amounts.
+    """
+    if not md or "<table" not in md.lower():
+        return md
+    start_bal, end_bal = _parse_statement_edge_balances(md)
+    if start_bal is None and end_bal is None:
+        return md
+
+    table_re = re.compile(r"<table\b[^>]*>.*?</table>", re.IGNORECASE | re.DOTALL)
+    tr_re = re.compile(r"<tr\b[^>]*>.*?</tr>", re.IGNORECASE | re.DOTALL)
+
+    tables = list(table_re.finditer(md))
+    row_refs = []
+    amounts = []
+    headers = {}
+
+    for ti, tm in enumerate(tables):
+        t = tm.group(0)
+        if re.search(r"(?:colspan|rowspan)\s*=", t, flags=re.IGNORECASE):
+            continue
+        trs = list(tr_re.finditer(t))
+        if len(trs) < 2:
+            continue
+        h_inner_m = re.search(r"<tr\b[^>]*>(.*)</tr>", trs[0].group(0), flags=re.IGNORECASE | re.DOTALL)
+        if not h_inner_m:
+            continue
+        h_cells = [c.group(0) for c in _CELL.finditer(h_inner_m.group(1))]
+        if not h_cells:
+            continue
+        htxt = [_cell_plain_text(c).strip().lower() for c in h_cells]
+        if any("balance" in h for h in htxt):
+            continue
+        date_idx = next((i for i, h in enumerate(htxt) if "date" in h), -1)
+        amt_idx = next((i for i, h in enumerate(htxt) if "amount" in h or "debit" in h or "credit" in h), -1)
+        if date_idx < 0 or amt_idx < 0:
+            continue
+        headers[ti] = (date_idx, amt_idx)
+        for ri, trm in enumerate(trs[1:], 1):
+            inner_m = re.search(r"<tr\b[^>]*>(.*)</tr>", trm.group(0), flags=re.IGNORECASE | re.DOTALL)
+            if not inner_m:
+                continue
+            cells = [c.group(0) for c in _CELL.finditer(inner_m.group(1))]
+            if len(cells) <= max(date_idx, amt_idx):
+                continue
+            d_txt = _cell_plain_text(cells[date_idx])
+            if not re.search(r"\d{1,2}/\d{1,2}(?:/\d{2,4})?", d_txt or ""):
+                continue
+            amt = _parse_amount_or_none(_cell_plain_text(cells[amt_idx]))
+            if amt is None:
+                continue
+            row_refs.append((ti, ri))
+            amounts.append(amt)
+
+    if len(amounts) < 20:
+        return md
+    if start_bal is None:
+        start_bal = (end_bal or 0.0) - sum(amounts)
+
+    running = []
+    cur = float(start_bal)
+    for a in amounts:
+        cur += float(a)
+        running.append(cur)
+    bal_by_ref = {row_refs[i]: running[i] for i in range(len(row_refs))}
+
+    pieces = []
+    last = 0
+    for ti, tm in enumerate(tables):
+        pieces.append(md[last : tm.start()])
+        t = tm.group(0)
+        if ti not in headers:
+            pieces.append(t)
+            last = tm.end()
+            continue
+        date_idx, amt_idx = headers[ti]
+        trs = list(tr_re.finditer(t))
+        out_t = []
+        cur2 = 0
+        for idx, trm in enumerate(trs):
+            out_t.append(t[cur2 : trm.start()])
+            tr_seg = trm.group(0)
+            inner_m = re.search(r"<tr\b[^>]*>(.*)</tr>", tr_seg, flags=re.IGNORECASE | re.DOTALL)
+            tr_open_m = re.search(r"<tr\b[^>]*>", tr_seg, flags=re.IGNORECASE)
+            tr_open = tr_open_m.group(0) if tr_open_m else "<tr>"
+            if not inner_m:
+                out_t.append(tr_seg)
+                cur2 = trm.end()
+                continue
+            inner = inner_m.group(1)
+            if idx == 0:
+                out_t.append(tr_open + inner + "<th>Balance</th></tr>")
+            else:
+                b = bal_by_ref.get((ti, idx))
+                if b is None:
+                    out_t.append(tr_seg)
+                else:
+                    out_t.append(tr_open + inner + f"<td>{b:,.2f}</td></tr>")
+            cur2 = trm.end()
+        out_t.append(t[cur2:])
+        pieces.append("".join(out_t))
+        last = tm.end()
+    pieces.append(md[last:])
+    return "".join(pieces)
+
+
 def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
     import pymupdf as fitz
     from PIL import Image
@@ -1047,6 +1329,9 @@ def run_ocr(uploaded_file):
         merged = "\n\n---page-separator---\n\n".join(all_pages) if all_pages else "(No content)"
         if merged and merged != "(No content)" and not merged.lstrip().startswith("Error:"):
             merged = stabilize_table_markup(merged, rounds=4)
+            merged = enrich_transaction_tables_with_daily_balances(merged)
+            merged = synthesize_running_balance_columns(merged)
+            merged = stabilize_table_markup(merged, rounds=2)
             merged = repair_thead_cell_semantics(merged)
             merged = strip_degenerate_html_tables(merged)
         return merged