validate.py

"""
Copyright (c) 2026 The Johns Hopkins University Applied Physics Laboratory LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0
"""


import argparse
import json
import sys
from typing import Annotated, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, Field, ValidationError


format_version = Literal["1.0"]


class Artifact(BaseModel):
    type: Literal["article", "press release", "paper", "figure"]
    text: str


class Problem(BaseModel):
    """
    Represents a problem with its claim text and artifacts
    """
    type: Literal["problem"]
    format_version: format_version
    problem_id: str
    problem_version: str
    claim: str
    artifacts: List[Artifact]
    domain: Literal["materials", "ai", "quantum"]
    subdomain: Optional[str] = Field(None, description="Valid subdomains depend on the selected domain")

    @staticmethod
    def valid_subdomains(domain: str) -> List[str]:
        subdomain_mapping = {
            "materials": ["alloys", "batteries", "semiconductors", "superconductors"],
            "ai": [],
            "quantum": []
        }
        return subdomain_mapping.get(domain, [])

    @classmethod
    def validate(cls, values):
        domain = values.get("domain")
        subdomain = values.get("subdomain")
        if domain and subdomain and subdomain not in cls.valid_subdomains(domain):
            raise ValueError(f"Invalid subdomain '{subdomain}' for domain '{domain}'. Allowed subdomains: {cls.valid_subdomains(domain)}")
        return values


LikertScore = Literal[-2, -1, 0, 1, 2]
Continuous = Annotated[float, Field(ge=0.0, le=1.0)]


class ExplanationItem(BaseModel):
    summary: Optional[str] = None
    text: str
    evidence: List[str]


class EvidenceItem(BaseModel):
    type: str
    source: Optional[str] = None
    summary: Optional[str] = None # descriptive summary of the evidence
    citation: Optional[str] = None # HTML encoded content from the source that is relevant such as images, text passages, code, etc.
    context: Optional[str] = None # any meta information about the evidence such as the approach used to find it, search terms, packages used etc


class Assessment(BaseModel):
    """
    Represents a claim assessment including a feasibility score and explanation with evidence.
    """
    type: Literal["assessment"]
    format_version: format_version
    problem_id: str
    problem_version: str
    team: str
    # unique id for all assessments in a submission
    run_id: str
    likert_score: LikertScore
    continuous_score: Continuous
    confidence: Continuous
    # time in seconds
    wall_clock_time: float
    explanation: List[ExplanationItem]
    evidence: Dict[str, EvidenceItem]


class GoldStandard(BaseModel):
    """
    Represents the authoritative definition of a scientific feasibility claim and explanation.

    A Problem instance for an evaluation can be created from a gold standard.
    This also contains metadata about author, tools, dimensions, etc.
    """
    type: Literal["gold standard"]
    format_version: format_version
    # problem id is usually an identified plus a counter (ex: dry_run_0001)
    problem_id: str
    # Set the version to "1.0" unless it is a revised claim.
    problem_version: str
    domain: Literal["materials", "ai", "quantum"]
    # ["alloys", "batteries", "semiconductors", "superconductors"]
    subdomain: str = Field(..., description="Valid subdomains depend on the selected domain")
    claim: str
    artifacts: Optional[List[Artifact]] = Field(default_factory=list)
    likert_score: LikertScore
    # a list of explanation text blocks with citations or an explanation gist as a string
    explanation: Union[str, List[ExplanationItem]]
    # dictionary of evidence name -> evidence item. If no evidence, set to empty dict.
    evidence: Dict[str, EvidenceItem]
    # the organization or subject matter expert who wrote the claim/explanation
    author: str
    # metadata about the claim such as tool use, evidence modality, dimensions
    tags: Optional[Dict[str, Union[str|int|bool|None]]] = Field(default_factory=dict)
    # comments by the author or a reviewer
    comments: Optional[List[str]] = Field(default_factory=list)


def validate_assessments_jsonl(filename: str):
    validate_jsonl(filename, Assessment)


def validate_problems_jsonl(filename: str):
    validate_jsonl(filename, Problem)


def validate_gold_standards_jsonl(filename: str):
    gold_standards = validate_jsonl(filename, GoldStandard)
    problem_ids = set()
    for gold_standard in gold_standards:
        if gold_standard.problem_id in problem_ids:
            print(f'ERROR: Duplicate gold standard ID "{gold_standard.problem_id}"')
        problem_ids.add(gold_standard.problem_id)


def validate_jsonl(filename: str, content_format: Union[Assessment, Problem, GoldStandard]):
    valid_count = 0
    invalid_count = 0
    inflated_content = []
    with open(filename, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f, start=1):
            try:
                data = json.loads(line)
                inflated_content.append(content_format(**data))
                valid_count += 1
            except (json.JSONDecodeError, ValidationError) as e:
                invalid_count += 1
                print(f"Line {i}: Invalid JSON or schema mismatch")
                print(e)
    print(f"\nValidation complete.")
    print(f"Valid lines: {valid_count}")
    print(f"Invalid lines: {invalid_count}")
    return inflated_content


def main():
    parser = argparse.ArgumentParser(
        description="Validate a JSONL file.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("-i", "--input", required=True, help="The JSONL file to validate")
    args = parser.parse_args()

    # check for problem or assessment
    with open(args.input, 'r', encoding='utf-8') as f:
        try:
            data = json.loads(f.readline())
            file_type = data.get("type")
        except json.JSONDecodeError as e:
            print(f"Invalid JSON on first line of {args.input}")
            sys.exit(-1)
    if file_type == "assessment":
        validate_assessments_jsonl(args.input)
    elif file_type == "problem":
        validate_problems_jsonl(args.input)
    elif file_type == "gold standard":
        validate_gold_standards_jsonl(args.input)
    else:
        print(f"Unknown type {file_type}")


if __name__ == "__main__":
    main()