File size: 6,628 Bytes
98bde72
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
"""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 = ""