Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import csv | |
| import json | |
| from pathlib import Path | |
| from data_agent_baseline.benchmark.schema import PublicTask | |
| def resolve_context_path(task: PublicTask, relative_path: str) -> Path: | |
| candidate = (task.context_dir / relative_path).resolve() | |
| context_root = task.context_dir.resolve() | |
| if context_root not in candidate.parents and candidate != context_root: | |
| raise ValueError(f"Path escapes context dir: {relative_path}") | |
| if not candidate.exists(): | |
| raise FileNotFoundError(f"Missing context asset: {relative_path}") | |
| return candidate | |
| def list_context_tree(task: PublicTask, *, max_depth: int = 4) -> dict[str, object]: | |
| entries: list[dict[str, object]] = [] | |
| def walk(path: Path, depth: int) -> None: | |
| if depth > max_depth: | |
| return | |
| for child in sorted(path.iterdir(), key=lambda item: (item.is_file(), item.name)): | |
| if child.name == "_auto.db": | |
| continue | |
| rel_path = child.relative_to(task.context_dir).as_posix() | |
| entries.append( | |
| { | |
| "path": rel_path, | |
| "kind": "dir" if child.is_dir() else "file", | |
| "size": child.stat().st_size if child.is_file() else None, | |
| } | |
| ) | |
| if child.is_dir(): | |
| walk(child, depth + 1) | |
| walk(task.context_dir, 1) | |
| return { | |
| "root": str(task.context_dir), | |
| "entries": entries, | |
| } | |
| def read_csv_preview(task: PublicTask, relative_path: str, *, max_rows: int = 20) -> dict[str, object]: | |
| path = resolve_context_path(task, relative_path) | |
| with path.open(newline="") as handle: | |
| reader = csv.reader(handle) | |
| rows = list(reader) | |
| if not rows: | |
| return { | |
| "path": relative_path, | |
| "columns": [], | |
| "rows": [], | |
| "row_count": 0, | |
| } | |
| header = rows[0] | |
| data_rows = rows[1:] | |
| return { | |
| "path": relative_path, | |
| "columns": header, | |
| "rows": data_rows[:max_rows], | |
| "row_count": len(data_rows), | |
| } | |
| def read_json_preview(task: PublicTask, relative_path: str, *, max_chars: int = 10000) -> dict[str, object]: | |
| path = resolve_context_path(task, relative_path) | |
| payload = json.loads(path.read_text()) | |
| # For {table, records} format (common in this benchmark), provide structured preview | |
| # showing ALL column names from the first record + sample records (minimum 2) | |
| if isinstance(payload, dict) and "records" in payload and isinstance(payload["records"], list): | |
| records = payload["records"] | |
| columns = list(records[0].keys()) if records else [] | |
| # Always show at least 2 sample records, even if it exceeds max_chars | |
| sample_count = max(2, min(len(records), 5)) | |
| sample_records = records[:sample_count] | |
| structured_preview = json.dumps( | |
| {"table": payload.get("table", ""), "columns": columns, | |
| "record_count": len(records), "sample_records": sample_records}, | |
| ensure_ascii=False, indent=2, | |
| ) | |
| # If more than 2 records, trim to min(10000, max_chars) | |
| if sample_count > 2 and len(structured_preview) > min(10000, max_chars): | |
| # Reduce sample_records to fit within max_chars | |
| for n in range(sample_count, 1, -1): | |
| sample_records = records[:n] | |
| structured_preview = json.dumps( | |
| {"table": payload.get("table", ""), "columns": columns, | |
| "record_count": len(records), "sample_records": sample_records}, | |
| ensure_ascii=False, indent=2, | |
| ) | |
| if len(structured_preview) <= min(10000, max_chars): | |
| break | |
| preview = structured_preview | |
| else: | |
| preview = json.dumps(payload, ensure_ascii=False, indent=2)[:min(10000, max_chars)] | |
| return { | |
| "path": relative_path, | |
| "preview": preview, | |
| "truncated": len(preview) >= max_chars, | |
| } | |
| def read_doc_preview(task: PublicTask, relative_path: str, *, max_chars: int = 5000) -> dict[str, object]: | |
| path = resolve_context_path(task, relative_path) | |
| text = path.read_text(errors="replace") | |
| return { | |
| "path": relative_path, | |
| "preview": text[:max_chars], | |
| "truncated": len(text) > max_chars, | |
| } | |