core/validator.py

from __future__ import annotations

import json
import re
from collections.abc import Callable

from pydantic import ValidationError

from core.schemas import SimulationResult

RetryFn = Callable[[str, str], str]


def extract_json(raw: str) -> str:
    stripped = raw.strip()
    fence_match = re.search(r"```(?:json)?\s*(.*?)\s*```", stripped, re.DOTALL | re.IGNORECASE)
    if fence_match:
        return fence_match.group(1).strip()
    start = stripped.find("{")
    end = stripped.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise json.JSONDecodeError("No JSON object found", stripped, 0)
    return stripped[start : end + 1]


def validate_result(
    raw_json: str,
    persona_id: str,
    retry_fn: RetryFn | None = None,
    max_retries: int = 2,
) -> SimulationResult | None:
    raw = raw_json
    for attempt in range(max_retries + 1):
        try:
            return SimulationResult.model_validate_json(extract_json(raw))
        except (json.JSONDecodeError, ValidationError) as exc:
            if retry_fn is None or attempt >= max_retries:
                return None
            raw = retry_fn(persona_id, str(exc))
    return None


def validate_swarm(
    raw_results: list[tuple[str, str]],
    retry_fn: RetryFn | None = None,
) -> tuple[list[SimulationResult], int]:
    valid: list[SimulationResult] = []
    failures = 0
    for persona_id, raw in raw_results:
        result = validate_result(raw, persona_id, retry_fn)
        if result is None:
            failures += 1
        else:
            valid.append(result)
    return valid, failures