Spaces:
Running on Zero
Running on Zero
| """Turn an uploaded file into text and rows. | |
| Deliberately format-agnostic and bank-agnostic: this module only gets the | |
| content *out* of the file. Deciding which column is a date and which is an | |
| amount is :mod:`finbot.ingest`'s job, and is done by the model so that any | |
| bank's export works without a hand-written layout per institution. | |
| Nothing here calls the network, so all of it is unit-testable offline. | |
| """ | |
| from __future__ import annotations | |
| import csv | |
| import io | |
| import logging | |
| from dataclasses import dataclass, field | |
| log = logging.getLogger(__name__) | |
| #: Guard rails. A personal statement is a few hundred rows; anything far past | |
| #: that is either a mistake or would blow the model's context and the user's | |
| #: token budget, so it is truncated with a visible note rather than silently. | |
| MAX_ROWS = 1500 | |
| MAX_TEXT_CHARS = 120_000 | |
| MAX_PDF_PAGES = 40 | |
| class ExtractedDocument: | |
| kind: str | |
| text: str = "" | |
| rows: list[list[str]] = field(default_factory=list) | |
| pages: int = 0 | |
| truncated: bool = False | |
| note: str = "" | |
| def is_empty(self) -> bool: | |
| return not self.text.strip() and not self.rows | |
| def _decode(data: bytes) -> str: | |
| for encoding in ("utf-8-sig", "utf-8", "cp1252", "latin-1"): | |
| try: | |
| return data.decode(encoding) | |
| except UnicodeDecodeError: | |
| continue | |
| return data.decode("utf-8", errors="replace") | |
| def _clean_row(row: list) -> list[str]: | |
| return [("" if c is None else str(c)).strip() for c in row] | |
| def _is_blank(row: list[str]) -> bool: | |
| return not any(cell for cell in row) | |
| def extract_csv(data: bytes) -> ExtractedDocument: | |
| text = _decode(data) | |
| sample = text[:8192] | |
| try: | |
| dialect = csv.Sniffer().sniff(sample, delimiters=",;\t|") | |
| except csv.Error: | |
| # Sniffing fails on single-column or irregular exports; comma is the | |
| # safe default and the model tolerates a wrong split better than a crash. | |
| dialect = csv.excel | |
| rows: list[list[str]] = [] | |
| truncated = False | |
| for row in csv.reader(io.StringIO(text), dialect): | |
| cleaned = _clean_row(row) | |
| if _is_blank(cleaned): | |
| continue | |
| rows.append(cleaned) | |
| if len(rows) >= MAX_ROWS: | |
| truncated = True | |
| break | |
| return ExtractedDocument(kind="csv", rows=rows, truncated=truncated) | |
| def extract_xlsx(data: bytes) -> ExtractedDocument: | |
| from openpyxl import load_workbook | |
| workbook = load_workbook(io.BytesIO(data), read_only=True, data_only=True) | |
| rows: list[list[str]] = [] | |
| truncated = False | |
| try: | |
| for sheet in workbook.worksheets: | |
| if truncated: | |
| break | |
| if len(workbook.worksheets) > 1: | |
| rows.append([f"# sheet: {sheet.title}"]) | |
| for row in sheet.iter_rows(values_only=True): | |
| cleaned = _clean_row(list(row)) | |
| if _is_blank(cleaned): | |
| continue | |
| rows.append(cleaned) | |
| if len(rows) >= MAX_ROWS: | |
| truncated = True | |
| break | |
| finally: | |
| workbook.close() | |
| return ExtractedDocument(kind="xlsx", rows=rows, truncated=truncated) | |
| def extract_pdf(data: bytes) -> ExtractedDocument: | |
| """Pull tables (preferred) and text out of a PDF statement. | |
| Scanned statements contain no text layer. That is detected and reported | |
| rather than silently returning nothing, because "0 transactions found" and | |
| "this file is a photograph of a statement" need very different user advice. | |
| """ | |
| import pdfplumber | |
| rows: list[list[str]] = [] | |
| text_parts: list[str] = [] | |
| truncated = False | |
| page_count = 0 | |
| with pdfplumber.open(io.BytesIO(data)) as pdf: | |
| page_count = len(pdf.pages) | |
| for index, page in enumerate(pdf.pages): | |
| if index >= MAX_PDF_PAGES: | |
| truncated = True | |
| break | |
| try: | |
| for table in page.extract_tables() or []: | |
| for row in table: | |
| cleaned = _clean_row(list(row)) | |
| if _is_blank(cleaned): | |
| continue | |
| rows.append(cleaned) | |
| if len(rows) >= MAX_ROWS: | |
| truncated = True | |
| break | |
| if truncated: | |
| break | |
| except Exception as exc: | |
| log.warning("Table extraction failed on page %d: %s", index + 1, exc) | |
| try: | |
| if page_text := page.extract_text(): | |
| text_parts.append(page_text) | |
| except Exception as exc: | |
| log.warning("Text extraction failed on page %d: %s", index + 1, exc) | |
| if truncated: | |
| break | |
| text = "\n".join(text_parts)[:MAX_TEXT_CHARS] | |
| doc = ExtractedDocument( | |
| kind="pdf", text=text, rows=rows, pages=page_count, truncated=truncated | |
| ) | |
| if doc.is_empty: | |
| doc.note = ( | |
| "This PDF has no selectable text, so it is almost certainly a scan " | |
| "or photo of a statement. Send it as a photo instead and I'll read " | |
| "it with the vision model, or export a CSV/XLSX from your bank." | |
| ) | |
| return doc | |
| def extract(data: bytes, filename: str = "", mime_type: str = "") -> ExtractedDocument: | |
| """Dispatch on filename extension, falling back to MIME type.""" | |
| name = (filename or "").lower() | |
| mime = (mime_type or "").lower() | |
| if name.endswith(".csv") or name.endswith(".tsv") or "csv" in mime: | |
| return extract_csv(data) | |
| if name.endswith((".xlsx", ".xlsm", ".xltx")) or "spreadsheet" in mime or "excel" in mime: | |
| return extract_xlsx(data) | |
| if name.endswith(".pdf") or "pdf" in mime: | |
| return extract_pdf(data) | |
| if name.endswith((".txt", ".md")) or mime.startswith("text/"): | |
| return ExtractedDocument(kind="text", text=_decode(data)[:MAX_TEXT_CHARS]) | |
| # .xls (old binary format) is not readable by openpyxl -- say so precisely. | |
| if name.endswith(".xls"): | |
| return ExtractedDocument( | |
| kind="unsupported", | |
| note=( | |
| "Old .xls files aren't supported. Open it and re-save as .xlsx " | |
| "or export CSV, then send it again." | |
| ), | |
| ) | |
| return ExtractedDocument( | |
| kind="unsupported", | |
| note=f"I can't read {filename or 'that file'}. Send a CSV, XLSX, PDF or a photo.", | |
| ) | |
| def rows_to_text(rows: list[list[str]], limit: int = MAX_ROWS) -> str: | |
| """Render rows as pipe-delimited lines for a model prompt.""" | |
| lines = [" | ".join(cell for cell in row) for row in rows[:limit]] | |
| return "\n".join(lines)[:MAX_TEXT_CHARS] | |