Blame for backend/media_analysis/data_export.py

"""Generic, safe extraction parser and downloadable CSV assembler.

A standardized transport envelope contains source-defined columns + compact row arrays.
There are NO hardcoded environmental engineering or chart fields.
"""
from __future__ import annotations
import base64
import csv
import io
import json
import math
import re

MAX_ROWS = 1000
MAX_COLUMNS = 100
MAX_CELL_CHARS = 10000
PROVENANCE = ["source_paper_id", "source_title", "source_doi", "source_paper_key", "source_media_label", "source_page"]


def _primitive(value):
    return value is None or isinstance(value, (str, bool, int)) or (isinstance(value, float) and math.isfinite(value))


def _valid_columns(columns):
    return (isinstance(columns, list) and 1 <= len(columns) <= MAX_COLUMNS
            and all(isinstance(x, str) and x.strip() and len(x) <= 160 for x in columns)
            and len(set(columns)) == len(columns))


def decode_structured_answer(answer: str):
    """Return (short answer, records, warning), for backward-compatible callers.

    New output is generic columns/rows. Legacy records are supported as a fallback.
    """
    s = (answer or "").strip()
    if s.startswith("```"):
        s = re.sub(r"^```(?:json)?\s*", "", s, flags=re.I)
        s = re.sub(r"\s*```$", "", s)
    try:
        obj = json.loads(s)
    except (ValueError, TypeError):
        return answer, [], "Incomplete or invalid model JSON; no CSV generated"
    if not isinstance(obj, dict):
        return answer, [], "Expected a JSON object"
    readable = obj.get("answer") if isinstance(obj.get("answer"), str) else ""
    columns, rows = obj.get("columns"), obj.get("rows")
    if columns is not None or rows is not None:
        if not _valid_columns(columns):
            return readable, [], "Invalid or duplicated column names"
        if not isinstance(rows, list) or len(rows) > MAX_ROWS:
            return readable, [], "Invalid row count"
        for row in rows:
            if (not isinstance(row, list) or len(row) != len(columns)
                    or any(not _primitive(v) or (isinstance(v, str) and len(v) > MAX_CELL_CHARS) for v in row)):
                return readable, [], "Mismatched columns/rows or unsupported cell type"
        return readable, [dict(zip(columns, row)) for row in rows], None
    legacy = obj.get("records")
    if isinstance(legacy, list) and len(legacy) <= MAX_ROWS and all(isinstance(r, dict) for r in legacy):
        if all(all(isinstance(k, str) and len(k) <= 160 and _primitive(v) for k,v in r.items()) for r in legacy):
            return readable, legacy, None
    return readable, [], "No valid columns/rows in model output"


def normalize_records(records):
    """Source-preserving general normalization. Never infer merged-cell content.

    Keep all source-defined column values unchanged; report no table-specific fields.
    """
    return [dict(rec) for rec in records], []


def _clean_csv_cell(value):
    if value is None:
        return ""
    if isinstance(value, str) and value[:1] in ("=", "+", "-", "@", "\t", "\r", "\n"):
        return "'" + value
    return value


def csv_attachment(records, sources):
    """Export arbitrary rectangular records to base64-encoded CSV for the authenticated API.

    Preserve source column order from model's `columns` list (dict insertion order).
    For several source items, source ID columns are populated only if provided
    explicitly by the model; never silently ascribe rows to the first source.
    """
    if not records:
        return None
    if len(records) > MAX_ROWS:
        raise ValueError("Too many rows")
    columns = list(dict.fromkeys(k for r in records for k in r.keys()))
    if not _valid_columns(columns):
        raise ValueError("Invalid CSV columns")
    for r in records:
        if any(not _primitive(v) for v in r.values()):
            raise ValueError("CSV data must contain only scalar values")
    extras = PROVENANCE if len(sources) == 1 else []
    names = columns + [x for x in extras if x not in columns]
    stream = io.StringIO(newline="")
    writer = csv.DictWriter(stream, fieldnames=names, extrasaction="ignore")
    writer.writeheader()
    src = sources[0] if len(sources) == 1 else {}
    provenance = dict(zip(PROVENANCE, [src.get(k) for k in ("paper_id","paper_title","doi","paper_key","label","page")]))
    for record in records:
        output = dict(record)
        for name in extras:
            if name not in output:
                output[name] = provenance[name]
        writer.writerow({key:_clean_csv_cell(output.get(key)) for key in names})
    data=stream.getvalue().encode("utf-8-sig")
    if len(data)>2_000_000:
        raise ValueError("CSV exceeds 2MB response limit")
    return {"filename":"waterscope_media_extracted.csv", "mime_type":"text/csv", "encoding":"base64",
            "content_base64":base64.b64encode(data).decode("ascii"), "row_count":len(records), "columns":names}