from __future__ import annotations
import argparse
import json
from typing import Any
import jsonschema
from common_json import load_answer_from_artifact, extract_first_json
def main():
ap = argparse.ArgumentParser(description="Validate model output JSON against a JSON Schema.")
ap.add_argument("--schema", required=True, help="Path to JSON schema file")
ap.add_argument("--pred", required=True, help="Model output artifact (.txt/.json or directory)")
ap.add_argument("--expect", choices=["object", "array", "any"], default="any", help="Expected top-level JSON type")
ap.add_argument("--out", default=None, help="Write report JSON here")
args = ap.parse_args()
schema = json.loads(open(args.schema, "r", encoding="utf-8").read())
raw, parsed = load_answer_from_artifact(args.pred)
if parsed is None:
try:
parsed = extract_first_json(raw)
except Exception:
parsed = None
report = {"parse_ok": 1 if parsed is not None else 0, "schema_ok": 0, "expected_type_ok": 0, "error": None}
if parsed is None:
report["error"] = "Could not parse JSON from model output"
else:
if args.expect == "any":
report["expected_type_ok"] = 1
elif args.expect == "object":
report["expected_type_ok"] = 1 if isinstance(parsed, dict) else 0
elif args.expect == "array":
report["expected_type_ok"] = 1 if isinstance(parsed, list) else 0
try:
jsonschema.validate(instance=parsed, schema=schema)
report["schema_ok"] = 1
except jsonschema.ValidationError as e:
report["schema_ok"] = 0
report["error"] = str(e).splitlines()[0]
except jsonschema.SchemaError as e:
report["schema_ok"] = 0
report["error"] = f"SchemaError: {str(e).splitlines()[0]}"
if args.out:
with open(args.out, "w", encoding="utf-8") as f:
json.dump(report, f, indent=2, ensure_ascii=False)
print(json.dumps(report, indent=2, ensure_ascii=False))
if __name__ == "__main__":
main()
