formats.py

"""JSON ↔ TOON (Tabular Object Object Notation) conversion.

Compact wire format for uniform object arrays: header once, rows as CSV.
For tabular tool results, ~30-60% smaller than JSON in token count.
"""

from __future__ import annotations

import csv
import io
import json
from typing import Any

from .tokens import TokenCounter


def to_toon(data: list[dict[str, Any]]) -> str:
    if not isinstance(data, list):
        raise ValueError("to_toon expects a list of dicts")
    if not data:
        return "records[0]{}:\n"
    if not all(isinstance(row, dict) for row in data):
        raise ValueError("to_toon requires every element to be a dict")
    fields = list(data[0].keys())
    for row in data[1:]:
        if list(row.keys()) != fields:
            raise ValueError("to_toon requires uniform keys across rows")
    for row in data:
        for v in row.values():
            if isinstance(v, (dict, list)):
                raise ValueError("to_toon does not support nested structures")

    out = io.StringIO()
    out.write(f"records[{len(data)}]{{{','.join(fields)}}}:\n")
    writer = csv.writer(out, lineterminator="\n")
    for row in data:
        writer.writerow([row[f] for f in fields])
    return out.getvalue()


def from_toon(text: str) -> list[dict[str, Any]]:
    lines = text.splitlines()
    if not lines:
        return []
    header = lines[0].strip()
    if not header.startswith("records["):
        raise ValueError("from_toon: missing records[N]{...}: header")
    try:
        n_part, rest = header.split("]", 1)
        n = int(n_part[len("records[") :])
        fields_part = rest.strip()
        if not fields_part.startswith("{") or not fields_part.endswith(":"):
            raise ValueError
        fields = (
            fields_part[1 : fields_part.index("}")].split(",")
            if fields_part != "{}:" else []
        )
    except (ValueError, IndexError) as e:
        raise ValueError(f"from_toon: malformed header: {header}") from e

    rows: list[dict[str, Any]] = []
    reader = csv.reader(lines[1 : 1 + n])
    for raw in reader:
        if not raw and not fields:
            rows.append({})
            continue
        rows.append({f: _coerce(v) for f, v in zip(fields, raw)})
    return rows


def _coerce(value: str) -> Any:
    if value == "":
        return ""
    try:
        return int(value)
    except ValueError:
        pass
    try:
        return float(value)
    except ValueError:
        pass
    if value.lower() in ("true", "false"):
        return value.lower() == "true"
    return value


def benchmark_format(data: list[dict[str, Any]], model: str = "gpt-4o-mini") -> dict[str, Any]:
    counter = TokenCounter()
    json_repr = json.dumps(data)
    toon_repr = to_toon(data)
    j = counter.count(json_repr, model)
    t = counter.count(toon_repr, model)
    delta_pct = ((j - t) / j * 100) if j else 0.0
    return {"json_tokens": j, "toon_tokens": t, "delta_pct": delta_pct}