Spaces:
Sleeping
Sleeping
| """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} | |