src/peppa/schema.py

"""Typed molecular identity, scientific requirements, and visible decisions.

Residues use one-based positions. Assay thresholds retain their original units.
Chemical identity includes modifications, bonds, and terminal chemistry.
"""
from __future__ import annotations
import hashlib
import json
from typing import Any, Literal
from pydantic import BaseModel, ConfigDict, Field, model_validator


def canonical(value: Any) -> str:
    """Serialize deterministically, rejecting NaN and infinite numbers."""
    if isinstance(value, BaseModel):
        value = value.model_dump(mode="json")
    return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)


def digest(value: Any) -> str:
    return hashlib.sha256(canonical(value).encode()).hexdigest()


class Record(BaseModel):
    model_config = ConfigDict(extra="forbid", validate_assignment=True, allow_inf_nan=False)


class Modification(Record):
    position: int = Field(ge=1)
    residue: str = Field(min_length=1, max_length=1)
    ccd: str = Field(min_length=1)


class Molecule(Record):
    sequence: str
    modifications: list[Modification] = Field(default_factory=list)
    # Each crosslink lists residue position and atom name at both ends.
    bonds: list[tuple[int, str, int, str]] = Field(default_factory=list)
    n_terminus: str = "free"
    c_terminus: str = "free"
    smiles: str | None = None

    @model_validator(mode="after")
    def chemistry(self):
        if not self.sequence or any(a not in "ACDEFGHIKLMNPQRSTVWY" for a in self.sequence):
            raise ValueError("sequence must contain canonical uppercase parent residues")
        seen = set()
        for m in self.modifications:
            if m.position > len(self.sequence) or self.sequence[m.position - 1] != m.residue:
                raise ValueError("modification position/residue does not match the parent sequence")
            if m.position in seen:
                raise ValueError("duplicate modification position")
            seen.add(m.position)
        for i, a, j, b in self.bonds:
            if not (1 <= i <= len(self.sequence) and 1 <= j <= len(self.sequence)):
                raise ValueError("bond position outside sequence")
            if (i, a) == (j, b):
                raise ValueError("a bond requires two different atoms")
        return self

    @property
    def identity(self) -> str:
        value = self.model_dump()
        value["modifications"] = sorted(value["modifications"], key=lambda m: (m["position"], m["ccd"]))
        value["bonds"] = sorted(tuple(sorted(((i, a), (j, b)))) for i, a, j, b in self.bonds)
        return digest(value)[:20]


class Target(Record):
    id: str
    accession: str
    molecule: Molecule
    species: str = "Homo sapiens"
    compartment: str
    structure_paths: list[str] = Field(default_factory=list)
    # Structural files must be matched to this construct and chemical identity.
    construct_start: int = Field(default=1, ge=1)


class Requirement(Record):
    endpoint: str
    unit: str
    direction: Literal["ge", "le"]
    threshold: float
    assay: str
    scale: float = Field(gt=0)
    required: bool = True


class DesignSpec(Record):
    episode_id: str
    task: Literal["binding", "ptm", "ternary", "direction"]
    objective: str
    targets: list[Target] = Field(min_length=1)
    countertargets: list[Target] = Field(default_factory=list)
    requirements: list[Requirement] = Field(min_length=1)
    administration: str
    peptide_format: Literal["linear", "head_to_tail", "modified"] = "linear"
    sequence_lengths: list[int] = Field(default_factory=lambda: [12, 18, 24])
    # These endpoint definitions remain fixed within the comparison.
    budgets: dict[str, float] = Field(default_factory=lambda: {"tool_calls": 64, "gpu_minutes": 480,
        "controller_calls": 40, "controller_tokens": 128000, "synthesis_slots": 24})
    seed: int = 2027

    @model_validator(mode="after")
    def scientific_contract(self):
        ids = [t.id for t in self.targets + self.countertargets]
        if len(ids) != len(set(ids)):
            raise ValueError("target IDs must be unique")
        if self.task == "ptm" and not self.countertargets:
            raise ValueError("a PTM task requires matched countertargets")
        if self.task == "ternary" and len(self.targets) != 2:
            raise ValueError("a ternary task requires two protein partners")
        if len({r.endpoint for r in self.requirements}) != len(self.requirements):
            raise ValueError("endpoint names must be unique")
        if not self.sequence_lengths or any(x < 2 for x in self.sequence_lengths):
            raise ValueError("invalid peptide length")
        if any(v < 0 for v in self.budgets.values()):
            raise ValueError("budgets must be nonnegative")
        return self


class Measurement(Record):
    candidate_id: str
    endpoint: str
    value: float
    unit: str
    source_id: str
    kind: Literal["prediction", "experiment"]
    lower: float | None = None
    upper: float | None = None
    assay: str
    replicate_id: str | None = None
    censor: Literal["none", "left", "right"] = "none"

    @model_validator(mode="after")
    def bounds(self):
        if self.lower is not None and self.lower > self.value:
            raise ValueError("lower bound exceeds estimate")
        if self.upper is not None and self.upper < self.value:
            raise ValueError("upper bound is below estimate")
        return self


class Candidate(Record):
    molecule: Molecule
    parent_ids: list[str] = Field(default_factory=list)
    generator: str
    revision: str

    @property
    def id(self):
        return self.molecule.identity


class Evidence(Record):
    id: str
    source_url: str
    source_version: str
    retrieved_at: str
    passage: str
    claim: str
    target_ids: list[str]
    species: str
    assay: str
    molecular_format: str
    direction: Literal["supports", "opposes", "unresolved"]
    evidence_type: Literal["experimental", "computational", "review"]


class Decision(Record):
    """An explicit action proposal with a concise scientific explanation."""
    tool: str
    arguments: dict[str, Any]
    hypothesis: str
    evidence_ids: list[str]
    decision_summary: str
    expected_observation: str
    stop: bool = False


class ToolResult(Record):
    candidates: list[Candidate] = Field(default_factory=list)
    measurements: list[Measurement] = Field(default_factory=list)
    evidence: list[Evidence] = Field(default_factory=list)
    artifacts: dict[str, str] = Field(default_factory=dict)
    message: str = ""