src/utils/json_repair.py

"""Small JSON object extraction helpers for model output."""

from __future__ import annotations

import json
from typing import Any


def parse_json_object(raw: str) -> dict[str, Any]:
    value = json.loads(_extract_json_object(raw))
    if not isinstance(value, dict):
        raise ValueError("Expected a JSON object.")
    return value


def _extract_json_object(raw: str) -> str:
    clean = _strip_markdown_fence(raw.strip())

    candidate = _scan_json_object(clean)
    if candidate is None:
        raise ValueError("No JSON object found.")
    return candidate


def _strip_markdown_fence(clean: str) -> str:
    if not clean.startswith("```"):
        return clean

    lines = clean.splitlines()
    if not lines:
        return clean

    if lines[0].strip().startswith("```"):
        lines = lines[1:]
    if lines and lines[-1].strip().startswith("```"):
        lines = lines[:-1]
    return "\n".join(lines).strip()


def _scan_json_object(clean: str) -> str | None:
    start = clean.find("{")
    if start == -1:
        return None

    stack: list[str] = []
    in_string = False
    escaped = False

    for index, char in enumerate(clean[start:], start=start):
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            continue

        if char == '"':
            in_string = True
        elif char == "{":
            stack.append("}")
        elif char == "[":
            stack.append("]")
        elif char in {"}", "]"}:
            if not stack or stack[-1] != char:
                return None
            stack.pop()
            if not stack:
                return clean[start : index + 1]

    if in_string or not stack:
        return None
    return clean[start:] + "".join(reversed(stack))