Formata

"""
Formata: LLM Schema Enforcer & Streaming JSON Repairer
"""
import json
import re
from typing import Dict, Any, Type
from pydantic import BaseModel, ValidationError

class EnforceResult(BaseModel):
    is_valid: bool
    repaired: bool
    payload: Dict[str, Any]
    error: str = ""

class FormataSchemaEnforcer:
    @staticmethod
    def repair_json_syntax(raw_text: str) -> str:
        # Strip markdown fences
        clean = re.sub(r"^```(json)?|```$", "", raw_text.strip(), flags=re.MULTILINE).strip()
        # Fix trailing commas in objects and arrays
        clean = re.sub(r",\s*([}\]])", r"\1", clean)
        # Close unclosed curly brace or brackets if truncated
        open_braces = clean.count("{") - clean.count("}")
        open_brackets = clean.count("[") - clean.count("]")
        clean += "}" * max(0, open_braces)
        clean += "]" * max(0, open_brackets)
        return clean

    def validate_or_repair(self, raw_stream: str, target_schema: Type[BaseModel]) -> EnforceResult:
        try:
            parsed = json.loads(raw_stream)
            model = target_schema.model_validate(parsed)
            return EnforceResult(is_valid=True, repaired=False, payload=model.model_dump())
        except (json.JSONDecodeError, ValidationError):
            repaired_text = self.repair_json_syntax(raw_stream)
            try:
                repaired_json = json.loads(repaired_text)
                model = target_schema.model_validate(repaired_json)
                return EnforceResult(is_valid=True, repaired=True, payload=model.model_dump())
            except Exception as e:
                return EnforceResult(is_valid=False, repaired=False, payload={}, error=str(e))