Spaces:
Sleeping
Sleeping
| """Export extracted data to structured formats: JSON, CSV, Excel.""" | |
| from __future__ import annotations | |
| import csv | |
| import io | |
| import json | |
| from app.schemas.documents import DocumentDetail | |
| def to_json(detail: DocumentDetail) -> bytes: | |
| payload = { | |
| "document": { | |
| "id": detail.id, | |
| "filename": detail.filename, | |
| "doc_type": detail.classification.doc_type if detail.classification else None, | |
| "num_pages": detail.num_pages, | |
| }, | |
| "classification": detail.classification.model_dump() if detail.classification else None, | |
| "summary": detail.summary, | |
| "extraction": detail.extraction.model_dump() if detail.extraction else None, | |
| "anomalies": [a.model_dump() for a in detail.anomalies], | |
| } | |
| return json.dumps(payload, indent=2, default=str).encode("utf-8") | |
| def _record_rows(extraction) -> tuple[list[str], list[list]]: | |
| """Flatten records into (header, rows) using the canonical key order.""" | |
| keys = list(extraction.record_keys) | |
| if not keys: # derive from the records themselves, preserving first-seen order | |
| seen = set() | |
| for rec in extraction.records: | |
| for f in rec.fields: | |
| if f.name not in seen: | |
| keys.append(f.name) | |
| seen.add(f.name) | |
| rows = [] | |
| for rec in extraction.records: | |
| m = {f.name: f.value for f in rec.fields} | |
| rows.append([m.get(k) for k in keys]) | |
| return keys, rows | |
| def to_csv(detail: DocumentDetail) -> bytes: | |
| buf = io.StringIO() | |
| w = csv.writer(buf) | |
| ex = detail.extraction | |
| # Multi-record docs export as a proper record table; otherwise field/value. | |
| if ex and ex.records: | |
| keys, rows = _record_rows(ex) | |
| w.writerow(keys) | |
| for row in rows: | |
| w.writerow(["" if c is None else c for c in row]) | |
| else: | |
| w.writerow(["field", "value", "confidence"]) | |
| if ex: | |
| for f in ex.fields: | |
| w.writerow([f.name, f.value, f"{f.confidence:.2f}"]) | |
| return buf.getvalue().encode("utf-8") | |
| def to_xlsx(detail: DocumentDetail) -> bytes: | |
| from openpyxl import Workbook | |
| wb = Workbook() | |
| ws = wb.active | |
| ws.title = "Fields" | |
| ws.append(["Field", "Value", "Confidence"]) | |
| ex = detail.extraction | |
| if ex: | |
| for f in ex.fields: | |
| ws.append([f.name, str(f.value) if f.value is not None else "", round(f.confidence, 2)]) | |
| # Repeated records get their own sheet, one row per record. | |
| if ex.records: | |
| keys, rows = _record_rows(ex) | |
| rec = wb.create_sheet(title=(ex.record_type or "Records")[:31]) | |
| rec.append(keys) | |
| for row in rows: | |
| rec.append([str(c) if c is not None else "" for c in row]) | |
| for ti, table in enumerate(ex.tables): | |
| tab = wb.create_sheet(title=(table.title or f"Table {ti+1}")[:31]) | |
| if table.columns: | |
| tab.append(table.columns) | |
| for row in table.rows: | |
| tab.append([str(c) if c is not None else "" for c in row]) | |
| out = io.BytesIO() | |
| wb.save(out) | |
| return out.getvalue() | |
| EXPORTERS = { | |
| "json": (to_json, "application/json", "json"), | |
| "csv": (to_csv, "text/csv", "csv"), | |
| "xlsx": (to_xlsx, | |
| "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "xlsx"), | |
| } | |