Spaces:
Sleeping
Sleeping
File size: 4,392 Bytes
d3d0e0e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | 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,
}
|