examples/parse_output.py

"""Parse the writer's first balanced JSON object without repairing its text.

This mirrors the parser used for the checkpoint's reported 4.3926 evaluation:
Markdown or prose around a complete JSON object is tolerated, but malformed or
truncated JSON is not repaired and is not retried.
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path
from typing import Any


CATEGORIES = ("content", "organization", "expression")


def extract_first_json_object(text: str) -> str | None:
    """Return the first balanced JSON-object substring in *text*."""

    start = text.find("{")
    if start == -1:
        return None

    depth = 0
    in_string = False
    escaped = False
    for index in range(start, len(text)):
        char = text[index]
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            continue

        if char == '"':
            in_string = True
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]
    return None


def coerce_score(value: object) -> int:
    """Return a 1..5 integer, applying historical half-up float rounding."""

    if isinstance(value, bool):
        raise ValueError(f"invalid boolean score: {value!r}")
    if isinstance(value, int):
        score = value
    elif isinstance(value, float) and math.isfinite(value):
        score = math.floor(value + 0.5)
    else:
        raise ValueError(f"score is not numeric: {value!r}")
    if not 1 <= score <= 5:
        raise ValueError(f"score is outside 1..5: {value!r}")
    return score


def parse_writer_output(text: str) -> dict[str, dict[str, Any]]:
    """Extract and validate the writer's nested content/organization/expression JSON."""

    candidate = extract_first_json_object(text)
    if candidate is None:
        raise ValueError("no balanced JSON object found")
    try:
        parsed = json.loads(candidate)
    except json.JSONDecodeError as error:
        raise ValueError(f"invalid JSON: {error}") from error
    if not isinstance(parsed, dict):
        raise ValueError("top-level JSON is not an object")

    validated: dict[str, dict[str, Any]] = {}
    for category in CATEGORIES:
        block = parsed.get(category)
        if not isinstance(block, dict):
            raise ValueError(f"{category}: expected an object")
        if "score" not in block or "rationale" not in block:
            raise ValueError(f"{category}: missing score/rationale")
        rationale = block["rationale"]
        if not isinstance(rationale, str) or not rationale.strip():
            raise ValueError(f"{category}: rationale must be a non-empty string")
        validated[category] = {
            "score": coerce_score(block["score"]),
            "rationale": rationale,
        }
    return validated


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "path",
        nargs="?",
        type=Path,
        help="raw model-output file; omit to read UTF-8 text from stdin",
    )
    args = parser.parse_args()
    raw = args.path.read_text(encoding="utf-8") if args.path else sys.stdin.read()
    try:
        parsed = parse_writer_output(raw)
    except ValueError as error:
        raise SystemExit(f"parse failure: {error}") from error
    print(json.dumps(parsed, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()