path: root/src/gap_pipeline/models.py

"""Validated schemas for every GAP pipeline boundary."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator


SCHEMA_VERSION = "gap-reference-v1"


class StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class CanonicalItem(StrictModel):
    item_id: str
    problem: str
    solution: str
    problem_type: Literal["proof", "calculation"] | str = "proof"
    variables: list[str] = Field(default_factory=list)
    parameters: list[str] = Field(default_factory=list)
    scientific_constants: list[str] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)

    @classmethod
    def from_public_record(cls, record: dict[str, Any]) -> "CanonicalItem":
        excluded = {
            "index",
            "question",
            "solution",
            "problem_type",
            "vars",
            "params",
            "sci_consts",
            "variants",
        }
        return cls(
            item_id=str(record["index"]),
            problem=str(record["question"]),
            solution=str(record["solution"]),
            problem_type=str(record.get("problem_type", "proof")),
            variables=list(record.get("vars", [])),
            parameters=list(record.get("params", [])),
            scientific_constants=list(record.get("sci_consts", [])),
            metadata={key: value for key, value in record.items() if key not in excluded},
        )


class ProofNode(StrictModel):
    node_id: str
    claim: str
    derivation: str
    dependencies: list[str] = Field(default_factory=list)
    source_span: str = ""
    mathematical_objects: list[str] = Field(default_factory=list)


class ProofDAG(StrictModel):
    nodes: list[ProofNode]
    terminal_node_id: str

    @model_validator(mode="after")
    def validate_graph(self) -> "ProofDAG":
        if not self.nodes:
            raise ValueError("proof DAG must contain at least one node")
        ids = [node.node_id for node in self.nodes]
        if len(ids) != len(set(ids)):
            raise ValueError("proof DAG node IDs must be unique")
        known = set(ids)
        if self.terminal_node_id not in known:
            raise ValueError("terminal node is absent from proof DAG")
        for node in self.nodes:
            if node.node_id in node.dependencies:
                raise ValueError(f"node {node.node_id} depends on itself")
            unknown = set(node.dependencies) - known
            if unknown:
                raise ValueError(
                    f"node {node.node_id} has unknown dependencies {sorted(unknown)}"
                )

        dependencies = {node.node_id: set(node.dependencies) for node in self.nodes}
        remaining = set(known)
        resolved: set[str] = set()
        while remaining:
            ready = {node_id for node_id in remaining if dependencies[node_id] <= resolved}
            if not ready:
                raise ValueError("proof graph contains a cycle")
            resolved.update(ready)
            remaining.difference_update(ready)

        ancestors = {self.terminal_node_id}
        frontier = [self.terminal_node_id]
        while frontier:
            node_id = frontier.pop()
            for dependency in dependencies[node_id]:
                if dependency not in ancestors:
                    ancestors.add(dependency)
                    frontier.append(dependency)
        disconnected = known - ancestors
        if disconnected:
            raise ValueError(
                "proof DAG contains nodes disconnected from the terminal "
                f"conclusion: {sorted(disconnected)}"
            )
        return self

    def topological_nodes(self) -> list[ProofNode]:
        by_id = {node.node_id: node for node in self.nodes}
        remaining = set(by_id)
        resolved: set[str] = set()
        ordered: list[ProofNode] = []
        while remaining:
            ready = sorted(
                node_id
                for node_id in remaining
                if set(by_id[node_id].dependencies) <= resolved
            )
            for node_id in ready:
                ordered.append(by_id[node_id])
                resolved.add(node_id)
                remaining.remove(node_id)
        return ordered

    def leaf_node_ids(self) -> list[str]:
        return sorted(node.node_id for node in self.nodes if not node.dependencies)


class MethodStep(StrictModel):
    node_id: str
    method_label: str
    input_roles: list[str] = Field(default_factory=list)
    output_role: str
    invariants: list[str] = Field(default_factory=list)


class MethodPlan(StrictModel):
    steps: list[MethodStep]
    terminal_node_id: str

    def validate_against(self, dag: ProofDAG) -> None:
        dag_order = [node.node_id for node in dag.topological_nodes()]
        plan_order = [step.node_id for step in self.steps]
        if plan_order != dag_order:
            raise ValueError(
                "method plan must contain every DAG node in topological order; "
                f"expected {dag_order}, received {plan_order}"
            )
        if self.terminal_node_id != dag.terminal_node_id:
            raise ValueError("method-plan terminal does not match DAG terminal")
        if any(not step.method_label.strip() for step in self.steps):
            raise ValueError("method labels must be non-empty")


class ReplacementSpec(StrictModel):
    target_node_id: str
    original_object: str
    replacement_object: str
    guard_conditions: list[str]
    guard_evidence: list[str]
    expected_downstream_changes: list[str]
    rationale: str

    def validate_against(self, dag: ProofDAG) -> None:
        if self.target_node_id not in dag.leaf_node_ids():
            raise ValueError(
                f"replacement target {self.target_node_id} is not a DAG leaf"
            )
        if self.original_object.strip() == self.replacement_object.strip():
            raise ValueError("replacement must genuinely change the leaf object")
        if not self.guard_conditions or not self.guard_evidence:
            raise ValueError("replacement needs explicit guards and guard evidence")


class NodeInstantiation(StrictModel):
    node_id: str
    method_label: str
    instantiated_claim: str
    instantiated_derivation: str
    dependencies: list[str] = Field(default_factory=list)


class DiffusedProof(StrictModel):
    node_instantiations: list[NodeInstantiation]
    regenerated_proof: str
    terminal_answer: str

    def validate_against(self, dag: ProofDAG, plan: MethodPlan) -> None:
        dag_order = [node.node_id for node in dag.topological_nodes()]
        actual_order = [node.node_id for node in self.node_instantiations]
        if actual_order != dag_order:
            raise ValueError("diffused proof node order differs from proof DAG")
        expected_labels = [step.method_label for step in plan.steps]
        actual_labels = [node.method_label for node in self.node_instantiations]
        if actual_labels != expected_labels:
            raise ValueError("diffused proof changes the method-label sequence")
        expected_dependencies = {
            node.node_id: node.dependencies for node in dag.topological_nodes()
        }
        for node in self.node_instantiations:
            if node.dependencies != expected_dependencies[node.node_id]:
                raise ValueError(
                    f"diffused proof changes dependencies for {node.node_id}"
                )


class KernelCandidate(StrictModel):
    problem: str
    proof: str
    terminal_answer: str
    node_instantiations: list[NodeInstantiation]
    replacement: ReplacementSpec

    def validate_against(self, dag: ProofDAG, plan: MethodPlan) -> None:
        DiffusedProof(
            node_instantiations=self.node_instantiations,
            regenerated_proof=self.proof,
            terminal_answer=self.terminal_answer,
        ).validate_against(dag, plan)
        self.replacement.validate_against(dag)


class JudgeVerdict(StrictModel):
    verdict: Literal["accept", "reject"]
    step_by_step_check: str
    blocking_issues: str = ""
    patch_suggestion: str = ""

    @model_validator(mode="after")
    def validate_verdict(self) -> "JudgeVerdict":
        if self.verdict == "accept":
            if self.blocking_issues.strip():
                raise ValueError("accept verdict cannot contain blocking issues")
        else:
            if not self.blocking_issues.strip():
                raise ValueError("reject verdict must identify a blocking issue")
            if not self.patch_suggestion.strip():
                raise ValueError("reject verdict must include a patch suggestion")
        return self


class IterationRecord(StrictModel):
    iteration: int
    candidate_sha256: str
    verdicts: list[JudgeVerdict]
    unanimous: bool
    pass_streak_after: int
    repaired: bool = False
    repaired_candidate_sha256: str | None = None


class KernelRunResult(StrictModel):
    item_id: str
    status: Literal["accepted", "rejected"]
    accepted_candidate: KernelCandidate | None = None
    iterations: list[IterationRecord]
    rejection_reason: str = ""
    schema_version: str = SCHEMA_VERSION

    @model_validator(mode="after")
    def validate_terminal_state(self) -> "KernelRunResult":
        if self.status == "accepted" and self.accepted_candidate is None:
            raise ValueError("accepted run must include its accepted candidate")
        if self.status == "rejected" and not self.rejection_reason:
            raise ValueError("rejected run must state a reason")
        return self


class ModelCallRecord(StrictModel):
    request_id: str
    model: str
    system_prompt: str
    user_prompt: str
    response_data: dict[str, Any]
    raw_text: str
    provider_response_id: str | None = None
    usage: dict[str, Any] = Field(default_factory=dict)
    created_at: str = Field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )


class StageArtifact(StrictModel):
    item_id: str
    stage: str
    payload_sha256: str
    payload: dict[str, Any]
    request_id: str | None = None
    schema_version: str = SCHEMA_VERSION
    created_at: str = Field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )